Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

2024/04/02 by David Raposo, Raposo, David, Sam Ritter +11 · 9 voices · 63 citations
Computer Science · #Natural Language Processing Techniques #cs.CL #cs.LG

paper · pdf · doi:10.48550/arxiv.2404.02258

openalex publication_date 2024/04/02 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens (k) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-k routing mechanism. Since k is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the k tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50% faster to step during post-training sampling.

Cited by

• Per-Token Fixed-Point Convergence in Depth-Recurrent Transformers
• Mobius Learning: Cyclic Depth Folding in Transformers
• Adaptive Compute in Latent World Models: When Depth Helps, Hurts, or Doesn't Matter
• Adaptive Depth Sparse Framework: Similarity-Driven Resource Allocation for Pre-Trained LLMs
• Adaptive Depth in Looped Transformers: Diagnosing Learned Halting Gates and Trajectory Readouts
• Dynamic Chunking for End-to-End Hierarchical Sequence Modeling
• Brain-Like Processing Pathways Form in Models With Heterogeneous Experts
• Learning When Not to Attend Globally
• MoE Pathfinder: Trajectory-driven Expert Pruning
• Dynamic Rebatching for Efficient Early-Exit Inference with DREX
• How Smoothing is N-simplicial Attention?
• Mixture of Attention Schemes (MoAS): Learning to Route Between MHA, GQA, and MQA
• VersatileFFN: Achieving Parameter Efficiency in LLMs via Adaptive Wide-and-Deep Reuse
• SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations
• MIDUS: Memory-Infused Depth Up-Scaling
• Subjective Depth and Timescale Transformers: Learning Where and When to Compute
• AdaPerceiver: Transformers with Adaptive Width, Depth, and Tokens
• Mixture of States: Routing Token-Level Dynamics for Multimodal Generation
• EarthSight: A Distributed Framework for Low-Latency Satellite Intelligence
• Think-at-Hard: Selective Latent Iterations to Improve Reasoning Language Models
• Teaching Pretrained Language Models to Think Deeper with Retrofitted Recurrence
• Route Experts by Sequence, not by Token
• Temporal Zoom Networks: Distance Regression and Continuous Depth for Efficient Action Localization
• Alleviating Forgetfulness of Linear Attention by Hybrid Sparse Attention and Contextualized Learnable Token Eviction
• Intent-Driven LLM Ensemble Planning for Flexible Multi-Robot Disassembly: Demonstration on EV Batteries
• Mirror Speculative Decoding: Breaking the Serial Barrier in LLM Inference
• Dr.LLM: Dynamic Layer Routing in LLMs
• Catch Your Breath: Adaptive Computation for Self-Paced Sequence Production
• DND: Boosting Large Language Models with Dynamic Nested Depth
• Bridging the Gap Between Latent and Explicit Reasoning with Looped Transformers
• Tapered Language Models
• TIDE: Token-Informed Depth Execution for Per-Token Early Exit in LLM Inference
• Dense2MoE: Restructuring Diffusion Transformer to MoE for Efficient Text-to-Image Generation
• Encode, Think, Decode: Scaling test-time reasoning with recursive latent thoughts
• Thoughtbubbles: an Unsupervised Method for Parallel Thinking in Latent Space
• Beyond Linear Probes: Dynamic Safety Monitoring for Language Models
• Two-Scale Latent Dynamics for Recurrent-Depth Transformers
• DynaNav: Dynamic Feature and Layer Selection for Efficient Visual Navigation
• WIDE: Boosting Adaptive LLM Inference via Token-level Dynamic Width Pruning
• Multi-Head Attention Residuals
• The Topological Trouble With Transformers
• TRACE: Unlocking Effective CXL Bandwidth via Lossless Compression and Precision Scaling
• CogVLA: Cognition-Aligned Vision-Language-Action Model via Instruction-Driven Routing & Sparsification
• Cognitive Workspace: Active Memory Management for LLMs -- An Empirical Study of Functional Infinite Context
• EA-ViT: Efficient Adaptation for Elastic Vision Transformer
• Variance-Based Pruning for Accelerating and Compressing Trained Networks
• Change of Thought: Adaptive Test-Time Computation
• Mono-InternVL-1.5: Towards Cheaper and Faster Monolithic Multimodal Large Language Models
• Preference-Driven Multi-Objective Combinatorial Optimization with Conditional Computation
• ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference
• Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
• OrthoRank: Token Selection via Sink Token Orthogonality for Efficient LLM inference
• Towards Distributed Neural Architectures
• Universal Retrieval for Multimodal Trajectory Modeling
• Learning to Skip the Middle Layers of Transformers
• Efficiency Robustness of Dynamic Deep Learning Systems
• Accurate Sublayer Pruning for Large Language Models by Exploiting Latency and Tunability Information
• AdaDecode: Accelerating LLM Decoding with Adaptive Layer Parallelism
• SkipGPT: Dynamic Layer Pruning Reinvented with Token Awareness and Module Decoupling
• Fast and Cost-effective Speculative Edge-Cloud Decoding with Early Exits
• System-1.5 Reasoning: Traversal in Language and Latent Spaces with Dynamic Shortcuts
• Efficient Data Driven Mixture-of-Expert Extraction from Trained Networks
• Chain-of-Model Learning for Language Model

Discussions

• Mixture-of-Depths: Dynamically allocating compute in transformers [hn, 281 points, 83 comments]
• Mixture-of-Depths: Dynamically allocating compute in transformer language models [hn, 5 points, 2 comments]
• Mixture-of-Depths: Dynamically allocating compute in transformers [hn, 4 points, 0 comments]
• This seems like a cool way to use an adaptive amount of compute per token. I speculate that models like these will have more faithful CoT since they don't get to do "extra" reasoning on easy tokens ht [bsky, 4 points, 1 comments]
• DeepMind: Mixture-of-Depths: Dynamically allocating compute in transformers [hn, 2 points, 1 comments]
• Dynamically allocating compute in transformer-based language models [hn, 2 points, 0 comments]
• Just remembering this from last year where if you give models telemetry about their own state they can learn to be efficient I’m surprised I haven’t seen more from this yet arxiv.org/abs/2404.02258 [bsky, 2 points, 0 comments]
• Adaptive Architectures? Reasoning is hot right now with OpenAI's O3. But that super costly to evaluate! Can we adapt the compute to not waste all our flops on less useful tokens? We should! arxiv.o [bsky, 1 points, 1 comments]
• Stochastic Depthの今どきバージョンみたいな感じで面白かった。 https://arxiv.org/abs/2404.02258 [bsky, 0 points, 0 comments]

Related

• Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• TwIST: Rigging the Lottery in Transformers with Independent Subnetwork Training
• Transformer^-1: Input-Adaptive Computation for Resource-Constrained Deployment
• Subjective Depth and Timescale Transformers: Learning Where and When to Compute