- The paper presents LYNX, a dynamic mechanism for selecting experts in MoE inference to reduce latency without significant accuracy loss.
- It employs a batch-aware strategy that adapts expert activation based on phase-specific sensitivity to balance high confidence and memory constraints.
- Evaluations on tasks like GSM8k and HumanEval show up to 1.55x speedup while preserving over 55% accuracy in challenging real-time scenarios.
LYNX: Enabling Efficient MoE Inference Through Dynamic Batch-Aware Expert Selection
The paper introduces "LYNX," a system designed to address significant inefficiencies that arise during the inference phase in Mixture-of-Experts (MoE) architectures when scaling LLMs. MoE models are increasingly being used to efficiently scale LLMs by activating only a subset of specialized 'experts' for any given input rather than utilizing the full model capacity. This selective computation leverages the principle that not all inputs require the entire neural network, thereby potentially reducing computational load during both training and, ideally, inference.
However, as highlighted, practical deployment environments pose a conflict between MoE's selective activation and the computational needs associated with batching multiple inference requests. During inference, especially the decode phase, batching necessitates activating every expert in the model even when the request does not require them, a requirement that negates selective computational advantages expected from MoEs. This tension is particularly acute during the decode phase, where memory bandwidth becomes the primary bottleneck.
LYNX addresses this challenge by implementing a dynamic, batch-aware expert selection mechanism. This system, informed by the inherent heterogeneity within MoE models, dynamically reduces the number of active experts during inference in a manner that minimizes accuracy degradation. The core insight driving this approach is that expert importance is not uniform; rather, it varies both across different tokens and between the prefill and decode phases of inference. While the system ensures a careful balance—it maintains necessary model accuracy while decreasing latency by up to 1.55 times—it eschews the need for a priori workload specific tuning or intrusive model modifications.
The paper reports that LYNX achieves a significant reduction in MoE inference latency with negligible accuracy loss across challenging tasks, specifically code generation (HumanEval) and mathematical reasoning (GSM8k). For instance, for compressive reasoning task GSM8k, LYNX reportedly maintains over 55% accuracy even when achieving a speedup of 1.5 times, whilst static pruning approaches see their accuracy decline to below 40%.
The systematic evaluation presented indicates that not all top-k experts are equally critical, a conclusion born from analyses showing heightened sensitivity to the primary (highest confidence) expert in relation to token assignments over secondary options. Additionally, the phase-specific sensitivity findings underscore that expert selection during the prefill phase is more crucial than during the decode phase, where computation redundancy compensates for any suboptimal mapping.
The design of LYNX is operationalized through two primary policy implementations: a latency-preserving policy and an accuracy-preserving policy. The latency-preserving policy prioritizes minimizing the latency impacts by adaptively dropping the least frequently activated experts while the accuracy-preserving policy focuses on maintaining primary expert assignments for high-confidence tokens.
Practically, LYNX has potential implications for AI systems requiring robust real-time inference capabilities without sacrificing model performance, particularly in memory-sensitive environments. The system holds promise for deployment scenarios where large scale MoE models are employed, contributing toward scalable, efficient LLM service delivery in practical, latency-constrained settings.
Theoretically, by elucidating varying importance distributions among experts and the operational temporal decoupling of these experts' importance between inference phases, LYNX opens pathways for further research into optimization strategies and expert allocation algorithms. The work encourages a reconsideration of how MoE models might be designed, particularly in terms of expert distribution and routing mechanics. Future studies could extend beyond the strategies interrogated within LYNX towards even more granular adaptive mechanisms, potentially informed by richer model introspection and on-the-fly model adaptation based on real-time data processing demands.