🤖 machine learning

Adaptive Inverted-Index Routing for Granular Mixtures-of-Experts

The paper introduces Adaptive Inverted-Index Routing for MoE (AIR-MoE), a two-stage, drop-in routing mechanism based on vector quantization that efficiently handles granular Mixture-of-Experts models by reducing routing costs while maintaining high performance without requiring structural changes to the model.

Original authors: Klaus-Rudolf Kladny, Maximilian Mordig, Bernhard Schölkopf, Michael Muehlebach

Published 2026-05-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Klaus-Rudolf Kladny, Maximilian Mordig, Bernhard Schölkopf, Michael Muehlebach

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Too Many Chefs" Problem

Imagine you are running a massive restaurant (a Large Language Model) with 65,000 different chefs (experts). Each chef is a tiny specialist who knows a little bit about everything.

In the old way of doing things, when a customer orders a dish (a token of text), the manager had to ask every single one of the 65,000 chefs, "Can you make this?" The manager would then pick the top 2 chefs to actually cook the meal.

The Problem: Asking 65,000 people takes forever. It's slow and wastes a huge amount of energy (computational power), even if you only use 2 chefs.

The "Granular" Solution:
Recent research suggests that having many tiny chefs is actually better than having a few giant ones. But this makes the "asking everyone" problem even worse. You have more chefs to ask, but you still only need a few.

The Solution: AIR-MoE (The Smart Librarian)

The authors propose a new system called AIR-MoE. Instead of asking every chef, they use a two-step "Smart Librarian" system inspired by how libraries organize books.

Step 1: The Coarse Shortlist (The Catalog)

Imagine the 65,000 chefs are organized into 1,000 different "bins" or "shelves" based on what they are good at. These bins are called codewords.

When a customer order comes in, the manager doesn't look at all 65,000 chefs.
They look at the order and quickly figure out which one bin it belongs to (e.g., "This is a French cooking order, so it goes in Bin #42").
Inside Bin #42, there is a pre-made list of the top 500 chefs who are best suited for French food.
The Magic: The manager only looks at these 500 chefs. They ignore the other 64,500 chefs entirely.

Step 2: The Fine Scoring (The Interview)

Now that the manager has narrowed it down to 500 chefs, they do a quick, precise interview with just those 500 to find the absolute top 2 to cook the meal.

Why this works: It's much faster to interview 500 people than 65,000. But because the "bins" were smartly organized, the top 2 chefs are almost certainly in that group of 500.

How It Learns (The "No-Brain" Librarian)

Here is the tricky part: How does the manager know which chefs go into which bin?

In many computer systems, the manager tries to learn this by guessing and getting graded by a teacher (using "gradients"). But in this system, the manager (the codebook) is a bit different.

The chefs and the customer orders are trained by the teacher (the main AI learning process).
The bins (the codebook) are updated separately using a simple, non-differentiable method called adaptive spherical k-means. Think of this as the librarian constantly rearranging the shelves based on what books are currently being checked out, without needing a teacher to tell them exactly how to move them.

Why Is This Better?

The paper claims three main things:

Speed vs. Quality: It finds the best chefs almost as well as asking everyone, but it uses significantly less energy (FLOPs). In their tests, it was up to 10% better at predicting text than other efficient methods, while using fewer resources.
No Rigid Rules: Previous methods forced chefs into fixed groups (like "French chefs only go in Group A"). AIR-MoE is flexible; a chef can be in multiple bins if they are good at many things. It doesn't force a rigid structure on the experts.
It Works: They proved mathematically that if the bins are organized well, the top chefs will almost always be in the shortlist. They also showed that this method prevents "dead chefs" (chefs who never get to cook), which is a common problem in these systems.

Summary Analogy

Old Way: You need to find the best 2 doctors for a specific illness. You call every doctor in the country to see who is available. (Too slow).
Other Efficient Ways: You only call doctors in one specific city or doctors who share the same last name. (Faster, but you might miss the best doctor who lives elsewhere or has a different name).
AIR-MoE: You use a smart directory. You look up your illness, and the directory instantly gives you a list of the top 500 doctors who specialize in that. You then pick the best 2 from that list. It's fast, flexible, and you rarely miss the best doctor.

The paper concludes that this "inverted index" approach (like a library catalog) is a powerful way to make huge AI models faster and smarter without breaking the bank on computing power.

Technical Summary: Adaptive Inverted-Index Routing for Granular Mixtures-of-Experts (AIR-MoE)

1. Problem Statement

Mixture-of-Experts (MoE) models have become a cornerstone of scalable Large Language Models (LLMs) by activating only a subset of experts per token. Recent empirical scaling laws suggest that performance improves with granular experts—using many small experts rather than a few large ones. However, this regime introduces a significant computational bottleneck: the routing cost. As the number of experts ( $E$ ) increases, the router must compute inner products between every token representation and every expert centroid to identify the top- $K$ experts. This overhead can dominate the total computation, rendering the training of highly granular models prohibitive.

Existing solutions to reduce routing costs typically impose structural constraints that limit model expressivity:

Hierarchical MoEs partition experts into fixed, non-overlapping groups, restricting routing to within those groups.
Product-Key/Structural Constraints (e.g., He, 2024) split expert centroids into components to enable exact retrieval but restrict the space of admissible expert representations (e.g., forcing a Cartesian product structure).
Fixed Hashing approaches lack learnability and have been shown to be fundamentally limited.

The challenge is to approximate the top- $K$ routing efficiently without imposing structural constraints on expert parameters or requiring full expert scoring.

