Scaling Embedding Layers in Language Models

Thank you for your supporting review! We appreciate your recognition of the novelty and contributions of our work. We also welcome any further comments you may have.

Thank you for your review. Please find our responses below, we’re happy to discuss further if needed. We’ve also included downstream evaluations in our response to Reviewer iJdC.

1. Missing references

We thank the reviewer for the insightful references. We have incorporated all the references in the following discussions, and added them to Related Work (Sec 5).

“Implicit $n$-gram patterns in transformers. Recent work analyzing the internal mechanisms of transformers has shown that these models often utilize implicit $n$-gram patterns for prediction (Geva et al., 2021; Geva et al., 2022; Voita et al., 2023). For instance, Chen et al., 2024 show that certain attention heads detect specific $n$-gram patterns, while MLPs can perform linguistic operations such as adding the “-ing” suffix. These findings underscore the importance of $n$-gram information in language modeling and offer a potential explanation for the effectiveness of SCONE. An interesting future direction is to 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 are one such scenario. Two phrases in different languages may refer to the same concept but correspond to different embeddings. Their embeddings should ideally be close. Recent work explores methods for learning transferable embeddings in cross-lingual settings (Artetxe et al., 2020; Chen et al., 2023). Another relevant example is scaling the embeddings for recommender systems (Chen et al., 2019; Liu et al., 2021), 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 a neural network, provides a complementary approach to mitigate sparsity issues.”

2. The focus on embedding offloading is somewhat narrow.

We agree that exploring offloading beyond input embeddings could further reduce inference costs and our work can be viewed as a first step. A key research challenge along this direction is deciding what should serve as keys in such a system. We have added the following discussion to the last section:

“An interesting future direction is to extend SCONE beyond short $n$-grams to include longer and frequent queries. A key challenge would be designing effective keys for such queries. Using raw text as keys may lead to low hit rates, as semantically similar queries often differ at the surface level. Alternatively, using semantic embeddings as keys would require discretization methods to map continuous embeddings to a set of keys that supports efficient indexing.”

3. Whether SCONE can outperform methods like MoE.

Comparing with MoE is indeed an interesting future work. That said, SCONE offers a key advantage: while both aim to scale model capacity under fixed inference FLOPS, MoE requires all expert weights to reside on accelerators, since any token might activate any expert. In contrast, SCONE is designed for fixed accelerators memory usage at inference.

4. Explain in simple terms how the speedup is achieved. (1B with SCONE v.s. ~2x inference FLOPS baseline)

The 1B SCONE-enabled model outperforms the 1.9B baseline. The 1.9B baseline requires ~2x inference FLOPS. Thank you for pointing this out; we’ve clarified this in the introduction.

5. Training sample efficiency of SCONE?

No, SCONE does not reduce sample efficiency. All models are trained on the same number of tokens, and as shown in Figure 13, the improvements are consistent throughout training.

6. Complexity of the counting algorithm & its runtime.

As noted in Section 3.1, there is an efficient implementation for the counting algorithm that requires $n-1$ linear passes over the corpus ($n$ is the max f-gram length). On 1T tokens using 8 processes and $n=5$, processing took about 10 hours. We've included this in the revision.

7. SCONE appears to bring less improvement to perplexity as model size increases.

As perplexity decreases, achieving the same absolute reduction becomes more difficult due to the nature of language modeling. In Figure 6, we used a linear scale on the y-axis, which can make improvements at lower perplexity appear smaller. We’ve uploaded a new figure (at this anonymous link) that uses a log y-axis (following recent scaling law studies, e.g., Figure 2 in Hoffmann et al., 2022). The improvements in log scale appear more consistent.

8. Some quotations are single commas.

We have changed the single comma quotations in line 567.

We thank the reviewer for the suggestions. Please find our response below.

1. The manuscript's quality is lacking.

We have carefully addressed the concerns regarding writing (see below). Please see the updated Figure 1 (and its caption), Figure 5, and Figure 6 at this anonymous link. We would like to clarify that these issues pertain to presentation and do not affect the validity of our main claims.

1.1 Figure 1: the line of 1.3B model and the x-label.

The 1.3B model corresponds to the column with inference FLOPS of $6.20\times 10^{12}$. We have revised the caption to improve clarity.

We have changed the x-label, and other occurrences of “inference-time” FLOPS in the paper, to inference FLOPS.

1.2 Line width not consistent. We have adjusted the width of all figures to be $0.9\times$ the text column width.

1.3 Figure 5: error bars. We have updated Figure 5 to include vertical lines in the error bars.

1.4 Figure 6: legend overlaps with the line. We have adjusted the figure to eliminate any overlap between the legend and the line.

2. The motivation is not clear. This architecture is not utilized by any open-weighted model. In addition, n-gram has been given up in the LLM era.

