Scaling Embedding Layers in Language Models
Authors: Da Yu
Edith Cohen
Badih Ghazi
Yangsibo Huang
Pritish Kamath
Ravi Kumar
Daogao Liu
Chiyuan Zhang
Venue: NeurIPS 2025 poster ()
Keywords: embedding layer scalability
contextualized token embeddings
scaling with fixed accelerator usage
Primary area: deep_learning
TL;DR: We introduce contextualized n-gram embeddings to extend input embedding layers, improving performance while maintaining fixed accelerator usage during inference.

Abstract: We propose SCONE (**S**calable, **C**ontextualized, **O**ffloaded, **N**-gram **E**mbedding), a new method for extending input embedding layers to enhance language model performance. To avoid increased decoding costs, SCONE retains the original vocabulary while introducing embeddings for a set of frequent $n$-grams. These embeddings provide contextualized representation for each input token and are learned with a separate model during training. After training, embeddings are precomputed and stored in off-accelerator memory; during inference, querying them has minimal impact on latency due to the low complexity of embedding lookups. SCONE enables two new scaling strategies: increasing the number of $n$-gram embeddings and scaling the model used to learn them, both while maintaining fixed accelerator usage  during inference (in terms of FLOPS and memory). We show that scaling both aspects enables a model with 1B accelerator-resident parameters to outperform a 1.9B-parameter baseline across diverse corpora, while using only about half the FLOPS and accelerator memory during inference.

Scaling Embedding Layers in Language Models
Da Yu?
Edith Cohen
Badih Ghazit
Yangsibo Huang
Pritish Kamath
Ravi Kumari
Daogao Liut
Chiyuan Zhang!
Abstract
We propose SCONE (Scalable, Contextualized, Offloaded, N-gram Embedding), a new method for extending input embedding layers to enhance language model performance. avoid increased decoding costs, SCONE retains the original vocabulary while introducing embeddings for a set of frequent n n-grams. These embeddings provide contextualized representation for each input token and are learned with separate model during training. After training, embeddings are precomputed and stored in off-accelerator memory; during inference, querying them has minimal impact on latency due to the low complexity of embedding lookups. SCONE enables two new scaling strategies: increasing the number R-gram embeddings and scaling the model used to learn them, both while main- taining fixed accelerator usage during inference (in terms of FLOPS and memory). We show that scaling both aspects enables a model with IB accelerator-resident parameters to outperform a 1.9B- -9B-parameter baseline across diverse corpora, while using only about half the FLOPS and accelerator memory during inference.
Introduction
Input embedding layers in language models map discrete tokens to continuous vector representa- tions [Mikolov et al.] 2013 Sennrich al., 2016) before passing them the subsequent layers. Since tokens are typically simple integer values, the mapping can be implemented purely memory fetch operations with no additional computation needed. This allows embedding layers to be offloaded from the limited accelerator memory to main memory or even to secondary storage, such as solid- state drives, with minimal impact on inference speed. This is highly desirable, as main memory and secondary storage are significantly more cost-effective than accelerator memory (e.g.. GPUs and TPUs) [MeCallum] 2024]. These advantages motivate us to explore methods for scaling up embedding layers.
However, scaling the embedding layer by simply increasing the token vocabulary size has limited benefits. A larger vocabulary also enlarges the output (logits) layer, whose weight matrix is tied to the vocabulary size and embedding dimension. Typically, predicting the next token requires computing logits over all tokens in the vocabulary. and prior work shows that the decoding cost becomes impractical once the vocabulary exceeds few hundred thousand tokens [Wang et al.) 2019 Zheng et al. 2021 Liang et al.) 2023 Tao et al. 2024 Dagan et al. 2024 Even if faster hardware or more effciint algorithms might partially offset this cost [Juulin al. [017 [Shim et al.) 2017), second problem remains: scaling the token vocabulary leads to a large number of tail rokens, which occur infrequently the training corpus. The embeddings of these tokens (both input and output) receive very few updates, resulting in lower-quality representations (Liao et al.) 2021 Dou et al. 2024). In Appendix DJ we train GPT-2 models with vocabulary sizes ranging from 32K to and observe performance degradation as the vocabulary size exceeds 512K. We attribute this degradation
Google. Correspondence to: (daywwork, edco, pritishk, chiyuan] agooaneeggool..comm
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Input
Contextualized Tokens (-grams)
Contextualized Embeddings
Tokens
Embeeding Transformer [-gram model)
Standard Tranaformer gmalin modill)
123
79 sia
123
32
deterees Equarra
32
142
60
142 mad
342
79 c..
142 oet
123
79
Contextualized Embedding Table
Procorpated Offeadeed Davig Kibrone
Figure/ 1: Illustration of SCONE with maximum 7-gram length of 3. The f-grams set of frequent n nt-grams selected using the method described in Section 3.1
to the increasing sparsity of updates per token as the vocabulary grows. Additionally, we observe linear increase in accelerator memory usage once the vocabulary size exceeds 1M.
Our contributions. In this paper, we propose SCONE, novel approach for scaling input embedding layers by learning them through a separate transformer model, referred to as the f-gram model. This model takes input a set of frequently occurring n--rams (called f-grams). which we select using an approach inspired by Byte-Pair Encoding-based tokenization (Section 3.1b. Crucially, the number of f-grams decoupled from the token vocabulary size, allowing us to build a separate f-gram input embedding table with up to billions of entries, without blowing up the vocabulary size. An overview of the proposed method is illustrated Figure I
During both training and inference, SCONE leverages the f-gram model to efficiently handle large embedding spaces without overwhelming accelerator resources. During training, the f-gram model learns to generate contextualized embeddings for each f-gram without requiring the instantiation of a massive embedding table. During inference, the output of the f-gram model can be precomputed to form the f- f-gram embedding layer. This embedding layer can then be offloaded from the accelerator, thereby eliminating the need for accelerator resources during inference.
SCONE introduces two novel scaling approaches for improving model performance: (i) increasing the number of cached f- f-gram embeddings and (ii) scaling up the f-gram model used to learn these embeddings. The first approach requires additional off-accelerator memory during inference, while the second demands greater accelerator resources during training. Importantly, both approaches preserve a fixed inference-time accelerator resource footprint, a property not supported by traditional scaling methods. Indeed, prior works (Jones] 202T Hoffmann et al. [022 have shown that scaling the training compute for a fixed model size beyond some compute-optimal threshold leads to diminishing returns. Therefore, the typical method for utilizing additional training compute to increase model size. However, directly scaling model size also increases FLOPS and accelerator memory requirements during inference. In contrast, SCONE leverages larger f-gram models effectively consume more accelerator resources during training but without increasing anference-time accelerator demands, offering novel scaling paradigm that previous studies have not explored.
There are important scenarios where maintaining a fixed inference-time accelerator footprint is especially valuable. In many deployments, where model is queried billions of times per day, inference costs during deployment can far exceed training costs. In such cases, even small increases in inference-time computation can lead substantial operational expenses. The recent emergence of test-time scaling techniques further highlights this trend [Jones] 2021 Snell et al. 22225) emphasizing situations where inference costs dominate the overall expenses of LLM deployment. Furthermore, applications impose strict limits on inference-time computation, leaving no margin for increased computational demands during deployment.
many
Our contributions can be summarized as follows:
A new scalable method, SCONE, to improve language models by expanding the input embeddings [eccionn 3) but requiring no additional accelerator resources at inference time.
Extensive experiments validate SCONE, analyzing key design choices and their impact on evaluation perplexity and accuracy on downstream tasks (Section 4).

---

15.5
Fanetine be
120 Eigrama gggam am madet
520
Baceate
120M fg atr mamaay/
320
creel
P-grame mode
2,95 Locel
Tergeeky 290
14.5
240
240
14.0
220
180
0.7
2-5
1.3 6.7
3.5 icc
13
Model
Figure 2: (Left) Perplexity on the OLMo Groeneveld et al. 2024 evaluation set. Model sizes along the zraxis indicate the number of parameters residing on the accelerator during inference. With 10M f-grams, the 1.3B model matches the performance of the 1 1.9B baseline; with 1B f-grams, the 1B model surpasses it. (Right) End-to-end token generation speed single A100 GPU. Storing f-gram embeddings in main memory adds negligible latency, while using NVMc storage introduces a minor slowdown without causing a bottleneck.
Our results show that SCONE significantly boosts model performance without introducing inference latency bottlenecks: Figure 2 highlights representative findings. Notably, a SCONE model with IB accelerator-resident parameters outperforms a 1.9B baseline that requires approximately 2x more inference FLOPS and accelerator memory.
2 Preliminaries
We focus on pre-training decoder-only language models with causal language modeling (Radford cet al., 2019]. We now introduce notations that will later help us formally describe the proposed SCONE method. For clarity, we omit details that are not essential to describing our method.
Decoder Transformer Model. Let Vioken denote the token vocabulary. The token embedding layer is parameterized by a function T Veoken + Rd, mapping each token to a d-dimensional embedding vector. We abstract the transformer itself function A (Rd) + Rd, which takes as input sequence of up to S token embeddings and outputs a single embedding vecto For completeness, we present the pseudocode of a basic next-token prediction model in Algorithm (Appendix C).
Efficient Indexing for Large Embedding Layers. Mappings from tokens or f-grams to embedding vectors can be implemented as key-value stores, enabling highly efficient lookup. Hash-basee based and tree-based data structures support lookup times that are constant or logarithmic in the number of entries, respectively. These data structures make it feasible, in principle, to offload large embedding layers from accelerators with minimal impact on latency. In practice, however, the traditional token embedding layer T is kept accelerator memory, as it is typically shared with the output (logits) layer, which requires fast access for matrix multiplications during next-word prediction. In contrast, the f-gram embedding layer introduced by SCONE is decoupled from the output layer and can therefore be offloaded, which is critical for maintaining a fixed accelerator memory footprint as the f-gram layer scales. In Seeetioo 4.3] we study two strategies for storing the f-gram embedding table. When stored in main system memory, the f-gram embedding table consists of a dense embedding matrix paired with hash dictionary that maps f-grams to matrix indices. When stored on NVMe solid-state drives, we use the Lightning Memory-Mapped Database (LMDB) Chu 201 map f- f-grams to their embeddings using B+ tree data structure.
to directly
3 SCONE Architecture
We propose to augment a standard transformer model with an additional f-gram embedding layer. Figure provides a high-level overview of the approach. We first construct a set Vi-gram S token :=
2= token consisting of frequently occurring n-grams of length up to K, which we refer to as f-grams. Throughout this paper, K denotes the maximum length of f-grams considered. To construct
Vn
AA decoder transformer typically produces sequence of embeddings for an input sequence. We define the output the final token embedding, used either for next-word prediction or as the embedding for f-gram.
3

---

Algorithm SCONE method F Ve-gram I.F.
T Vioken I
token embedding layer.
Vigram S sokon' set of E-grams.
TRAINING: f-gram transformer model
At.gram (BR)S + Rd
INFERENCE: f-gram embedding layer b F:: Vigrram + IRd.
Input: A sequence (01,-.., Tm) ) E V token of tokens.
Output: Embeddings (ei...., em) E (R"")".
for I...., do
j <- smallest j' < such that , ai) E Vagram if such j exists, otherwise
if i then
e, <- T(o))
else
T(o.)) training. es f- (F(o)...,,0.)
inference.
return (e.)...,em)
Vi-gram. an efficient implementation that requires only K 1 linear scans over the training corpus, as detailed in /eection 3.1 the formal construction is described in [Algorithm 3 Appendix C).
Next, we define the SCONE method, which maps a given sequence of tokens to sequence of embeddings. SCONE behaves differently during training and inference, as described in Algorithm During training, it is parameterized by an f-gram transformer model At-gram (12 ) - R"
which takes an f-gram input and uses the output embedding of the final token as the embedding for that f-gram. During inference, in contrast, it is parameterized by an f-gram embedding layer F: Vi-gram - Rdd key-value store that directly maps each f-gram its precomputed embedding. This embedding layer is implemented by caching the outputs of Ar.ramm for all f-grams in Vi Vi-gram and storing them in off-accelerator memory.
The embeddings produced by SCONE are passed to a standard transformer model Amains referred the main model. This is followed by prediction head D Rd A Avizem Together, these components form the end-to-end process for next-word prediction with SCONE. We present the full description in Algorithm 4 Appendix C).
In the rest of this section, we will discuss the motivation behind the design chocies of SCONE and provide further implementation details.
3.1 BPE-Style Discovery of f-grams
The construction of Vi-gram- outlined in Algo- rithm 3 Appendix C), can be implemented ei- ficiently with K- 1 linear scans over the train- ing corpus. We perform
1600
2-grams
for each n E K), starting with 2-grams. In subse- quent scans, impose a minimum frequency threshold of 5 to reduce memory usage. At the (n 1)th scan, the set of n-grams from the
3-grams
1200
4-grams
of uniquu nigrams
1000
5-grams
6-grams
a00
600
previous scan allows us to skip any (n + 1)- gram candidates that cannot meet the minimum threshold. Specifically, if an (n + 1)-gram sur- passes the threshold, its nufuffix prefix must appear at least many times. Figure shoows how the number of unique n-grams (up grams) grows as the training corpus scales from a few billion to one trillion tokens. Finally, found n-grams (for E [2, K]) are ranked by frequency, and the top ones are selected to com- prise Vi Virgaam
Number of training takens tin billions)
Figure 3: Number of unique 2- to 6-grams ap- pearing at least five times. We uniformly sample tokenized sequences from Dolma (Soldaini et al.) 2024 to vary the corpus size.

