PaperlayerView exact revision on arXiv

Source-linked AI summary

Low-Bit N-Gram Embeddings for Language Models

Stephen Chung, Wenyu Du

arXiv:2610.31646v1cs.CLcs.AI

TL;DR

N-gram embedding tables can add model capacity with little additional computation, but their memory footprint can be substantial. This paper compresses their embeddings using low-bit representations and studies quantization and bit allocation. Four-bit tables retain nearly all of the tables’ validation-loss improvement while saving memory, and saved capacity can also be reinvested in larger tables.

  • Problem

    N-gram tables improve language-model performance with little additional computation, but storing them can require substantial memory, motivating ways to retain their benefits with a smaller footprint.

    Source: §1 · p. 1

  • Method

    The paper stores n-gram embedding values at low precision and compares quantization methods and bit allocations across tables.

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

  • Results

    Four-bit tables save approximately 75% of table memory while preserving more than 99% of the improvement provided by full-precision tables; DeepSeek-V4.1-Flash and Qwen3.8-Flash-Next save 47–55% with little change in perplexity and downstream performance.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Four times as many two-bit embeddings achieve 122–126% of the original tables’ improvement while using 46–48% less table memory.

    Source: §1 · p. 2

Abstract

from arXiv · show

Many recent open-weight language models use n-gram embedding tables to improve performance. These tables store learned representations of short token sequences, such as pairs or triples of consecutive tokens, which the model retrieves during processing. However, they can require substantial memory: DeepSeek-V4.1-Flash, for example, allocates 196 billion parameters to its Engram n-gram memory. We therefore investigate whether this memory cost can be reduced while preserving the performance improvement that the tables provide. We propose storing n-gram embeddings at low precision and systematically investigate different quantization methods and bit allocations across tables. Across nanochat-, Llama-, and GPT-2-family models trained from scratch, four-bit tables save approximately 75% of table memory while preserving more than 99% of the validation-loss improvement from adding the tables. Using fewer bits leads to a larger drop in performance, although we find that a short fine-tuning phase recovers a significant part of this loss. In addition, reinvesting the saved memory in larger tables can improve performance beyond that of the original full-precision tables while still using less memory. These table memory savings also extend to large existing models: on DeepSeek-V4.1-Flash and Qwen3.8-Flash-Next, low-bit tables save 47-55% of table memory, with perplexity increases of less than 1.2% and largely preserved downstream task performance. These findings establish low-bit precision as a practical way to reduce the memory requirements of n-gram embeddings and use the saved capacity to improve model performance.

1 Introduction

N-gram tables add learned representations of short token sequences with little additional computation, but their memory cost can be substantial. This work studies low-bit compression and bit allocation to retain their benefits, recover losses, and use saved memory for larger tables.

1 Introduction

  • N-gram tables store learned representations of short token sequences and provide model capacity with little additional computation, but their memory cost can be substantial.DeepSeek-V4.1-Flash’s Engram memory contains 196 billion parameters; table hashing lets multiple n-grams share embeddings, while smaller tables increase that sharing.Source: §1 · p. 1
  • The proposed approach stores n-gram embedding values at low precision, with quantization methods and bit allocations studied to balance table memory and performance.The introduction identifies evenly spaced, normal-quantile, and learned Lloyd–Max codebook values as alternatives.Source: §1 · pp. 1–2
  • Three-bit n-gram tables save 80% of table memory while retaining approximately 97% of the validation-loss improvement from adding tables across three model architectures.The study also reports substantial table-memory savings for DeepSeek-V4.1-Flash and Qwen3.8-Flash-Next.Source: §1 · p. 2
  • The study compares codebook construction, scale fitting, and precision allocation across bigram and trigram tables.The evaluated models include nanochat-, Llama-, and GPT-2-family models trained from scratch, as well as DeepSeek-V4.1-Flash and Qwen3.8-Flash-Next.Source: §1 · p. 2
  • A short fine-tuning phase substantially reduces the performance loss from low-bit compression.This is presented as making more aggressive compression effective.Source: §1 · p. 2
  • Four times as many embeddings at two bits achieve 122–126% of the original tables’ improvement while using 46–48% less table memory.The authors present reinvesting saved memory in larger tables as a way to improve performance beyond the original full-precision tables.Source: §1 · p. 2

2 Method

The model uses four n-gram tables to retrieve embeddings that supplement the current token embedding. The method compresses table values into low-bit indices and shared codebooks with per-embedding scales, while comparing quantization and scale choices and allowing mixed precision across n-gram orders.

