M2R2: Mixture of Multi-Rate Residuals for Efficient Transformer Inference

TL;DR

M2R2 introduces a dynamic residual modulation framework for transformers that improves inference speed and maintains quality by adjusting residual velocity based on token complexity, outperforming existing methods.

Contribution

The paper proposes M2R2, a novel method that dynamically modulates residual velocity in transformers, enhancing inference efficiency over static or distance-based approaches.

Findings

Achieves up to 2.8x speedup on reasoning tasks like MT-Bench.

Outperforms state-of-the-art distance-based residual strategies.

Reduces expert-switching and accelerates decoding in MoE architectures.

Abstract

Residual transformations enhance the representational depth and expressive power of large language models (LLMs). However, applying static residual transformations across all tokens in auto-regressive generation leads to a suboptimal trade-off between inference efficiency and generation fidelity. Existing methods, including Early Exiting, Skip Decoding, and Mixture-of-Depth address this by modulating the residual transformation based on token-level complexity. Nevertheless, these approaches predominantly consider the distance traversed by tokens through the model layers, neglecting the underlying velocity of residual evolution. We introduce Mixture of Multi-rate Residuals (M2R2), a framework that dynamically modulates residual velocity to improve early alignment, enhancing inference efficiency. Evaluations on reasoning oriented tasks such as Koala, Self-Instruct, WizardLM, and MT-Bench show M2R2 surpasses state-of-the-art distance-based strategies, balancing generation quality and speedup. In self-speculative decoding setup, M2R2 achieves up to 2.8x speedups on MT-Bench, outperforming methods like 2-model speculative decoding, Medusa, LookAhead Decoding, and DEED. In Mixture-of-Experts (MoE) architectures, integrating early residual alignment with ahead-of-time expert loading into high-bandwidth memory (HBM) accelerates decoding, reduces expert-switching bottlenecks, and achieves a 2.9x speedup, making it highly effective in resource-constrained environments.