---

Our procedure for counting and ranking n-grams is analogous to continuing the training of a BPE tokenizer on an existing vocabulary. In each BPE iteration [Gage] 1994 Sennrich al., 2016]. the frequencies of all token pairs (2-grams) are counted, and the most frequent pair is merged to form new token, expanding the vocabulary by one. However, merging and recounting pairs repeatedly to obtain a large number of f-grams is prohibitively expensive for large training corpora. Instead, we simply collect and sort all n-grams up to a small constant K.
3.2 Learning f-gram Embeddings with Af.rrmm
We motivate the use of At-gram by contrasting it with the alternative of directly backpropagating gradients to large embedding table. The direct approach fails to exploit dependencies between n-grams, leading to fewer updates per embedding. We observed this by pretraining GPT-2 models with vocabulary sizes ranging from 32K to 2M. As vocabulary size increases, embedding updates become sparser, eventually degrading performance. For example, when training on 100M tokens, 97.6% of tokens in 32K vocabulary receive more than 100 updates, compared to only 7.3% in a 2M vocabulary (Appendix D). This sparsity makes it difficult to train a large embedding table through direct gradient updates. SCONE addresses this by parameterizing emheddings with f-gram transformer Ai-gram avoiding the sparse update problem.
Additionally, Ar-gram removes the need to instantiate a full embedding table during training, requirement that would otherwise strain accelerator memory. This because, unlike inference, where word prediction is largely sequential, training parallelizes computation across the sequence dimension, demanding much higher token throughput. Moreover, during training, embeddings must be updated frequently, whereas inference only requires read access. Together, these factors make difficult to offload the embedding layer from accelerators during training. As a result, avoiding full-table instantiation is crucial for scaling the emhedding layer extremely large sizes.
SCONE jointly trains the Argram model with the main model Amain and the token embedding layer T. This overcomes the sparse updates issue but also introduces additional compute costs. For each w E Vi-gram- the computation the same as that of processing a short sequence of length [w| through a standard transformer. Since |س) is small constant, the primary overhead comes from the feed-forward layer. In our experiments (Section 4), we account for this overhead in one of two ways: (1) by training baseline models long enough to reach ceae-conveeeence a fixed model size, ensuring that additional training compute would yield minimal gains; or (2) by reducing the number of training tokens for models using SCONE so that their total training FLOPS match those of the baselines.
During inference, the f-gram embedding layer F can be precomputed and stored a lookup table, offloaded to system main memory or secondary storage while still permitting efficient retrieval. Meanwhile, the token embedding layer T remains on the accelerator for decoding. In(Section 4.3) we evaluate the query latency and space usage of the f-gram embedding layer under various configurations. We show that the latency is not a bottleneck for language model inference and the space costs are low due to the use of relatively inexpensive system memory and solid-state drives.
Experimental Evaluation
In this section, we evaluate SCONE in pre-training settings. Section 4.1, we assess SCONE on uPT-2-sized models to study various design choices, and (Seectn 4.2 extend the evaluation to large-scale pre-training scenarios involving trillions of tokens. Finally, in Section 4.3 we analyze the inference and storage costs during deployment.
For completeness, we also evaluate SCONE in post-training settings by applying it during the SFT stage of recent Qwen3 models (Yang et al. 2025]; these results, presented in Appendiix E.3] show that SCONE remains effective post-training as well.
4.1 Experiments with GPT-2
We analyze three key hyperparameters: (i) the maximum f-gram length K in Vr-gram- (ii) the number
of f-grams used, VV..ram and (iii) the Ar Argram model size. We use the released GPT-2 tokenizer, which has Vooken| = 50,257, and train on the WebText dataset (Peterson et al. 2019 The tokenized corpus contains 9B training tokens, from which extract f-grams using the method in Section 3.1
5

---

Mude: deme
228-2290
20
pffgpmm urbedingg
Figure 4: Evaluation perplexity on WebText (left) and WikiText- 103 (right) as function of VVigraml- Model sizes in the legend are number of parameters residing on the accelerator during inference. Dashed lines and leftmost stars show baseline performance.
We consider three main model sizes with 76M, 340M, and 510M non-embedding parameters. In- cluding the token embedding layer, the total parameter counts increase to 128M, 419M, and 589M, respectively. These models are either trained using only the token embedding layer as a baseline or with an additional Ar-gram when SCONE is applied. We train all models for 80B tokens, roughly twice the number of training tokens used in Radford et al. [2018], to ensure that the baseline models approach convergence. For evaluation, we use the validation split of WebText and WikiTText 103 [Merity al.] (207)), one of the largest downstream datasets in Radford al. [2019]. Additional implementation details are provided in Appendix F.I
4.1.1 Varying the Maximum f-gram Length
We study the impact of varying the maximum f-gram length K in Vi-gram- We vary K from 2 to while keeping the total number of f-grams fixed at 20M. As K increases, the frequency cutoff, which the minimum number of times an n-gram appears in the training corpus to be included Vr g-ram- also increases. The cutoff is 7 when K and 108 when K = 8.
For each value of K, we evaluate (i) the model's perplexity and (ii) the average length of matched f-grams WikiText-103. 103. The results are shown in Figure 5] We find that evaluation perplexity and
4, after which it plateaus with some fluctuations. similar trend is observed for the average match length, defined as the average length of the f-grams matched by SCONE for each token in the evaluation set. The average match length also increases between K = and K = 4 and then stabilizes. This trend is likely because longer f- f-grams rarer after frequency- based ranking. As a result, even when K is larger, most selected n-grams remain short. Additionally, longer f-grams from the training corpus are less likely to match sequences in downstream data. Experiments on the WebTeex validation split (Appendix E.T) show a similar trend, although the average match length continues to increase slightly longer, plateauing around K = 6.
rises between
Considering these findings, for the experiments the remainder this paper, we set the maximum f-gram length to K 5 unless stated otherwise.
4.1.2 Varying the Number of f-grams
We observe consistent improvements in language modeling performance as we scale up VV.gram| |. To implement Ar Araram, we replicate the baseline model architecture but remove the token embedding layer. This results in the size of At.gram matches the baseline model's non-embedding parameters.
Figure 4 shows the evaluation perplexity as |Vigraml increases from 512K to 100M. On the WebText validation split, the perplexity decreases consistently as the number f-gram embeddings increases. Similarly, on 1ikiText-103, the perplexity generally decreases with more f-gram embeddings, though minor fluctuations are observed.
In Figure 4 we include three additional baselines where the non-embedding parameters of the three main models are doubled, resulting models with 204M 759M, and 1099M parameters for the original I 128M, 419M, and 589M models, respectively. This ensures that the total parameter count of each baseline matches the training-time parameter count when SCONE applied. With 100M f-gram embeddings, the 419M and 589M models trained with SCONE match or surpass the performance

---

27
2.8
Perplexity
matched length
Model size=120M
Evel perpiexitt on wastext-10
26
2.6
Model si2e
Perplciity Willexx-115
32
2.4
Model size- size=759M
30
24-
28
2.0
23
26
1.8
24
1.6
2
5 Max n nigram size
200
400
600
800
1000
1200
1400
1600
Egram model size (M)
Figure Effect of the maximum f-gram length perplexity and matched length. Perplexity decreases as the maximum length increases from to 4, then plateaus with minor fluctuations. Similarly, the average matched length stabilizes after length 4.
Figure 6: Evaluation perplexity improves as the size of Ar-gram grows. Model sizes the legend are main model sizes, including the token em- bedding layer. Dashed lines and stars on the left represent baselines without an Argrram model.
Table Zero-shot evaluation accuracy on downstream tasks. Models with f-gram embedding layers show clear improvements over the baseline models.
Model
PIQA
HellaSwag
ARC-E
ARC-C
CSQA
MMLU
Avg.
OLMo-1B
73.6
60.9
69.5
31.8
48.7
37.6
$3.7 (+0)
OLMo-1.9B
75.3
65.9
74.2
36.8
49.7
38.6
56.8 (+3.1)
OLMo-1B OE-12.8M [Huang al. 2025
73.7
62.7
70.3
32.1
49.9
37.8
54.4 (+0.7)
OLMo- 1B 10M I-grams
74.0
63.6
70.4
32.1
49.9
39.3
54.9 (+1.2)
0LMo-1.3B 10M I-grams
75.0
65,5
75.3
36,4
49.9
38.5
56.8 (+3.1)
OLMo-IB 1B f-grams
75.3
67.1
72.5
36.4
50.8
39.9
57.0 (+3.3)
of the 759M and 1099M baselines, respectively, despite using only half many non-embedding parameters during inference.
4.1.3 Varying the Size of the Ar-gram Model
We observe that, for fixed |Vi-graml, scaling up the At-rram model size provides a new way to improve language modeling performance. We vary the model size by changing the number of layers in the main model architecture. For each Amain model size, we evaluate four Ar model sizes: 0.5x, 1x, 2x, and 3x the non-embedding parameters of the main model. We set |Vt.gram to be 100M. Figure presents the evaluation perplexity on Wikitext- 103. The observations on WebText validation split are similar, and we present the results in Appendix E.
The results Figure 6 show that the perplexity generally decreases as the At-gram model size increases, although the improvements become smaller as the model size grows larger. For instance, with the 419M main model, a 100M Af-gram model improves the perplexity from 26.1 to 23.4, outperforming the 589M baseline (24.7) by a clear margin. Further scaling of the At gram model to 1020M (resulting in 1439M total parameters during training) lowers the perplexity to 22.1, which is slightly higher than the 1099M baseline (21.9). This suggests that scaling up the Argramm model initially yields better scaling curve, beyond a certain size, becomes less optimal compared to directly scaling up Amain- However, scaling Amain also increases accelerator usage during inference, whereas scaling At-gram does not, since it is replaced with an off-accelerator lookup table at inference time. This highlights method a novel way to leverage additional training compute while maintaining fixed accelerator usage during inference.
4.2 Scaling Up the Training Corpus
After exploring several design choices for SCONE, we now evaluate its performance in large-scale pretraining. Our implementation builds on the open-source OLMo codehase [Groeneveld et 2024 licensed under Apache 2.0. Additional implementation details are provided in Appendix F.2

---

Downstream Tasks. We report zero-shot accuracy on six standard downstream benchmarks: MMLU- var, Hellaswag, ARC-Challenge, ARC-Easy, CommonsenseQA (CSQA), and PIQA. In the main text, we focus on presenting downstream accuracy results under our primary training setting. Additional results, including perplexity evaluations, training curves, and further SCONE configurations under alternative settings, are provided in Appendix E.2. Notably, perplexity trends closely correlate with downstream performance.
Baselines. We compare SCONE against three baselines: OLMo- IB, OLMo-1.9D. and a concurrent method, the over-tokenize transformer [Huang et al. 2025]. Model configurations for OLMo-I IB and OLMo- 1.9 are detailed in Appendix F.2 For the over- wer-tokenized transformer, we adopt the best-performing OE- 12.8M variant, which introduces an additional input embedding layer comprising 12.8M embedding vectors, applied on top of the OLMo- 1B model. Further discussion of this method can be found in [Section All baseline models are trained for 1T tokens.
SCONE Configuration. We apply SCONE with two f-gram embedding layer sizes: 10M and 1B f-grams, respectively. The cutoff frequencies are 21,956 for 10M f-grams and 70 for IB f-grams. SCONE is integrated into OLMo-IB and OLMo-1.33 models, with an Ar-rram model size of 1.8B parameters (matching the architecture of OLMo-1.9B but excluding the token embedding layer). Models equipped with SCONE are trained for 500B tokens, half the number of tokens used by the baselines, account for the additional training cost of the Ar-gram model and ensure comparable total training FLOPS. In Appendix E.2 we also explore smaller Af-gram model of 0.6B parameters, where we observe consistent improvements.
presents the downstream accuracy results. Models with SCONE demonstrate clear improvements over the baselines. Specifically, with 10M f-grams, OLMo-1 1.3B achieves parity with the OLMo-1.99 1.B baseline while using approximately 32% less inference FLOPS and accelerator memory. With IB f-grams, OLMo- IB slightly outperforms the OLMo- .99B baseline while requiring approximately 48% less inference FLOPS and accelerator memory. Compared to the over-tokenized transformer, OLMo- IB + 10M f-grams surpasses OLMo- IB OB-12.8M, which we attribute the additional capacity introduced by the Araramm model during training.
Results. Table 1
4.3 Space Usage and Query Latency
We show that the latency of the f-gram embedding layer does not constitute bottleneck during inference, and that system memory and solid-state storage are relatively inexpensive. In Appendix E.4 we also provide a unified table summarizing the key metrics: (I) GPU memory, (2) CPU memory, (3) disk usage, and (4) FLOPS/iatency for training and inference.
Table 2: Space usage of the f-gram embedding layer F, along with unit cost for memory and NVMe solid-state drives [McCallum] 2024).
n-grams
System memory (GB)
Solid state drive (GB)
107
41.4
77.3
1
413.6
766.8
10"
(does not fit)
7665.4
Price (per GB)
2 USD
0.1 USD
We experiment with |Viggraml being 10M, 100M, and IB with embedding dimension of d = 2048 and 16-bit precision per floating point value. Experiments were conducted on workstation with 64 Intel Xeon CPU cores and 512 GB of memory. Space and latency were measured for both in-memory and on-disk storage. In memory, embeddings are stored
single matrix with hash dictionary
mapping f-grams to indices, while on- on-disk storage uses the Lightning Memory-Mappe Database [Chul [011] to directly store f-gram and embedding pairs NVMe solid-state drives.
Table summarizes the space usage for both storage methods. In both cases, the space required increases linearly with the number of embedding vectors. The 10M and 100M f-gram embedding layers are able to fit within main memory, with the 10M layer requiring 41.4 GB. For on-disk storage, there is additional overhead as the same 10M layer occupies 77.3 GB storage.
8

---

System memory
Solid-stake stake drive
0.08
1 10
1 10
0.07
x 10
2.5
1 10
0.06
x 10
Latency (ms)
Latency (ms)
2.0
0.05
0.04
1.5
0.03
1.0
0.02
0.5
0.01
0.00
0.0
Batch size
Batch
Figure 7: Amortized per-token query latency (ms), averaged over 100,000 batches. The latency spike from batch size 2 when reading from system memory due to batch operator overhead, which is less pronounced for solid-state drives.
shows the latency of retrieving embeddings at different batch sizes. Latency is measured as the end-to-end time from loading batch of tokens having the f-gram embeddings ready on the GPU. Before each test, the CPU cache is cleared. Up to four queries per token are performed to find the longest matching n-gram (for a maximum n-gram length of K = 5). For in- me-memory storage, sequential queries are sufficient since they are not the bottleneck, whereas on-disk storage uses parallel queries to the database. At a batch size of 1, retrieval from 10M f-gram embedding layer on NVMe takes 1.1 ms, increasing to 2.3 ms for IB-layer: both well below the latency threshold for LLM inference (typical commercial APIs generate at 100 tokens/sec, or 10 ms/token AntificialAnlys 2025)). Larger batch sizes further improve efficiency: at batch size 16, amortized per-token latency drops to 0.5 ms. In-memory access is much faster: for a 100M f-gram layer at batch size 16, per-token latency is only 0.017 ms. We also report end-to-end generation speed in Figure 2] which aligns with these latency trends: when embeddings are stored in main memory, there negligible impact on throughput; when stored NVMe drives, throughput slightly decreases but no major bottleneck arises.
Figure 7
5 Related Work
Scaling of Embedding Layers. Prior research on the scalability of embedding layers has primarily focused on token embeddings, where the vocabulary is shared with the language model's decoding head. For instance, [Tao et al. [2024) show that larger models benefit from larger vocabularies,
reflecting trends where vocabulary sizes have grown from tens of thousands (Devlin et al. 2019 Radford et al.) 2019 hundreds of thousands of tokens [Google, 2024] Adler et al.) 2024| Dubey et al. 2024 Liu et al. 2024]. However, even for the largest models, the optimal vocabulary sizes predicted by Tao al. 2024] remain much smaller relative model size. For instance, a vocabulary 216K tokens for 70B-parameter model. This relatively small token embedding layer limits the potential for expanding model capacity through token embeddings alone. To address this, Roy et al.) [2022) propose decoupling input and output embeddings by introducing additional embedding layer for bi-grams. Building on this idea, we also decouple the input and decoding embedding layers, but crucially, we parameterize the additional embedding table with neural network during training. This approach allows us to efficiently scale the additional input embedding layer and introduces a new scaling strategy: scaling the At Arggrmm model that learns to generate the additional embeddings.
Concurrently, [Hang et al. 22025] propose the over-tokenized transformer, which also decouples input and decoding embeddings by introducing an additional input embedding layer for n-grams. They observe similar performance gains as the input embedding size increases. Unlike our approach, which selectively retains only frequent n-grams, they hash n-grams into fixed number of embeddings to manage the vast n-gram space, a strategy that likely also mitigates sparse updates (Appendix D). A key distinction is that their additional embedding layer must be fully instantiated and resident on accelerators during training, leading to significant memory challenges despite tensor sharding. Moreover, through the design of Ar.gram SCONE expands model capacity not only by enlarging the embedding table but also by introducing a scalable Ar Ar.rram model for learning the embeddings. We compare SCONE with the over-tokenized transformer in Section 4.2 Results suggest that while