2. Methodology: AIR-MoE

The authors propose Adaptive Inverted-Index Routing for MoE (AIR-MoE), a routing architecture inspired by inverted-file indices used in information retrieval and approximate nearest neighbor search. The method approximates the true top- $K$ experts via a two-stage coarse-to-fine retrieval process, avoiding full scoring of all experts.

2.1 Two-Stage Retrieval

For each input token $h$ , the routing proceeds in parallel:

Coarse Shortlisting: The token is mapped to the closest codeword $c_g$ in a learned codebook $\mathcal{C} = \{c_1, \dots, c_G\}$ using cosine similarity. Each codeword $c_g$ is associated with a pre-computed shortlist $L_g$ of $M$ candidate experts. These shortlists are essentially "posting lists" in an inverted index, storing the indices of experts with the highest inner products relative to the codeword centroid.
Fine Scoring: The router computes exact inner products only between the token $h$ and the experts within the retrieved shortlist $L_g$ . The top- $K$ experts are then selected from this reduced candidate set.

This hierarchy ensures that the set of active experts $T$ is a subset of the shortlist $L$ , which is a subset of the full expert set $[E]$ .

2.2 Training and Optimization

A central challenge is that the assignment of tokens to codewords is non-differentiable. AIR-MoE addresses this via a bi-level optimization strategy:

Gradient-Based Training: Token representations and expert centroids are updated via standard backpropagation from the downstream language modeling loss.
Gradient-Free Codebook Training: The codebook $\mathcal{C}$ $C$ and its associated shortlists are updated independently using an adaptive spherical k-means algorithm. This algorithm:
- Uses an Exponential Moving Average (EMA) to track the distribution of token representations.
- Operates on the unit sphere to handle high-dimensional cosine geometry effectively.
- Includes a reinitialization mechanism for "dead" codewords (those with low assignment counts) to prevent collapse.
- Updates shortlists $L_g$ only after optimizer steps, ensuring the index remains consistent with the current state of expert centroids.

Crucially, AIR-MoE does not use the straight-through estimation trick; gradients do not flow through the quantization step to the codebook.

2.3 Theoretical Guarantees

The paper derives a lower bound on the Mass Recall (the total probability mass of the true top- $K$ experts retained in the shortlist). The bound depends on:

The quantization error $\epsilon(h) = \|h - c(h)\|$ : Smaller error leads to better retention.
The shortlist mass $\rho_M(c)$ : The probability mass captured by the top- $M$ experts relative to the codeword.
Theoretically, AIR-MoE retains most routing mass when tokens are well-quantized and the shortlist size $M$ is sufficiently large.

3. Key Contributions

AIR-MoE Architecture: An inverted-index-inspired routing mechanism that approximates top- $K$ routing for granular MoEs without imposing structural constraints (e.g., grouping or parametric restrictions) on expert representations.
Bi-Level Optimization: A training strategy that decouples the gradient-free adaptation of the codebook (via adaptive spherical k-means) from the gradient-based training of model parameters, avoiding the need for straight-through estimation.
Theoretical Analysis: A derived lower bound on mass recall, providing intuition on the conditions under which the approximation preserves routing quality.
Empirical Performance: Demonstration that AIR-MoE achieves a superior perplexity–FLOPs trade-off compared to existing efficient granular routing methods.

4. Experimental Results

The authors evaluated AIR-MoE on WikiText-103, C5, and OpenWebText2 across small, medium, and large model scales (up to ~0.45B parameters) with 65,536 experts.

Performance vs. Baselines: AIR-MoE consistently outperformed existing efficient granular baselines, including PEER (He, 2024) and Hierarchical MoE.
- It achieved perplexity improvements of up to 10% over the PEER baseline at similar FLOP costs.
- While standard granular routing (scoring all experts) achieved the lowest perplexity, it incurred prohibitive computational costs (e.g., ~467.5P FLOPs vs. 324.1P for AIR-MoE on WikiText-103 Small).
Efficiency: AIR-MoE reduced training FLOPs significantly compared to exact granular routing while maintaining competitive perplexity.
Ablation Studies:
- Removing centroid normalization or using Euclidean distance instead of spherical k-means slightly degraded performance.
- Using a static codebook (non-adaptive) resulted in the worst performance, confirming the necessity of the adaptive codebook.
- Expert-choice gating led to high overlap but severe expert starvation (78.6% dead experts), indicating that overlap alone is not a sufficient metric for routing quality.
Expert Usage: Unlike hierarchical or PEER routers which impose rigid structures leading to routing artifacts, AIR-MoE allowed for flexible expert usage without structural constraints.

5. Significance and Claims

The paper claims that AIR-MoE provides an effective mechanism for routing in sparse mixture-of-experts architectures, specifically enabling the use of granular experts without the computational penalty of full routing.

Drop-in Replacement: The method requires no modifications to the underlying model architecture or loss function, serving as a drop-in replacement for standard routers.
Flexibility: By avoiding structural constraints on expert parameters, AIR-MoE allows the model to learn token-expert dependencies more flexibly than methods requiring fixed groupings or product-key structures.
Scalability: The approach scales favorably in large effective-batch regimes because the coarse shortlisting cost is amortized and independent of the number of routed tokens.

The authors acknowledge limitations, noting that while FLOP reductions are significant, actual wall-clock speedups depend on hardware support and optimized implementations to mitigate memory traffic and indexing overheads inherent in granular MoEs. They also note that the current implementation does not employ dimensionality reduction, leaving that as an orthogonal direction for future work.