Extreme Compression of Per-Layer Embeddings via Cross-Layer Latent Vector Quantization

Paper Recent language models keep increasing the amount of information packed into a single vocabulary entry. Per-layer embeddings (PLE), first introduced in Gemma 3n, carried over to Gemma 4: instead of a single input-layer vector per vocabulary item, each decoder layer has a different learned vector. This post examines OPTIC, which shrinks this table by about 100× while retaining most of the original performance, starting from the problem setting and proceeding through a small example, the method, results, and limitations.

OPTIC’s core ideas are two. One is to store a single code per token and have each layer unfold that code with its own decoder; the other is to fit the codes and decoders not by “how faithfully the stored vectors are reconstructed” but by “how much difference appears in the output path the model actually uses.”

Why PLE is a problem

A typical input embedding turns one token into one vector. For the token "cat", it simply retrieves the single vector that the token points to in the vocabulary dictionary. PLE adds a layer axis on top of this. Even for the same token, the vector used at layer 1 differs from the vector used at layer 42, and each layer retrieves its vector from its own stage’s residual connection. The residual connection is the pathway that adds a term to the hidden state as it passes through each layer.

Storage grows immediately. E4B has a vocabulary of 262,144, 42 decoder layers using PLE, and a vector width of 256 dimensions, so storing a separate vector for each token-layer pair results in 11 million token-layer vectors (§1, §2.2). In bfloat16 that is about 5.64GB (§2.2). In the same checkpoint, PLE accounts for 35.2% of parameters (abstract). In the smaller E2B it is 45.8% (§1).

Illustrative example: let’s see what a single token actually stores and does (E4B).

Compressing this table by about 100× leaves just over 5 bytes per vector (§1). Existing quantization methods spend those few bytes on “a code per vector.” With 11 million vectors each having its own code, the codes become far too short and reconstruction quality drops sharply.

How OPTIC stores and reconstructs

The problem is clear. 5 bytes is not enough to properly hold a single vector. The PLE table has two directions, the token axis and the layer axis, and OPTIC shares both directions at once (§2.2).

• Token direction: instead of each token storing its 42 vectors separately, each token stores one code.
• Layer direction: all tokens in a layer share a single affine decoder (matrix multiply plus addition). The size of this decoder is independent of vocabulary size.

Concretely, token t stores M indices. Each index points to a 4-dimensional codeword in its own codebook, and concatenating the M of them yields a latent vector z_t that summarizes the token. Then layer ℓ unfolds z_t into a 256-dim row with its own affine decoder.

In the ~100× setting for E4B, M=161 and 256 codewords (K=256) are used. A single token needs to store only 161 indices, i.e. 161 bytes, reducing 21,504 bytes of 42 vectors to 161 bytes (§2.2). In exchange, the codebook (about 330KB) is shared fully, and the decoder weights (~13.8MB) and biases of the 42 layers are shared by all tokens in the respective layer. Total is 56.4MB, 99.94× versus the original 5.64GB (§2.2, Table 16).

Reading the figure above as a state transition: the input is a single ID for token t. The stored state is the M indices placed in vocabulary row t (e.g., 161 bytes). Each index q_{tm} selects one row (highlighted cell) of its own codebook C_m, and concatenating the 161 selected codewords computes the latent vector z_t (length 644). At this point the layer does not intervene yet. Then layer ℓ applies its own affine weight A_ℓ and bias b_ℓ to compute the 256-dim row ê_{tℓ}, and that row leads to the PLE residual of layer ℓ. Cells of the same color indicate the same layer, and the dashed ellipsis indicates omitted layers and codebooks. Importantly, the original PLE is used only in the “fitting step” and not for actual storage and reconstruction (original Figure 1).

The key point is that the cost of storing indices is now proportional only to vocabulary size V and not to V × L. Conventional row-wise quantization must store an index set per vector, which multiplies by the number of layers, whereas OPTIC removes this layer factor (§2.2, Equation 3).

Why fit on the “output path”

Sharing storage alone is not the end. PLE vectors are not used on their own; at each layer they pass through the gate and output projection and are added to the hidden state. Specifically, the PLE row is multiplied elementwise by a gate derived from the layer’s hidden state, then projected to the hidden width by the output projection matrix W_out (§2.3).

Therefore, what matters more directly to model quality is not how similar the reconstructed vector is “in coordinates” to the original, but how much difference remains after passing through the gate and output projection. OPTIC fits the codes and decoders to minimize the error in this projected output rather than in the stored coordinates (§2.3, Equation 4–5).

