AIR-MoE routes with vector-quantized two-stage inverted indexing, improving the perplexity-FLOPs trade-off at over a million tiny experts
Synopsis
The work introduces AIR-MoE, an inverted-index-inspired two-stage MoE routing architecture based on vector quantization: a first stage assigns tokens to VQ codewords to build a candidate expert set, and a second stage computes exact routing scores restricted to that shortlist, approximating true top-K routing while avoiding full expert scoring and imposing no structural constraints on expert parameters; the authors also provide a lower bound on the mass recall achieved by AIR-MoE, and empirically improve the perplexity-FLOPs trade-off over existing MoE routers in regimes with up to more than a million tiny experts while achieving consistent perplexity improvements over the best existing baseline.
Figure 1: Overview of AIR-MoE . The method consists of two parts: The coarse shortlisting stage uses vector quantization to select a codeword (blue) that stores a pre-computed expert shortlist L L that references specific expert centroids (green) and is updated after each optimizer step. The fine scoring stage takes the shortlisted expert weights and scores them exactly. Notably, the codebook is learned via gradient-free optimization (gear symbol) and only token representations and expert centroids are trained using the downstream gradient (dashed orange) without straight-through estimation trick.
arXivInterpretation
Introduces AIR-MoE two-stage routing: tokens are first assigned to VQ codewords via vector quantization to construct a candidate expert set, then exact scoring is performed within that shortlist. Relative to standard routing that scores all experts, this coarse-shortlist-plus-fine-score procedure approximates true top-K routing while avoiding the cost of full expert scoring. The abstract describes the architecture and the two-stage procedure and states that it approximates true top-K routing; implementation details and experimental setup are not expanded in the abstract.
AIR-MoE serves as a drop-in replacement for standard routers, requiring no modifications to model architecture or loss function and imposing no structural constraints on expert parameters. Unlike prior work that imposes structural constraints on expert parameters, this method replaces the router while keeping architecture and loss unchanged. The abstract explicitly states 'drop-in replacement,' 'no modifications to the model architecture or loss function,' and 'no structural constraints on expert parameters.'
Provides a lower bound on the mass recall achieved by AIR-MoE, offering insight into its inner workings. Adds a theoretical characterization of routing recall beyond empirical results. The abstract states it will 'provide a lower bound on the mass recall achieved by AIR-MoE that yields insights into the inner workings,' without giving the bound's form or value.
In granular regimes with up to more than a million tiny experts, AIR-MoE improves the perplexity-FLOPs trade-off and achieves consistent perplexity improvements over the best existing baseline. Addresses the bottleneck of rising routing cost in granular experts and reports empirical gains at extreme expert counts. The abstract reports the scale of 'up to more than a million tiny experts' and 'consistent perplexity improvements over the best existing baseline,' without specific perplexity values, FLOPs figures, or dataset names.
Perspective
The work targets the routing component of sparse MoE and applies to granular regimes with very large expert counts (the abstract's 'up to more than a million tiny experts'), used as a drop-in replacement for standard routers without changing model architecture or loss function. Its value lies in preventing routing overhead from canceling the performance gains of granular experts, making it most relevant to MoE practitioners attentive to compute and memory budgets in training and inference. The mass-recall lower bound described in the abstract offers an analytical lens on why the method works.
The abstract does not give specific perplexity or FLOPs values, baseline names, datasets, or model scales, nor does it state the derivation conditions or tightness of the mass-recall lower bound; how hyperparameters such as the number of VQ codewords and shortlist size affect recall and final performance, and the magnitude and statistical stability of the 'consistent perplexity improvements,' need confirmation in the main text. The abstract also does not indicate whether the method applies equally when expert counts are small or non-granular.