Research map / Architectures and efficiency / Attention and transformers
N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation
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 N-vium (Quadrivium), a mixture-of-exits transformer whose next-token distribution is a learned token-adaptive mixture over intermediate exit heads, enabling exact early sampling and KV-cache recovery via piggybacked deferred computation, with up to 57.9% decode latency speedup at matched perplexity.
Problem
Autoregressive LLM decoding is sequential through all layers; early-exit methods reduce FLOPs per token via approximations that degrade quality, skip KV caches, and need hand-tuned confidence thresholds.
Approach
Define the next-token distribution as a learned mixture over multiple exit heads with router-determined weights, sample exactly via per-junction Bernoulli routing, and defer upper-layer KV computation to be batched with later tokens ("piggybacking"), trading sequential depth for GPU parallelism rather than reducing FLOPs.
Key results
- Largest model (1.5B, 4-vium) reaches 57.9% wall-clock latency-per-token speedup over a parameter- and data-matched dense baseline at no perplexity cost.
- Across depth scaling (6 to 48 layers at width 768), speedups range from 33.3% to 58.2% with perplexity differences reported as relative improvements over baseline.
- Routing analysis shows the mixture perplexity lies well below each individual exit's perplexity, indicating exits contribute complementary roles.
Surprising findings
- Retrofitting a pretrained Llama-3.2-1B with the mixture loss leads to router collapse, with all tokens routed to the final head.
Limitations
- Batched inference must be adapted since tokens may exit at different junctions, similar to early-exit and MoD challenges.
- Requires additional loss terms and tuning of the compute penalty beta for speedup.
- Evidence limited to models up to 1.5B parameters.
Experimental setup
Datasets: C4 Tulu 3 SFT mixture
Benchmarks: SpecBench MMLU ARC-Easy HellaSwag PIQA OpenBookQA WinoGrande
Metrics: relative perplexity difference (%) latency per token speedup (%) accuracy
Baselines: CALM* LayerSkip parameter-matched dense transformer
Builds on: CALM LayerSkip speculative decoding Mixture of Depths LLaMA 2
Closest papers by meaning
- Ladder-Residual: Parallelism-Aware Architecture for Accelerating Large Model Inference with Communication Overlapping (ICML 2025 · similarity 0.85)
- Learning Rate Transfer for Hybrid Transformer-SSM Architectures (NeurIPS 2026 · similarity 0.85)
- Ultra-Sparse Memory Network (ICLR 2025 · similarity 0.84)
- A General and Efficient Training for Transformer via Token Expansion (CVPR 2024 · similarity 0.84)
- Linear Transformers Implicitly Discover Unified Numerical Algorithms (NeurIPS 2025 · similarity 0.84)
- SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention (NeurIPS 2024 · similarity 0.84)
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.
