PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiacheng Liu, Yaxin Luo, Xiaohan Zhao, Peng Sun, Xinyi Shang, Jiacheng Cui, Yonina C. Eldar, Zhiqiang Shen

arXiv:2610.31309v1cs.LG

TL;DR

Per-Layer Embeddings (PLE) are a substantial compression target in Gemma 4, where high compression rates are useful. OPTIC shares token codes across layers and fits them to the PLE output path; at about 100× compression, it retains 97.8% of BF16 E4B performance across seven benchmarks, versus 66.3% for the strongest conventional baseline.

  • Problem

    PLE accounts for 35.2% of Gemma 4 E4B parameters, making extreme compression useful as embedding capacity expands.

    Source: §1 · p. 1

  • Method

    OPTIC stores one code tuple per token, decodes it with a layer-specific affine decoder, and fits codes and decoders to minimize error at the projection output.

    Source: §1 · p. 1

  • Results

    At about 100× PLE compression, OPTIC retains 97.8% on BF16 E4B across seven benchmarks, versus 66.3% for the strongest of six conventional baselines.

    Source: §1 · p. 2, §8 · p. 8

  • Takeaways & Limitations

    At about 100× PLE compression, BF16 E4B serving uses 29.7% less GPU memory with at most 0.40% added wall time at a fixed key–value cache size.

    Source: §1 · p. 2, §8 · p. 8

    The study evaluates Gemma 4 E2B and E4B, and its main comparisons use one fixed fitting seed, so paired-bootstrap intervals reflect benchmark sampling uncertainty.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Gemma 4’s Per-Layer Embeddings make up a substantial share of model parameters, while extreme compression leaves conventional quantizers only a few bytes per token-layer vector. OPTIC shares token codes across layers and fits them to the PLE output path, retaining more benchmark performance than conventional baselines at high compression.

1 Introduction

  • PLE accounts for 45.8% of Gemma 4 E2B parameters and 35.2% of E4B parameters, making it a substantial storage-reduction target.Source: §1 · p. 1
  • At near 100× compression, E4B’s 11 million token-layer vectors of width 256 leave about five bytes per vector for conventional quantizers.Source: §1 · p. 1
  • OPTIC stores one code tuple per token, decoded by a layer-specific affine decoder, so codes are shared across layers while decoders preserve layer differences.Decoder storage does not grow with vocabulary size.Source: §1 · p. 1
  • OPTIC fits codes and decoders to reconstruction error at the PLE projection output after the gate, rather than only to stored-vector error.Source: §1 · p. 1
  • At about 100× PLE compression, OPTIC retains 97.8% on BF16 E4B across seven benchmarks, versus 66.3% for the strongest of six conventional baselines.At 594× and 1,002× compression, OPTIC retains 90.5% and 84.1%; on W4A16 E4B it retains 96.9% at about 100× compression.Source: §1 · pp. 1–2
  • On an RTX 4090 with a fixed key–value cache size, compact PLE lowers GPU memory for BF16 E4B serving by 29.7% and adds at most 0.40% wall time.Source: §1 · p. 2

2 Method

The method compresses the pretrained PLE tensor under a byte budget using token-shared codes and layer-specific decoders. It fits reconstructions along the gated projection path, alternates continuous and discrete updates, and also describes two additional compact representations.

2.1 Problem Setup

  • The setup treats each token-layer entry of the pretrained PLE tensor as a width-d row and compresses it while keeping remaining model parameters fixed.A byte budget constrains the compact representation.Source: §2.1 · p. 2

2.2 Shared Token Codes

  • OPTIC assigns one code tuple to each token across layers, while a shared affine decoder per layer reconstructs that layer’s row.Sharing removes the layer factor from token-code storage; decoder storage grows with layers but not vocabulary size.Source: §2.2 · p. 2

2.3 Output-Path Fitting

  • The fitting target is reconstruction error at the gated PLE branch’s projection output, with fitting-record gates held fixed.The figure depicts the shared latent reconstructed through layer-specific affine decoders; the original PLE is used only for fitting.Source: §2.3 · pp. 2–3, Figure 1 · p. 3

