N-Gram Memories at Two Bits: Quantizing Lookup Tables in Large Language Models

AI summary

Overview

Research area: Model compression and efficient inference for large language models, specifically the quantization of conditional memory — the large n-gram hash lookup tables that augment modern LLMs.

Technical level: Intermediate. The motivation and results are accessible to anyone who knows what an embedding table is; the encoding machinery (trellis-coded quantization, Gaussian source modeling) will be most familiar to readers with a background in quantization.

Scope: The paper proposes and evaluates Conditional Program (CP), a compressed representation of n-gram lookup tables that stores each row as a short trellis code plus a local scale, and measures its effect on task accuracy, output-distribution fidelity, lookup latency, and GPU serving capacity across three models.

What This Paper Is About

Conditional memory increases a model's capacity by letting it look up learned representations from a huge table instead of computing over all parameters, but the entire table must still be stored — DeepSeek's Engram tables alone occupy 202.7 GB. Standard LLM weight-quantization methods assume dense computation, and compressing these tables like ordinary weights can also destroy the property that makes them useful: each row must be retrievable independently, without decoding unrelated data. The paper asks whether such tables can be stored far more compactly while reconstructing only the handful of rows any given input actually requests.

Key Contributions

1. Conditional Program (CP): a compressed table representation that represents each row by a short code and a local scale, decoded through a reconstruction rule shared across the whole table, so storage is traded for reconstruction work confined to the few rows an input selects.
2. A concrete instantiation using a shared Gaussian trellis (HYB bitshift trellis with tail-biting, 16-bit state, two values per step), per-row norm restoration rather than least-squares scaling, and independently addressable row records with a fused fetch-and-decode kernel.
3. An evaluation across three n-gram memory models (DeepSeek-V4.1-Flash, Qwen3.8-Flash-Next, LongCat-Flash-Lite) on 40 tasks in five categories, plus six-domain output-KL measurements, headline storage reductions of roughly four- to eightfold at two bits, and 41–63% lower output KL than 2-bit round-to-nearest.
4. A serving study covering lookup-layout microbenchmarks (host DRAM and HBM) and multi-GPU workload capacity, showing that smaller tables permit larger decode batches and prefill sizes.

Main Findings

• Storage: At two bits, CP compresses the original lookup tables by roughly four- to eightfold. For DeepSeek the table falls from 202.7 GB to 50.7 GB at 2 bits (the abstract rounds this to 203 GB → 51 GB), and to 38.4 GB at 1.5 bits; Qwen goes from 102.4 GB to 13.4 GB and 10.2 GB, and LongCat from 62.8 GB to 8.1 GB and 6.1 GB.

• Output fidelity: At two bits, CP lowers output KL by 41–63% relative to 2-bit round-to-nearest, and a further 3–9% relative to QTIP, with the latter gain requiring an additional bf16 scale per row. At 2.5 bits CP approaches the fidelity of 3-bit RTN with less storage, and at higher rates it beats four-bit floating-point formats (NVFP4/MXFP4) at smaller size.

• Downstream accuracy: Overall scores remain close to the originals — DeepSeek 79.8 original versus 79.9 (CP 2-bit) and 79.5 (CP 1.5-bit); Qwen 75.6 versus 75.5 and 75.4; LongCat 68.2 versus 67.3 and 67.0. Pooling all 40 DeepSeek tasks by item, accuracy changes by −0.07 points with CP, −0.34 with 2-bit RTN, and −12.0 with the table removed. The authors caution that losses are task- and model-dependent.

• Weak spots at the lowest rate: On DeepSeek, the 1.5-bit code scores below QTIP 2.0 on the knowledge group (difference −0.80 percentage points, 95% interval [−1.23, −0.38]), with a PopQA difference of −1.50 [−2.14, −0.84].

• Learned structure exists but does not pay off: Sixteen fitted directions on the 4,096 most-read held-out rows capture 19.1% and 21.1% of energy for DeepSeek layers 1 and 14, 33.3% for Qwen, and 10.3% for LongCat, against a Gaussian reference of 6.25% (10.0% for Qwen). Even so, at a matched 2.10 bpp the shared trellis reaches KL 0.0198 versus 0.0238 for a fitted 256×4 codebook and 0.0226 for a 65,536×8 codebook; a mean-plus-trellis variant reaches 0.0193, and mean plus four principal components plus trellis costs 2.50 bpp for 0.0188, while the trellis reference at 2.50 bpp is 0.0144.

• Norm restoration beats minimum-error scaling: With identical codes and storage, restoring each row's original norm gives lower KL than a row-wise least-squares scale on all three models (DeepSeek 0.0401 vs 0.0404, Qwen 0.0327 vs 0.0349, LongCat 0.0426 vs 0.0462). Within-row rotation shows no statistically resolved benefit, and frequency-based bit allocation did not consistently generalize across text domains — so CP uses no rotation and uniform bits.

