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

AI summary

Overview

Research area: Machine learning — model compression and efficient inference for large language models, specifically compressing the large per-layer embedding (PLE) tables that recent models add alongside their Transformer computation.

Technical level: Advanced. The paper assumes familiarity with vector quantization (product quantization, residual/additive quantization), PCA, and Transformer internals such as gated residual branches and output projections.

Scope: The paper introduces OPTIC, a compression scheme that stores one shared code tuple per vocabulary token and decodes it into a separate embedding for each of the model's layers, and it evaluates that scheme on Gemma 4 E2B and E4B at compression ratios from roughly 100× to over 1,000×.

What This Paper Is About

Recent language models scale embedding capacity as a separate axis from computation: Gemma 4's Per-Layer Embeddings give every vocabulary item a distinct learned vector at every decoder layer, and the paper states that PLE accounts for 35.2% of E4B's parameters and 45.8% of E2B's. Extreme compression of that table leaves only a few bytes — or less than one byte — per token–layer vector, and conventional quantizers, which spend a separate code on each vector, degrade sharply at that budget. The goal is a representation that keeps model quality near the uncompressed checkpoint while shrinking PLE storage by two to three orders of magnitude and still allowing the model to be served directly from the compressed form.

Key Contributions

1. A cross-layer sharing scheme for PLE. OPTIC (Output-Path-fitted, Token-shared Indexed Codes) stores one tuple of codebook indices per token rather than one per token–layer pair, and reconstructs each layer's row with a layer-specific affine decoder. This removes the layer factor from the code-storage term, while decoder storage grows with the number of layers but not with vocabulary size.

2. An output-path fitting objective. Instead of minimizing error in the stored coordinates, OPTIC fits its codes and decoders to minimize the error at the PLE branch's projection output, after the gate and output projection (with gates and projections frozen). An aggregated per-layer metric makes this tractable for code reassignment over the whole vocabulary.

3. A full fitting procedure. An alternating optimization recipe: output-path-weighted PCA initialization, OPQ-style eigenvalue allocation to assign latent coordinates to codebook groups, K-means codebook initialization, AdamW continuous updates, and discrete code reassignment with handling for unused codewords.

4. A frequency-tiered variant and a compact inference path. A variant giving more code groups to frequent tokens and storing special-token rows exactly in bfloat16, plus a serving implementation that decodes only the rows of the current tokens rather than materializing the full PLE tensor.

Main Findings

• Large quality advantage at ~100× compression: On BF16 E4B, OPTIC retains an aggregate 97.8% of the uncompressed model's seven-benchmark performance, versus 66.3% for product quantization, the strongest of six conventional baselines — a 31.5 percentage-point lead. On the official W4A16 release, OPTIC retains 96.9% versus 82.8% for the AQLM Euclidean beam encoder, a 14.0-point lead.

• Holding up at much higher ratios: At 594× and 1,002× PLE compression of BF16 E4B, OPTIC retains 90.5% and 84.1%; on W4A16 E4B it retains 87.7% and 81.8% at the same nominal tiers.

• Row-wise output-aware compression degrades faster: Output-path-weighted PCA-PQ, the authors' row-wise method, retains 93.3% near 100× on BF16 E4B, 25.1% near 500× and zero near 1,000×, while OPTIC stays above 84%.

• Smaller model results: On E2B, OPTIC retains 89.1% on BF16 and 91.5% on W4A16; the frequency-tiered variant gives the highest E2B retention at 90.0% (BF16) and 93.0% (W4A16).

• Sharing codes across layers is what matters: With nearly equal storage and the same fitting recipe, independent latent codes for each token–layer vector retain only 79.3% on BF16 E4B (86.8% on W4A16), versus 97.8% and 96.9% for OPTIC. Output-path-fitted additive codes retain 90.0% on BF16 E4B and 94.4% on W4A16 E4B.

• Decoder rank matters: Refitting only decoder weights on fixed codes, capping each layer decoder at rank 20 drops retention from 95.09% (± 0.48) to 70.55% (± 1.18) over fit seeds 0, 1 and 2.