2.4 Alternating Optimization

  • Alternating optimization updates codebooks and decoders continuously, then reassigns token codes using a surrogate error metric with layer gate information averaged across fitting records.Initialization uses weighted PCA, OPQ eigenvalue allocation, and within-group K-means.Source: §2.4 · p. 3

2.5 Additional Compact Representations

  • Output-path-weighted PCA-PQ keeps separate codes per token-layer row, while frequency-tiered OPTIC gives frequent tokens additional code groups and uses shared codewords for others.The frequency-tiered variant also uses INT8 decoder weights and stores special-token rows exactly in bfloat16.Source: §2.5 · p. 4

3 Experimental Setup

The experiments compare OPTIC and output-path-weighted PCA-PQ with six quantization baselines under matched resident-byte accounting. Calibration combines natural and synthetic token records, and evaluation uses seven benchmarks with a capped, equal-weight geometric-mean retention score.

3.1 Baselines

  • Six adapted baselines, including PQ, RVQ, OPQ-P, AQ, AQLM, and INT8 randomized low-rank SVD, are compared with output-path-weighted PCA-PQ and OPTIC under matched resident-byte accounting.The fitting and representation rules are specified in Appendix A.2.Source: §3.1 · p. 4

3.2 Calibration Data

  • Calibration combines 1,536 natural-text sequences with synthetic records covering every vocabulary token, including repeated special-token records.Natural records use frozen-model gates; synthetic records use all-ones gates, and evaluation examples are excluded from calibration.Source: §3.2 · p. 4

3.3 Evaluation Protocol

  • Seven benchmarks cover reasoning, world knowledge, instruction following, and vision, with methods sharing backbone weights, prompts, generation settings, and scorers.Retention is the equal-weight geometric mean of capped per-task score ratios; capping prevents above-reference gains on one task from offsetting losses on another.Source: §3.3 · pp. 4–5
  • W4A16 compact PLE artifacts are fitted from matched unquantized checkpoints and applied to official releases whose PLE lookup remains in bfloat16.The source checkpoints come from the official quantization-aware-training pipeline.Source: §3.3 · p. 5

4 Compression Results

Near 100× compression, OPTIC leads E4B retention on both BF16 and W4A16; at more extreme compression, it retains substantial performance and its advantage over row-wise output-path-weighted PCA-PQ widens.

4.1 Compression and Quality

  • 97.8% BF16 and 96.9% W4A16 E4B retention near 100× are OPTIC’s highest results, leading the strongest baseline by 31.5 and 14.0 percentage points, respectively.The frequency-tiered OPTIC variant leads E2B with 90.0% BF16 and 93.0% W4A16 retention.Source: §4.1 · p. 5
  • On BF16 E4B, OPTIC retains 92.8% GSM8K accuracy versus 92.9% for uncompressed PLE and 73.5% for PQ, and 41.7% IFBench versus 44.0% and 16.0%.The aggregate advantage also includes gains in mathematical reasoning, instruction following, and visual question answering.Source: §4.1 · p. 5
  • 67.5% InfographicVQA accuracy for OPTIC compares with 69.7% for uncompressed PLE and 38.4% for PQ.PIQA changes little for either compressed method.Source: §4.1 · p. 5
  • 95% paired-bootstrap intervals for E4B remain separated between OPTIC and PQ in BF16 and W4A16, and intervals remain separated in both E2B settings.For BF16 E4B, intervals are 95.5–99.0% for OPTIC and 60.6–70.5% for PQ; for W4A16 E4B, they are 93.9–97.8% and 76.8–87.1%.Source: §4.1 · p. 5

4.2 Compression Range

  • 90.5% and 84.1% BF16 E4B retention at 594× and 1,002× compression contrasts with output-path-weighted PCA-PQ’s 25.1% near 500× and zero near 1,000×.OPTIC remains above 84% as fewer bytes represent each PLE vector.Source: §4.2 · p. 5

5 Analysis and Ablations