Our motivation is to reduce inference FLOPS by trading computation for RAM or SSD, which are significantly cheaper and more abundant than accelerators. A natural target for this is the embedding layer, due to its inherent lookup-based structure. While we agree that the NLP community has largely moved away from traditional n-gram methods, recent work showed that $n$-grams still play important roles in transformers (due to space limit, please see the references in reviewer LBJ6’s review). We hope our work offers a new perspective on how n-grams can be leveraged to improve the efficiency of LLMs.

3. … it is known to all that the 7B model is always trained by 1T tokens

In lines 88–93, our intention is to highlight that increasing the vocabulary size can eventually degrade model performance. We use the update counts over 100M tokens to help explain this phenomenon. We have revised the manuscript to clarify this point:

“In Appendix C, we train GPT-2 models for 80B tokens with vocabulary sizes ranging from 32K to 2M, and observe performance degradation as the vocabulary size exceeds 512K. This degradation may be attributed to the increasing sparsity of updates per token as the vocabulary grows. …”

With 1T training tokens, the absolute number of updates each embedding receives will increase. However, the relative sparsity remains: larger vocabularies still result in fewer updates per embedding on average.

4. Lines 391 and line 396, I feel confused that +10M Vf-gram (0.6B Af-gram) has different performance over c4-en, books, etc. … In addition, where is the 1.3B model mentioned in the caption of Table 2?

Although lines 391 and 396 use the same f-gram configuration, they correspond to different main model sizes. Specifically, the results from lines 390 to 394 are based on the 1B model, while lines 395 to 399 correspond to the 1.3B model.

To clarify this distinction, we have added the following to the caption of Table 2:

“We train three baseline models of sizes 1B, 1.3B, and 1.9B. For the 1B and 1.3B baseline models, we apply our SCONE method with four different configurations, and present the results directly below each corresponding baseline model.”

5. Could you evaluate your pre-trained model over downstream tasks?

Yes, we evaluated the zero-shot performance of our models on PIQA, MMLU, HellaSwag, ARC-Easy, ARC-Challenge, and CommonsenseQA, following the recent implementation in the OLMO codebase. Due to space constraints, please see our response to reviewer iJdC for the full results. The downstream evaluation outcomes align with the perplexity trends.

6. … deploying this model requires substantial memory or disk space.

The configurations in Table 1 use large f-gram embedding layers to demonstrate that even the most resource-intensive setup explored in the paper remains feasible in certain server-based settings.

As shown in Figure 6, SCONE offers clear improvements even with much smaller embedding sizes. For example, with a 512K f-gram embedding layer—approximately 20× smaller than the 10M setting in Table 1—the perplexity of the 589M base model improves from 18.1 to 16.8 on WebText.

7. The memory and disk requirements for the Gemma model will surge significantly.

For a given embedding dimension, the storage usage of the f-gram embedding layers is only determined by the number of f-grams, which is a configurable parameter independent of the vocabulary size. Therefore, please note that a larger vocabulary does not lead to higher memory or disk requirements compared to smaller vocabularies.

We thank the reviewer for the comments. Please find our response below.

1. The evaluation metric (perplexity) is not a good metric. Can you add numbers for standard LLM evals like MMLU, hellaswag and so on?

While perplexity is a commonly used metric for evaluating language models, we acknowledge that additional downstream evaluations can further strengthen our work. In response, we have incorporated zero-shot evaluations on standard benchmarks including PIQA, MMLU, HellaSwag, ARC-Easy, ARC-Challenge, and CommonsenseQA, following the recent implementation in the OLMO codebase. The results are shown below. Thank you for suggesting; we will add the benchmark results in the revision.

Applying SCONE, i.e., adding f-gram embeddings, does not increase inference FLOPS and requires only off-accelerator storage. Notably, the 1.9B baseline incurs roughly twice the inference FLOPS of the 1B model. The downstream evaluation results align with the perplexity trends and further reinforce our main claims.

2. Can you clarify how the system works during inference? If I decode 3 tokens from the base vocabulary and find that they match an n-gram from the f-gram model, should the model recompute the k,v-cache for these three tokens with the embeddings from the matched n-gram?

To clarify, SCONE does not require recomputing the (k,v) cache during inference. Instead, it fetches the f-gram embedding for each decoded token individually. For example, if the current context is [t0, t1] and the newly decoded token is t2, and we find a match for [t0, t1, t2] in the f-gram embedding layer, we use the f-gram embedding of [t0, t1, t2] as the input embedding for t2 in the main model. This f-gram embedding of [t0, t1, t2] corresponds to the embedding of the last token embedding in the output of the f-gram model (precomputed for [t0, t1, t2]) and only involves simple lookups during inference. If it helps, we can add this example in the revision.