• Selective lookup is fast: Replaying 4,096-token DeepSeek requests (48 rows per token across two Engram layers), CP completes lookups in 0.190 ms versus 0.221 ms for a QTIP tile-window decoder in HBM, and 6.54 ms versus 10.78 ms when fetching from pinned host DRAM over PCIe.

• Serving: On four H200 GPUs, table memory is 50.7 GB per GPU in FP8, 27.6 GB in NVFP4, and 12.7 GB with CP 2-bit. An 8×4k-token prefill drops from 64.66 s to 13.27 s, and 64k-token prefill from 59.81 s (NVFP4) to 18.23 s, while decode times stay similar. On eight GPUs, CP frees 19.0 GB per GPU and raises the largest successful decode batch from 1,024 to at least 2,048 sequences and prefill from 96k to at least 128k tokens.

• Speed caveat: When GPU memory is already sufficient, complete inference steps show little speed improvement in the authors' implementation; the benefit is capacity and memory-pressure relief rather than raw latency.

Methodology in Plain English

The starting point is a table of learned rows. Each row is split into two things: its direction (the shape of the vector) and its size (its length). The size is cheap to store as one bf16 number per row, so the hard part is compressing the direction.

The authors argue that a normalized row resembles a draw from a standard Gaussian: each coordinate accumulates many small, weakly dependent learning contributions, so a central-limit argument gives roughly Gaussian coordinate statistics. If that is true, you do not need to learn a codebook from your specific table — you can use one fixed, shared quantizer designed for Gaussian data.

They use a trellis quantizer, which encodes a vector as a path through a small state machine rather than rounding each coordinate independently, choosing the path that minimizes reconstruction error. Every row in the table shares this same trellis and the same small lookup table; all a row stores is its path symbols plus a single scale. That scale is set to restore the row's original length exactly, rather than to minimize squared error, because the experiments show length matters more to the model than per-coordinate accuracy.

Because each row's code and scale sit in their own independently addressable record, inference works like the original lookup: hash the input n-grams to addresses, fetch only those records, decode them with a fused kernel, rescale, and hand the vectors to the unchanged readout. Nothing is retrained — the address function, backbone, and readout stay fixed. The authors compare this against round-to-nearest quantization, NVFP4 and MXFP4, a QTIP-style trellis adapted to lookup tables with a single table-wide scale and within-row rotation, and fitted vector-quantization codebooks, measuring both task scores on 40 benchmarks and KL divergence from the original model's output distribution over six text domains.

Why This Matters

Compression of conditional memory has been treated as a footnote to weight quantization, even though these tables can be hundreds of gigabytes and are accessed sparsely rather than densely. This paper makes the case that the sparse access pattern changes the design problem: a good codec must not only shrink the table but preserve per-row addressability, and the metric that matters is downstream behavior rather than reconstruction error. The finding that norm restoration beats minimum-MSE scaling is a concrete instance of that gap.

Real-world applications:

• Serving large-memory models on fixed GPU fleets — smaller tables mean larger batches and longer prompts fit on the same hardware, which the four- and eight-GPU experiments illustrate directly.
• Host-memory or SSD-backed memory tiers — the row-wise layout cuts PCIe lookup time from 10.78 ms to 6.54 ms, which matters when the table cannot live on the accelerator at all.
• On-device or edge deployment of retrieval-augmented and n-gram-augmented models where storage, not compute, is the binding constraint.
• Cheap archiving and distribution of trained checkpoints whose memory tables dominate the parameter count.

Industry relevance: the work targets inference economics rather than training. Practitioners running memory-augmented models under memory pressure get a drop-in, retraining-free compression scheme with measured serving numbers, and the negative results (no benefit from rotation, no consistent gain from frequency-based bit allocation) save others from re-testing them.

Future Directions

• Extending beyond hashed n-gram tables. The authors explicitly limit their claims to post-training compression of three hashed-table models

Authors’ abstract

Conditional memory expands language-model capacity through selective lookup, but the full embedding table must remain available for every possible input. We introduce Conditional Program (CP), a compressed representation that reconstructs only the rows an input requests. CP represents each row by a short code and a local scale, with a reconstruction program shared across the table. We instantiate this representation with a Gaussian trellis, norm restoration and independently decodable row records. CP's two-bit code lowers output KL by 41--63\% relative to 2-bit round-to-nearest. Analysis across three trained n-gram memory families shows that Gaussian-like coordinate statistics coexist with learned relationships among rows; at matched storage budgets, the shared trellis outperforms the fitted codebooks and low-rank residual variants we test. For DeepSeek, it reduces table storage from 203GB to 51GB, leaving room for larger workloads when the table is kept on GPU.