The analysis examines PLE structure, cross-layer code sharing, and fitting components. It finds that token-specific variation matters, shared codes retain more performance than independent codes at nearly equal storage, and output-path fitting improves over Euclidean fitting.

5 Analysis and Ablations

  • 97.8% retention near 100× on BF16 E4B; OPTIC’s advantage over the shown baselines remains large at higher compression ratios.Figure 2 reports rate–retention curves normalized to uncompressed PLE.Source: Figure 2 · p. 6

5.1 PLE Structure

  • Layer means account for 8.1% of PLE’s squared Frobenius norm, but replacing rows with means drives seven-task retention to zero; 90% variation capture requires 212–227 stored-coordinate directions versus 1–205 under the output-path metric.The output-path direction count has a median of 158, indicating substantial differences across layers.Source: §5.1 · p. 6

5.2 Code Sharing

  • At nearly equal storage, OPTIC retains 97.8% with BF16 and 96.9% with W4A16, versus 79.3% and 86.8% for independent latent codes; rank-20 decoder refits retain 70.55% versus 95.09% without the rank limit.OPTIC shares each token’s code tuple across layers; the rank control keeps indices, codebooks, and biases fixed and refits decoder weights.Source: §5.2 · pp. 6–7

5.3 Fitting Components

  • The context-dependent fitting objective gains 1.98 percentage points in paired mean retention over Euclidean fitting and 0.31 points over diagonal weighting.Each loss is normalized by its value for an all-zero reconstruction, with codes and initialization fixed for 5,000 continuous updates.Source: §5.3 · p. 7, Table 5 · p. 7
  • Continuous fitting supplies most of the optimization gain, while code reassignment improves all three paired fits by 0.65–1.14 percentage points.The comparison starts from the common Euclidean initialization.Source: §5.3 · p. 7, Table 6 · p. 7

6 Compact Inference

The serving experiment decodes compact PLE directly from OPTIC codes on an RTX 4090 with a fixed 1 GiB key–value cache. Compact lookup adds at most 0.40% wall time for BF16 and 0.69% for W4A16 across eight workloads.

6 Compact Inference

  • 5.64 GB to 56.4 MB: resident PLE storage falls for E4B with compact inference.The reported GPU-memory measurement uses an RTX 4090, a fixed 1 GiB key–value cache, and averages two engine initializations.Source: Table 7 · p. 8
  • At most 0.40% added wall time with BF16 and 0.69% with W4A16 across eight prefill and generation workloads.Compact PLE keeps codes, codebooks, and layer decoders on the GPU and decodes current-token rows at each step.Source: §6 · p. 8

7 Related Work

OPTIC adapts discrete embedding compression to PLE’s additional layer axis by sharing token codes while retaining layer-specific affine decoders. Its fitting objective follows the gated PLE output path rather than a linear weight matrix.

7 Related Work

  • OPTIC shares a token’s latent codes across layers, unlike embedding-compression methods for ordinary token embeddings such as CARVQ’s grouped residual codes and nonlinear corrective adaptor.PLE assigns a different embedding to each vocabulary item at each layer; multi-codebook discrete codes for word embeddings are described by Shu and Nakayama (2018).Source: §7 · p. 8
  • OPTIC combines discrete token codes with a shared latent representation and layer-specific affine decoders, unlike low-rank factorization’s continuous coordinates and shared basis.The related-work discussion also distinguishes product, optimized product, residual, and additive quantizers.Source: §7 · p. 8
  • Unlike output-aware weight-compression methods GPTQ, AWQ, and AQLM, OPTIC’s loss follows the gated PLE output path; PLE compression complements backbone weight quantization.The paper demonstrates the combination in BF16 and W4A16 experiments.Source: §7 · p. 8
  • The comparisons include the AQLM Euclidean beam encoder and output-path-fitted additive codes, with PLE adaptations and storage budgets given in the appendices.Source: §7 · p. 8

8 Conclusion

OPTIC shares one code tuple per token across layers and fits codes and decoders at the PLE projection output. It retains substantial performance at high PLE compression and reduces inference GPU memory with little added wall time.

