Research map / Architectures and efficiency / Attention and transformers
A General and Efficient Training for Transformer via Token Expansion
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 Token Expansion (ToE), a parameter-free "initialization-expansion-merging" token growth scheme that trains ViTs on a progressively growing subset of tokens to accelerate training ~1.3x without accuracy loss while keeping the original hyper-parameters, architecture, and strategy.
Problem
Training Vision Transformers requires very large compute, and existing training-acceleration methods either lose accuracy, are not universal across models, or break training consistency (hyper-parameters, architecture, or extra strategies like EMA).
Approach
Apply a token growth schedule after the first Transformer block: start from a uniformly sampled seed token set, grow it over training stages by greedily/parallel adding the tokens with the largest cosine distance to the selected set (widest feature-distribution expansion), and merge the remaining tokens into their nearest selected tokens, so all parameters are still trained and the final model is
Key results
- ToE with r1=0.5 keeps DeiT-tiny at 72.6% top-1 (+0.4 over the 72.2% baseline) with 1.24x practical training speedup.
- On DeiT-small, ToE r1=0.5 matches the baseline (79.8%) at 1.22x practical speedup, while NetworkExpansion loses about 1.0 point.
- Integrated with EfficientTrain, ToE r1=0.6 improves DeiT-tiny to 73.5% (+1.0) with 1.63x practical speedup.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply ToE r1=0.5 | DeiT-tiny · ImageNet-1K | Top-1 accuracy: 72.2 -> 72.6 (+0.4) | Helped |
| Apply ToE r1=0.5 | DeiT-small · ImageNet-1K | Top-1 accuracy: 79.8 -> 79.8 (+0.0) | No clear effect |
| Apply ToE r1=0.5 | DeiT-base · ImageNet-1K | Top-1 accuracy: 81.8 -> 81.6 (-0.2) | Hurt |
| Apply ToE r1=0.4 | LV-ViT-T · ImageNet-1K | Top-1 accuracy: 79.1 -> 79.4 (+0.3) | Helped |
| Apply ToE r1=0.4 | LV-ViT-S and LV-ViT-M · ImageNet-1K | Top-1 accuracy: 83.3 (+0.0); 84.1 (+0.0) | No clear effect |
| Add ToE r1=0.6 | EfficientTrain with DeiT-tiny · ImageNet-1K | Top-1 accuracy: 73.3 -> 73.5 | Helped |
| Add ToE r1=0.6 | EfficientTrain with DeiT-small · ImageNet-1K | Top-1 accuracy: 80.4 -> 80.4 | No clear effect |
| Use ToE-pretrained weights | DeiT-tiny with baseline fine-tuning · CIFAR-10 | Top-1 accuracy: 98.07 -> 98.19 (+0.12) | Helped |
| Set r0/r1=1/3 instead of the default (ablation) | ToE on DeiT-tiny · ImageNet-1K | Top-1 accuracy: 72.6 -> 72.3 | Hurt |
| Set r0/r1=2/3 instead of the default (ablation) | ToE on DeiT-tiny · ImageNet-1K | Top-1 accuracy: 72.6 -> 72.5 | Hurt |
| Set training stages Ng=2 instead of 3 (ablation) | ToE on DeiT-tiny · ImageNet-1K | Top-1 accuracy: 72.6 -> 72.4 | Hurt |
| Set training stages Ng=4 instead of 3 (ablation) | ToE on DeiT-tiny · ImageNet-1K | Top-1 accuracy: 72.6 -> 72.5 | Hurt |
Surprising findings
- Reducing token redundancy early in training sometimes improves final accuracy over full-token training (e.g., DeiT-tiny +0.4 and CIFAR-100 fine-tuning gains).
Limitations
- Speedup is modest (about 1.2-1.3x practical) compared with the theoretical FLOP reduction.
- The method is applied only at the first block's output; the gains from later blocks were weaker in ablations.
- Some hyper-parameter-relaxed results (ToE Hyper) break the consistency that the method emphasizes.
Experimental setup
Datasets: ImageNet-1K CIFAR-10 CIFAR-100
Benchmarks: ImageNet-1K CIFAR-10 CIFAR-100
Metrics: Top-1 accuracy GFLOPs total GPU hours practical acceleration rate
Baselines: DeiT baseline S2ViTE ToMe NetworkExpansion StackBERT EfficientTrain LV-ViT baseline
Builds on: ToMe NetworkExpansion EfficientTrain DeiT
Closest papers by meaning
- Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer (AAAI 2022 · similarity 0.90)
- Slicing Vision Transformer for Flexible Inference (NeurIPS 2024 · similarity 0.87)
- Linearly Decomposing and Recomposing Vision Transformers for Diverse-Scale Models (NeurIPS 2024 · similarity 0.87)
- Vision Transformers with Self-Distilled Registers (NeurIPS 2025 · similarity 0.87)
- Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalization (NeurIPS 2024 · similarity 0.87)
- Vision Transformers Need Registers (ICLR 2024 · 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.