Mixture-of-Recursions Boosts LLM Efficiency, Delivering Faster Inference and Lower Costs

Researchers at KAIST AI and Mila have developed a new Transformer architecture called Mixture-of-Recursions (MoR), designed to enhance the efficiency of large language models (LLMs) by combining parameter sharing with adaptive computation. This innovation addresses the mounting challenges of training and deploying LLMs, which often require extensive computational resources and memory, making them less accessible to organizations beyond hyperscale data centers. The Scaling Challenges of LLMs The remarkable capabilities of modern LLMs are closely linked to their size. However, as these models grow larger, their memory footprints and computational demands become increasingly burdensome. This has spurred a quest for more efficient designs. Two primary methods have been explored: parameter sharing and adaptive computation. Parameter sharing techniques, such as layer tying, reduce the total number of unique parameters by reusing weights across different parts of the model, cutting computational complexity. Adaptive computation methods, like early exiting, adjust the model to use only the necessary resources during inference, optimizing performance by dynamically allocating compute based on input complexity. Despite these advancements, creating a unified framework that leverages both approaches effectively has remained a challenge until now. How Mixture-of-Recursions (MoR) Works MoR addresses the high computational demands of LLMs by integrating parameter sharing with adaptive computation. It builds on the concept of Recursive Transformers, which apply a set of shared layers multiple times instead of using a deep stack of unique layers. This design allows for increased computation without expanding the model's size. MoR introduces two key enhancements: 1. Lightweight Router: This component intelligently assigns a specific recursion depth to each token, similar to the routing mechanism in Mixture-of-Experts (MoE) models. However, MoR's "experts" are different recursion depths, enabling the model to dynamically decide how many times a shared block of layers should be applied based on the token's complexity. This ensures that computation is used only where it is most needed, reducing wasted cycles on simple inputs. 2. Efficient Key-Value (KV) Caching: Traditional KV caching, which stores information from previous tokens to speed up generation, can become a memory bottleneck in recursive models. MoR implements a "recursion-wise" KV caching mechanism that selectively stores and retrieves key-value pairs for active tokens at each recursion step. This targeted caching reduces memory traffic and improves throughput, enhancing the model's overall efficiency. MoR in Action To validate the framework, the researchers trained MoR models with varying parameter counts, from 135 million to 1.7 billion, and compared them against vanilla and standard recursivebaseline models on validation loss and few-shot accuracy benchmarks. The results were promising: Accuracy and Efficiency: An MoR model achieved 43.1% average few-shot accuracy, surpassing a vanilla baseline model's 42.3% while using almost 50% fewer parameters. Reduced Training Time: When trained on the same amount of data, MoR models reduced training time by 19% and cut peak memory usage by 25% compared to vanilla models. Scalability: At smaller scales (135 million parameters), MoR slightly underperformed the vanilla model, but this gap closed rapidly as the model size increased. For models with over 360 million parameters, MoR matched or exceeded the performance of standard Transformers, particularly on lower compute budgets. Inference Throughput: MoR's design dramatically boosts inference throughput, with one configuration achieving a 2.06x speedup over the vanilla baseline. This improvement could translate into substantial operational cost savings for companies. Practical Path for Enterprise Adoption The research team's practical insights are valuable for enterprise adoption. While the paper showcases results from models trained from scratch, uptraining existing open-source models could be a more cost-effective approach for businesses. According to Sangmin Bae, a co-author and PhD student at KAIST, this method allows for simultaneous inference on more samples, processing a higher number of tokens and handling longer context windows efficiently. MoR also introduces new architectural "knobs" for developers, enabling fine-tuning of the balance between performance and efficiency. For simpler tasks, using models with more recursion steps can offer greater flexibility, while for more complex scenarios, fewer steps might be optimal. The optimal settings will depend on the specific deployment environment, prompting teams to explore the trade-offs based on MoR's principles. Future Potential MoR is not limited to text processing; it is modality-agnostic. This versatility could lead to significant efficiency gains in handling video, audio, and other complex data types. By dynamically adjusting the processing depth for each segment of these data streams, MoR could unlock further cost savings and performance improvements, making large-scale AI more accessible for a broader range of enterprise applications. Industry Evaluation and Company Profiles The introduction of MoR is seen as a significant breakthrough in the field of AI, particularly for organizations looking to leverage LLMs without the prohibitive costs typically associated with training and deployment. Experts praise the architecture for its innovative combination of parameter sharing and adaptive computation, which could redefine the landscape of AI efficiency. KAIST AI and Mila, the institutions behind MoR, are renowned for their contributions to AI research. KAIST AI is part of the Korea Advanced Institute of Science and Technology, a leading global university in science and technology, while Mila is a leading AI research institute in Montreal, Canada, known for its work in deep learning and reinforcement learning. The collaboration between these two institutions highlights the growing international effort to advance AI technologies and make them more practical for real-world applications. By addressing the memory and computational inefficiencies of LLMs, MoR offers a viable path for businesses to achieve high levels of AI performance with significantly reduced resource consumption. This could democratize access to powerful AI models, fostering innovation and efficiency across various industries.