---

both SCONE and the over-tokenne wer-tokenized transformer show clear improvements over the baseline, SCONE offers better gains owing the additional capacity provided by Argram
Mixture of Lookup Experts. A concurrent work by Jie et al. 2025 proposes discretizing the inputs to Mixture-of-Experts (MoE) layer so that, at inference time, it can be implemented as simple lookup table. They use token embeddings from the input embedding layer as keys for the lookup. This design enables training with larger MoE layer to improve performance, while having negligible impact on FLOPs and accelerator memory usage during inference. Their approach shares a similar insight with our design of the Af-gram model: constructing neural network module with discretized inputs that can be switched to an efficient lookup table at inference. However, a key difference is that we introduce method for scaling the number of keys by constructing the set Vi-gram- This enables an additional axis of scaling, namely expanding the size of the lookup table at inference. It also allows for fully utilizing the increased model capacity during training. Without this extension, small number of discretized keys would lead to diminishing returns when scaling the parameters of the discretized module during training.
We discuss additional related work in Appendix due to space constraints.
6 Conclusion
We introduce SCONE, a scalable approach for generating r-gram contextualized embeddings for each input token. These embeddings are learned during training and cached in off afceleceerrtorr storage for inference. SCONE enables two new strategies for scaling language models under fixed inference-ting accelerator memory and FLOPS budgets, making particularly useful for reducing serving costs and supporting latency-sensitive applications.
We discuss limitations and promising directions for future work in Appendix A
Acknowledgments
The authors would like to thank Andrew Tomkins for his helpful feedback on an early draft.
References
Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H Anh, Pallah Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, et al. Nemotron-4 340b technical report. arXiv:2406.11700 2024.
Mikel Artetxe, Sebastian Ruder, and Dani Yogatama. On the cross-lingual transferability of monolingual representations. In ACL, pages 4623 4637, 2020.
ArtificialAnlys. Independent analysis of Al models and API providers, 2025. artiificialanalysis.. ai// Accessed: 2025 1025-01-07.
URL https://
Vincent-Pierre Berges, Barlas Oğuz, Daniel Haziza, Wen-tau Yih, Luke Zettlemoyer, and Gargi Gosh. Memory layers at scale. arXiv:2412.09766 2024.
Yihong Chen, Bei Chen, Xiangnan He, Chen Gao, Yong Li, Jian-Guang Lou, and Yue Wang. Aopt: Learn to regularize recommender models in finer levels. In KDD, pages 978-986, 2019.
Yihong Chen, Kelly Marchisio, Roherta Raileanu, David Adelani, Pontus Lars Erik Saito Stenetorp, Sebastian Riedel, and Mikel Artetxe, Improving language plasticity via pretraining with active forgetting. In NeurlPS, pages 31543-31557, 2023.
Yihong Chen, Xiangxiang Xu, Yao Lu, Pontus Stenetorp, and Luca Franceschi. Jet expansions of residual computation. arXiv:2410.06024, 2024.
Dokook Choe, Rami Al-Rfou, Mandy Guo, Hecyoung Lee, and Noah Constant, Bridging the gap for tokenizer- free language models. arXiv: 10322, 2019.
Howard Chu. Lightnnng mappod database, 2011. URL http:://www: Imdb. tech/ 2025-01-23.
Accessed:
10

---

Gautier Dagan, Gabriel Synnaeve. and Baptiste Roziere. Getting the most out of your tokenizer for pre- vre-trainin and domain adaptation. In ICML, 2024.
DeepSpeed. DeepSpeed, 2024. URL
Accessed: 2025-01-19,
Bjöm Deiseroth, Manuel Brack, Patrick Schramowski, Kristian Kersting, and Samuel Weinbach. T-FREE: Subword tokenizer-free generative LLMs via sparse representations for memory-cfficient embeddings. In EMNLP, 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv:/810.04800 2019.
Longxu Dou, Qian Liu, Guangtan Zeng, Jia Guo, Jiahui Zhou, Xin Mao, Ziqi Jin, Wei Lu, and Min Lin. Sailor: Open language models for south-Fast Asia. In EMNLP, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama 3 herd of models. arXiv:2407.21783, 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. JMLR, 2022.
Philip Gage. A new algorithm for data compression. The Users Journal, 1994.
Mor Geva, Roci Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. In EMNLP, pages 5484 2021.
Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In EMNLP, pages 30-45, 2022.
Google. Gemma 2: Improving open language models at practical size. arXiv:2408.001/8 2024.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinncy, Oyvind Tafjord, A. Jha, Hamish Ivison, lan Magnusson, Yizhong Wang. Shane Arora, David Atkinson, Russell Authur, Khyathi Raghavi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Daniel Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew E. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Schwenk, Saurabh Shah, Will Smith, Emma Strubell, Nishant Subramani, Mitchell Wortsman, Pradeep Dasigi, Nathan Lambert, Kyle Richardson, Luke Zettlemoyer, Jesse Dodge, Kyle Lo. Luca Soldaini, Noah A. Smith, and Hanna Hajishirzi. OLMo: accelerating the science of language models. In ACL, pages 15789- 15809, 2024.
Prakhar Gupta, Matteo Pagliardini, and Martin Jaggi. Better word embeddings by disentangling contextual n-gram information. In NAACI-HLT,, pages 933-939, 2019.
Xu Owen He. Mixture of a million experts. arXiv:2407.04153 2024.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. 220: 15556, 2022.
Hongzhi Huang, Defa Zhu, Banggu Wu, Yutao Zeng. Wang, Qiyang Min, and Xun Zhou. Over-tokenized transformer: Vocabulary is generally worth scaling. In ICMI. 2025.
Yuzhen Huang, Jinghan Zhang, Zifei Shan, and Junxian He. Compression represents intelligence linearly. In COLM, 2024.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv:2401, 00408,, 2024.
Shibo Jie, Yehui Tang, Han, Yitong Li, Duyu Tang, Zhi-Hong Deng, and Yunhe Wang. Mixture lookup experts. In JCML, 2025.
Jeff Johnson, Matthijs Douze, and Hervé Jégou, Billion-scale similarity search with GPUs, IEEE Trans. Big Data, 2019.
Andy Jones. Scaling scaling laws with board games. arXi::2104.03113. 2021.
Armand Joulin, Moustapha Cissé, David Grangier, Hervé Jégou, et al. Efficient softmax approximation for GPUs. In ICML, pages 1302 1310, 2017.

---

