Research map / Architectures and efficiency / Pruning and efficient tuning
Transformer as Linear Expansion of Learngene
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 TLEG, which learns a compact learngene (two Transformer layers' worth of parameters) via a linearly expanded auxiliary Transformer trained with soft distillation from LeViT-384, then linearly expands it to initialize Transformers of varying depths.
Problem
Efficiently producing and initializing Transformers of varying depths for diverse resource constraints without storing and pre-training a separate model for each scale.
Approach
Observes that the layer-position-to-weight relationship in well-trained ViTs is approximately linear, so each layer's weights are parameterized as theta_B + (l-1)/L * theta_A; only the two shared modules are learned (via distillation from an ancestry model) and then linearly expanded to initialize descendant models of any depth.
Key results
- On ImageNet-1K, 6-layer-to-12-layer Des-S/Des-B initialized from learngene match or exceed scratch training (e.g., Des-B 12-layer 79.9% vs 78.6% scratch) at about 2x lower training cost.
- On downstream transfer with Des-S, TLEG beats Mini-Init by 7.66% on CIFAR-100 and 7.53% on Mini-ImageNet while storing 3.9M vs 11.0M parameters for initialization.
- For Des-B with 5 depths on CIFAR-100, TLEG stores 14.7M parameters versus 285.7M for Pre-Fin(U), a 19x reduction, with comparable accuracy (86.52-88.34% vs 87.01-88.62%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Initialize Des-Ti from a linearly expanded learngene instead of training from scratch | Des-Ti · ImageNet-1K, 3 layers | Top-1 accuracy (%): 45.0 -> 46.6 (+1.6) | Helped |
| Initialize Des-S from a linearly expanded learngene instead of training from scratch | Des-S · ImageNet-1K, 6 layers | Top-1 accuracy (%): 68.7 -> 69.5 (+0.8) | Helped |
| Initialize Des-B from a linearly expanded learngene instead of training from scratch | Des-B · ImageNet-1K, 12 layers | Top-1 accuracy (%): 78.6 -> 79.9 (+1.3) | Helped |
| Use TLEG rather than Mini-Init to initialize Des-S | Des-S · Mi-INet | Top-1 accuracy (%): 73.39 -> 80.92 | Helped |
| Use TLEG rather than Mini-Init to initialize Des-S | Des-S · Ti-INet | Top-1 accuracy (%): 64.56 -> 71.32 | Helped |
| Use TLEG rather than Mini-Init to initialize Des-S | Des-S · C-100 | Top-1 accuracy (%): 75.98 -> 83.64 | Helped |
| Use TLEG rather than Pre-Fin(U) to initialize Des-S | Des-S · C-100 | Top-1 accuracy (%): 84.43 -> 83.64 | Hurt |
| Include linear expansion of LN along with MSA and MLP (ablation) | TLEG Aux-S / Des-S · C-100, 6 layers | Top-1 accuracy (%): 82.03 -> 83.64 | Helped |
| Expand MSA, MLP, and LN rather than MSA only (ablation) | TLEG Aux-S / Des-S · C-100, 6 layers | Top-1 accuracy (%): 79.16 -> 83.64 | Helped |
| Expand MSA, MLP, and LN rather than MLP only (ablation) | TLEG Aux-S / Des-S · C-100, 6 layers | Top-1 accuracy (%): 77.26 -> 83.64 | Helped |
| Initialize all six layers rather than only layers 1–3 or 4–6 (ablation) | TLEG Des-Ti · C-100, 6 layers | Top-1 accuracy (%): 74.29 -> 78.66; 72.00 -> 78.66 | Helped |
| Initialize all six layers rather than only layers 1–3 or 4–6 (ablation) | TLEG Des-S · C-100, 6 layers | Top-1 accuracy (%): 74.88 -> 83.64; 74.33 -> 83.64 | Helped |
Surprising findings
- Partial linear expansion (from the 3rd layer onward) slightly outperformed linear expansion on all layers for Des-Ti on CIFAR-100 (79.02% vs 78.66%) and Des-S (84.13% vs 83.64%).
Limitations
- Relies on an approximate linear fit derived from PCA of a single well-trained ViT-B; the linearity assumption is empirically motivated but not theoretically justified.
- Des-Nets are still fine-tuned, and many variants of linear expansion strategies are left for future work.
- Compares against a limited set of initialization baselines; some results (e.g., Des-Ti on ImageNet-1K with 9 and 12 layers) show marginal gains.
Experimental setup
Datasets: ImageNet-1K iNaturalist 2019 Mini-ImageNet Tiny-ImageNet CIFAR-10 CIFAR-100 Food-101
Benchmarks: ImageNet-1K iNat 2019 Mini-ImageNet Tiny-ImageNet CIFAR-10 CIFAR-100 Food-101
Metrics: Top-1 accuracy Top-5 accuracy FLOPs Params S-Params
Baselines: Scratch Pre-Fin(U) He-LG Share-Init Mini-Init
Builds on: Learngene DeiT MiniViT Knowledge distillation
Closest papers by meaning
- Initializing Variable-sized Vision Transformers from Learngene with Learnable Transformation (NeurIPS 2024 · similarity 0.90)
- Vision Transformers as Probabilistic Expansion from Learngene (ICML 2024 · similarity 0.90)
- Cluster-Learngene: Inheriting Adaptive Clusters for Vision Transformers (NeurIPS 2024 · similarity 0.87)
- Linearly Decomposing and Recomposing Vision Transformers for Diverse-Scale Models (NeurIPS 2024 · similarity 0.87)
- Trainable Transformer in Transformer (ICML 2024 · similarity 0.86)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.85)
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.