The scaling relationship of compute and performance for language models (LMs) has been extensively studied; more recently, models are being designed to scale memory use, as well. In particular, adding embeddings for n-grams scales memory separately from compute use, because each pass only processes individual embedding rows, which are much smaller than the storage used for the whole embedding table. However, existing methods train embeddings from scratch, which needs expensive GPU memory for fast reading and writing. Meanwhile, the ordinary memory of the host system is cheap and underutilized by these training methods. To scale on host memory, we introduce COUNT (Counting Offline Unlocks N-gram Transformers), which builds n-gram embeddings without adding learned parameters by storing precomputed n-gram counts on host. We use the corpus frequencies to weight together the LM's own unembeddings of the next tokens most often following the current n-gram history. At a fixed compute budget, COUNT can be added as an additional input embedding for pretraining 1B-parameter LMs to improve validation loss as much as an LM with 57% more compute. If GPU memory is limited, using COUNT to let backbone parameters fill GPU memory gives better performance than splitting memory with additional learned n-gram embeddings. Beyond efficiency, COUNT allows switching in a different table of n-gram counts from a domain-specific corpus at inference time to adapt to a new task without any training.
Figures & tables
Figure 1: Left: COUNT uses corpus n-gram frequencies to score continuations and gather rows of the unembedding matrix into an additional input embedding, ft . The count score S(w∣h) can be stupid backoff ( Brants et al., 2007 ) , or any other frequency score based on count(h,w)/count(h) . Unlike n-gram embeddings learned from scratch, COUNT’s additional storage can be offloaded to host RAM at training time. Right: COUNT reaches validation loss equivalent to a standard transformer trained with 46–89% more compute. It does this without increasing compute by instead scaling n-gram statistics on otherwise unused host RAM by up to 797 GB. Solid curves are the compute-optimal fits of each method’s scaling law. Each violet guide shows the extra compute required to make a standard transformer reach the loss of an optimally trained COUNT run. Figure 2 shows task loss and downstream accuracy.
Figure 2: COUNT improves performance without increasing compute or learned parameters. By storing up to 797GB of n-gram frequencies, COUNT allows scaling on host RAM. Left: task loss against training compute at token to parameter ratios D/N∈{10,20,40} . The solid curves minimize fitted scaling laws for each arm, showing compute optimal performance ( Hoffmann et al., 2022 ) , and violet guides show the additional compute required to make a standard transformer reach the loss of an optimally trained COUNT run. Right: downstream accuracy for the same runs.
Figure 3: For a given budget of total learned parameters (a hardware-agnostic measure that drives GPU memory footprint) COUNT delivers best LM loss across a range of backbone model scales from 60M–1B, and downstream accuracy (bottom row) follows the same trend with more noise. Each column uses the indicated token to parameter ratio D/N∈{10,20,40} relative to that run’s own backbone.
Figure 4: The count-based continuation scores in COUNT can be cheaply fit on new domains and swapped in at inference time with no training. We find in-domain task loss improves on all domains when interpolating 1:3 between a domain-specific table and the domain general table used in pretraining this 1B parameter COUNT checkpoint. Accuracy results are more mixed but still increase on average for math and law domains. Each point is a result on one subtask at one training budget. Count-based late-fusion methods also support domain adaptation (Figure 10 ); COUNT combines this capability with the stronger pretrained improvements of early fusion.
Figure 5: The n-gram table used in COUNT can be replaced at inference time with more compressed versions. This allows the method to work on host-RAM constrained deployments with a graceful and predictable performance degradation.
Figure 6: COUNT retains most of the training and inference throughput of a standard transformer. Its throughput is similar to learned n-gram embeddings, and further optimizations would not have far to go to recover both methods to full standard transformer throughput. These measurements use the 1B configurations for each method, and other configurations with smaller tables and lower J further improve COUNT’s throughput, see Appendix B.2 .
Appendix figures & tables22 assets
Supplementary material from the paper’s appendix.
Appendix
Ladder label
Layers
Hidden width
Heads
Feed-forward width
60M
8
384
8
1,536
100M
12
512
8
2,048
190M
12
768
12
3,072
370M
16
1,024
16
4,096
600M
16
1,280
16
5,120
760M
16
1,536
16
6,144
Appendix
Table 1: OLMo-core model ladder configurations. The feed-forward width is the intermediate width of each transformer block.
Setting
Value
Tokenizer
Dolma 3 tokenizer ( allenai/dolma2-tokenizer )
Vocabulary
100,278 tokens; embedding allocation padded to 100,352 rows
Sequence length
8,192 tokens
Position encoding
Rotary position embeddings with base 500,000
Precision
Bfloat16 model computation; float32 gradient reduction
Optimizer
AdamW with anomalous-update skipping ( SkipStepAdamW ); β1=0.9 ; β2=0.99 for 60M and 100M and 0.95 for 190M–1B; ϵ=10−8
Appendix
Table 2: Training configuration. N is the backbone parameter count defined above.
Figure 7: The size of our n-gram tables depends on several configurations. History and n-gram occurrence thresholds limit the number of rows in the table. Each order of n-gram is its own set of rows in the table, with orders four through six accounting for most of the rows. Finally lowering K to store fewer top continuations reduces the table width.
Domain
Size (GB)
K
History-count threshold
Math
83.398
32
4
Biomedical
81.723
32
45
Law
80.151
32
8
Appendix
Table 3: Domain-specific tables used in Figure 4 . Sizes are decimal GB with 8-bit scores.
Figure 8: Composition of the task-loss (BPB) and accuracy (ACC) aggregates by qualitative task category. Each eligible configuration contributes once: 74 configurations for BPB and 41 for accuracy. Subject subsets, splits, prompt formats, and language variants are counted separately. Categories describe task families rather than isolated skills.
Compiled lookup; lookup caching; fused embedding; KV caching; CUDA graphs
Appendix
Table 6: Decode TPS experiment configuration. Prompt processing and graph capture are excluded from the timed generation steps.
Figure 9: Generalist inference-only late fusion compared with COUNT and the standard transformer.
Figure 10: Domain adaptation with inference-only late fusion and COUNT.
Figure 11: Pretraining versus inference-time late fusion at 190M.
J=6 17.1 GB
J=1 17.1 GB
J=6 797.5 GB
J=1 797.5 GB
Standard transformer
55,836
55,836
55,836
55,836
Learned n-gram
54,890 (98.3%)
54,890 (98.3%)
54,890 (98.3%)
54,890 (98.3%)
COUNT
52,127 (93.4%)
54,585 (97.8%)
52,008 (93.1%)
54,532 (97.7%)
Appendix
Table 7: Training throughput at batch 32 on eight H200s by COUNT lookup ( J ) and table size.
Batch 1 GPU-resident
Batch 8 GPU-resident
Batch 1 Host-resident
Batch 8 Host-resident
Standard transformer
116,017
122,389
116,017
122,389
Learned n-gram
115,181 (99.3%)
122,091 (99.8%)
111,403 (96.0%)
117,402 (95.9%)
COUNT
—
—
111,110 (95.8%)
116,931 (95.5%)
Appendix
Table 8: Prefill throughput on one H200 by batch size and n-gram storage placement.
Batch 1 GPU-resident
Batch 256 GPU-resident
Batch 1 Host-resident
Batch 256 Host-resident
Standard transformer
340.2
19,216
340.2
19,216
Learned n-gram
290.8 (85.5%)
18,795 (97.8%)
319.3 (93.9%)
18,448 (96.0%)
COUNT
—
—
326.3 (95.9%)
18,586 (96.7%)
Appendix
Table 9: Decode throughput on one H200 with CUDA graphs by batch size and n-gram storage placement.
Workload
Batch
Standard transformer
Learned n-gram
COUNT
GPU-resident n-gram storage
Training
32
60.8 (61.7)
63.4 (64.6)
—
Prefill
1
4.4 (6.1)
6.4 (12.2)
—
Prefill
8
14.4 (20.1)
16.4 (21.9)
—
Decode
1
4.0 (4.9)
6.0 (6.9)
—
Decode
256
47.4 (110.8)
49.4 (115.1)
—
Appendix
Table 10: Peak GPU memory in GiB per executing GPU for each throughput measurement, as allocated (reserved) high-water marks of the PyTorch allocator. Training is sharded across eight H200s; inference runs on one H200. The learned n-gram baseline’s GPU-resident table holds 2.16 GB.
Figure 12: Full domain-table interpolation grid for 1B COUNT at D/N=10,20,40 . Each metric is change relative to its own checkpoint’s original-table score.
Figure 13: Inference- K sweep for the 1B COUNT model with its original 797.487 GB table fixed.
Figure 14: Inference- J sweep for the 1B COUNT model with its original 797.487 GB table and K=8 fixed.
Figure 15: 190M COUNT trained with nmax=3 versus nmax=6 maximum n-gram order with same 17.1 GB table, K=8 , and embedding amplification.
Figure 16: 190M COUNT with and without embedding amplification (17.1 GB table, K=8 , J=6 ).
Figure 17: Embedding amplification across standard transformer scales at D/N=10,20,40 .
Figure 18: 1B COUNT with an unembedding or token-embedding basis.
Transformers lack a native lookup mechanism, requiring repeated dense computation to recognize and reuse local static patterns. Lngram v1 introduces tokenizer-independent conditional memory through discrete latent n-gram addressing, but its memory capacity is coupled with the backbone width, limiting scalability due to high parameter and activation costs. We propose Lngram v2, which decouples the number of routes, memory dimension, and backbone width, and introduces a context-aware grouped-query attention readout to scale memory capacity independently. A zero-value Sink and counterfactual surrogate gradients further improve readout selectivity and routing trainability while preserving hard discrete addressing. Experiments across vision--language models (VLMs) of different scales show consistent improvements, including successful scaling to a 30B-parameter model. Compared with Lngram v1, Lngram v2 substantially reduces both total and activated memory parameters while maintaining or improving language modeling performance. Further analysis shows that its discrete IDs preserve substantial semantic structure of continuous hidden states, enabling semantic recovery from IDs alone and stable ID--semantic associations across datasets. These results establish Lngram v2 as an efficient and scalable latent conditional memory mechanism whose discrete addresses also provide a structured interface for analyzing internal model representations.
Yunao Zheng, Bin Wen, Xiaojie Wang +12
Beijing University of Posts and Telecommunications. · Kuaishou Technology.
Scaling conditional memory offers a promising way to increase language-model capacity, but existing methods such as Engram learn large memory tables from scratch during pre-training, making memory scaling expensive and sometimes ineffective. We propose Memory Grafting, a conditional memory scaling method that utilizes frozen hidden states from a grafting model as conditional n-gram memory. Given frequent local n-grams, we run the grafting model offline, store final-token hidden representations as memory values, and let the recipient model retrieve them through exact longest-match suffix lookup. Retrieved memories are adapted by lightweight projections and gates, while a hash-based Engram fallback preserves coverage for unmatched contexts. Since the grafting model is only run offline and exact lookup has expected O(1) complexity with respect to memory-bank size, Memory Grafting expands external latent capacity with limited training and inference overhead. Experiments under matched recipient architectures and pre-training budgets show that Memory Grafting improves over both MoE and vanilla Engram baselines. In the 2.8B-scale setting, it improves the average benchmark score from 51.95 for MoE and 52.43 for vanilla Engram to 53.86. In the 0.92B-scale setting, all grafting-model variants improve over the baselines, with Qwen3.5-35B-A3B giving the strongest gains. These results suggest that pretrained models can serve as reusable constructors of external latent memory, providing a practical step toward scaling future language models beyond trainable parameters alone.
Large token-indexed lookup tables provide a compute-decoupled scaling path, but their practical gains are often limited by poor parameter efficiency and rapid memory growth. We attribute these limitations to Zipfian under-training of the long tail, heterogeneous demand across layers, and "slot collapse" that produces redundant embeddings. To address this, we propose X-GRAM, a frequency-aware dynamic token-injection framework. X-GRAM employs hybrid hashing and alias mixing to compress the tail while preserving head capacity, and refines retrieved vectors via normalized SwiGLU ShortConv to extract diverse local n-gram features. These signals are integrated into attention value streams and inter-layer residuals using depth-aware gating, effectively aligning static memory with dynamic context. This design introduces a memory-centric scaling axis that decouples model capacity from FLOPs. Extensive evaluations at the 0.73B and 1.15B scales show that X-GRAM improves average accuracy by as much as 4.4 points over the vanilla backbone and 3.2 points over strong retrieval baselines, while using substantially smaller tables in the 50% configuration. Overall, by decoupling capacity from compute through efficient memory management, X-GRAM offers a scalable and practical paradigm for future memory-augmented architectures. Code aviliable in https://github.com/Longyichen/X-gram.
Yilong Chen, Yanxi Xie, Zitian Gao +10
Institute of Information Engineering, Chinese Academy of Sciences · School of Cyber Security, University of Chinese Academy of Sciences · IQuest Research +1