• The output-path objective helps, and most of it comes from diagonal weighting: With fixed codes and 5,000 continuous updates, context-dependent output-path fitting reaches 94.74% (± 0.59), versus 92.76% (± 1.16) for Euclidean fitting and 94.44% (± 0.94) for diagonal output-path weighting — gains of 1.98 and 0.31 percentage points.

• Both optimization stages contribute: From a common Euclidean initialization, initialization alone gives 63.60% (± 3.68), continuous updates only give 94.67% (± 0.53), and continuous updates plus 18 reassignment sweeps give 95.49% (± 0.40).

• Per-task behavior on BF16 E4B: OPTIC preserves GSM8K accuracy at 92.8% versus 92.9% uncompressed and 73.5% for PQ; IFBench at 41.7% versus 44.0% uncompressed and 16.0% for PQ; InfographicVQA at 67.5% versus 69.7% uncompressed and 38.4% for PQ. The paper states PIQA changes little for either compressed method.

• Confidence intervals separate the methods: Paired-bootstrap intervals for E4B BF16 are 95.5–99.0% for OPTIC and 60.6–70.5% for PQ; for E4B W4A16, 93.9–97.8% versus 76.8–87.1% for the AQLM encoder. E2B intervals are 89.1 [85.8, 91.8] and 91.5 [89.3, 93.3] for OPTIC.

• Real memory savings at serving time: On one RTX 4090 with a fixed 1 GiB key–value cache and vLLM, BF16 E4B drops from 18.57 GB to 13.05 GB (29.7% saved) and W4A16 from 12.95 GB to 7.37 GB (43.1% saved); resident PLE storage falls from 5.64 GB to 56.4 MB. Compact lookup adds at most 0.40% wall time with BF16 and 0.69% with W4A16.

• Layer means alone are insufficient: The per-layer vocabulary means account for 8.1% of the table's squared Frobenius norm, but replacing every vector with its layer mean lowers GSM8K from 92.9% to 43.5%, GPQA Diamond from 56.9% to 23.0%, and drops AIME 2026 to zero.

• Token-specific variation is high-dimensional in stored coordinates but less so at the output: Capturing 90% of a layer's token-specific variation needs 212–227 of 256 principal directions in stored coordinates, but 1–205 directions (median 158) under the output-path metric, with substantial variation across the 42 layers.

Methodology in Plain English

The PLE table is a three-dimensional array: one vector of width 256 for every combination of vocabulary item and decoder layer. For E4B that is 11 million vectors across a vocabulary of 262,144 tokens and 42 layers, occupying 5,637,144,576 bytes in bfloat16.

Rather than compress each of those 11 million vectors separately, the authors exploit the fact that the table has two indices. They give each token a single compact code: a short list of indices into small codebooks, whose selected entries are concatenated into one latent vector of width Z. Each layer then gets its own small affine map (a matrix plus a bias) that turns that shared latent into that layer's 256-dimensional embedding. Because a layer's decoder is shared by all tokens, its cost does not grow with vocabulary size; because a token's code is shared by all layers, the code payload does not grow with layer count. In the ~100× E4B configuration, Z = 644, there are 161 codebook groups with 256 codewords each and codeword width four, totaling 56,404,992 bytes — a 99.94× ratio.

Two design choices make this work at extreme ratios. First, the codes and decoders are not fit to reproduce the stored vectors as such. The model only ever uses a PLE vector by multiplying it element-wise with a gate computed from the hidden state, projecting it through a frozen output matrix, and adding it as a residual. So the fitting loss measures the error at that projection output, with gates fixed. Second, training alternates between gradient updates to the codebooks and decoders (codes held fixed) and discrete reassignment of each token's codebook indices. Because reassignment must cover the whole vocabulary, including tokens barely seen in calibration text, the gate information is averaged into one per-layer matrix, giving a weighted squared-error objective that can be evaluated for every token; indices are updated one group at a time, and unused codewords are replaced using the worst-fitting tokens. Initialization comes from PCA of the per-token concatenated layer vectors, scaled by those per-layer matrices, with coordinates assigned to codebook groups via OPQ's eigenvalue allocation and codewords set by K-means.

