Channel Selection for Mixture-of-Experts Decoding Needs No Calibration Data

AI summary

Overview

Research area: Efficient inference for Mixture-of-Experts (MoE) large language models, specifically sparse channel selection inside routed experts during decoding.

Technical level: Intermediate — the paper is readable with a basic grasp of transformer MoE layers, but the hardware-sector accounting and kernel details assume some systems background.

Scope (one sentence): The paper argues that the calibrated thresholds used by nearly all MoE channel selectors are unnecessary and unreliable, shows that top-k selection needs no calibration data, and introduces a calibration-free "reader" built purely from model weights that keeps 98.1% of full-model accuracy at 17.8% of routed-expert reads while running MoE blocks up to 2.69× faster than optimized vLLM.

What This Paper Is About

At the small batch sizes used for interactive chat, MoE decoding is limited by the cost of pulling routed-expert weights from GPU memory, and most of the intermediate "channels" in an expert barely affect any single token. Existing selectors thus predict which channels matter and load only those, generally deciding with a cutoff threshold calibrated on offline data. The authors show this calibration step is both unnecessary and actively misleading, and they build a selector that needs no calibration data at all.

Key Contributions

1. A demonstration that calibrated thresholds do not control read budget. Because channel scores vary per token and per input, the number of channels passing a fixed cutoff drifts; selectors with calibrated thresholds read below their target in all 72 benchmark runs the authors made, with Prox calibrated for 18% actually reading 12.9% on average and as little as 9.9%.

2. A controlled keep-rule comparison holding the channel score fixed. Each selector is run twice on its own score — once with its calibrated threshold, once with top-k — at the budget the threshold was calibrated for. Top-k matches or beats the threshold in 40 of 43 comparisons at 13% and 18% budgets, including all 23 on generation tasks.

3. A reader built entirely from model weights. It screens channels with a rotated ternary sketch (about 7% of the BF16 checkpoint), computes top-k exactly, and recovers skipped channels through a rank-32 map of the down projections (about 0.5% extra memory). Nothing is fitted to the target model.

4. Triton kernels with measured speedups. MoE blocks run 1.60× to 1.90× faster than optimized vLLM at batch one and up to 2.69× at larger batches, on the same three models, with reads counted in 32-byte hardware sectors and verified against a GPU profiler.

Main Findings

• Read drift is systematic, not noise. Selectors with calibrated thresholds read less than their budget in all 72 benchmark runs, across Prox and the authors' reader at targets of about 13%, 18%, 28% and 38%. Prox at an 18% target reads 12.9% on average, and roughly 1.7 percentage points of that gap is built into Prox's published budget-splitting rule once reads are counted in sectors.

• Counts fall as answers grow. Around 6,000 tokens into an answer, Qwen3-Thinking keeps 71% of its calibrated channel count, Gemma 4 keeps 92%, and Qwen3.5 keeps 88%. Even on held-out chat, the middle 90% of Qwen3-Thinking tokens keep 9.0% to 16.2% of channels, and on benchmark answers the median falls to 10.0%.

• Kept channels track earlier context, not token difficulty. Placing a 256-token window from position 6,000 directly after the question makes those same tokens keep 57% more channels. Kept counts correlate with the full model's next-token loss at only −0.29 to +0.14 (Spearman).

• Top-k ties or wins on the same score. Of 43 paired comparisons, the top-k version ties in 30, leads in 10, and trails in 3 (each time by at most 2.8 points). Run at the reads the threshold version actually used — a check that favors the threshold — top-k still matches or beats it in 37 of 41 comparisons.

• Thresholds cause repetition-loop collapse on Gemma 4. At the 13% budget, top-k versions of Prox and its output-scaled variant beat their threshold versions by 38 to 45 points. Prox with an 18% threshold reads 13.3% and scores 11.7; with top-k at nearly the same reads (13.5%) it scores 56.3, and 145 rather than 674 of 792 answers hit the token limit. Neither a 4% floor on channels nor calibration on Gemma 4's own reasoning prevents it.

• Scoring matters more than the keep rule. Switching Prox from threshold to top-k changes KL divergence by only −17% to −3% on held-out OLMoE documents; rescaling the same score to output units divides KL by 1.3 to 1.5. Across four scores and four budgets, KL falls as kept channels carry more of the routed output's energy (Spearman −0.96).

• The reader is the most accurate selector tested at low budgets. At 12.8% of reads it reaches 94.5% mean accuracy (86.9% worst) over nine model–benchmark pairs; at 17.8% it reaches 98.1% mean (95.0% worst). Prox as published at a 13% target reads 8.5% and averages 51.5%, with a 0.0% worst case. WINA as published, which already uses top-k, averages only 28.4% at 20% of reads — evidence that top-k cannot rescue a weak score.

• Speedups are real but partial. MoE blocks run 1.90×, 1.71× and 1.60× faster at batch one on Qwen3-Thinking, Gemma 4 and Qwen3.5, and 2.39×, 2.34× and 2.51× at batch four. Whole models decode 1.15× to 1.36× faster at batch one and up to 1.88× at batch eight or sixteen, falling to 1.00× to 1.24× at batch 128.

Methodology in Plain English

The authors separate a channel selector into two parts: a score that guesses how much each channel would contribute, and a keep rule that decides which channels to actually compute. Most prior work uses a calibrated threshold as the keep rule; the alternative is top-k, which simply takes each token's k highest-scoring channels.