Here g_{iℓ} is the gate at calibration position i and layer ℓ, τ is a fixed PLE scale, and ⊙ is elementwise multiplication. Gates from the training records are collected from the pre-compression model and used fixed. Gates for synthetic records are filled with 1 (§3.2). The objective is the mean of this projected output error ‖Δr‖² over positions and layers (§2.3, Equation 5).

The fitting procedure alternates two kinds of updates (§2.4, Algorithm 1). One updates the codebooks and decoders continuously (with AdamW) while fixing the codes; the other reassigns each token’s index to the codeword with the smallest error while fixing the codebooks. Reassignment covers all tokens in the vocabulary. To handle rare tokens that do not appear often in natural-language calibration, it uses a per-layer averaged gate metric G_ℓ as a surrogate objective (§2.4, Equation 6–7). Initialization concatenates per-token vectors, subtracts the vocabulary mean, and builds the latent vectors and decoders via PCA scaled by the output-path weights, then starts with an OPQ-style assignment that splits coordinates into codebook groups and runs K-means within each group.

Results: quality preserved, memory reduced

Quality first. The paper reports the geometric mean of “retention relative to the original” across 7 tasks (PIQA, GSM8K, GPQA Diamond, AIME 2026, MMLU-Redux, IFBench, InfographicVQA). Each task’s retention is capped at 100%, so over-performance on one task cannot mask losses on another. If any single task scores 0, the geometric mean becomes 0 (§3.3, Equation 8).

Question: at ~100× compression, how well do existing quantization methods and OPTIC hold up? Below are the E4B results (Table 1).

On BF16 E4B, OPTIC leads the strongest existing method (PQ) by 31.5 percentage points, and on W4A16 it leads the strongest method (AQLM) by 14.0 percentage points (§4.1). On E2B, OPTIC retains 89.1% on BF16 and 91.5% on W4A16, while the frequency-tiered variant is highest at 90.0% and 93.0% (§4.1, Table 1). The frequency-tiered variant gives more code groups to frequent tokens, uses learned shared codewords per group for the remaining tokens, and preserves exactly the PLE rows of special tokens in bfloat16 (§2.5, Appendix A.3).

Looking at actual task scores shows where the difference comes from. On BF16 E4B, GSM8K is 92.8% for OPTIC versus 92.9% original and 73.5% for PQ. IFBench is 41.7% for OPTIC versus 44.0% original and 16.0% for PQ, and InfographicVQA is 67.5% for OPTIC versus 69.7% original and 38.4% for PQ (§4.1, Table 12). The drop for existing methods is especially large in mathematical reasoning, instruction following, and visual question answering. Bootstrap 95% confidence intervals also separate. On BF16 E4B, OPTIC is 95.5–99.0% and PQ is 60.6–70.5%. On W4A16 E4B, OPTIC is 93.9–97.8% and AQLM is 76.8–87.1% (Table 2). These intervals reflect benchmark sampling uncertainty.

OPTIC holds up even at higher compression rates. The curve below plots retention (y-axis) against compression rate (x-axis, log scale).

Dashed lines are existing methods (PQ, OPQ-P), solid lines are the authors’ method. OPTIC retains 97.8% at ~100×, and still holds 90.5% and 84.1% at 594× and 1002×. W4A16 E4B at the same points is 87.7% and 81.8% (§4.2, Table 16). In contrast, output-path weighted PCA-PQ, which improves only row-wise, is good at 93.3% near 100× but falls to 25.1% near 500× and reaches 0 near 1000× (§4.2). The fewer bytes left per vector, the larger the benefit of sharing codes at the token level.

Kernel speed and system performance are measured separately

Beyond quality, how memory and time behave in actual serving is a separate measurement. The authors compared original PLE and ~100× compressed PLE in vLLM on a single RTX 4090 under identical settings (fixed 1GiB KV cache) (§6).

Memory drops substantially from a system perspective. For BF16 E4B, process GPU memory falls from 18.57GB to 13.05GB, a 29.7% reduction, and for W4A16 it falls from 12.95GB to 7.37GB, a 43.1% reduction (Table 7). This is because the PLE resident on the GPU shrinks from 5.64GB to 56.4MB.

Time must be read from both a kernel and a system perspective. In single-lookup measurements where the table is in host memory, dense lookup is faster for a single token (compressed lookup is 0.6×), but from 32 tokens compressed lookup becomes faster, reaching 12.1× at 8192 (Table 11). In terms of total serving time, compressed PLE is only up to 0.40% slower on BF16 and up to 0.69% slower on W4A16 (Table 9). Compressed lookup must perform extra multiplications to unfold codes with the decoder, so under the same cache budget the conclusion is “memory drops a lot, time stays about the same.”

