# Scaling Embedding Layers in Language Models

arXiv: https://arxiv.org/abs/2502.01637 · HTML: https://mlsanity.com/paper/2502.01637/

## Key facts

- Cited by 20 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 74 references that are in the map.
- First posted on arXiv 3 Feb 2025.

## Builds on (most cited first, 30 of 57)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [Deep contextualized word representations](https://mlsanity.com/paper/1802.05365/) (2018)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing](https://mlsanity.com/paper/1808.06226/) (2018)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates](https://mlsanity.com/paper/1804.10959/) (2018)
- [OLMo: Accelerating the Science of Language Models](https://mlsanity.com/paper/2402.00838/) (2024)
- [On the Cross-lingual Transferability of Monolingual Representations](https://mlsanity.com/paper/1910.11856/) (2019)
- [Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space](https://mlsanity.com/paper/2203.14680/) (2022)
- [ByT5: Towards a token-free future with pre-trained byte-to-byte models](https://mlsanity.com/paper/2105.13626/) (2021)
- [Dolma: an Open Corpus of Three Trillion Tokens for Language Model Pretraining Research](https://mlsanity.com/paper/2402.00159/) (2024)
- [Memory Networks](https://mlsanity.com/paper/1410.3916/) (2014)
- [MEGABYTE: Predicting Million-byte Sequences with Multiscale Transformers](https://mlsanity.com/paper/2305.07185/) (2023)
- [Byte Latent Transformer: Patches Scale Better Than Tokens](https://mlsanity.com/paper/2412.09871/) (2024)
- [Nemotron-4 340B Technical Report](https://mlsanity.com/paper/2406.11704/) (2024)

## Built on by (most cited first, 20 of 20)

- [Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models](https://mlsanity.com/paper/2601.07372/) (2026)
- [On the Design of Qwen3.8-Next Architecture: Evaluation, Efficiency, and Training Stability](https://mlsanity.com/paper/2608.30320/) (2026)
- [Exploiting Vocabulary Frequency Imbalance in Language Model Pre-training](https://mlsanity.com/paper/2508.15390/) (2025)
- [Lossless Token Sequence Compression via Meta-Tokens](https://mlsanity.com/paper/2506.00307/) (2025)
- [DynaSpec: Context-aware Dynamic Speculative Sampling for Large-Vocabulary Language Models](https://mlsanity.com/paper/2510.13847/) (2025)
- [JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation](https://mlsanity.com/paper/2602.00800/) (2026)
- [Proxy Compression for Language Modeling](https://mlsanity.com/paper/2602.04289/) (2026)
- [Tensorizing Engram: Sharing Latents Across N-Gram Embeddings is Beneficial in LLMs](https://mlsanity.com/paper/2606.08347/) (2026)
- [VocabTailor: Dynamic Vocabulary Selection for Downstream Tasks in Small Language Models](https://mlsanity.com/paper/2508.15229/) (2025)
- [L3: Large Lookup Layers](https://mlsanity.com/paper/2601.21461/) (2026)
- [Mixture of Lookup Key-Value Experts](https://mlsanity.com/paper/2512.09723/) (2025)
- [Beyond N-gram: Data-Aware X-GRAM Extraction for Efficient Embedding Parameter Scaling](https://mlsanity.com/paper/2604.21724/) (2026)
- [Scratchpad Patching: Decoupling Compute from Patch Size in Byte-Level Language Models](https://mlsanity.com/paper/2605.09630/) (2026)
- [NGM: A Plug-and-Play Training-Free Memory Module for LLMs](https://mlsanity.com/paper/2605.16893/) (2026)
- [Dynamic Injection of Entity Knowledge into Dense Retrievers](https://mlsanity.com/paper/2507.03922/) (2025)
- [Compressing Sequences in the Latent Embedding Space: K-Token Merging for Large Language Models](https://mlsanity.com/paper/2604.15153/) (2026)
- [MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup](https://mlsanity.com/paper/2609.15126/) (2026)
- [Large language models in medical time series analysis](https://mlsanity.com/paper/2609.22262/) (2026)
- [Forking: Sudden Overfitting Under Replay](https://mlsanity.com/paper/2610.00394/) (2026)
- [The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation](https://mlsanity.com/paper/2610.20456/) (2026)

## Abstract (arXiv)

We propose SCONE (Scalable, Contextualized, Offloaded, N-gram Embedding), 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.