MoRE：通过硬件感知低秩路由扩展混合专家

专家混合 (MoE) 层是前沿语言模型的核心，最近的架构推动了更多、更小的专家的发展。在这种情况下，标准线性路由器成为瓶颈：在具有 $M$ 专家和隐藏维度 $h$ 的情况下，一旦 $M$ 很大，其每个词元成本 $θ(Mh)$ 将在 MoE 层中占据主导地位。我们引入了 MoRE（Rank-reduced-routed Experts 的混合），它将路由器权重矩阵进行秩为r的分解，并将路由成本降低到 $O((h + M)r)$。我们证明，当活跃专家的数量固定时，$M$ 中的对数排名足以满足路由表达能力，并且对于精度因子来说是必要的。我们还证明对数排名可以保持高斯记忆模型中的负载平衡，并且对合成电话簿任务的训练表明低秩不会损害记忆。在匹配的有效 FLOP 中，因式分解允许多出 $θ(h/r)$ 的专家。为了节省挂钟时间，我们在推理时设计了一个融合的 Triton 内核，避免了 HBM 上昂贵的内存操作。根据经验，MoRE 在预训练后提高了电话簿任务的记忆和知识密集型问答基准的性能，同时匹配推理能力。代码可在 https://github.com/Matheart/MoRE_code. 获取