Yoon Kim, Yacine Jernite, David Sontag, and Alexander Rush. Character-aware neural language models. In AAAI, 2016.
Taku Kudo. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. In EMNLP (Demonstraation) pages 66-71, 2018a. 8a.
Taku Kudo. Subword regularization: Improving neural network translation models with multiple subword candidates. In ACL., pages 66-75, 2018b.
Woosuk Kwon, Zhuohan Li. Siyuan Zhuang. Ying Sheng, Lianmin Zheng. Cody Hao Yu, Joseph E. Gon- zalez, Hao Zhang, and lon Stoica. Efficient memory management for large language model serving with PagedAttention. SOSP, 2023.
Guillaume Lample, Alexandre Sablayrolles, Marc' Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou. Large memory layers with product keys. In NIPS, 2019.
Dmitry Lepikhin. HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. (GS)hard: Scaling giant models with conditional computation and automatic sharding. In ICLR, 2021.
Davis Liang. Hila Gonen, Yuning Mao, Rui Hou, Naman Goyal, Marjan Ghazvininejad, Luke Zettlemoyer, and Madian Khabsa. XLM-V: Overcoming the vocabulary bottleneck in multilingual masked language models. In EMNLP, 2023.
Xianwen Liao, Yongzhong Huang, Changfu Wei, Chenhao Zhang, Yongqing Deng, and Ke Yi. Efficient estimate of low-frequency words' embeddings based the dictionary: case study Chinese. Applied Sciences, 2021.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang. Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv:2412. 19437, 2024.
Siyi Liu, Chen Gao, Yihong Chen, Depeng Jin, and Yong Li. Learnalee emhedding sizes for recommender systems. In ICLR, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019.
John C. McCallum. Price and performance changes of computer technology with time, 2024. URL https net.iii.ett/idex... htm Accessed: 2025-01-20
Bryan McCann, James Bradbury, Caiming Xiong, and Richard Socher. Learned in translation: Contextualized word vectors. NIPS, 2017.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In ICLR, 2017.
Meta. Large concept models: Language modeling in sentence representation space. arXiv:2412.08821 2024.
Tomás Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. In ICLR (Workshop), 2013.
Artidoro Pagnoni, Ram Pasunuru, Pedro Rodriguez, John Nguyen, Benjamin Muller, Margaret Li, Chunting Zhou, Lili Yu, Jason Weston, Luke Zettlemoyer, et Byte latent transformer: Patches scale better than tokens. 21X 2024.
Matthew E. Peters, Mark Neumann, Mohit lyyer, Mau Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations. In NAACL-HILT, 2018.
Joshua Peterson, Stephan Meylan, and David Bourgin. Openwehtext, 2019. URI htttps::t//:/ com/ sysetterson//peexebbeexx Accessed: 2025-01-19.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding with unsupervised learning. Technical report, OpenAI, 2018
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Aurko Roy, Rohan Anil, Guangda Lai, Benjamin Lee, Jeffrey Zhao, Shuyuan Zhang. Shibo Wang, Ye Zhang, Shen Wu, Rigel Swavely, et al. N-grammer: Augmenting transformers with latent n grgrmms, arXivv2207.06366 2022.
12

---

Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In ACZ, 2016.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gates gated mixturre-f-expertt layer. In ICLR, 2017.
Kyuhong Shim, Minjae Lee, Iksoo Choi, Yoonho Boo, and Wonyong Sung. SVD-softmax: Fast softmax approximation large vocabulary neural networks. In NIPS, 2017.
Charlie Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling L.I. test-time compute optimally can be effective than scaling model parameters. In ICLR. 2025.
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Magnusson, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew Peters, Abhilasha Ravichander, Kyle Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani. Oyvind Tafjord, Evan Walsh, Luke Zettlemoyer. Noah Smith. Hannaneh Hajishirzi. Iz Beltagy. Dirk Groeneveld, Jesse Dodge, and Kylc Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research. In ACI,, 2024.
Sainbayar Sukhbaatar, Jason Weston, Rob Fergus, et al. End-to-end memory networks. NIPS, 2015.
Chaofan Tao, Qian Liu, Longxu Dou, Niklas Muennighoff, Zhongwei Wan, Ping Luo, Min Lin, and Ngai Wong. Scaling laws with vocabulary: Larger models deserve larger vocabularies. In Neur/PS, 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv:2307. 092888 2023.
Elena Voita, Javier Ferrando, and Christoforos Nalmpantis, Neurons in large language models: Dead, n-gram, positional. In ACL (Findings), pages 1288 1301, 2024.
Hai Wang. Dian Yu, Kai Sun, Jianshu Chen, and Dong Yu. Improving pre-trained multilingual model with vocabulary expansion. CoNLL, 2019.
Junxiong Wang, Tushaar Gangavarapu, Jing Nathan Yan, and Alexander M Rush. MambaByte: Token-free selective state space model. In COLM, 2024.
Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. NIPS, 2015.
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoe V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao. Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah. Melvin Johnson, Xiaobing Liu, Lukasz Kaiser. Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young. Jason Smith, Jason Riesa, Alex Rudnick. Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google's 's neural machine translation system: Bridging the gap between human and machine translation. orXiv:1609.. 08/ 44, 2016.
Linting Xue, Aditya Barua, Noah Constant, Rami AI-Rfou Sharan Narang, Mihir Kale, Adam Roberts, and Colin Raffel. Byt5: Towards token-free future with pre- pre-trained byte-to-byte models. TACL, 2022.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint 2rXiv:2555.00935 2025.
Lili Yu, Dániel Simig. Colin Flaherty. Armen Aghajanyan. Luke Zettlemoyer. and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers. NewrIPS, 2023.
Bo Zheng, Li Dong, Shaohan Huang, Saksham Singhal, Wanxiang Che, Ting Liu, Xia Song, and Furu Wei. Allocating large vocabulary capacity for cross lonsull language model pre pre-training. In EMNLP, 2021.
13

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: All claims made are supported by detailed experimental results. Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the
results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We discuss the limitations of our work in Appendix and also reference them in Section
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: We do not present new theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
14

---

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yss]
Justification: We present implementation details required to reproduce all results in both Section 4 Appendix F.I. and Appendix F.2 We will release our code after the reviewing process.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e-g.. in the case of a large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a way to access this model for reproducing the results or a way to reproduce the model (e.g.. with an open-source dataset or instructions for how to construct the dataset).
model (e.g., large language model), then there should either be
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-soure models, it may be that access to the model is limited in some way (c.g.. to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
All experiments in this work build well-known opensource framework and datasets.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurIPS code and data submission guidelines lattps://nips.cccppulll
guides/Cood Coseee//ooSSSubissse. ssionpolccyy for more details.
While we encourage the release code and data, we understand that this might not be possible,
"No" is acceptable answer. Papers cannot be rejected simply for not including code. unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines
guidess/ooeeuubeeee wtt.ddddddnssssoooooooooooo onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
15

---

The assumptions made should be given (e.g., Normally distributed errors). The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why. At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
Experimental setting/details Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary understand the results? Answer: [Yes]
Justification: We present training and evaluation details both Section Appendix F.1 and Appendix F.2
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or supplemental material. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We made efforts include statistical information in some of our experiments, such as in Figure 7 and Figure 5 However, for the large-scale pre-training experiments, we did not conduct formal statistical significance testing due the prohibitive computational cost. Nevertheless, we observe consistent performance improvements across different settings. Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1 I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (c.g. negative error rates). error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [Yes]
Justification: We describe the hardware used in this work in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well as estimate the total compute.
16

---

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't 't make
into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, every respect, with the NeurIPS Code of Ethics httpps:/neurr: http://p.....p..tttthhoooooooooa
Answer:
[Ye]]
Justification: We confirm we have reviewed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there a special consideration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The goal of our work to advance the field of language models. Our work does not introduce new potential societal consequences beyond those already associated with LLMs, and therefore we do not believe any specific impacts need to be highlighted.
Guidelines:
The answer NA means that there is socictal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g.. gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse. mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e-g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
17

---

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We've ensured that all assets used in this work are properly licensed and have provided references and links their sources.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwpprswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's
creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects Question: For crowdsoureing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
18

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required. Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer our LLM policy htte::///aaarrgecccco/ooorrrrrrr for what should or should not be described.
19

---

A Limitations and Future Work
promising direction for future work is extending SCONE beyond short n-grams by enabling caching for longer queries. A central challenge lies in designing effective keys for such queries. Using raw text as keys may result in low cache hit rates, as semantically similar queries often vary at the surface level. Alternatively, using semantic embeddings as keys would require discretization techniques to map continuous embeddings into discrete key space that supports efficient indexing.
One limitation of the current study is that we evaluate SCONE only on models with up to 3B parameters at training time. This constraint is primarily due to hardware limitations, which restrict our ability to conduct large-scale pretraining on larger models. Although the model scales we tested widely used in real-world applications, exploring the performance of SCONE at larger scales presents an exciting direction for future research. We believe applying SCONE to larger models would be highly beneficial to the community and leave this exploration for future work. Notably, our experiments in Seccionn 4 provide encouraging evidence, as the benefits of SCONE are consistent across all model sizes we tested.
B Additional Related Work
Contextualized Word Embeddings. Words can have different meanings depending on context. Prior work has incorporated context into word embeddings, either from the entire sequence [McCann et al.. 2017 Peters et al. 2018) or short n-grams [Gupta et al. [211)]. before applying them to down- stream tasks. Modern language models inherently use contextualized token embeddings, leveraging attention mechanisms. In this study, we extend the embedding layer to include contextualized f-gram embeddings for each token. A key novelty that our approach allows embeddings to be precomputed and offloaded from accelerators, providing contextual embeddings for each token without increasing FLOPS and accelerator memory usage at inference.
Tokenization in Language Models. Our method assumes predefined vocabulary from a trained tokenizer. Several popular algorithms exist for training tokenizers (Sennrich al., [0166] Wu et al.. 2016 Kudo 20] In this work, we use a BPE tokenizer, following prior seminal works [Radford] jt al. 2019 Touvron et al.] 2023]. However, our method not tied to any specific tokenization algorithm and can be applied seamlessly to others.
Tokenization-freee -ree language models have also been widely explored [Kim al. 2016 Choe et al.
2019 Xue et al. [2022] Yu et al. 2023 Wang et al. [2024] Deiseroth et al.] 2024 Meta 2024 Pagnoni al. 2024]. While we have not tested our method on tokenization-free models, we believe our core
idea introducing an off-accelerator embedding layer by precomputing embeddings for frequent
input patttmastteetteee applicable.
Mixture-of-Experti (MoE) and Memory Layers. MoE and memory layers are two established approaches for scaling language models within fixed FLOPS budget.
MoE layers replace traditional feedforward layers with multiple parallel "experts," activating only one few per token using lightweight routing mechanism [Shazeer al. 2017 Lepikhin et al.) 2021 Fedus et al.] [022]] Jiang et al., 2024] He [2024]. This allows the model to scale by increasing the number of experts without increasing the computational cost per token. However, all experts must reside on the accelerator, resulting in significantly higher memory usage.
Memory layers, on the other hand, store large collections of embeddings (continuous vectors) and retrieve the nearest neighbors during the forward pass via (approximate) similarity search [Weston et al., 2015 Sukhbaatar et al., 2015. Lample et al.) 2019 Berges et al., 2024]. These retrieved embeddings enhance the model's 's capabilities without adding much to the FLOPS budget. Despite
improvements in similarity search techniques [Lample et al. 2019 Johnson al. 2019], memory layers still require storing the embeddings on the accelerator, making memory demands impractically high at larger scales (Berges et al.] 2024]. Furthermore, because embeddings are typically updated via backpropagation, memory layers introduce additional challenges related to sparse updates as the memory size grows.
While MoF. layers, memory layers, and our proposed method SCONE all maintain fixed inference FLOPS, key advantage of SCONE is it also maintains fixed accelerator memory usage at inference.
20

---