However, situations where memory is the real bottleneck are different. Under the same GPU budget, there are cases where dense PLE either fails to start because it does not fit in the KV cache or processes requests sequentially, while compressed PLE cuts completion time by up to 2.1× (§B.3, Table 10). In other words, the goal is not saving serving time per se; the gain appears when memory becomes latency.

Why this works

Start with the analysis showing that layer averaging alone is insufficient. Replacing each vector wholesale with its layer’s vocabulary mean drops GSM8K from 92.9% to 43.5%, GPQA from 56.9% to 23.0%, and AIME to 0, even though the mean vector accounts for only 8.1% of the table’s total squared norm (§5.1, Appendix E.3). In other words, the information that differs per token is the core of performance.

So is that per-token variation concentrated in low dimensions? The next figure answers this question.

The top panel shows stored coordinates and the bottom panel shows coordinates after output projection; each curve (layer) shows how many principal components explain the per-token variation (x-axis: number of principal components, y-axis: cumulative explained ratio). In stored coordinates, 212–227 directions are needed to explain 90%. After output projection, this drops to 1–205 per layer (median 158) (§5.1, Figure 3). This gap is the basis for OPTIC’s choice to “fit on the output path.” After the gate and projection, important variation is concentrated substantially toward lower dimensions, a signal that a low latent width can be tolerated to some degree.

There are also experiments that isolate and validate the two designs. One is the effect of code sharing. If each token-layer vector gets an independent latent code and the same fitting procedure is applied, E4B BF16 retention drops from 97.8% to 79.3% at nearly the same storage budget (Table 3). The caveat that this difference mixes the effect of sharing itself with the effect of reallocating bytes between codes and decoders is revisited in the limitations section. Decoder freedom also matters. If only the decoder is refit by least squares with indices, codebooks, and biases fixed while limiting rank to 20 per layer, retention drops from 95.09% to 70.55% (§5.2, Table 4).

The other is the effect of output-path fitting. With codes and initialization fixed and only the objective changed, the context-dependent output-path objective gives 1.98 percentage points higher retention on average than Euclidean error and 0.31 percentage points higher than the diagonal approximation (§5.3, Table 5). Also, continuous updates produce most of the gain (63.60% with initialization only, 94.67% with continuous updates), and index reassignment adds 0.65–1.14 percentage points on top (§5.3, Table 6).

Limitations and applicability

To decide whether to use OPTIC in practice, one must separate the limitations the authors state from other conditions.

Start with the limitations the authors state explicitly. Evaluation is limited to Gemma 4 E2B/E4B, and the main comparison uses a single fixed fitting seed, so confidence intervals reflect only benchmark sampling uncertainty (Limitations). The six baselines are reconstructions of the original papers’ methods adapted to PLE, and cannot be read directly against results from their original settings (Limitations). The code-sharing ablation also mixes in the effect of reallocating bytes between codes and decoders, so it does not measure only the “pure sharing effect” (Limitations). Calibration adds synthetic records (gate = 1) covering the full vocabulary to 786,432 natural-language positions (§3.2), so tokens that rarely appear in natural language are fitted with synthetic records alone. The output-path loss is also an approximation that measures error only at each layer’s projected output with gates fixed (Limitations).

From a reviewer’s perspective, three things are worth noting. First, the method is specialized for embeddings with a “token×layer” structure like PLE, so the benefit does not carry over directly to ordinary single input embeddings, which lack the layer axis for code sharing. Second, retention is a geometric mean over 7 tasks, and any single task scoring 0 makes the whole thing 0, so if a low-sample task collapses the metric swings sharply. One should check individual scores to see whether a particular task collapses under extreme compression (1000×), but the paper does not report per-task scores for extreme compression (Tables 12 and 13 contain only the 100× condition; those numbers are not reported in the original). Third, the serving gains above depend on a memory-bottleneck condition, so there is less to expect in environments without it.

In conclusion, OPTIC is promising for settings where “the embedding table makes up a large part of the model and GPU memory is constrained,” especially when serving PLE-family models with many concurrent requests. It can also be combined with W4A16 weight-only quantization. The perspective of fitting where the model actually consumes the vectors rather than where they are stored shows that embedding capacity and compressed representation can be designed together from the start (Conclusion).

Figures in this post are taken from the original arXiv:2610.31309 (CC BY 4.0). Only size and format were changed.