What if your language model could “think harder” only when needed—without blowing up your GPU bill? Meet Mixture-of-Recursions (MoR), Google DeepMind has released a new Transformers architecture that combines the best of adaptive compute and parameter sharing. Here’s what you need to know.
The Transformer Dilemma: Power vs. Price
We all love what massive Transformers can do—stunning few-shot learning, tricky reasoning, code completion, etc. But let’s be real: scaling up to billions of parameters means you need serious hardware and cash, whether you’re training or deploying.
Two Popular Paths to Efficiency (But Both Are Halfway Solutions)
- Parameter Sharing: Reuse the same weights across layers (think: Universal Transformers, recurrent models). Saves memory, but doesn’t solve compute waste.
- Adaptive Computation: Use early-exit or routing, so “easy” tokens don’t go through all layers. Saves computation, but still requires lots of unique parameters.
But what if you could get the best of both worlds?
Enter Mixture-of-Recursions (MoR): Deep, Smart, and Lean
MoR, introduced by Sangmin Bae and collaborators, is a new Transformer architecture that dynamically decides how many “thinking steps” each token deserves, while reusing the same set of layers again and again.
Picture this:
- All tokens start together, but as the model processes, a “router” decides which tokens need another pass through the shared block of layers.
- “Simple” tokens exit early, “hard” tokens get more recursive attention.
- At every step, only the “active” tokens are part of the expensive attention computation.
- Key-Value (KV) caching is optimized: only store the states you’ll actually need, further saving memory!
It’s like a Transformer with a built-in “focus engine”—more brainpower where it’s needed, less wasted everywhere else.
How Does It Work? (Spoiler: It’s Actually Elegant)
- Recursive Layer Block: Instead of stacking dozens of unique layers, MoR reuses a small stack (e.g., 3 layers) recursively. This slashes parameter count.
- Token-Level Routing: A lightweight router is trained to assign recursion depths per token. Some tokens take 1 pass, others 2, others 3+, depending on their “difficulty.”
- Efficient KV Caching: Only the active tokens at each recursion depth contribute to the key-value cache, reducing memory and compute.
- Routing Variants: The authors explore “expert-choice” (layers pick which tokens to keep) and “token-choice” (each token picks its own path). Both have trade-offs, but “expert-choice” with an auxiliary loss works best.
Why Should You Care?
- Smaller Models, Same or Better Results: MoR models with half or a third the parameter count match—or sometimes beat—vanilla Transformers on validation loss and few-shot tasks.
- Cheaper Training: For a fixed compute budget, MoR trains on more tokens and achieves better scores.
- Faster Inference: Thanks to dynamic depth and efficient batching, MoR can double throughput compared to regular Transformers on the same hardware.
- Memory Footprint: Parameter and KV cache sharing means more room for longer contexts or bigger batch sizes.
Real-World Impact
- Deploy Large-Language-Model Power Without LLM Costs: For startups, researchers, and tinkerers, MoR means you can do more with less.
- Flexible Computation: You can “turn up” the thinking depth for more difficult tasks at inference, without retraining.
- Open Source: The code is available! Check it out on GitHub.
Curious to dive in?
Read the full paper here or check out the official GitHub repo.
#transformers #machinelearning #deeplearning #llm #nlp #ai #research #devjournal
Top comments (0)