Research map / Architectures and efficiency / Attention and transformers
Vision Transformers as Probabilistic Expansion from 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 PEG, which samples MSA heads and FFN layers from a pretrained DeiT ancestry model via Gaussian probabilistic mixtures into a compact learngene, then expands it with learnable non-linear mappings to initialize elastic-scale ViT descendants.
Problem
Pretrained ViTs assume a one-size-fits-all model, and reusing them at full scale for downstream tasks with varying resource constraints is inelastic and risks negative transfer; prior learngene methods select integral layers and manually stack them with random layers.
Approach
Sample MSA heads and FFN layers from every layer of an ancestry ViT as Gaussian-weighted mixtures (learngene), then expand them with learnable nonlinear mappings to initialize descendant models of arbitrary width and depth, with theoretical Gaussian-mixture approximation guarantees.
Key results
- On downstream tasks, PEG with a 6-layer Small-scale descendant beats Auto-Learngene by 7.81% on iNaturalist-2019 and outperforms From-Scratch and Weight-Multi across datasets (e.g., Flowers Tiny 87.85 vs 68.82 From-Scratch).
- Base-scale PEG descendants reach 3.6x training speedup over From-Scratch on Swin-Base and about 1.7x training cost reduction on DeiT-Base versus 100-epoch from-scratch training.
- With only 25% of ImageNet-1K training data, PEG on DeiT-Base reaches 76.3 vs 65.7 for From-Scratch (+10.6%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the Gaussian mixture with a Uniform mixture (ablation) | PEG · CIFAR-100 | Top-1 accuracy: 83.59 -> 83.02 | Hurt |
| Replace the Gaussian mixture with a Beta mixture (ablation) | PEG · CIFAR-100 | Top-1 accuracy: 83.59 -> 82.13 | Hurt |
| Replace the Gaussian mixture with a Poisson mixture (ablation) | PEG · CIFAR-100 | Top-1 accuracy: 83.59 -> 82.26 | Hurt |
| Add MSAQ inheritance (ablation) | PEG learngene sampling · CIFAR-100 | Top-1 accuracy: 75.18 -> 83.59 | Helped |
| Add MSAK inheritance (ablation) | PEG learngene sampling · CIFAR-100 | Top-1 accuracy: 75.24 -> 83.59 | Helped |
| Add MSAV inheritance (ablation) | PEG learngene sampling · CIFAR-100 | Top-1 accuracy: 74.74 -> 83.59 | Helped |
| Add FFN inheritance (ablation) | PEG learngene sampling · CIFAR-100 | Top-1 accuracy: 78.85 -> 83.59 | Helped |
| Add FFN inheritance to complete MSA inheritance (ablation) | PEG learngene sampling · CIFAR-100 | Top-1 accuracy: 79.02 -> 83.59 | Helped |
| Use PEG initialization instead of From-Scratch training | DeiT-Base · ImageNet-1K, 50% training data | Top-1 accuracy: 74.7 -> 77.1 | Helped |
| Use PEG initialization instead of From-Scratch training | DeiT-Base · ImageNet-1K, 25% training data | Top-1 accuracy: 65.7 -> 76.3 | Helped |
| Use PEG initialization instead of From-Scratch training | Swin-Base · ImageNet-1K, 50% training data | Top-1 accuracy: 76.2 -> 79.7 | Helped |
| Use PEG initialization instead of From-Scratch training | Swin-Base · ImageNet-1K, 25% training data | Top-1 accuracy: 68.1 -> 79.5 | Helped |
Surprising findings
- PEG with Tiny-scale descendants lags behind fully-trained From-Scratch models on ImageNet-1K, while Small and Base scale descendants outperform From-Scratch.
- On Flowers with Tiny-scale descendants, PEG outperforms full Fine-tuning of the pretrained DeiT.
Limitations
- Tiny-scale descendant models underperform fully trained From-Scratch models on ImageNet-1K.
- Theoretical guarantees are stated for approximating parameter distributions with Gaussian mixtures; the Gaussian mixture weights are drawn randomly rather than learned during sampling.
Experimental setup
Datasets: ImageNet-1K Oxford Flowers CUB-200-2011 Stanford Cars CIFAR-10 CIFAR-100 Food101 iNaturalist-2019
Benchmarks: ImageNet-1K Oxford Flowers CUB-200-2011 CIFAR-10 CIFAR-100 Food-101 iNaturalist-2019 Stanford Cars
Metrics: Top-1 accuracy training speedup FLOPs parameter count
Baselines: From-Scratch Fine-tuning Heur-Learngene Weight-Multi Auto-Learngene
Builds on: Learngene DeiT Auto-Learngene Swin Transformer
Closest papers by meaning
- Initializing Variable-sized Vision Transformers from Learngene with Learnable Transformation (NeurIPS 2024 · similarity 0.91)
- Cluster-Learngene: Inheriting Adaptive Clusters for Vision Transformers (NeurIPS 2024 · similarity 0.90)
- Transformer as Linear Expansion of Learngene (AAAI 2024 · similarity 0.90)
- Linearly Decomposing and Recomposing Vision Transformers for Diverse-Scale Models (NeurIPS 2024 · similarity 0.89)
- Elastic ViTs from Pretrained Models without Retraining (NeurIPS 2025 · similarity 0.88)
- Probabilistic Conceptual Explainers: Trustworthy Conceptual Explanations for Vision Foundation Models (ICML 2024 · similarity 0.86)
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.