8 Conclusion

  • At about 100× PLE compression, OPTIC retains 97.8% on BF16 E4B and 96.9% on W4A16 E4B; independent latent codes retain 79.3% on BF16 E4B with nearly equal storage.At 594× and 1002× compression, BF16 E4B retention is 90.5% and 84.1%, respectively.Source: §8 · p. 8
  • With a fixed key–value cache size, BF16 E4B served directly from OPTIC codes uses 29.7% less GPU memory with at most 0.40% added wall time.Source: §8 · p. 8
  • The results motivate designing embedding capacity and compressed representations together with deployment memory as a constraint; extending the combination to other embedding structures is a promising direction.Source: §8 · p. 8

Limitations

The evaluation covers Gemma 4 E2B and E4B, with main comparisons based on one fitting seed. Several controlled comparisons have narrower interpretations because they hold specific initialization or fitting conditions fixed.

Limitations

  • The study evaluates PLE only in Gemma 4 E2B and E4B, and main comparisons use one fixed fitting seed.The paired-bootstrap intervals describe benchmark sampling uncertainty; three-fit comparisons and controls describe variation under a common Euclidean initialization.Source: Limitations · p. 9
  • At a matched byte budget, the code-sharing ablation combines the effect of sharing with changes to the code–decoder byte allocation.The fitting recipe is held fixed in that comparison.Source: Limitations · p. 9
  • The fixed-code objective ablation concerns continuous fitting, and its output-path loss measures projection-output error with gates held fixed.The baselines are adaptations of their original methods to PLE, and a single low-count task can strongly affect retention.Source: Limitations · p. 9

Ethical Considerations

The paper reports resource-related ethical considerations and uses its appendix to specify notation, storage accounting, baselines, representations, fitting, and evaluation protocols. The implementation and evaluation details define the settings under which the reported comparisons are made.

Ethical Considerations

  • The authors note that compression can broaden access to model research while also lowering the cost of harmful applications, and disclose AI assistance for polishing and code.The authors retain responsibility for results, references, implementation, and writing.Source: Ethical Considerations · p. 9

Appendix

  • The appendix organizes details on representations and storage, implementation and serving, additional results, output-path-weighted PCA-PQ, and controlled comparisons.These sections supplement the main methods and results.Source: Appendix · p. 12

A.1 Notation and Storage

  • For E2B and E4B, the PLE tensor has V times L_a token–layer rows; resident-size accounting includes persistent representation tensors, and W4A16 uses the same PLE reference size as BF16.Benchmark evaluation materializes reconstructed PLE, whereas serving computes only requested rows.Source: §A.1 · p. 12

A.2 Baselines

  • Six established methods are adapted to PLE with matched complete resident-state budgets, including vector quantizers and INT8 randomized low-rank SVD.The baselines include PQ, RVQ, OPQ-P, AQ, AQLM's Euclidean beam encoder, and SVD.Source: §A.2 · p. 12

A.3 Proposed Representations

  • The proposed representations include output-path-weighted PCA-PQ and OPTIC, whose token-code assignment term scales with vocabulary size rather than vocabulary times layers.Frequency-tiered OPTIC adds latent groups for frequent tokens while using shared codewords for other tokens and retaining selected special-token rows exactly.Source: §A.3 · p. 12

B Implementation and Protocol

B.1 Fitting Procedure

  • OPTIC fits use output-path-weighted PCA initialization, K-means codebook initialization, and a fixed 5,000-step AdamW schedule followed by the reported final-update artifact.Two code-reassignment sweeps follow each scheduled block; unused codewords are reinitialized from high-error tokens and their optimizer moments reset.Source: §B.1 · p. 13

B.2 Models and Evaluation

  • The evaluation crosses E2B and E4B with BF16 and official W4A16 releases, using shared weights, prompts, generation settings, and scorers across methods.For W4A16, fitting uses the corresponding unquantized checkpoint from the official quantization-aware-training pipeline; the seven tasks span reasoning, knowledge, instruction following, and vision.Source: §B.2 · p. 13

