DEV Community

Seth Wheeler
Seth Wheeler

Posted on Originally published at sethwheeler.dev

The Embedding Table Was 72% of the Model

An earlier experiment in this series had established something slightly deflating about a small transformer: quantising the whole network to int8 is free, and the bytes you save are better spent on count tables than on network precision. That is a useful result and it invites a sharper question, which is where the bytes actually were.

tok.weight (embedding)   1,024,000 params   71.7%
everything else            404,992 params   28.3%
Enter fullscreen mode Exit fullscreen mode

Nearly three quarters of that "neural network" is a lookup table. So the dial worth turning is not the model's precision, it is the embedding's precision, quantised independently of the transformer body. The research repo is not public, so what follows comes from the experiment's own sweep and its recorded results rather than from anything you can clone.

Two axes, and only one of them is the obvious one

The sweep varies two things. Width is int8, int4, int3 and int2 on the embedding, with the body held at int8. Granularity is one scale for all 8,000 rows, which is what the earlier experiment used, against one scale per token.

Per-row scales cost V extra fp16 values, which is 16 KB, under 2% of the table. They should matter more for an embedding than for a weight matrix. Token embedding norms differ by orders of magnitude between frequent and rare words, so a single tensor-wide scale spends most of its levels on a handful of outlier rows; the rest of the vocabulary shares a coarse grid.

Everything is reported on the same evaluation as the earlier frontier work, with the count table and document cache mixed in at their tuned weights, so the numbers land on a comparable curve. No retraining anywhere.

embedding / body MB network alone mix top-1 OOD top-1
fp32 / fp32 5.72 0.351 0.384 0.185
int8 per-tensor 1.43 0.350 0.381 0.185
int8 per-row 1.44 0.350 0.381 0.185
int4 per-tensor 0.92 0.325 0.366 0.185
int4 per-row 0.93 0.347 0.384 0.185
int3 per-row 0.80 0.313 0.371 0.177
int2 per-row 0.68 0.110 0.276 0.143

Granularity is the entire reason int4 works

int4 per-tensor loses 2.5 points of network top-1, from 0.350 down to 0.325. int4 per-row loses 0.3, landing at 0.347; those 16 KB of extra scales, 1.7% more bytes, recover 2.2 of the 2.5 points.

At int8 the two granularities are identical, both 0.350. The distinction does not exist until quantisation levels get scarce, which is why it is easy to miss. You test granularity at the width where it does not matter; you find no effect, and drop it from the sweep.

That is the part worth taking away. A sweep over width alone would have measured int4 at 0.325, concluded that four bits is where this model breaks, and stopped at int8. The conclusion would have been wrong and it would have looked well-supported, because the width axis really does show a cliff; it is just not at four bits. The cliff is between four and three: 0.371 then 0.276 in the mixture. At four bits with per-row scales there is nothing to pay. The mixture scores 0.384, which is fp32's 0.384 exactly, at 0.93 MB against 5.72 MB. A 6.2x reduction for zero measurable loss.

Two axes where one only reveals itself in the tail of the other is not a rare shape, and the only defence I know is to sweep the cross product when it is cheap. Here it was seven runs.

The budget arithmetic, done honestly

Those mixture numbers use the full count table, about 25 MB, so none of the rows above describe a small system. The follow-up redoes the allocation properly: int4-per-row as the network, and the largest pruned count table that fits in whatever the budget leaves.

budget system actual MB in top-1 OOD top-1
0.75 MB counts only 0.68 0.289 0.136
0.75 MB int4 hybrid network does not fit
1.0 MB counts only 0.68 0.289 0.136
1.0 MB int4 + pruned counts 0.96 0.356 0.175
1.5 MB counts only 1.28 0.305 0.140
1.5 MB int4 + pruned counts 1.28 0.366 0.175
2.0 MB counts only 1.28 0.305 0.140
2.0 MB int4 + pruned counts 1.61 0.372 0.179

The earlier experiment had reported that at 1 MB nothing but counts fits, giving 0.289. A real 1 MB system now scores 0.356, which is 6.7 points better in distribution and 3.9 better out of it. That 0.96 MB system also beats the full-precision 5.72 MB network outright, 0.356 against 0.351, at six times fewer bytes.

I want to be careful about what that comparison is and is not. It is not "quantisation improved the model." It is that the bytes freed by quantising bought count table and cache, and those were worth more than the precision they replaced. The network alone got slightly worse (0.351 to 0.347). The system got better.

The thing I did not expect

The best 1 MB configuration prunes the count table down to contexts seen at least 400 times, which is a 24 KB table, and weights the document cache at 0.40.

At that budget, almost all of the non-neural contribution is coming from statistics of the document currently being processed, which cost nothing to store because they are derived from text that is already in memory. An earlier finding in this series was that the cache earns a higher weight the tighter the budget gets. This sharpens it: at the tightest budgets the cache is the only count model you can afford, and it is enough.

Out of distribution the story is flatter and it makes the same point from the other side. OOD top-1 sits at 0.185 all the way down through int4 per-tensor, and only moves at int3. Out of distribution the mixture's accuracy is dominated by counts and cache, so the network's precision barely registers, which is consistent with everything else in this series: the cheap components matter most exactly where the trained one is weakest.

What this does not show

Quantisation shrinks a table that was trained at full width. A factorised embedding attacks the same 72% by never materialising the wide table at all, and could plausibly go below int4-per-row's 0.93 MB floor, but it requires training, so it is a separate experiment. The open question there is whether factorisation and 4-bit quantisation compose, or whether a low-rank bottleneck has already removed the redundancy that per-row scaling was exploiting. I would guess the latter. Guessing is all that is on offer until it is run.

Top comments (0)