N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation

Abstract: Improving the inference efficiency of autoregressive transformers typically means reducing FLOPs per token, usually through approximations that degrade model quality. We introduce N-vium, a mixture-of-exits transformer that partially parallelizes computation across depth on standard hardware, increasing effective FLOPs per second rather than minimizing compute per token. N-vium attaches prediction heads at multiple depths and defines the next-token distribution as a learned mixture over these exits, with token-adaptive routing. This formulation strictly generalizes the standard transformer, which is recovered exactly when routing assigns zero mass to all intermediate heads. Sampling from the mixture is exact, and complete KV caches are recovered by deferring the upper-layer computation and batching it with later tokens. We pretrain N-vium at scales up to 1.5B parameters. Our largest model reaches 57.9% wall-clock speedup over a parameter- and data-matched standard transformer at no perplexity cost.

• The paper demonstrates a mixture-of-exits transformer that uses adaptive early exits to achieve 57.9% wall-clock speedup without sacrificing perplexity.
• It introduces a router MLP to compute token-adaptive exit probabilities, enabling exact sampling and efficient KV cache recovery via piggybacking.
• Empirical results across various scales and hardware show that N-vium maintains quality comparable to dense models while outperforming prior early-exit methods.

N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation

Motivation and Problem Setting

Autoregressive transformers for language modeling typically process tokens sequentially through all decoder layers, coupling hidden state computation for KV cache with next-token prediction. Empirical evidence using logit lens analysis demonstrates that, for many tokens, output distributions stabilize well before the final layer. Existing early exit methods attach prediction heads at intermediate layers to exploit this convergence, emitting predictions before full depth. However, these approaches treat intermediate predictions as approximations, rely on confidence thresholds, discard upper-layer KV computation for speed, and suffer from train-test mismatch, leading to compromised quality and inexact sampling.

N-vium Architecture and Method

N-vium introduces a mixture-of-exits framework strictly generalizing the standard transformer. Prediction heads are attached at multiple depths, each acting as a sampling junction. At every junction $k$, a router MLP computes a token-adaptive exit probability $w_k(h_t^{(\ell_k)})$ from the current hidden state. The next-token distribution is a convex combination $\pi_{\mathrm{mix}} = \sum_{k=1}^{N} p_k \pi_k$, where $p_k$ is the probability of exiting at junction $k$, computed recursively via Bernoulli sampling.

If a token exits early, remaining layers are processed later—batched in parallel with future tokens—to recover full KV caches, termed piggybacking. All exits share a projection head, with lightweight adapter MLPs for per-junction specialization. Sampling at inference is exact and matches the trained $\pi_{\mathrm{mix}}$. At training, N-vium minimizes cross-entropy over $\pi_{\mathrm{mix}}$ and incorporates a compute penalty weighted by $\beta$ to incentivize early exits.

Critically, N-vium does not reduce per-token FLOPs but increases utilization via parallelization. This approach decouples next-token prediction from KV state buildup, utilizing hardware more efficiently.

Theoretical Properties

The mixture-of-exits formulation ensures that the token-wise loss is bounded between the loss of the earliest and latest head, and an oracle router can always achieve at least the performance of the best individual head. Unlike early exit methods, N-vium guarantees exact sampling and full cache recovery.

Empirical Results

N-vium models were pretrained up to 1.5B parameters on C4 with the LLaMA2 backbone. The largest 1.5B model achieved 57.9% wall-clock speedup over a parameter- and data-matched dense baseline with no perplexity degradation ($\Delta = +0.55\%$). Depth scaling from 6 to 48 layers and width scaling up to 1536 consistently yielded speedups (up to 53%-58%) with perplexity differences remaining negligible. Intermediate heads exhibited distinctly higher perplexity than the joint mixture, validating learned routing rather than collapse.

Comparisons with CALM and LayerSkip baselines show that N-vium achieves quality-preserving speedup, whereas CALM incurs significant perplexity degradation and LayerSkip’s gains do not match the speed-quality Pareto attained by N-vium. Benchmark transfer via supervised fine-tuning on Tulu3 and downstream tasks (MMLU, ARC-Easy, HellaSwag, etc.) confirms that N-vium models remain steerable and retain downstream performance comparable to dense baselines. Portability of the speedup was demonstrated across A100, RTX A6000, RTX 3090, with consistent acceleration; only older Turing hardware shows lower gains due to missing bfloat16 support.

Practical and Theoretical Implications

N-vium’s approach to maximizing effective FLOPs/s instead of minimizing per-token computation fundamentally alters inference optimization, directly addressing hardware utilization limits in transformer generation. The exact early sampling and deferred cache computation introduce theoretical guarantees absent in prior adaptive/early exit architectures. The technique is lightweight—parameter overhead is negligible at scale—and compatible with standard transformer variants, speculative decoding, and RL fine-tuning. The tradeoff is slightly increased training FLOPs per token, but inference efficiency dominates lifecycle costs as scale increases.

Implementing batched inference with N-vium and hyperparameter tuning (especially $\beta$) require further study, but initial evidence indicates it scales favorably.

Future Directions

N-vium establishes a new axis for transformer inference acceleration, orthogonal to speculative decoding and mixture-of-experts methods. Future work could integrate N-vium with reinforcement learning, broader backbone variants, or extend piggybacking strategies for even denser parallelization. Robustified batched inference or application to massive model scales ($w_k(h_t^{(\ell_k)})$010B) are natural next steps.

Conclusion

N-vium reframes inference acceleration as a function of computation scheduling, not computation avoidance, while guaranteeing exact token sampling and cache consistency. Comprehensive empirical evidence supports its claims: significant speedup with no quality loss up to billion-scale, compatibility with downstream fine-tuning, and hardware generality. This methodology defines a principled foundation for inference-efficient transformers, addressing theoretical and practical gaps in existing early-exit and adaptive-depth computation paradigms (2605.13190).