Algorithm 2 Basic Next-Word Prediction Model M,,A.
Parameters:
T Vioken + R token embedding layer,
A: (Rd)ST > IR: transformer model, where 7 is the maximum sequence length,
D Rd A Vesken prediction head.
Input: (o1,..., Tm) Vioken for S T.
Output: Probability distribution over next token 0m+1-
for i = 1,..., do
E: T(.)) Input embedding per token.
Cout + A(e em): Output embedding. return ((euut)
Algorithm 3 Constructing a set of f-grams Vi.raam
Parameters: S: desired size of Vi-gram-
Input: ((o.,...,or)" token sequences from training set.
Output: Vi-gram v=(2,k],
set of f-grams of size S.
for 2..... K do
for w := (o (o' V token do
C f the number of times w appears in all sequences f(o1....,oy") or) Let w1, ... be list of elements of U V ties arbitrarily.
token sorted such that Cwn, 2 Cwn 2 ... breaking
return {wi...., ws}: set of f- f-grams of size S.
By focusing on the input embedding layer, SCONE ensures that the computational overhead remains O(1) during inference and enables offloading the additional parameters off-accelerator with negligible impact end-to-end latency.
Implicit n--ram Patterns in Transformers. Recent work analyzing the internal mechanisms of transformers has shown that these models often utilize implicit n- R-gram patterns for prediction [Geva
et al., 2021 2022, Voita et al.] 2024). For instance, Chen et al. [2024] demonstrate that certain attention heads can detect specific n-gram patterns, while MLPs can perform linguistic operations such adding the "-ing" suffix. These findings underscore the importance of r--grmm information in language modeling and offer a potential explanation for the effectiveness of SCONE. An interesting future direction is examine how SCONE'S f-gram embeddings interact with the transformer's implicit n-gram patterns.
Embedding Sparsity in Multilingual Applications and Recommender Systems. This work focuses on a common setting for training LLMs: language modeling on large-scale text corpora, primarily in English. However, scaling embedding layers presents challenges beyond this context, particularly due to frequency-related performance degradation caused by sparsity. Multilingual applications
such scenario. Two phrases in different languages may refer to the same concept but correspond to different embedding vectors. Their embeddings should ideally be close. Recent work has explored methods for learning transferable embeddings in cross- cross-lingual settings [Artetxe] let al. 2020 systems [Chen al. 2019 Liu et al. 2022]. where embeddings often dominate the model's parameter count due to the high cardinality of user or item categories. For both scenarios, SCONE 's strategy, i.e., parameterizing large embedding tables using neural network, provides complementary approach to help mitigate sparsity issues.
Chen Chen et al. [023]. Another relevant example scaling the embeddings for recommender
C Additional Algorithms
21

---

Algorithm 4 Next-word prediction with SCONE Mr, Vegraom Atigram I Amain.
Parameters:
T: <Vooken - RRd token embedding layer, Vi-gram S yla.k), set of f-grams,
TRAINING: f-gram transformer model
D Ai-gram (IR)) + Rd
INFERENCE: f-gram embedding layer
b DFF Vi.gram - Rd.
A.aain : (R4)ST + RR. main transformer model.
D Rdd > Vioten" Prediction head.
Input: (01,-.., Tm) V Vtoken for S T. where T is the maximum sequence length.
Output: Probability distribution over next token om+l.
(el...., em) t Fy,, Vegram.Aeaan (o Tm) (Algorithm
Cout + AmAiin(ee,,..... em)
return D(eout)
we discuss simple next-word prediction model, MA,AD consisting of token embedding layer T, a transformer model A, and a prediction head D. This model takes a token sequence (...... Jm). with each token from the token vocabulary Vioken, and produces a probability distribution for the next token. We provide the pseudocode for MT,A,D in Algorithmm 2]
In Section 2)
In Section introduce algorithm for constructing grams. We present the pseudocode for the construction process in Algorithm 3] Although Algorithm 3 clearly illustrates the procedure, it is too expensive implement in practice. In /Section 3.1|
of f-grams given a target number of f-
describe an efficient implementation that requires only K linear scans over the pre-training corpus, where K is hyperparameter controlling the maximum length of f-grams.
In Algorithm tualized f- --gram embeddings. Next, describe the end-to-end next-word prediction process using SCONE (Algorithm 4p. Specifically, the process, denoted as MT Aexem IF. .F.Auiia takes an input sequence (0 .m) E Vooken and produces a distribution over the next token 6m+1- Note that in Allorithh 4) f-gram embeddings generated with A gram during training and retrieved from a lookup table F during inference.
in Section we present the pseudocode for SCONE'S process of generating contex-
D Challenges of Scaling Vocabulary Size in Embedding Layers
Scaling the token vocabulary size is the most straightforward way enlarge embedding layer, but we find that larger token vocabularies degrade performance beyond a certain threshold and significantly increase accelerator usage during decoding. We pre-train GPT-2 models [Radford et al., 2019] with three sizes of non-embedding parameters: 85M (small), 302M (medium), and 708M (large) the WebText dataset (Peterson al. 2019]. testing six vocabulary sizes ranging from 32,768 to 2,097,152. The tokenizers are trained using the BPE algorithm [Gage] 1994 (Sennrich et al. 2016]. We follow the implementation in Tao al. [2024), which allows token merges across word boundaries. Each model is trained on 80B tokens. Since larger vocabularies produce fewer tokens for the same dataset, they effectively enable models process
data. Additional implementation details such as training hyperparameters are provided in Appendix F.
Figure g presents the average bits per character (BPC) the WebText validation set. We report BPC instead of cross-entropy loss because the latter sensitive to vocabulary size, with larger vocabularies typically producing higher losses. BPC, by contrast, is common vocabulary-insensitiee metric for comparing models trained with different tokenizers [Huang et al. 2024). We observe that BPC for all three models initially improves with larger vocabulary sizes but eventually deteriorates.
Figure 9 shows the percentages tokens that receive more than a given number of updates over 100M training tokens. In standard embedding layers, gradients are directly backpropagated to the embedding vectors. With a fixed number of training tokens, larger vocabularies lead to fewer updates per token. For vocabulary size of 2,097,152, only 7.3% of tokens receive more than 100 updates,
22

---

GP12 small CPT-2 meldum large
1.18
0.98
0.46
0.94
Figure 8: BPC of three model sizes on the validation set (lower is better). For all three model sizes, BPC initially improves as vocabulary size increases but eventually deteriorates.
Poreentte 28: Caane
32K
128K.
512K
1024K
Trachei count
Figure 9: Percentages of tokens (y-axis) that receive more than given number of updates (I-axis), measured over 100M training tokens. As the vocabulary size increases, tokens receive increasingly sparse updates.
compared to 97.6% for vocabulary size of 32,768. This suggests that the performance drop for larger vocabularies may stem from sparse updates to per-token embedding vectors.
GPT-2 small
GPT-2 small
2500
GPT-2 meidum
GPT-2 meidum
Emaxebing perameeeen (nt 2e14
large
large
6RU (08)
1900
500
Vocabalery size
Vocabulary
Figure 10: Number of embedding layer parameters stored on GPU and the corresponding memory usage. For large vocabulary sizes, the memory usage increases linearly with the vocabulary size.
In addition to performance degradation, increasing the vocabulary size significantly raises accelerator usage during the inference stage. This is because predicting the next token involves running linear layer and softmax operation across the entire vocabulary to identify the closest embedding. Figure 10 illustrates that both the number of embedding layer parameters stored on the GPU and the GPU memory cost increase linearly with vocabulary size. These costs are measured using a batch size of 1, sequence length of 1024, and 16-bit precision.
23

---

Perplexity
Aug. matched length
Model size-
Model sizo-415M
1.0
Model 5tc0- B5M
Model 1tte 24M
perpledity on
204
2.6
Aappor wea
Model sizem] DO0M
177
2.4
16
15
1.8
Max n grom 5ize
400
1200 1400 1600
fgram model sioe (M)
Figure 11: Effect of the maximum f-gram length in Vi-grams evaluated on the WebText validation split.
Figure 12: Evaluation perplexity on WebText as function of the size of A Argram
E
Additional Experiments
E.1
More Results for Training on WebText
Varying Maximum f-gram Length. In Sectionn 4.1.1 we discuss the impact of varying the maximum f-gram length in Vi.gram and present results on Wikitext-103 We observe that a relatively small maximum length is sufficient, as long as is not too small, otherwise, the number of available n-grams for ranking becomes too limited. Here, in Figure 11 we show the corresponding results on WebTeet, which exhibit similar trends. The left y-axis represents the evaluation loss (averaged over three seeds), with the leftmost star indicating baseline performance. The right y-axis shows the average length matched f-grams. As the maximum size increases, the loss initially decreases but then plateaus with some fluctuations. Meanwhile, the matched length rises initially before stabilizing for larger values.
Varying Afgram Model Size. In Section 4.1.3 we discuss the impact of varying the size of Aigrram on evaluation perplexity for Wikitext- 103. We find that increasing the model size leads further performance improvements for a fixed VV-graml- In Figure 12 we present the results on WebText, which show similar trend. Model sizes in the legend correspond to inference time sizes on accelerators. Dashed lines and stars on the left represent baseline performance. The evaluation perplexity improves as the size of AF grows.
E.2 Additional Results for Training on the OLMo Corpus
We present training curves and perplexity evaluations for models trained on the OLMo corpus. For SCONE, we introduce an additional Af-gram model size of 0.6B, alongside the 1.8B At-rram model discussed in [Section 4.2) Due to the large number of experiments, all models in this section are trained for 2008 tokens, constrained by computational resources. This differs from Section 4.2 where SCONF-enabled models are trained for 500B tokens and baseline models for IT tokens.
Duusiinn mode aDE 0.78
Daucd no modd JBE 23
Dusi ins modul succ 1.22
33.0
wllee groass 88 gam mode)
#ioe granss 681 gam madee)
33
1106 gr 48T mode)
1OR1 groest d81 gan mode
-10e gramect gon moded
granee model
27.5
fgnin
27.
modil
275
nodel
23.0
20:0
20.0
3755
175
25
52
183
125
350
375
270
23
57
75
120
125
150
3.5
200
25
30
75
1.00
150
375
200
Figure 13: Average perplexity on the OLMo evaluation mixture throughout training. Models with SCONE enabled converge later, indicating stronger capacity, and achieve better perplexity.
24

---

