适用于细粒度专家混合的自适应倒排索引路由

专家混合 (MoE) 模型通过仅激活每个词元的一部分专家来实现可扩展的 Transformer 架构。最近的证据表明，随着专家数量的增加，即许多小专家而不是少数大专家，性能会得到提高。然而，这种机制大大增加了路由成本，这可能会主导计算。我们引入了 MoE 自适应倒排索引路由 (AIR-MoE)，这是一种基于矢量量化 (VQ) 的倒排索引启发路由架构。在第一阶段，AIR-MoE 通过将词元分配给 VQ 码字来执行粗略筛选，以构建候选专家集。在第二阶段，精细评分计算仅限于此候选列表的精确路由分数。这个两阶段过程近似于真正的 top-k 路由，同时避免了完整的专家评分，并且与之前的工作相比，对专家参数没有施加结构限制。 AIR-MoE 可直接替代标准路由器，无需修改模型架构或损失函数。我们进一步提供了 AIR-MoE 实现的大规模召回的下限，从而深入了解其内部运作。根据经验，我们证明，与精细 MoE 设置中的现有路由方法相比，AIR-MoE 实现了更高的性能。