To isolate the keep rule from scoring error, they take existing selectors (Prox, TEAL, and others) and run each one twice on its identical score — once with its calibrated threshold, once with top-k — at the same target budget. This is the key experimental design choice, because earlier studies that reached opposite conclusions often changed the score along with the keep rule or tested dense models on short tasks.

They also count cost carefully. Instead of counting "channels read," they count bytes fetched in 32-byte GPU sectors, charging every touched sector in full, and they validate the count with a GPU profiler (18.3% measured versus 17.8% counted at batch one on Qwen3-Thinking).

For their own reader, they build everything at load time from the weights: a randomized Hadamard rotation in blocks of 128 coordinates to spread out outliers, ternary codes (−1, 0, 1) for the rotated gate and up rows, one scale and one quantization-error term per block, and a rank-32 singular-value map of the down-projection columns for recovery. At decode time the reader rotates the hidden state once, scores all candidate channels by estimated output magnitude, loads and exactly computes the top-k, and adds back the skipped channels' estimated contribution through the low-rank map.

They evaluate on Qwen3-30B-A3B-Thinking-2507, Gemma-4-26B-A4B-it and Qwen3.5-35B-A3B in vLLM 0.29.0, on GPQA-Diamond, IFEval and MMLU-Pro, plus OLMoE for likelihood and KL-divergence analysis.

Why This Matters

Impact on research. The paper resolves a live disagreement in the literature: six prior studies compared thresholds against top-k and split roughly evenly, but five tested dense models on mostly short tasks, two changed the score along with the rule, and none ran top-k at the reads a threshold actually used. This paper holds the score fixed, tests MoE models on long chain-of-thought generation, and reports both the target-budget and realized-read comparisons — and finds no model where the threshold version wins. It also documents a failure mode (repetition-loop collapse under calibrated thresholds) that perplexity and short likelihood benchmarks do not expose.

Real-world applications:

• On-device and edge chat assistants, where batch size is one and decoding latency is dominated by weight loading from memory.
• Low-latency interactive coding or reasoning assistants, where 32,768-token reasoning traces are common and must not degenerate into repetition.
• Cost reduction in serving fleets, since the paper reports that at 100 tokens per second per user, one GPU serves 3.3 to 4.3 times as many tokens.
• Memory-constrained deployment of MoE models, where the modest 7% sketch overhead is a cheaper trade than keeping full expert weights hot.

Industry relevance. The reader is a drop-in addition to vLLM, requires no retraining, fits nothing to the target model, and ships as Triton kernels plus vLLM patches in a public repository. The "no calibration data" property removes a deployment step and a source of silent failure that operators currently cannot observe in advance, since the realized read budget under a threshold is known only after the fact.

Future Directions

• Batch-wide channel selection. Per-token selection stops helping as batches grow, because the GPU must load every channel any token in the batch selects. Selecting jointly across a batch, so the gain survives past batch 16, is explicitly left as future work.
• Closing the gap between saved reads and saved time. The MoE block moves only 18.3% of the full model's memory traffic yet runs 1.90× faster rather than roughly 5×, because fixed per-step costs such as kernel launches do not shrink. Reducing that overhead is open.
• Generalizing the partial-screen setting. The share of coordinates the partial screen reads (80) was chosen on OLMoE and checked only on Gemma 4, and the kernels were tuned for one budget on one GPU type.
• Understanding the collapse cliff. The authors hypothesize that once a sparse model starts repeating, it rarely stops, so accuracy falls sharply past a model-dependent error level — but they do not measure where that level is, which is what a deployer would need to choose a safe budget.

Target Audience

Systems and inference-optimization researchers working on MoE serving; practitioners deploying large MoE models in latency-sensitive, small-batch settings; and anyone currently using Prox, TEAL or WINA-style calibrated channel selection who wants to know why their effective read budget does not match their target and why long reasoning outputs sometimes loop. Readers interested in the broader calibration-versus-data-free debate in sparse inference will also find the controlled keep-rule comparison the most useful part.

Authors’ abstract

Mixture-of-Experts (MoE) models are bottlenecked by loading routed-expert weights from memory during interactive decoding at small batch sizes. Because most intermediate channels in a routed expert contribute little to any single token, channel selectors predict which channels matter and read only their weights, typically keeping channels above a threshold calibrated on offline data. We show that channel selection needs no calibration data. First, a calibrated threshold does not fix how much a selector reads, because the number of channels passing it varies across inputs; Prox with thresholds calibrated for an 18% read budget actually reads only 12.9% on average across our benchmarks. Second, top-k keeps each token's k highest-scoring channels and needs no calibration data. At the calibrated budget, a selector's top-k version matches or beats its threshold version on the same score in 40 of 43 comparisons at 13% and 18% budgets, including all 23 on generation tasks. The top-k versions also avoid the repetition-loop collapse that calibrated thresholds cause on Gemma 4. Calibrating a threshold thus provides neither of the two things selectors use it for: a known read budget and a consistent accuracy gain at that budget. We furthermore introduce a reader built entirely from the model's weights, which screens channels with a ternary sketch and recovers the skipped ones with a low-rank map. Across three post-trained MoE models, our reader retains 98.1% of full-model accuracy at 17.8% of routed-expert reads, with Triton kernels running MoE blocks 1.60 to 1.90x faster than optimized vLLM at batch one and up to 2.69x at larger batches on the same three models.