Early-exiting LLMs that don't break the batch or the KV cache

What it does

Mixture-of-Recursions (MoR) is a Llama-based transformer that lets individual tokens bail out of layers early instead of marching through the full stack. A trained router assigns each token a dynamic recursion depth on the fly, while shared parameters let tokens at different depths stay in the same batch. The authors claim up to 2× inference throughput at similar accuracy to standard transformers, plus lower training FLOPs and memory.

The interesting bit

The clever part is treating early-exiting and batched inference as one problem rather than two. MoR uses a recursion-wise KV caching strategy that selectively stores pairs to avoid the “missing KV cache” trap, and its expert-choice or token-choice routers scatter and gather tokens through shared recursive blocks so nothing idles waiting for slower neighbors.

Key highlights

• Built by modifying LlamaForCausalLM, with two routing flavors: expert-choice and token-choice.
• Trains end-to-end on FineWeb-Edu (SmolLM-Corpus) using 4× H100 or A100 GPUs via DeepSpeed or Accelerate.
• Evaluates with lm-evaluation-harness for few-shot tasks and tracks validation loss for scaling-law analysis.
• The codebase includes an automated YAML config generator for pretraining and evaluation scripts.
• Recursion-wise KV caching is designed to resolve the missing-cache problem while optimizing memory usage.

Caveats

• The current KV-sharing implementation computes all sequences against a shared cache and extracts outputs by masking, which the authors note is a simplification that leaves “further optimizations” on the table.
• Input packing currently allows attention across unrelated documents, and the authors flag FlexAttention and grouped_mm as future optimizations rather than implemented features.
• FSDP, Tensor Parallelism, and Pipeline Parallelism are explicitly left for future work, so scaling beyond the current 4-GPU setup is untested.

Verdict

Worth a look if you’re researching adaptive computation, early-exiting, or efficient LLM inference architectures. Production engineers seeking drop-in speedups should probably wait for the promised optimizations to land.