2 Method

  • Four tables—two for bigrams and two for trigrams—each retrieve one embedding per current token; concatenating the four embeddings and adding them to the token embedding forms the model input.Each retrieved embedding has dimension d/4, so their concatenation matches the token embedding dimension.Source: §2 · pp. 2–3
  • Each embedding component is stored as a b-bit index into a shared codebook of up to 2^b values, and reconstruction multiplies the selected value by a 16-bit per-embedding scale.For p = 128, two-bit precision reduces per-embedding memory by approximately 87%, from 2,048 to 272 bits.Source: §2 · p. 3
  • Mixed bit widths across bigram and trigram tables allocate more bits to the tables more sensitive to quantization under a fixed table-memory budget.The two n-gram orders represent different context patterns and may need different precision to preserve their contributions.Source: §2 · p. 3
  • The method compares uniform, zero-inclusive uniform, normal-quantile, and Lloyd–Max codebooks, with Lloyd–Max values fitted to normalized trained embeddings.Fitting assigns components to their nearest codebook values and replaces each value with the mean of its assigned components.Source: §2 · p. 3
  • Embedding scales are chosen by matching maximum magnitudes, using root-mean-square magnitude, or minimizing reconstruction mean squared error.The fitted codebook and scales are stored with the low-bit indices, so inference retrieves embeddings without repeating fitting.Source: §2 · p. 3

3 Experiments

The experiments evaluate low-bit n-gram tables across models trained from scratch and two large existing models, then test fine-tuning and larger tables as ways to recover or improve performance. Across these settings, compression can substantially reduce table memory while preserving much of the tables’ benefit, and reinvestment can exceed the original tables’ improvement.

3.1 Experimental setup

  • The study trains five models from scratch with four n-gram tables by default, using two bigram and two trigram tables with about 262k embeddings each.Experiments use ClimbMix and three random seeds; performance is measured against both full-precision-table and no-table models.Source: §3.1 · p. 4

3.2 Table memory savings and model performance

  • Four-bit tables save 74–75% of table memory while retaining more than 99% of the original improvement; two-bit tables retain 85.9–91.0%.The comparison covers default-size tables and selected models with 1M embeddings per table.Source: §3.2 · p. 4, Figure 2 · p. 5

3.3 Codebook choices and mixed precision

  • With error-fitted scales, three codebooks retain 86.36–86.52% of improvement, versus 58.04% for uniform codebooks with maximum-based scaling.Lloyd–Max with RMS scaling retains 86.19%; it is used as the default in later experiments.Source: §3.3 · p. 4, Figure 3 · p. 5
  • Across five models, four-bit bigrams and three-bit trigrams retain 98.9–99.4% of improvement while saving 77–78% of table memory.At equal total bit counts, assigning more bits to bigrams consistently retains more improvement than the reverse allocation.Source: §3.3 · p. 4, Figure 3 · p. 5

3.4 Evaluation on large existing models

  • Across three held-out datasets, selected settings keep perplexity changes below 1.2%; downstream accuracy falls by 0.93% for DeepSeek and 0.27% for Qwen.Figure 5 reports low-bit minus original accuracy for DeepSeek 4/5/4 and Qwen 4/3 settings.Source: §3.4 · p. 6, Figure 5 · p. 7

3.5 Recovering performance with low-bit training

  • After-compression fine-tuning raises retained improvement to 65.3% at one bit and 92.0% at two bits; one-bit training from the start retains 77.0%.The comparison tests whether recovery comes from adapting to low-bit embeddings rather than simply receiving more training.Source: §3.5 · p. 7

3.6 Reinvesting memory in larger tables

  • At 1M embeddings per table and two-bit precision, tables use 52–54% of baseline memory while providing 22–26% more improvement than the original 16-bit tables.Fine-tuning raises the two-bit improvement to 129–145% at the same memory cost.Source: §3.6 · pp. 7–8, Figure 7 · p. 8

4 Related work

Prior work explores n-gram representations through hashing, gating, factorization, and external memory, while embedding compression work includes quantization and compositional representations. This paper compares fewer embeddings with fewer bits per embedding for already-shared hashed embeddings.

N-gram and lookup memory in language models

  • The Over-Tokenized Transformer hashes multi-gram inputs into large input tables, while Deep Sparse Embedding places hashed n-gram memory inside the network with gating and short convolutions.Other approaches include factorized embeddings, frequency-aware hashing, precomputed contextualized embeddings, product-key memories, and external datastores.Source: §4 · p. 9