B.3 Compact Serving

Compact PLE is served directly from resident codes, codebooks, and decoders, with requested rows decoded during inference. The supplementary sections report additional compression comparisons, uncertainty, a row-wise output-path-weighted baseline, and controlled fitting comparisons.

B.3 Compact Serving

  • At smaller GPU budgets, compact PLE can permit serving when dense PLE leaves no room for a key–value cache; at larger budgets, compact serving generally matches or exceeds dense serving, with stated workload exceptions.For dense W4A16 E4B at 16 GiB, compact PLE finishes 8 and 16 requests 1.20× faster; long-context retrieval remains correct through 130,000 tokens for compact BF16 E4B.Source: §B.3 · pp. 14–15, Table 10 · p. 15

C Additional Results

  • The additional-results section reports the seven constituent benchmark scores and nominal-100× retention, with further tables extending comparisons across compression rates and weight formats.Source: C Additional Results · p. 15

C.1 Original BF16

  • The original-BF16 results include a nominal-100× comparison for both model scales and extend the E4B comparison across practical compression rates.Source: §C.1 · p. 15

C.2 Official W4A16

  • The official W4A16 results compare three baselines with two OPTIC variants, while the full nominal-100× comparison is reported elsewhere in the paper.Source: §C.2 · p. 15

C.3 Sampling Uncertainty

  • The 95% percentile intervals use paired benchmark resampling and represent item-sampling uncertainty for the fixed model, artifact, and evaluation configuration.Resampling groups GPQA permutations and AIME attempts by their original question.Source: §C.3 · p. 15

D Output-Path-Weighted PCA-PQ

  • Output-path-weighted PCA-PQ retains one code tuple per token–layer row but changes the reconstruction metric using the local PLE projection geometry.Its procedure fits a shared PQ bank and row assignments, selects residual corrections using calibration-token counts, and fits a post-quantization layer-coordinate scale.Source: D Output-Path-Weighted PCA-PQ · pp. 15–17

E Controlled Comparisons

  • Controlled comparisons use matched resident bytes and a common mixed calibration policy, with fitting seeds 0, 1, and 2 for the representation comparison.The E4B comparison covers shared, independent latent, and additive coding representations.Source: E Controlled Comparisons · p. 17

E.1 Objective and Optimization

  • The full output-path objective gains 1.98 ± 1.19 percentage points over Euclidean fitting and 0.31 ± 0.45 points over diagonal weighting; periodic reassignment adds a mean 0.83-point retention difference over continuous fitting with fixed codes.The objective comparison holds initial codes and continuous parameters fixed, while the optimization comparison contrasts initialization, fixed-code fitting, and periodic reassignment.Source: §E.1 · pp. 17–18

E.2 Evaluation Details

Evaluation uses matched model settings and fixed task protocols across the paper's benchmark suite. The PLE analysis finds that layer means alone do not preserve performance, while projection-output structure is more concentrated than stored-coordinate structure.

E.2 Evaluation Details

  • The evaluation uses checkpoint-matched comparisons and fixed task protocols, including benchmark-specific prompts, scoring rules, and attempt populations.The seven-task evaluation includes reasoning, knowledge, instruction-following, and vision benchmarks.Source: §E.2 · pp. 18–19
  • At about 100× BF16 E4B compression, OPTIC answers 92.0% of Belebele English questions and 89.1% of Simplified Chinese questions, versus 90.9% and 89.8% with uncompressed PLE.Source: §E.2 · p. 19

E.3 PLE Structure and Decoder-Rank Control

  • Layer means account for 8.10% of squared PLE-table Frobenius norm, and replacing vectors with those means yields zero seven-task retention.At projection output, the leading 20 eigenvalues hold 25.0–98.2% of layer trace, compared with 10.5–24.9% in stored coordinates.Source: §E.3 · p. 20
  • The decoder-rank control keeps indices, codebooks, and biases fixed while refitting each layer decoder under a rank constraint at the same resident byte size.Source: §E.3 · p. 20

Related papers

Loading 2610.31309v1…