Title: Scaling Embedding Layers in Language Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: embedding layer scalability

contextualized token embeddings

scaling with fixed accelerator usage
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nsee cons.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-20"}

comments: 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	the scaling law with regards to the number of n-gram embeddings and the size of the model that used to learn n-gram embeddings.

STRENGTHS AND WEAKNESSES:
Strength:
1. The paper is well-written and presents detailed results for scaling experiments.


Weaknesses:
1. 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?
2. 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.
3. 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?

QUESTIONS:
See Strengths And Weaknesses

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n**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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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,\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nPlease refer to my weakness part.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}