Compressing embedding tables

  • Embedding-table compression includes four-bit codebooks, compositional embeddings, and mixed embedding dimensions; word embeddings have also been replaced by learned compositional codes.These methods span recommendation-system tables and word embeddings.Source: §4 · p. 9
  • Because hashed embeddings are already shared, the paper directly compares using fewer embeddings with using fewer bits per embedding.Related approaches also include hashing for parameter sharing and n-gram representations over discrete latent codes or bytes.Source: §4 · p. 9

5 Conclusion

The paper studies low-bit n-gram tables as a way to reduce memory while preserving performance, and finds that saved capacity can also support better-performing, larger tables.

5 Conclusion

  • Four-bit tables save approximately 75% of table memory while preserving more than 99% of the improvement from full-precision tables across models trained from scratch.Compression also saves 47–55% of table memory in DeepSeek-V4.1-Flash and Qwen3.8-Flash-Next, with little change in perplexity and downstream performance.Source: §5 · p. 9
  • The reported results indicate that, beyond reducing n-gram table memory use, reinvesting the savings in additional embeddings can improve performance over the original full-precision tables.Source: §5 · p. 9
  • At lower precision, brief fine-tuning can recover much of the lost performance, and mixed precision lets different n-gram tables use different bit levels.The paper also reports that choosing embedding scales matters more than choosing the codebook.Source: §5 · p. 9

6 AI Use

The work builds on a low-precision n-gram-table finding from an autonomous agent on rStation. The authors reproduced and extended the experiments outside the platform and revised the manuscript for broader accessibility.

6 AI Use

  • An autonomous agent, given the open-ended task of improving language-model pre-training, found that n-gram tables could use low precision with minimal performance loss.The agent ran on rStation, a multi-agent research platform.Source: §6 · p. 9
  • The authors reproduced and extended the experiments outside rStation using additional models and datasets.They considered the agent’s finding potentially useful to the broader research community.Source: §6 · p. 9
  • The authors revised the manuscript to make the findings accessible to researchers.Source: §6 · p. 9

A Detailed methods

The methods hash bigrams and trigrams to table addresses, compare three shared-codebook construction approaches, and select per-embedding scales to reconstruct the learned vectors.

A.1 N-gram hashing and lookup

  • Bigram and trigram keys combine the current token with preceding tokens using base V + 1, assign missing predecessors V, and map keys to table addresses by modulo.The default base table has 262,141 entries; the 1M-embedding experiments use 1,048,573.Source: §A.1 · p. 13

A.2 Codebook construction

  • The codebook options are uniform values, normal-quantile values, and Lloyd–Max values fitted to normalized embedding components.Uniform* includes zero; Lloyd–Max initializes values at sample quantiles and iteratively assigns components to their nearest value and replaces each value with its assigned components’ mean, repeating 60 times.Source: §A.2 · p. 13

A.3 Embedding scales

  • For each embedding, maximum-based scaling matches largest magnitudes, RMS scaling uses overall magnitude, and error-fitted scaling selects the candidate with lowest reconstruction mean squared error.Error-fitted scaling tests the RMS scale and maximum-based scales multiplied by α values from 1.00 to 0.30, with BF16 rounding and reconstruction.Source: §A.3 · p. 14

B.1 Models and training settings

The experiments use nanochat-, GPT-2-family, and Llama-family models with a shared n-gram lookup design. They train and evaluate on packed ClimbMix sequences, using held-out validation data for final scoring.

B.1 Models and training settings

  • The five from-scratch models span nanochat, GPT-2-family, and Llama-family architectures, use an 8,192-token vocabulary, and share the same n-gram lookup design.Models with and without n-gram tables use the same training setup, and table embeddings are initialized to zero.Source: §B.1 · p. 14, Table 1 · p. 14, Table 2 · p. 15

B.2 Training and validation data

  • Training uses shuffled ClimbMix packed into 2,049-token sequences, yielding 2,048 prediction targets per sequence; validation uses 10,240 held-out sequences scored identically across models.The validation set contains 20.97M prediction targets, with duplicate documents and documents also present in training or development removed.Source: §B.2 · p. 15

B.3 DeepSeek and Qwen evaluation

The appendix details how low-bit embeddings are trained, evaluated, and compared across table sizes, quantization choices, and mixed-precision allocations. At approximately equal memory cost, assigning more bits to bigrams consistently preserves more improvement than the reverse allocation across five models and two table sizes.