Calibration mixes natural text — 1,536 sequences of 512 tokens (786,432 positions) from FineWeb, FineWeb-Edu, OpenWebMath and FineWeb2 — with one synthetic record per vocabulary token using all-ones gates, so the full vocabulary is covered. Evaluation is a capped, equal-weight geometric mean of per-task score ratios across seven benchmarks (PIQA, GSM8K, GPQA Diamond, AIME 2026, MMLU-Redux, IFBench, InfographicVQA), with capping so gains on one task cannot offset losses on another; a single zero task score makes the whole retention zero.

Why This Matters

The paper shows that a large, structurally regular parameter block can be compressed by two to three orders of magnitude with only a few points of benchmark loss, provided the compressor exploits the structure of the table (the token × layer grid) and the structure of how the model consumes it (a gated, projected residual branch). It also shows PLE compression stacks with 4-bit weight quantization rather than competing with it, since the W4A16 retention figures remain high.

Real-world applications:

• On-device and edge deployment: a 29.7% (BF16) or 43.1% (W4A16) reduction in GPU memory at a fixed key–value cache size on a single RTX 4090 makes larger checkpoints fit on consumer or single-GPU hardware.
• Serving many models per machine: shrinking resident PLE from 5.64 GB to 56.4 MB frees memory that can be spent on key–value cache, more concurrent requests, or additional models.
• Distributing models over constrained networks: a compressed embedding table is far cheaper to ship and store, which matters for models whose embedding capacity, not their compute, dominates size.
• Research on scaling embedding capacity: if large embedding tables are cheap to store, the trade-off between parameters spent on embeddings versus on Transformer computation changes, which the authors argue should be designed jointly rather than after the fact.

Industry relevance: the results are directly aimed at inference economics — memory footprint and wall-clock overhead on standard serving stacks (vLLM) — and apply to a publicly released model family in both its original BF16 and officially quantized W4A16 forms, which is the format most deployments actually use.

Future Directions

• Extending the scheme beyond Gemma 4's PLE: the study covers only E2B and E4B; the conclusion names extending the combination of PLE compression and backbone quantization to other embedding structures as a promising direction.
• Co-designing embedding capacity with compression: the authors argue for designing embedding capacity and compressed representations together, using deployment memory as a constraint rather than compressing a table designed without that constraint.
• Better decoders at fixed code storage: the rank-20 decoder control (70.55% versus 95.09%) indicates decoder capacity is a live lever; how to allocate bytes between codes and per-layer decoders at a fixed budget is left open.
• Stronger fitting objectives: the output-path loss holds gates fixed and measures error only at each layer's projection output; richer or adaptive objectives, and comparisons across multiple fitting seeds (the main comparisons use seed 0), are natural follow-ups.

Target Audience

Researchers and engineers working on model compression, quantization, and LLM inference systems; practitioners deploying Gemma-family or similar models under memory constraints; and readers interested in vector-quantization methods (product, residual, additive quantization) applied to unusual tensor structures. The paper is written for an advanced audience comfortable with quantization mathematics and Transformer architecture, though the core idea — share one code across layers, decode per layer, and fit at the point where the model actually uses the vector — is accessible to intermediate readers.

Authors’ abstract

Scaling embedding capacity alongside Transformer computation is an emerging trend in language models. This trend makes high compression rates increasingly useful. Hence, we study extreme compression of Per-Layer Embeddings (PLE), which follow this trend and account for 35.2% of the parameters of Gemma 4 E4B. At such rates, conventional quantizers give each token-layer vector its own code of only a few bytes. We instead develop OPTIC (Output-Path-fitted, Token-shared Indexed Codes), which stores one quantized code tuple per token and decodes it at every layer with a layer-specific affine decoder. Codes and decoders are fitted to the PLE output path rather than to the stored vectors. At about $100\times$ PLE compression of bfloat16 (BF16) E4B, OPTIC retains an aggregate 97.8% of the uncompressed model's performance across seven benchmarks, whereas product quantization, the strongest of six conventional baselines, retains 66.3%. At $594\times$ and $1,002\times$ PLE compression of BF16 E4B, OPTIC retains 90.5% and 84.1%. On the official E4B release with 4-bit linear weights and 16-bit activations (W4A16), OPTIC retains 96.9% at about $100\times$ PLE compression. During inference at about $100\times$ PLE compression, BF16 E4B uses 29.7% less GPU memory than with uncompressed PLE at the same key-value cache size, while running at most 0.40% slower.