Research map / Architectures and efficiency / Attention and transformers
TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters
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 TokenFormer, a Transformer where all linear projections are replaced by token-parameter attention over learnable key-value parameter tokens, enabling incremental parameter scaling (124M to 1.4B) with reuse of smaller pretrained models.
Problem
Transformer scaling requires architectural changes (e.g., channel dimensions) that force retraining from scratch, making growing foundation models increasingly costly.
Approach
Treat model parameters as tokens: each linear projection becomes a cross-attention (Pattention) layer where input tokens are queries and learnable key-value parameter tokens are keys and values. New key-value pairs can be appended to scale the model without changing hidden dimension, with zero-initialized keys preserving the existing output; a GeLU-after-L2-norm replaces softmax.
Key results
- Progressively scaling from 124M to 1.4B with an extra 30B tokens reaches perplexity 11.77 vs 11.63 for a 1.4B Transformer trained from scratch on 300B tokens, at roughly one-tenth of the incremental budget (3x cheaper cumulatively, per the paper).
- TokenFormer-1.5B reaches 59.3 average zero-shot accuracy over eight tasks, versus 55.2 for Pythia-1.3B and 59.7 for Mamba-1.4B.
- Ours-B/16 reaches 82.5% ImageNet top-1 with 109M params, matching ViT-B/16 (MAE) at 82.3%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| progressively reuse smaller Tokenformer parameters when scaling to 1.4B, with 30B training tokens | Tokenformer · OpenWebText validation, 1.4B parameters | perplexity: 13.34 -> 11.77 | Helped |
| use Tokenformer parameter reusing rather than scratch training | 1.4B model · OpenWebText validation, 1.4B parameters | perplexity: 11.63 vs. 11.60 | No clear effect |
| replace the exponential nonlinearity with GeLU in Pattention (ablation) | Tokenformer Pattention · ImageNet classification | Top-1 accuracy: +2.1 points | Helped |
| replace L1 normalization with L2 normalization in Pattention (ablation) | Tokenformer Pattention with GeLU · ImageNet classification | Top-1 accuracy: +0.8 points | Helped |
| remove learnable layer-normalization weight (ablation) | Tokenformer · ImageNet classification | Top-1 accuracy: 82.6 -> 82.5 | No clear effect |
| remove learnable layer-normalization weight and bias (ablation) | Tokenformer · ImageNet classification | Top-1 accuracy: 82.6 -> 82.5 | No clear effect |
| progressively reuse smaller Tokenformer parameters at 354M | Tokenformer · OpenWebText validation, 354M parameters, 30B tokens | perplexity: 15.97 -> 14.02 | Helped |
| progressively reuse smaller Tokenformer parameters at 757M | Tokenformer · OpenWebText validation, 757M parameters, 30B tokens | perplexity: 13.78 -> 12.59 | Helped |
Surprising findings
- Replacing softmax with GeLU after L2 normalization in the token-parameter attention improved ImageNet accuracy (79.6 to 82.5), contrary to the usual softmax choice.
- Non-parametric layer norm (removing learnable gamma/beta) cost essentially nothing (82.6 vs 82.5 top-1).
Limitations
- Token-parameter computation scales linearly with parameter count, so parameter growth increases inference cost.
- Experiments are limited to language modeling up to 1.5B and ImageNet classification; no large-scale multimodal evaluation.
- Compute was reported on TPU v4 hardware with unspecified chip count.
Experimental setup
Datasets: OpenWebText Pile EnWik8 ImageNet-1K
Benchmarks: Pile LAMBADA HellaSwag PIQA ARC-Easy ARC-Challenge WinoGrande ImageNet-1K
Metrics: perplexity zero-shot accuracy top-1 accuracy training FLOPs
Baselines: Pythia Mamba OPT GPT-Neo ViT DeiT Net2Net HyperCloning RWKV
Builds on: Transformer LoRA Net2Net Pythia
Closest papers by meaning
- Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.89)
- Dynamical Properties of Tokens in Self-Attention and Effects of Positional Encoding (NeurIPS 2025 · similarity 0.88)
- Distilling Sequential Computation in Transformer Language Models (NeurIPS 2026 · similarity 0.87)
- Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.87)
- Looking Beyond the Top-1: Transformers Determine Top Tokens in Order (ICML 2025 · similarity 0.87)
- Pool Me Wisely: On the Effect of Pooling in Transformer-Based Models (NeurIPS 2025 · similarity 0.87)
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.