B.3 DeepSeek and Qwen evaluation

  • The DeepSeek and Qwen evaluation uses Lloyd–Max codebooks with error-fitted scales, comparing table precisions against each model’s original eight-bit floating-point tables.Perplexity is evaluated on WikiText-2; downstream accuracy is measured over the full task splits with paired bootstrap confidence intervals.Source: §B.3 · p. 15, Table 7 · p. 22

B.4 Training and fine-tuning with low-bit embeddings

  • Training retains BF16 embeddings for updates, treats quantization as an identity operation in backpropagation, and refits scales each pass and codebooks on a progressively sparser schedule.The codebook is refitted every update for the first 100 updates, then every 25 updates.Source: §B.4 · p. 15
  • Fine-tuning uses unseen sequences for the 262k-table experiments, while the 1M-table experiments use new sequences for smaller models and pre-training sequences for the 322M model.The 262k-table procedure uses 381 updates of 32 sequences, with all weights trainable and a learning rate set to 0.1 times the original peak.Source: §B.4 · p. 16

B.5 Computational resources

  • Reported memory savings count indices, one 16-bit scale per embedding, and shared codebook values, relative to the original tables rather than total model memory.The reference tables are 16-bit in the from-scratch experiments and eight-bit for DeepSeek and Qwen.Source: §B.5 · p. 16

C.1 Compression and table size

  • Tables 3 and 4 report validation loss, loss change, and improvement for 262k and 1M embeddings per table, using 262k 16-bit tables as the reference.The table summaries report means and sample standard deviations over three seeds.Source: §C.1 · p. 17, Table 3 · p. 17, Table 4 · p. 18

C.2 Codebook and scale choices

  • The codebook and scale comparison extends to other models and includes uniform-codebook variants.Table 5 reports validation loss in BPB averaged over three seeds.Source: §C.2 · p. 19, Table 5 · p. 19

C.3 Mixed precision

  • At approximately equal memory cost, assigning more bits to bigrams consistently preserves more improvement than the reverse allocation across five models and both table sizes tested.The complete sweep covers two-, three-, and four-bit bigram/trigram pairs for 262k and 1M embeddings per table.Source: §C.3 · p. 20, Figure 8 · p. 20

C.4 DeepSeek and Qwen

  • DeepSeek and Qwen are evaluated against their original FP8 tables on WikiText-2, with held-out perplexity and downstream accuracy reported for selected allocations.DeepSeek uses bigram, trigram, and four-gram tables; Qwen uses bigram and trigram tables. The held-out comparisons include WikiText-2, Penn Treebank, C4, HellaSwag, ARC-Challenge, and WinoGrande.Source: §B.3 · p. 15, §C.4 · p. 20, Table 6 · p. 22

C.5 Training and fine-tuning with low-bit embeddings

The appendix reports training-sequence and larger-table fine-tuning results, then compares component-wise hashing with conventional whole-embedding hashing. Component-wise methods improve average two-bit results across three architectures, while same-column hashing also shows gains in most tested settings.

C.5 Training and fine-tuning with low-bit embeddings

  • Table 8 covers four training sequences on the 31M Llama-family model, while Table 9 reports fine-tuning of 1M-embedding tables before or after compression.The tables report improvement and table memory relative to the original 262k-embedding, 16-bit tables; fine-tuning comparisons use equally fine-tuned references.Source: §C.5 · p. 23, Table 8 · p. 23, Table 9 · p. 23

Experimental setup.

Across three architectures, component-wise hashing improves average two-bit results over conventional hashing, and same-column hashing also improves mean performance in most tested settings. The comparisons use matched training setups and evaluate both 16-bit and two-bit tables.

Experimental setup.

  • The three methods are compared on nanochat 26M, GPT-2-family 30M, and Llama-family 31M using the same setup, three seeds, and both 16-bit and two-bit tables.Two-bit tables are quantized during pre-training with Lloyd–Max codebooks and error-fitted scales.Source: Experimental setup · p. 24

Discussion

  • Component-wise methods retain 93.6–100.3% of conventional 16-bit-table improvement at two bits, compared with 87.0–93.7% for conventional two-bit hashing across three architectures.Gains are larger on GPT-2 and Llama than on nanochat.Source: Experimental setup · p. 24, Table 10 · p. 24
  • Same-column hashing improves mean performance over conventional hashing in five of six settings, but consistency and the proposed collision explanation require further investigation.The authors suggest component-wise hashing may reduce interference from whole-embedding collisions.Source: Experimental setup · p. 24

Related papers

Loading 2610.31646v1…