Training Curves. Figure 13 shows the evaluation perplexity curves for OLMo-0.7B, OLMo- IB, and OLMo-1.3B throughout training. The curves indicate that models trained with SCONE converge more slowly, suggesting that SCONE effectively increases model capacity. Furthermore, both increasing the number of f-grams and enlarging the A gram model size enhance model capacity.
Table 3: SCONE consistently improves perplexity (lower is better) across all evaluation corpora. We train three baseline models models, apply SCONE using four different configurations and present the results directly below each corresponding baseline.
with sizes of 1B, 1.3B, and 1.9B parameters. For the 1B and 1.3B baseline
Model size
o4-en
books
common-craw]
pes2u
reddit
stack
wiki
ice
m2de-s2re
pile
14 11kites-1033
Averaige
IB baselime
16813
21,570
16.752
11,682
22.612
3.360
14,453
15.281
27.900
10.429
16,053
16,082
+10M Veggann (0.6R Auguun
16.087
20.963
16,099
11,270
21.797
13.777
14.979
26.
10.128
15.371
15:499
+10M Vegran (1.8 Al Acguas)
15.727
20.429
15:473
11.124
21 388
3
13,454
14,709
25.785
9.956
15.104
15.125
Vegpan Ausuaal
15.684
13.543
+18 Vegaa (1 (1.8 Augrnn)
15.158
19,680
14857
10.761
20.757
3,133
12,964
14.230
24.95%
9.553
14,354
14.581
1,3 baseline
15,994
20.157
15.921
11,148
21.644
3,248
13.721
14,651
26.583
9.927
15.143
15.264
+10M V/ goe (0.68 A !
15,509
19.816
15.407
10.857
222
3,192
13,260
14.372
25.450
9.757
14,616
14.844
+10M Viggon (1.00 A A
15.193
19.557
14,995
10,795
20.735
3,171
13,071
14.272
25.258
9.674
14,438
14,654
Vegpan A
13,073
+18 ya (1 18 At
14,03
18.996
14:541
10,502
20 206
3.085
12.637
13.971
24.533
9.357
13.971
14.245
1.91 baseline
15,270
19,017
15.184
10.719
20.752
3,163
13.119
14,095
25,461
9.570
14,229
14,598
Perplexity Evaluation. Figure presents perplexity results the OLMo evaluation mixture which covers 11 diverse corpora, including web crawl data, literature, online forums, scientific writing, coding, and more. TTblle 3 details the performance breakdown by corpus. Results show that increasing both |Vtgraam and the size of Af-gram consistently improves performance across all datasets.
Additionally, Figure Seguuee2repore reports token generation speeds measured using the vLLM framework [K Kwon et al. 2023] with a batch size of 1. Even with large |Vr.gaml|, embedding retrieval remains efficient and is not bottleneck for inference.
As a representative example, in the 1B model variant, the baseline achieves an average perplexity of 16.082. Setting Vi-gramm 10M improves perplexity to 15.459 with a 0.6B Ar-gram model and to 15. with a 1.8B At-gram model- the latter outperforming the 1.3B baseline (15.284). Further increasing |Vi-gram| to IB improves perplexity to 15.187 (0.6B Ar-gram) and 14.581 (1.8B At-gram,, surpassing the 1.9B baseline (14.598) while requiring only about half the FLOPS and accelerator memory at inference time.
E.3 Apply SCONE in Post-training
We apply SCONE to supervised fine-tuning of Qwen3-4B-barl using the open-rl framework and the open-r htttt////iuuu-o--ooougge while Qwen3-8B-base or Swen3-14B-bbae are used the f-gram models, We set the number of f-grams to 10M and follow the training hyperparameters in open-rl. Table Joolprcssperes the resulting ScoNE-enabled models with the Qwen3-4B baseline in terms of both accuracy and decoding latency.
datase(s this setup, Qwen3-4B serves as the main model,
Decoding Latency (per-token) (ms)
AIME 2024
LiveCodeBench v4_v5 pass@1
Model
pass @1
Qwen3-4B-base
45.3
30.8
10.05
SCONE-4B (8B f-gram model)
48.3 (+3.0)
34.5 (+3.7)
10.13
SCONE-4B (14B f-gram model)
51.6 (+6.3)
36.3 (+5.5)
10.13
Table Performance comparison of Qwen3-4B variants on AIME 2024 and LiveCodeBench v4/v5 with decoding latency.
yam
http:///ihhhbbb.c////ggggggggggg
25

---

These results show that SCONE consistently improves performance over the baseline, with larger f-gram models yielding greater gains, all while maintaining similar inference latency.
E,4 Summary of Comparison on Computational Resources
In Table 5]
summarize the key metrics of SCONE for the three model variants evaluated in
Section 4.2 The metrics include (1) GPU memory, (2) CPU memory, (3) disk usage, and (4) taken with context length of 2048 and a batch size of 4 on a single 1100 80 GB GPU. The three settings are: (1) the 1.9B baseline model, (2) SCONE 1.3B, with 1.8B f-gram model and 10M cached f-gram embeddings, and (3) SCONE IB, with a 1.8B f-gram model and IB cached f-gram embeddings.
FLOPS/latency for training and inference. All measurements
Peak GPU Memory (GB)
CPU Memory Overhead (GB)
Disk Usage (TB)
Decoding Latency (per-token) (ms)
Training FLOPS (per-seq.) (x10
Model Variant
1.9B baseline
8.38 5.60
N/A
N/A
6,45
2.73
SCONE-1.3B
41.76
N/A
4.83
4.94
(10M [-grams)
SCONE-IB
4.45
7.67
5.57
(IB [-grams)
Table 5: Comparison of SCONE model variants on memory, storage, latency, and compute efficiency.
Both memory and disk usage are reported for inference. The decoding latency is averaged over one thousand decoding steps. As shown in /Section 4.2] all three models achieve similar downstream performance. Compared to the 1.9B baseline. SCONE-enabled models significantly reduce GPU memory usage and decoding latency, at the cost of increased CPU memory or disk storage.
Implementation Details
We provide additional implementation details below. Most of our experiments are conducted on 4 8xH100 nodes, while some experiments are conducted on 16xA100 nodes.
While f-gram lookup is efficient for inference, it creates a bottleneck during training since at training time transformer models process all token positions in parallel. This leads GPU idle time when fetching the longest matching f-gram on the fly. To remove this bottleneck, after we construct the of f-grams (Vi-gram), we pre-scan the training sequences tag the longest matching length for each token. During training, we can then directly retrieve the corresponding f-gram for forward computation with the At-gram model.
For the At model, we use an absolute position embedding layer where the maximum position equals the longest n-gram in Vi Vi.gram- Within each batch, all f-grams padded the longest n-gram length that batch. We train all models with the 6 precision.
F.1 WebText
For pre-training on WebText [Peterson et al. 2019). J. we follow Radford 2019 and the batch
size and sequence length to 512 and 1024, respectively. Radford et al. [2099] do not specify the
number of training tokens or optimizer details. We train the models for 80B tokens, roughly doubling
the count in Radford et al. [2018). For optimization, we use AdamW [Loshchilov and Hutter] 2019 with weight decay
Following Hoffmann et al. [2022), we set the maximum learning rate to 2 x 10-1 and apply cosine learning rate scheduler. We list the model configurations in Table 6
26

---

Parameters (M)
model
ffw size
n layers
128
1024
4096
6
204
1024
4096
12
491
1536
6144
12
759
1536
6144
24
589
1536
6144
18
1099
1536
6144
36
Table 6: Baseline model configurations for pre-training on WebText. For constructing the f-gram
model (At-gram). we vary the number of layers in the 128M, 491M, and 589M variants and discard the token embedding layer.
Parameters (M)
d model
ffw size
layers
711
2048
8192
12
1014
2048
8192
18
1316
2048
8192
24
1920
2048
8192
36
Table Model configurations for pre- pre-training on the OLMo corpus. To construct the f-gram model
(Ar.grm)) we use the 71 1M and 1920M variants, excluding the token embedding layers.
F.2 OLMo Tokenized Training Corpus
For pre-training on the OLMo tokenized training corpus, we follow the optimizer settings for the IB variant in Groeneveld et al. 2024 All models sequence length of 2048. We use DeepSpeed [DeepSpeed] (2224] with ZeRO stage that partitions the optimizer state across GPUs to reduce GPU
memory usage. We list the model configurations in Table
27


Reviews:
[{"review_id": "Y7lWyjL2o6", "reviewer": "Reviewer_NngF", "summary": "This paper introduces SCONE, a novel method for improving language models by scaling the input embedding layer without increasing inference-time accelerator costs. The core idea is to augment the standard token embeddings with contextualized embeddings for a large set of frequent n-grams (termed \"f-grams\"). To avoid the challenges of training a massive, static embedding table, SCONE employs a separate \"f-gram model\" during training to learn and generate these embeddings dynamically. After training, these f-gram embeddings are precomputed and stored in an off-accelerator lookup table (e.g., in main memory or on an NVMe drive), which is queried efficiently during inference. This design decouples the input embedding size from the main model's vocabulary and output layer, sidestepping the computational bottlenecks of traditional vocabulary scaling.", "questions": "* Interaction with Mixture-of-Experts (MoE) Architectures: The paper positions SCONE as an alternative scaling approach to MoE, highlighting SCONE's advantage in maintaining a fixed accelerator memory footprint. However, given that MoE is a prevalent architecture for state-of-the-art large language models, it would be valuable to explore the interaction between these two methods. Could SCONE and MoE be used together, and if so, are their benefits complementary or do they overlap? For instance, would applying SCONE to an MoE model yield additive performance gains, or would one method's benefits diminish the other's? An experimental analysis or at least a discussion on this potential synergy would be highly relevant.\n* Domain Adaptation and Continual Learning: The current approach relies on a static f-gram set defined at the start of pre-training. A key practical question is how the model adapts when moved to a new domain (e.g., legal or medical texts). A discussion or experiment on how to efficiently integrate new, high-frequency n-grams from a new domain would significantly strengthen the paper's contribution to real-world applicability.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strength:\n\n* The idea of training an f-gram model is an good approach to address the challenges of training large-scale, sparse vocabularies.\n* The paper is well-written, clear, and easy to understand.\n* The work demonstrates two effective scaling dimensions: increasing the size of the f-gram vocabulary and scaling up the f-gram model.\n\nWeakness:\n\n* Clarity on Training and Pre-computation Costs: While the authors make a commendable effort to create a fair comparison by reducing the number of training tokens for SCONE models to approximate equivalent training FLOPS (as mentioned in Section 4.2), the comparison between the SCONE-enhanced 1B model and the 1.9B baseline could be further clarified. The training dynamics of a smaller model trained for more steps versus a larger, more complex architecture trained for fewer steps can be different. A more direct comparison under a fixed total training time or cost budget would help establish whether the architectural complexity of SCONE offers a definitive advantage over simply training a larger, monolithic model, which offers a simpler architecture. Furthermore, the paper could benefit from a more detailed analysis of the pre-computation costs. For instance, quantifying the time and computational resources required for the 1.8B f-gram model to generate the 1B embeddings would provide a more complete picture of the total resource footprint of the SCONE methodology.\n* Scalability Concerns in High-Concurrency Deployment: The paper provides a valuable analysis of the query latency for the offloaded f-gram embedding table (Section 4.3), demonstrating its feasibility in single-instance, batched inference scenarios. However, the experiments are conducted on a single workstation with batch sizes up to 16. It would strengthen the paper to discuss the potential engineering challenges and performance implications in a large-scale, production deployment. In a real-world serving environment with high concurrency (i.e., many simultaneous, independent user requests), the contention for system memory bandwidth or NVMe I/O from these frequent lookups could potentially become a bottleneck. A discussion on how the system would perform under such high-throughput, multi-tenant conditions would add significant practical depth to the evaluation.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review! We have incorporated your suggestions to improve our manuscript. Please find our responses below.\n\n**1. Clarity on Training and Pre-computation Costs**\n\n**1.1** The comparison between the SCONE-enhanced 1B model and the 1.9B baseline could be further clarified. \n\nWe compare the training dynamics of the 1.9B baseline with the SCONE-enhanced 1B model (which uses a 1.8B f-gram model and 1B cached f-gram embeddings).  Since we cannot share images via external links, we summarize the results in the following tables. The first table shows the training FLOPs per sequence (for SCONE averaged across sequences), and the second shows evaluation perplexity at comparable training FLOPs.\n\n**Training FLOPS / sequence**\n\n| SCONE 1B (1.8B f-gram model) | 1.9B baseline |\n|----------|----------|\n| $5.57\\times 10^{13}$| $2.73\\times 10^{13}$ |\n\n**Evaluation perplexity vs. training FLOPs**\n\n| | $\\sim 2.66\\times 10^{21}$  | $\\sim 5.32\\times 10^{21}$ | $\\sim 7.98\\times 10^{21}$ | $\\sim 1.06\\times 10^{22} $\n|----------|----------|----------|----------|----------|\n| 1.9B baseline | 12.680 | 11.246  |  10.570 | 10.085 |\n| SCONE 1B (1.8B f-gram model) | 12.598 | 11.256 | 10.524 | 10.033 | \n\nThe SCONE-enhanced 1B model achieves performance comparable to the 1.9B baseline when matched on training FLOPs, while requiring roughly half the FLOPs at inference. We will add the corresponding plots to the revision.\n\n\n**1.2** The paper could benefit from a more detailed analysis of the pre-computation costs.\n\nThe FLOPs and processing time for pre-computation on a node with 8 A100 GPUs are shown below:\n\n\n| Total FLOPS | \tProcessing time |\n|----------|----------|\n| $\\sim 1.836\\times 10^{17}$| 2 hour 35 mins |\n\nThe cost of pre-computing 1 billion f‑gram embeddings (requiring forward passes over 1B short n‑grams) is negligible compared to pre-training, which involves processing 1T tokens over full sequence lengths. We will add a discussion of these pre-computation costs to the revision.\n\n\n**2. A discussion on how the system would perform under such high-throughput, multi-tenant conditions would add significant practical depth to the evaluation.**\n\nThank you for the insightful question. While a batch size of 16 is already larger than what is typically used for model inference in production (since accelerator memory usage grows linearly with batch size and decoded sequence length), we agree that real-world deployments often operate in high-throughput, multi-tenant environments. In such scenarios, a common approach is to replicate the serving instances so that each replica handles at most the target batch size of concurrent queries. For SCONE, this implies that the f-gram embedding table would need to be replicated across physical devices. However, we note that this requirement is not unique to SCONE: traditional LLM serving pipelines also need to replicate model parameters across servers to handle high concurrency.\n\n\n**3. Could SCONE and MoE be used together, and if so, are their benefits complementary or do they overlap?**\n\nWhile SCONE and MoE can easily be combined since SCONE only modifies the input embeddings of the main model, we acknowledge that our current experiments do not include results with MoE models. This is primarily because pre-training large MoE models requires substantial computational resources, which is beyond the scope of our current work. We do, however, discuss the conceptual connections between SCONE and MoE in Appendix B of our submission, and we will revise to explicitly address this point in the Limitations section as well.\n\nTo address the question raised by Reviewer 5vFb, we have conducted additional experiments using the recent Qwen3 4B/8B/14B models at the **post-training stage**, which is considerably less resource-intensive than pre-training. Although these Qwen models are not MoE, we believe the results provide **partial evidence for the broader applicability of SCONE across different model architectures**.\nSpecifically, we apply SCONE during supervised fine-tuning (SFT) on [Qwen3-4B-base](https://huggingface.co/Qwen/Qwen3-4B-Base), following the [open-r1](https://github.com/huggingface/open-r1) framework, while using Qwen3-8B-base or Qwen3-14B-base as the f-gram model. These experiments use the [open-r1/Mixture-of-Thoughts](https://huggingface.co/datasets/open-r1/Mixture-of-Thoughts) SFT dataset, with 10M f-grams and the same training hyperparameters as open-r1. The table below summarizes the results, comparing SCONE-enhanced models against the Qwen3-4B baseline in terms of accuracy and inference latency:\n\n\n\n|   |  AIME 2024 pass@1  |  LiveCodeBench v4_v5 pass@1 | Decoding Latency at inference (per-token) | \n|----------|----------|----------|----------|\n| Qwen3-4B-base Baseline| 45.3 | 30.8 | 10.05 ms | \n|Qwen3-4B SCONE  (Qwen3-8B f-gram model) | 48.3 (+3.0) | 34.5 (+3.7) | 10.13 ms | \n|Qwen3-4B SCONE  (Qwen3-14B f-gram model) | 51.6 (+6.3)  | 36.3 (+5.5)  | 10.13 ms | \n\nThese results demonstrate that SCONE remains effective when applied to state-of-the-art architectures, with negligible impact on decoding latency.\n\n**4. A discussion or experiment on how to efficiently integrate new, high-frequency n-grams from a new domain would significantly strengthen the paper's contribution to real-world applicability.**\n\nWe construct the f-gram set from a large and diverse pre-training corpus. The zero-shot evaluation results in Section 4.2 (covering math, coding, factuality QA, etc.) indicate that this f-gram set generalizes well across a wide range of tasks. That said, we agree that it would be valuable for the f‑gram set and its embeddings to adapt to new domains that are not presented in pre-training data, and we will include this point in the Limitations section. To partially address this concern, the new post-training results suggest that SCONE can also be applied after pre‑training, allowing the f‑gram embedding table to be updated without repeating the entire pre‑training process."}, {"review_id": "HrTNTgdohG", "reviewer": "Reviewer_rTKx", "summary": "This paper introduces SCONE, a novel and practical method for scaling language models by enhancing the input embedding layer. Instead of simply enlarging the token vocabulary, which has known drawbacks, the authors propose augmenting it with embeddings for frequently occurring n-grams (f-grams). These f-gram embeddings are learned by a separate, smaller transformer model during training. Critically, for inference, these embeddings are precomputed and offloaded to cheaper, more abundant resources like system memory or SSDs. This design allows the model to leverage richer, contextualized input representations to improve performance without incurring additional computational (FLOPS) or memory costs on the accelerator during inference. The experiments convincingly demonstrate that this approach enables smaller models to outperform significantly larger baselines,", "questions": "Please refer to my weakness part.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper's core idea is highly innovative. Much of the research on LLM architecture focuses on attention and FFN layers, while the input embedding layer is often overlooked.  SCONE provides a fresh perspective by demonstrating that this seemingly simple component can be a powerful lever for scaling. The approach of decoupling the learning of contextual n-grams (during training) from their application (a simple lookup during inference) is both clever and effective.\n\n- The experimental validation is thorough and convincing.  The authors conduct systematic experiments across multiple model sizes and datasets, including well-designed ablation studies on key hyperparameters like the number of f-grams and the f-gram model size. The initial reliance on perplexity was a concern, but the authors' rebuttal comprehensively addressed this by providing extensive results on standard downstream benchmarks, which strongly reinforce their claims. The paper is well-written, and the core concept is simple and easy to grasp. \n\n**Weaknesses:**\n- The paper focuses exclusively on offloading computations related to input n-grams. Transformers implicitly learn to process n-gram patterns in their attention and feedforward layers.  A fascinating future direction would be to investigate if other frequent computational patterns within the main model could also be identified, precomputed, and offloaded to further improve efficiency.\n- While the empirical results are compelling, a deeper theoretical analysis of the trade-offs would be beneficial. For instance, quantifying the expressiveness added by the f-gram embeddings versus the computational cost of the main model could provide a more formal understanding of SCONE's scaling properties.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their insightful review. Please find our responses below. We would be happy to address any further questions or clarifications if needed.\n\n**1. A fascinating future direction would be to investigate if other frequent computational patterns within the main model could also be identified, precomputed, and offloaded to further improve efficiency.**\n\nWe completely agree that identifying and caching other frequent computational patterns within the main model is a fascinating direction for future work. In Appendix A of our submission, we discuss two interesting challenges in extending caching beyond n-gram embeddings: (1) Using raw text as cache keys can lead to low hit rates as the sequence length increases, and (2) using semantic embeddings as keys would require carefully designed discretization techniques to map continuous embeddings into a discrete key space that also supports efficient indexing.\n\n**2. A deeper theoretical analysis of the trade-offs would be beneficial. For instance, quantifying the expressiveness added by the f-gram embeddings versus the computational cost of the main model could provide a more formal understanding of SCONE's scaling properties.**\n\nWhile the primary focus of this work is on proposing a new method and conducting an empirical study, we agree that a deeper theoretical analysis would be a valuable direction for future research. We will acknowledge this in the “Limitations and Future Work” section of the revision.  However, please note that providing a formal theoretical characterization of modern LLMs has been challenging, even without SCONE. \n\n**3. New results on post-training of Qwen3 models**\n\nIn addition to addressing the reviewer’s comments, we conducted new experiments during the rebuttal period to demonstrate that SCONE can also be applied at the post-training stage of the latest Qwen3 models. Although these experiments are not directly related to your specific questions, we share them here to further illustrate the broad applicability of SCONE.\n\nSpecifically, we applied SCONE to supervised fine-tuning (SFT) of [Qwen3-4B-base](https://huggingface.co/Qwen/Qwen3-4B-Base) using the [open-r1](https://github.com/huggingface/open-r1) framework and the [open-r1/Mixture-of-Thoughts](https://huggingface.co/datasets/open-r1/Mixture-of-Thoughts) dataset. In this setup, Qwen3-4B serves as the main model, while Qwen3-8B-base or Qwen3-14B-base are used as the f-gram models. We set the number of f-grams to 10M and use the same training hyperparameters as open-r1. The table below compares the resulting SCONE-enhanced models with the Qwen3-4B baseline in terms of both accuracy and decoding latency.\n\n|   |  AIME 2024 pass@1  |  LiveCodeBench v4_v5 pass@1 | Decoding Latency at inference (per-token) | \n|----------|----------|----------|----------|\n| Qwen3-4B-base Baseline| 45.3 | 30.8 | 10.05 ms | \n|Qwen3-4B SCONE (Qwen3-8B f-gram model) | 48.3 (+3.0) | 34.5 (+3.7) | 10.13 ms | \n|Qwen3-4B SCONE (Qwen3-14B f-gram model) | 51.6 (+6.3)  | 36.3 (+5.5)  | 10.13 ms | \n\nThese results show that SCONE consistently improves performance over the baseline, with larger f-gram models yielding greater gains, all while maintaining similar inference latency."}, {"review_id": "BE7HHUO5ev", "reviewer": "Reviewer_yxCy", "summary": "This paper focuses on scaling embedding tables for large language models (LLMs). The authors propose a solution that (1) enables effective scaling of embedding tables and (2) is reasonably practical for both training and inference. Experimental results are strong, demonstrating improved performance when trained with equivalent resources and served using fewer resources.", "questions": "**Main Questions:**\n\n1. Regarding Figure 7, could you provide specific latency metrics such as p50, p90, and p99? Usually I imagine for embedding table lookup, training is fine since you can do pipelining but it is on the critical path for inference. See also the paper I mention below. \n2. Can you share e2e latency numbers for Figure 7?\n\n**Clarifying Questions:**\n\n1. What are the dimensions of the embedding tables?\n2. I am not familiar with the Lightning Memory-Mapped Database. Does it utilize CPU memory for caching? Do you have any data on cache miss behavior?\n3. On page 8, you mention that on-disk storage incurs an 86% memory overhead. Do you have insights into the reasons behind this?\n\n**Minor Suggestions / Nits:**\n\n1. I understand why you don't want to offload the \"embedding table\" for f-grams in training since you don't want to materialize it and you still want to write to it. But that reminds me of \"Toward 100TB Recommendation Models with Embedding Offloading\" which handles it by cache (roughly 10% of total table) and pipelining. \n2. Page 3 \"During training, it is parameterized by ...\" this sentence is pretty confusing to me.\n3. Page 5 \"... instantiate a full embedding table during training\" would be much better if you add this is the f-gram table.", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n\n* The paper is extremely well-written, with clear motivation and thorough explanations that proactively address potential questions.\n* The literature review is comprehensive.\n* The ideas presented are natural, and the design choices are well-justified.\n* The solution is co-designed for both training and inference.\n\n**Weaknesses:**\n\n* The main tradeoffs (GPU memory, disk usage, latency) for training and inference are scattered throughout the paper. The authors should consider summarizing key metrics in a table—potentially in the Appendix—including (1) GPU memory, (2) CPU memory, (3) disk usage, and (4) latency for both training and inference. \n* The sell of the paper is \"more training resources usage but equal inference resources usage for better performance\", with the assumption that (1) inference usage is way more than training resources and (2) allocating additional resources to training in this manner is the most effective way to improve performance. While I agree with (1) and somewhat (2), it remains to see if this method is truly the best way to spend the extra training resources.", "was_revised": true, "final_justification": "Same score.", "rebuttal": "Thank you for your review and suggestions! We have incorporated them to improve our manuscript. Please find our responses below.\n\n**1. The authors should consider summarizing key metrics in a table—potentially in the Appendix—including (1) GPU memory, (2) CPU memory, (3) disk usage, and (4) latency for both training and inference.**\n\nThank you for the suggestion.  The following table summarizes the computational resource usage for the three model variants evaluated in Section 4.2. All measurements are taken with a context length of 2048 and a batch size of 4 on a single A100 80 GB GPU. The three settings are: \n(1) the 1.9B baseline model, \n(2) SCONE 1.3B, with a 1.8B f‑gram model and 10M cached f‑gram embeddings, and \n(3) SCONE 1B , with a 1.8B f‑gram model and 1B cached f‑gram embeddings.\n\n| Model Variant | Peak GPU Memory   |  CPU Memory Overhead | Disk Usage | Decoding latency (per-token)  | Training FLOPS (per-seq.) |\n|----------|----------|----------|----------|----------|----------|\n| 1.9B baseline |  8.38 GB | N/A | N/A | 6.45 ms | $2.73\\times 10^{13}$ |\n| SCONE 1.3B (10M f-grams) | 5.60 GB  | 41.76 GB  | N/A | 4.83 ms  | $4.94\\times 10^{13}$ |\n| SCONE 1B (1B f-grams) | 4.45 GB | N/A | 7.67 TB | 4.90 ms   | $5.57\\times 10^{13}$ |\n\nBoth memory and disk usage are reported for inference. The decoding latency is averaged over one thousand decoding steps. As shown in Section 4.2, all three models achieve similar downstream performance. Compared to the 1.9B baseline, SCONE-enabled models significantly reduce GPU memory usage and decoding latency, at the cost of increased CPU memory or disk storage.  As you suggest, we will add this summary table in the revision.  \n\n\n**2. … with the assumption that (1) inference usage is way more than training resources and (2) allocating additional resources to training in this manner is the most effective way to improve performance. While I agree with (1) and somewhat (2), it remains to see if this method is truly the best way to spend the extra training resources.**\n\nWhile in Section 4.2 we keep the total training FLOPs for SCONE roughly the same as for the 1.9B baseline to show that SCONE provides a better use of training resources, we acknowledge that our experiments do not yet explore broader model families, such as significantly larger models. Therefore, whether this approach is always the most effective way to allocate additional training resources remains an open question and is an interesting direction for future work. We will include a discussion of this point in the ‘Limitations and Future Work’ section.\n\nThat said, as discussed in Section 1, there are important scenarios where training resources could be relatively abundant but inference resources are strictly constrained. SCONE is particularly well-suited for such settings.  Two such scenarios are 1) deploying models on edge devices with limited compute, and 2) latency-sensitive applications where inference FLOPs cannot exceed some threshold.\n\n**3. Regarding Figure 7, could you provide specific latency metrics such as p50, p90, and p99?** \n\nWe summarize the specific latency matrics in the table below.\n\n**Setup.** We report the largest configurations from Fig. 7: 1) 100 M f‑gram embeddings held in RAM, 2)  1 B f‑gram embeddings stored on NVMe.\nThe numbers are per‑token latencies averaged over 100, 000 batches and are expressed in milliseconds (ms). For batch sizes larger than 1, the total batch time is divided by the number of tokens, i.e., the latency is amortized over all tokens in the same batch.\n\n|   |  Batchsize = 1  |  2 | 4 | 8 | 16 |\n|----------|----------|----------|----------|----------|----------|\n| p50 (RAM)|  0.016 | 0.058 | 0.046 | 0.027 | 0.017 |\n| p90  (RAM)| 0.021 | 0.081 | 0.065 | 0.034 | 0.021 |\n| p99 (RAM) | 0.026 | 0.107 | 0.076 | 0.050 | 0.032 |\n| p50 (disk) | 0.783  | 1.105 | 1.158 | 1.029 |0.783 |\n| p90 (disk) | 5.438 | 3.094 | 2.288 | 1.670 | 1.214 |\n| p99 (disk) |  6.675 | 3.946 | 2.968 | 2.211 | 1.522 |\n\n**Observations.** \n - In‑memory (100 M). The distribution is tight: p50 essentially matches the mean reported in Fig. 7, and p99 is within 2 $\\times$ p50.\n\n- On-disk (1B). Variance is higher: p99 can be up to  $\\sim$ 8.5  $\\times$ p50 for batchsize = 1. Increasing the batch size amortizes the look‑up cost across tokens, shrinking the p99/p50 ratio to  $\\sim$ 2.8 $\\times$ at batch = 16.\n\nIn the revision, we will extend Figure 7 with a percentile plot (p50 / p90 / p99) numbers in the Appendix.  \n\n**4. Can you share e2e latency numbers for Figure 7?**\n\nIn the submission, we report the end-to-end decoding speed in the right panel of Figure 2. As shown in our response to the first question, a 1B model with a 1B f-gram embedding table stored on an NVMe drive achieves a per-token decoding latency of 4.90 ms (approximately 200 tokens per second). This is still significantly faster than the 1.9B baseline, which has an average latency of 6.45 ms, while delivering comparable downstream performance.\n\n\n**5. Clarifying Questions**\n\n**5.1** What are the dimensions of the embedding tables?\n\nThe default embedding dimension is 2048, except in Section 4.1, where we conduct ablation experiments using GPT with a smaller embedding dimension of 1024.\n\n**5.2** Does Lightning Memory-Mapped Database utilize CPU memory for caching? Do you have any data on cache miss behavior?\n\nLMDB does not add an extra cache. The database file is memory‑mapped as every page on disk is assigned a position in the process’s virtual memory address (on a 64‑bit process the space is vast enough to cover terabytes). The memory mapping process itself does not load any data into RAM yet. As LMDB code touches a page, the CPU triggers a page fault, the kernel pulls that page from disk into the OS page cache, and execution continues as if the page had always been in memory.\nBecause our embedding database is many times larger than the machine’s physical RAM, most look‑ups require pages that are not in the cache, so the latencies we report largely reflect disk‑bound access. \n\n**5.3** On page 8, you mention that on-disk storage incurs an 86% memory overhead. Do you have insights into the reasons behind this?\n\nThe ~86% larger on‑disk footprint is primarily due to internal fragmentation from fixed‑size pages. Unlike the contiguous in‑memory layout, the on‑disk format is partitioned into fixed blocks, many of which are only partially filled: especially under small, frequent updates. Beyond that, each page carries metadata (e.g., headers, checksums). These factors collectively account for most of the observed overhead.  We are happy to add this description in the revision.  \n\n\n**6. Suggestions**\n\n**6.1** Materialize the f-gram embedding table at training and handle it by cache and pipelining.\n\nThank you for the pointer. We will discuss this related work in the revision. As you noted, if the f‑gram embedding table were to be fully materialized during training, techniques such as caching and pipelining could help mitigate the latency issue. However, beyond latency, a key challenge with materializing the f‑gram embedding table is the problem of sparse updates. Due to the heavy-tail distribution of natural language, many f‑gram embeddings would receive only a handful of gradient updates per epoch. As we show in Appendix D of the submission, this sparse update issue can eventually degrade model performance. In contrast, when f‑gram representations are generated through a neural network, the parameter updates are shared across all f‑grams, which alleviates this problem.\n\n\n**6.2** Page 3\"During training, it is parameterized by ...\". This sentence is confusing.\n\nWe have revised the corresponding sentence to:\n\n“During training, instead of storing an explicit embedding table for f-grams, we generate their embeddings on the fly. A separate transformer model, $\\mathcal{A}_{\\text{f-gram}}$, takes an f-gram as input, and the embedding of its final token is used as the representation of that f-gram.”  Hope this is clearer.  \n\n**6.3** Page 5 \"... instantiate a full embedding table during training\" would be much better if you add this is the f-gram table.\n\nWe have revised accordingly."}, {"review_id": "KDZtqxyawV", "reviewer": "Reviewer_5vFb", "summary": "This work introduced SCONE, which extends input embedding layers to decrease the decoding costs. SCONE adds additional embeddings for a set of frequent n-grams. These newly added embeddings are learned with a separate model during training. During inference, these embeddings are pre-computed and stored in memory. The experiments demonstrate\tthe scaling law with regards to the number of n-gram embeddings and the size of the model that used to learn n-gram embeddings.", "questions": "See Strengths And Weaknesses", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The paper is well-written and presents detailed results for scaling experiments.\n\n\nWeaknesses:\n1. Lack of experiments for larger model. The scaling experiments are mainly conducted with 76M/340M/510M parameters, which is limited in scale. How about the performance of SCONE applied for larger model, e.g., Qwen2.5-3B/7B?\n2. The evaluation on downstream tasks only focus on language modeling tasks. Including more non-generation tasks, e.g., math or code tasks, could enhance the effectiveness of the proposed method, since these tasks usually require the model to generate very long context.\n3. The proposed method seems quit similar to a variant of prefix tuning. Could the authors discuss the difference between prefix tuning and SCONE? Is it possible to use SCONE for post-training stage of pre-trained LLM?", "was_revised": true, "final_justification": "The authors' rebuttal has addressed some of my concerns, in terms of including math and code tasks. I would like to raise my score to 4.", "rebuttal": "We thank the reviewer for their review and suggestions. Following the reviewer’s suggestion, we have added new results applying SCONE during the post-training stage on larger models. Please find our detailed response below.\n\n**1. The scaling experiments are mainly conducted with 76M/340M/510M parameters, which is limited in scale. How about the performance of SCONE applied for larger model, e.g., Qwen2.5-3B/7B?**\n\nWhile the scaling experiments in Section 4.1 are based on GPT-2 sized models, we also include experiments with OLMo-1B in Section 4.2, which is a commonly used efficient language model, and compare against baselines up to 1.9B parameters. As we acknowledge in the Limitations section (Appendix A), pre-training larger models requires substantial computational resources and is beyond the scope of this work.\n\nNevertheless, as suggested in the third question, we conducted additional post-training experiments on the latest Qwen3 models, since post-training is much less resource-intensive than pre-training. The results of these experiments can be found in our response to item 4 below. We hope these additional results can showcase the effectiveness of SCONE in larger-model settings.\n\n**2. Including more non-generation tasks, e.g., math or code tasks, could enhance the effectiveness of the proposed method, since these tasks usually require the model to generate very long context.**\n\nIn Section 4.2 of our submission, we report the performance on the MMLU tests used in the OLMo repo, which includes math and coding related problems. In our new experiments for post-training with Qwen3 models, we also report results on two additional benchmarks: AIME (math) and LiveCodeBench (code).\n\n**3. Could the authors discuss the difference between prefix tuning and SCONE?**\n\n[Prefix tuning](https://arxiv.org/abs/2101.00190) is a parameter-efficient fine-tuning approach that optimizes a small set of continuous vectors prepended to the model’s input context. While both SCONE and prefix tuning involve learning additional embeddings for a language model, SCONE differs in two key ways: (1) SCONE learns embeddings for a predefined set of n-grams, which can be cached and reused at inference time; and (2) SCONE generates these embeddings using a separate model, which can be scaled up during training to improve performance. Importantly, as the n-gram embeddings can be cached, scaling this separate model does not increase accelerator usage at inference.  We will highlight these points better in the revision.  \n\n\n**4. Is it possible to use SCONE for the post-training stage of pre-trained LLM?**\n\nYes.  We conducted additional experiments to demonstrate that SCONE can indeed be applied during the post-training stage. Specifically, we use SCONE with supervised fine-tuning (SFT) to post-train [Qwen3-4B-base](https://huggingface.co/Qwen/Qwen3-4B-Base), following the [open-r1](https://github.com/huggingface/open-r1) framework. The SFT dataset is [open-r1/Mixture-of-Thoughts](https://huggingface.co/datasets/open-r1/Mixture-of-Thoughts). In our setup, Qwen3-4B serves as the main model, while Qwen3-8B-base or Qwen3-14B-base are used as the f-gram models. We set the number of f-grams to 10M and follow the same training hyperparameters as open-r1. The table below presents the new results, comparing our SCONE-based model with the Qwen3-4B baseline in terms of both performance and latency.\n\n\n|   |  AIME 2024 pass@1  |  LiveCodeBench v4_v5 pass@1 | Decoding Latency at inference (per-token) | \n|----------|----------|----------|----------|\n| Qwen3-4B-base Baseline| 45.3 | 30.8 | 10.05 ms | \n|Qwen3-4B SCONE (Qwen3-8B f-gram model) | 48.3 (+3.0) | 34.5 (+3.7) | 10.13 ms | \n|Qwen3-4B SCONE (Qwen3-14B f-gram model) | 51.6 (+6.3)  | 36.3 (+5.5)  | 10.13 ms | \n\nThese results show that incorporating 8B and 14B f-gram models yields clear improvements on both AIME and LiveCodeBench, with larger f-gram models providing greater gains. Importantly, scaling up the f-gram model does not increase inference latency because the f-gram embeddings are pre-computed and cached.\n\nThank you for suggesting the post-training direction (and suggesting to use Qwen).  We will include these new results in the revision."}, {"review_id": "3OWeKMCXgE", "reviewer": "Reviewer_H6nF", "summary": "This paper proposes enhancing LLM performance by incorporating n-gram embeddings. To control the size of the n-gram embeddings, the authors suggest building embeddings only for high-frequency n-grams. They demonstrate that scaling the size of the n-gram model and increasing the number of n-grams both effectively reduce the perplexity of LLMs. A key advantage of this approach is that n-gram embeddings are context-free, allowing them to be read directly without recomputation. However, a major drawback is the substantial storage overhead—since the total number of global n-grams scales with the fourth power of the vocabulary size, the method can only accommodate a limited number of n-grams in practice.", "questions": "see cons.", "limitations": "Yes.", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Pros**:\n\nNo recomputation required: The n-gram embeddings are context-free, meaning they can be directly retrieved without recomputation, unlike context-dependent representations. This leads to low inference overhead.\n\nEffective performance improvement: Experiments show that increasing both the number and dimensionality of n-gram embeddings significantly reduces the perplexity (PPL) of LLMs, indicating the method effectively enhances language modeling capability.\n\nClear writing and logic: The paper is well-written with clear logic, making it easy to read and understand.\n\n**Cons**:\n\nLarge storage overhead: The storage cost is substantial. Loading the n-gram embeddings during the prefill phase incurs significant overhead, although the cost is more manageable during decoding.\n\nNot a game changer: The improvement in perplexity, while present, is relatively limited and may not be sufficient to position the method as a major advancement.\n\nLimited coverage: It is not feasible to compute embeddings for all n-grams; only a limited number can be supported. As a result, the method may offer limited benefits for low-frequency or rare tasks.\n\nLack of architectural generalization: The method's generality across architectures is not explored. Demonstrating its effectiveness on alternative models such as Mamba or GLA could significantly strengthen the paper's claims.", "was_revised": true, "final_justification": "I think after adding the rebuttal experiments, the paper should be accepted.", "rebuttal": "Thank you for the review and comments! Please see our response below.\n\n**1. Loading the n-gram embeddings during the prefill phase incurs significant overhead, although the cost is more manageable during decoding.**\n\nWe thank the reviewer for raising this point. We agree that during the prefill phase, substantially more n‑gram embeddings must be accessed compared to the decoding phase, particularly for long-context tasks. However, techniques such as chunked prefill are commonly used to improve efficiency. This is because processing a few hundred (or even just a few dozen) tokens in parallel is typically enough to fully saturate the accelerator. Under these conditions, SCONE can load hundreds of token embeddings without adding noticeable latency beyond what is shown in Figure 7. When the n‑gram embedding table is stored in CPU memory, this loading can be handled efficiently with batched operations. If the embedding tables are stored on disk, the solution is to increase the number of parallel read processes. We will incorporate this discussion in the revision.  \n\n\n**2. The improvement in perplexity, while present, is relatively limited.**\n \nIn addition to the perplexity evaluation, our submission also reports downstream task results in Table 1 (Section 4.2). These results show that the best SCONE‑1B variant outperforms a 1.9B baseline across multiple tasks. While the improvements are apparently modest, we believe they are nonetheless meaningful and demonstrate the effectiveness of our approach.\n\n\n**3. Only a limited number of n-grams can be supported. As a result, the method may offer limited benefits for low-frequency or rare tasks.**\n\nWhile the zero-shot evaluation results in Section 4.2 (covering math, coding, factuality QA, and other tasks) show that the f-gram set built from a diverse pre-training corpus can generalize well across a broad range of tasks, we acknowledge that f-gram coverage may be limited for extremely rare corpora.\n\nTo partially address this concern, as discussed in our response to Reviewer 5vFb, we show that SCONE can also be applied at the **post-training** stage. This allows the **f-gram embedding table to be updated on a new corpus without having to repeat the entire pre-training process**.\n\nSpecifically, we applied SCONE to the supervised fine-tuning (SFT) of [Qwen3-4B-base](https://huggingface.co/Qwen/Qwen3-4B-Base) using [open-r1](https://github.com/huggingface/open-r1) and the [open-r1/Mixture-of-Thoughts](https://huggingface.co/datasets/open-r1/Mixture-of-Thoughts) dataset. In our setup, Qwen3-4B serves as the main model, and Qwen3-8B-base or Qwen3-14B-base are used as the f-gram models. We set the number of f-grams to 10M and follow the training hyperparameters in open-r1. The table below compares the resulting SCONE-enhanced models with the Qwen3-4B baseline in terms of both accuracy and decoding latency.\n\n|   |  AIME 2024 pass@1  |  LiveCodeBench v4_v5 pass@1 | Decoding Latency at inference (per-token) | \n|----------|----------|----------|----------|\n| Qwen3-4B-base Baseline| 45.3 | 30.8 | 10.05 ms | \n|Qwen3-4B SCONE (Qwen3-8B f-gram model) | 48.3 (+3.0) | 34.5 (+3.7) | 10.13 ms | \n|Qwen3-4B SCONE (Qwen3-14B f-gram model) | 51.6 (+6.3)  | 36.3 (+5.5)  | 10.13 ms | \n\nThese results indicate that both the f-grams and the f-gram embedding table can be updated during post-training, without requiring a full re-run of pre-training.\n\n**4. The method's generality across architectures is not explored.**\n\nWhile the transformer architecture we used is a common choice in the literature, we agree that evaluating SCONE on alternative architectures such as Mamba or GLA would be valuable. Unfortunately, due to limitations in computational resources and human effort, such an extensive study is beyond the scope of the current work. In the revision, we will include this point in the Limitations section and pose it as an interesting research direction to pursue."}]

Author rebuttal:

Meta-review: The paper received strong evaluations. It is well-written, has clear motivation, and the design choices are well-justified. I recommend that the authors include additional experiments conducted during the rebuttal and the limitations of their work (which they presented in the final remarks) in the camera-ready version of the paper.
Decision: Accept (poster)