AIR-MoE 用向量量化两阶段路由，在百万级细粒度专家下改善困惑度-FLOPs 权衡
核心概要
该工作提出 AIR-MoE，一种基于向量量化的倒排索引式两阶段 MoE 路由架构：第一阶段用 VQ 码字对 token 做粗筛选构建候选专家集，第二阶段仅在候选集内计算精确路由分数，从而在避免全专家打分的同时近似真实 top-K 路由，且不对专家参数施加结构约束；作者还给出 AIR-MoE 质量召回的下界，并在多达百万以上微小专家的场景中相对现有 MoE 路由器改善困惑度-FLOPs 权衡、相对最佳基线取得一致的困惑度提升。
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.
arXiv深度剖析
提出 AIR-MoE 两阶段路由：先用向量量化把 token 分配到 VQ 码字以构建候选专家集，再在候选集内做精确打分。 相对直接对所有专家打分的标准路由，该流程以粗筛选加细打分近似真实 top-K 路由，从而避免全专家打分带来的路由开销。 摘要给出架构描述与两阶段流程说明，并称其近似真实 top-K 路由；具体实现细节与实验设置未在摘要中展开。
AIR-MoE 可作为标准路由器的即插即用替代，无需修改模型架构或损失函数，且不对专家参数施加结构约束。 与对专家参数施加结构约束的先前工作不同，该方法在保持架构与损失不变的前提下替换路由器。 摘要明确陈述“drop-in replacement”“no modifications to the model architecture or loss function”以及“no structural constraints on expert parameters”。
给出 AIR-MoE 所达到的质量召回（mass recall）下界，用以理解其内部工作机制。 在经验结果之外补充了关于路由召回性质的理论刻画。 摘要称“provide a lower bound on the mass recall achieved by AIR-MoE that yields insights into the inner workings”，未给出下界的具体形式或数值。
在多达百万以上微小专家的细粒度场景中，AIR-MoE 改善困惑度-FLOPs 权衡，并相对最佳现有基线取得一致的困惑度提升。 针对细粒度专家导致路由成本上升这一瓶颈，在极端专家数量规模上给出经验改进。 摘要报告“up to more than a million tiny experts”的规模与“consistent perplexity improvements over the best existing baseline”，未给出具体困惑度数值、FLOPs 数值或数据集名称。
启示与展望
该工作面向稀疏 MoE 的路由环节，适用于专家数量极大（摘要所述“up to more than a million tiny experts”）的细粒度场景，并以标准路由器即插即用替代的形式使用，不改变模型架构与损失函数。其价值在于让细粒度专家带来的性能收益不被路由开销抵消，因此对训练与推理中关注计算与内存预算的 MoE 实践者最为相关。摘要所述的质量召回下界为理解该方法为何有效提供分析视角。
摘要未给出困惑度与 FLOPs 的具体数值、对比基线名称、数据集与模型规模，也未说明质量召回下界的推导条件与紧致程度；两阶段路由中 VQ 码字数量、候选集大小等超参数如何影响召回与最终性能，以及“consistent perplexity improvements”的幅度与统计稳定性，均需在正文中确认。此外，摘要未说明该方法在专家数量较少或非细粒度场景下是否同样适用。