OmniMoE：通过编排原子级专家实现大规模高效混合专家模型 OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

OmniMoE的方法分为三个紧密耦合的步骤：第一步是原子专家定义与动态专家组装（DEA）。原子专家定义为最小可路由计算单元，由输入向量 $w_{in}^i \in \mathbb{R}^d$ 和输出向量 $w_{out}^i \in \mathbb{R}^d$ 参数化，计算为 $E_i(x) = \sigma(x {w_{in}^i}^\top) w_{out}^i$（使用SWIGLU激活）。所有N个原子专家的参数存储在两个全局参数矩阵 $W = [w_{in}^0, \ldots, w_{in}^{N-1}]^\top$ 和 $V = [w_{out}^0, \ldots, w_{out}^{N-1}]^\top$ 中。DEA的检索步骤通过索引 $I_x$ 从全局矩阵中获取token局部参数块 $w_x = W[I_x] \in \mathbb{R}^{K \times d}$，组装步骤计算 $y = (g_x \odot \sigma(x w_x^\top)) v_x + \text{MLP}(x)$。第二步是笛卡尔积路由器设计。将专家id $n$ 视为二维坐标 $(i,j)$ 在 $N_r \times N_c$ 网格上，用两个投影矩阵 $W_r, W_c$ 计算行和列的logits：$s_r = x W_r, s_c = x W_c$，然后通过LogSoftmax得到对数概率 $p_r, p_c$，隐式定义评分矩阵 $S_{ij} = p_r[i] + p_c[j]$。通过并行Top-K选择（将隐式网格分块分配给GPU线程块）提取全局Top-K专家索引。第三步是专家中心调度。首先收集路由计算任务（M = L × K个任务），然后按专家ID排序并分组为大小B的连续组，通过层次排序（主键为组ID，次键为token ID）改善局部性，最后对每个活跃组执行Grouped GEMM：$O_q = (G_q \odot \sigma(X_q W_q^\top)) V_q$。

OmniMoE的技术新颖性体现在四个维度。首先，原子专家的定义是全新的——将专家粒度从FFN级（数百到数千参数）推至向量级（2d参数），这不仅是粒度的细化，更是计算范式的转变：从静态的全连接计算变为动态的参数组装。这种设计使得每个激活参数都为特定token计算，实现了极端的参数效率。其次，笛卡尔积路由器的隐式评分机制是创新的核心。现有方法（包括PKM）都需要显式计算或存储N维评分向量，而OmniMoE通过因子化假设 $p(i,j|x) \approx p_r(i|x) \cdot p_c(j|x)$，在对数空间中将联合分布转化为可加形式 $S_{ij} = p_r[i] + p_c[j]$，无需显式构造N维评分矩阵即可进行Top-K选择。第三，专家中心调度反转了MoE执行的基本范式。传统MoE（无论粗粒度还是细粒度）都采用token中心的执行方式——每个token独立获取其选择的专家参数，导致分散的内存访问。OmniMoE将执行顺序反转为专家中心——先按专家分组，再在组内按token排序，将分散的随机访问转化为连续的流式读取。最后，共享密集MLP与路由分支的异构设计也是新颖的——不同于DeepSeekMoE等将共享专家作为独立的FFN实例，OmniMoE将共享MLP作为通用语义处理的稳定骨干，让路由分支专注于长尾知识检索，形成互补的分工。

OmniMoE存在几个值得关注的弱点。首先，原子专家的表达能力受限——每个原子专家仅由一对向量参数化，计算为简单的内积加激活，缺乏深度非线性变换能力。虽然DEA通过组合多个原子专家缓解了这一问题，但对于需要复杂推理的任务，单层的原子专家组合可能不如深层FFN。改进方向可以探索多层原子专家组合或引入轻量级的跨原子专家交互机制。其次，笛卡尔积路由器的因子化假设限制了路由灵活性——独立的行和列投影无法捕捉行列之间的交互，可能导致某些专家组合被系统性忽略。可以考虑引入低秩交互项或注意力机制来增强行列间的依赖建模。第三，专家中心调度在动态批处理场景下面临挑战——当不同请求的活跃专家集合差异很大时，分组效率可能下降，导致GPU利用率波动。改进方向包括动态调整分组大小或引入工作窃取机制。最后，当前设计缺乏显式的负载均衡机制——虽然消融实验显示专家使用率为100%，但这可能受益于特定的训练设置，在更极端的专家规模下可能出现负载不均衡。

基于OmniMoE的成果，可以延伸出多个有前景的研究方向。作者提出的未来工作包括：将OmniMoE扩展到更大规模的模型（数十亿到数千亿参数）并验证分布式训练场景下的扩展性，以及在指令微调和RLHF后评估下游任务性能。此外，可以探索将原子专家设计应用于其他模态（如视觉、多模态）的MoE架构，验证其通用性。另一个有趣的方向是动态专家组装的可解释性——通过分析哪些原子专家被频繁共激活，可以理解模型内部的知识组织方式。还可以探索自适应粒度的MoE，即根据输入的复杂度动态调整激活的专家粒度（简单token激活少量粗粒度专家，复杂token激活更多细粒度专家）。最后，将OmniMoE的笛卡尔积路由思想应用于其他需要大规模索引的场景（如推荐系统中的候选生成、检索增强生成中的知识检索）也值得探索。

OmniMoE在复现性方面表现良好。作者已将代码开源在 https://github.com/flash-algo/omni-moe，这为社区复现和扩展提供了便利。训练数据使用公开的SmolLMCorpus（40B tokens），涵盖Web、教科书、代码和数学领域，使用NeoX tokenizer（128,256词表）。训练超参数遵循已发表的缩放律和Chinchilla计算最优性协议，使用AdamW优化器和WSD学习率调度器。所有实验在NVIDIA PyTorch容器中运行，使用Hugging Face Transformers框架。推理/评估在单节点8× NVIDIA A100 GPU上进行。然而，复现的主要挑战在于算力需求——6.4B-A1.7B模型的预训练需要大量GPU小时，且论文未提供精确的训练时间和成本估计。此外，专家中心调度的实现涉及复杂的GPU内核优化，可能需要特定的硬件支持（如Tensor Core）和CUDA编程经验。对于资源有限的研究者，可以从小规模实验（如280M-A80M配置）开始验证核心思想。