Research map / Architectures and efficiency / RNNs and normalization
Scaling Embedding Layers in Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SCONE, which augments the input embedding layer with contextualized embeddings for frequent n-grams (f-grams) learned by a separate f-gram transformer during training, then precomputed and offloaded off-accelerator for inference.
Problem
Scaling language model capacity usually increases inference FLOPS and accelerator memory; enlarging token vocabularies degrades performance (sparse tail-token updates) and inflates the output layer cost.
Approach
Keep the original vocabulary but add embeddings for a set of frequent n-grams (up to length K) produced by a separate transformer (the f-gram model) during training; at inference these are precomputed into an offloaded lookup table, so the accelerator footprint stays fixed while the number of cached f-grams and the f-gram model size can be scaled.
Key results
- With 1B f-grams, OLMo-1B + SCONE (1.8B f-gram model) averages 57.0 zero-shot accuracy vs 56.8 for OLMo-1.9B, using about 48% less inference FLOPS and accelerator memory.
- With 10M f-grams, OLMo-1.3B + SCONE reaches 56.8 average accuracy, matching OLMo-1.9B with about 32% less inference FLOPS and accelerator memory.
- In SFT on Qwen3-4B-base, SCONE with a 14B f-gram model improves AIME 2024 pass@1 from 45.3 to 51.6 and LiveCodeBench from 30.8 to 36.3 at similar per-token decoding latency.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add 10M f-grams | OLMo-1B · Average downstream tasks | accuracy: 53.7 -> 54.9 | Helped |
| add 1B f-grams | OLMo-1B · Average downstream tasks | accuracy: 53.7 -> 57.0 | Helped |
| add SCONE with 10M f-grams | OLMo-1B · PIQA | accuracy: 73.6 -> 74.0 | Helped |
| add 10M f-grams using a 0.6B f-gram model | OLMo-1B · OLMo evaluation mixture | average perplexity: 16.082 -> 15.459 | Helped |
| increase f-gram model size from 0.6B to 1.8B with 10M f-grams (ablation) | OLMo-1B + SCONE · OLMo evaluation mixture | average perplexity: 15.459 -> 15.125 | Helped |
| increase f-gram count from 10M to 1B with a 1.8B f-gram model (ablation) | OLMo-1B + SCONE · OLMo evaluation mixture | average perplexity: 15.125 -> 14.581 | Helped |
| increase maximum f-gram length from K=2 toward K=4 (ablation) | SCONE on WikiText-103 · WikiText-103 | perplexity | Mixed |
| increase maximum f-gram length beyond K=4 (ablation) | SCONE on WikiText-103 · WikiText-103 | perplexity | No clear effect |
| increase f-gram count from 512K to 100M (ablation) | SCONE on WebText · WebText validation | perplexity | Helped |
| increase f-gram count (ablation) | SCONE on WikiText-103 · WikiText-103 | perplexity | Helped |
| scale the f-gram model from 170M to 1020M parameters (ablation) | SCONE with 419M main model on WikiText-103 · WikiText-103 | perplexity: 26.1 to 23.4; 22.1 | Helped |
| add the OE-12.8M input embedding layer | OLMo-1B · Average downstream tasks | accuracy: 53.7 -> 54.4 | Helped |
Surprising findings
- Scaling the token vocabulary beyond roughly 512K degrades GPT-2 performance, with sparse per-token updates (7.3% of 2M-vocab tokens get over 100 updates vs 97.6% for 32K).
Limitations
- Only evaluated on models up to 3B parameters at training time due to hardware limits.
- Only short n-grams (up to K=5) are cached; extending to longer queries needs new keys, which is unsolved.
- Training compute increases due to the f-gram model, so token budgets were reduced to match FLOPS.
Experimental setup
Datasets: WebText WikiText-103 Dolma OLMo pretraining corpus open-r1/Mixture-of-Thoughts
Benchmarks: MMLU-var HellaSwag ARC-Challenge ARC-Easy CommonsenseQA PIQA WikiText-103 WebText validation AIME 2024 LiveCodeBench v4_v5
Metrics: perplexity zero-shot accuracy pass@1 decoding latency bits per character
Baselines: OLMo-1B OLMo-1.9B Over-tokenized transformer (OE-12.8M) GPT-2 baselines with larger non-embedding parameters
Builds on: Over-tokenized transformer Byte-Pair Encoding OLMo
Closest papers by meaning
- EmbedLLM: Learning Compact Representations of Large Language Models (ICLR 2025 · similarity 0.90)
- Layer by Layer: Uncovering Hidden Representations in Language Models (ICML 2025 · similarity 0.89)
- DEPT: Decoupled Embeddings for Pre-training Language Models (ICLR 2025 · similarity 0.89)
- Leviathan: Decoupling Input and Output Representations in Language Models (NeurIPS 2026 · similarity 0.89)
- Scaling Inference-Efficient Language Models (ICML 2025 · similarity 0.89)
- MoME: Mixture-of-Memory Embeddings for Context-Aware Sparse Lookup (NeurIPS 2026 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.