Research map / Architectures and efficiency / Attention and transformers
Adaptive Computation Modules: Granular Conditional Computation for Efficient Inference
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.
Introduces Adaptive Computation Modules (ACMs), per-token width-adaptive blocks of nested learners with a gate, and a distillation-based procedure to convert pre-trained transformers (ViT-B, XLS-R-300M) into ACMized models with lower FLOPs.
Problem
Transformer layers use full width for every token, although many tokens need only a fraction of capacity; existing conditional computation methods are coarse (whole-input exits, token dropping, MoE routing) or hard to run efficiently on GPUs.
Approach
A module of ordered learners whose outputs are summed progressively; a per-token gate (Gumbel-Softmax) chooses how many learners to run, so computation width varies per token and per layer. Pre-trained models are converted by distilling each replaced block into learners, pre-training gates with heuristic labels, then end-to-end fine-tuning with budget and diversity losses.
Key results
- On ImageNet-1k, ACMized ViT-B gives favorable accuracy across budgets, with the largest advantage below 12.5 GFLOPs versus A-ViT, MoEfication and Zero Time Waste.
- On CommonVoice-es with XLS-R-300M, ACMs achieve lower WER than MoEfication at every evaluated computational budget.
- Wall-clock latency of an ACM layer on an A100 has negligible overhead over a static MLP and scales linearly with the average number of executed learners.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ACMs instead of MoEfication | XLS-R-300M on CommonVoice · CommonVoice-es | WER | Helped |
| Use ACMs instead of the tested conditional-computation baselines | ViT-B · ImageNet-1k | performance-efficiency trade-off | Helped |
| Use ACMs instead of the tested conditional-computation baselines below 12.5 GFLOPs | ViT-B · ImageNet-1k | performance-efficiency trade-off | Helped |
| Add the entropy loss to budget loss (ablation) | ACMized ViT-B · ImageNet-1k validation | within-image computation diversity | Helped |
| Add the entropy loss to budget loss (ablation) | ACMized ViT-B · ImageNet-1k validation | between-image computation diversity | No clear effect |
| Use the diversity loss instead of the entropy loss (ablation) | ACMized ViT-B · ImageNet-1k validation | between-image computation diversity | Helped |
| Use the diversity loss instead of the entropy loss (ablation) | ACMized ViT-B · within-image computation diversity | within-image computation diversity | Hurt |
| Add both entropy and diversity losses alongside budget loss (ablation) | ACMized ViT-B · ImageNet-1k validation | computation allocation diversity | Helped |
| Increase the average number of executed learners | ACM layer on A100 · 128-sample batch | latency | Hurt |
Surprising findings
- Without the auxiliary diversity losses, gates collapse to choosing the same number of learners for every token, making the model effectively static.
Limitations
- Only MLP and MHA projection modules replaced; self-attention unchanged.
- Evaluated only on ViT-B image classification and XLS-R-300M speech recognition.
- A-ViT and ACMs require multiple training runs at different budget hyperparameters.
Experimental setup
Datasets: ImageNet-1k CommonVoice
Benchmarks: ImageNet-1k CommonVoice-es
Metrics: top-1 accuracy GFLOPs word error rate (WER)
Baselines: A-ViT MoEfication Zero Time Waste
Builds on: Mixture-of-Experts Gumbel-Softmax A-ViT MoEfication Zero Time Waste
Closest papers by meaning
- Resource-Efficient Transformer Pruning for Finetuning of Large Models (CVPR 2024 · similarity 0.87)
- Constrained Belief Updates Explain Geometric Structures in Transformer Representations (ICML 2025 · similarity 0.85)
- Separations in the Representational Capabilities of Transformers and Recurrent Architectures (NeurIPS 2024 · similarity 0.85)
- Adaptive Rank Allocation: Speeding Up Modern Transformers with RaNA Adapters (ICLR 2025 · similarity 0.85)
- Efficient Large Language Model Inference with Neural Block Linearization (NeurIPS 2025 · similarity 0.85)
- Transformer-Squared: Self-adaptive LLMs (ICLR 2025 · 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.