DOT-MoE: 面向MoE化的可微最优传输

TL;DR DOT-MoE 将密集 FFN 分解为 MoE 建模为可微最优传输问题，通过 Sinkhorn 迭代与 STE 端到端学习神经元分配与路由，在减少 50% 活跃参数的同时保留 90% 的密集模型性能。

问题

问题背景

大语言模型（LLM）通过参数扩展持续提升性能，但高推理成本成为部署瓶颈。Mixture of Experts（MoE） 架构通过稀疏激活解耦模型容量与计算量，在保持模型规模的同时降低推理 FLOPs，因此受到广泛关注。然而，从零训练 MoE 存在不稳定、超参数敏感、计算开销大等问题，促使研究者探索密集模型到 MoE 的转换（MoEfication），将预训练密集模型重构成稀疏混合专家结构。

现有方法局限

当前 MoEfication 方法主要依赖启发式神经元聚类或随机分割来划分前馈网络（FFN）层为多个专家。

• 启发式聚类（如基于权重相似度的 K-Means）仅利用参数结构信息，并未直接考虑层输出重建质量，导致转换后模型与原始输出分布偏差较大。
• 随机分割简单高效，但缺乏对专家功能分化的指导，容易造成专家冗余或负载失衡。
• 更根本的问题在于，这些方法将神经元分配与后续的 token 路由割裂处理：分配阶段只关注局部统计量，路由策略则在后训练中独立优化，无法保证端到端的最优匹配。这种两层分离设计使得最终性能难以逼近原始密集模型。

为什么这个问题难且重要

MoEfication 的核心难点在于：

1. 离散组合优化：神经元到专家的分配是一个需要满足容量约束的离散指派问题，直接求解计算复杂度极高，且不可微分，难以与下游任务联合优化。
2. 负载均衡与路由协同：推理时 token 路由必须与专家分配对齐，任何不均衡都会导致部分专家过载、部分闲置，加剧延迟或精度损失。 业界日益重视稀疏推理的经济性——无论是云服务按 token 计费，还是端侧部署对内存和时延的严格限制，低活跃参数比的高性能模型都有巨大需求。若能高效实现密集到稀疏的迁移，可复用庞大密集模型的预训练知识，大幅降低从零训练稀疏模型的成本。

行业类比：这类似于移动端模型压缩中，将大型图像分类网络转化为动态条件卷积，每个输入仅激活部分卷积核，只是此处面对的是语言模型的万亿级 token 流，对均衡性和重建精度的要求更高。

核心洞察

• 将稠密模型向MoE的转换形式化为**可微分最优传输（DOT）** 问题，通过Sinkhorn-Knopp迭代动态平衡专家容量，使神经元到专家的分配不再是启发式聚类或随机分割，而是全局代价最小化的高精度匹配。相比现有方法，DOT-MoE在分配阶段就显式建模了负载均衡约束，并借助可微分层实现端到端优化，有效避免了因静态聚类导致的局部次优解，从而在显著降低参数量的同时更完整地保留原始模型表现。
• 通过**Straight-Through Estimator (STE)** 将离散的神经元分配与连续的token路由联合学习，使整个转换过程可训练且可微分，打破了传统MoEfication中“先固定分配、后训练路由”的两阶段割裂。这允许分配策略直接感知下游任务信号，而非仅拟合权重相似度，因此输出的稀疏模型在推理时能以更少激活参数达到更优的性能保持，在工程上提供了一条更具普适性的稠密模型稀疏化路径。

方法

DOT-MoE 将稠密 FFN 层的专家化（MoEfication） 建模为一个端到端可微的流程，核心输入是预训练 Transformer 中 FFN 的权重矩阵，输出是满足容量约束的稀疏 MoE 模型。

问题形式化

给定一个稠密 FFN（通常包含两个线性层和一个激活函数），DOT-MoE 把第一个线性层的输出神经元（或中间特征）划分为 K 个专家，并同时学习一个路由网络，使得每个 token 只激活少量专家。该方法将这一分解定义为一个平衡最优传输问题：将每个神经元视为供应方，每个专家视为需求方，每个专家的容量（参数总量或神经元数）被严格约束为均等。通过最小化神经元特征与专家原型之间的传输代价，得到软分配矩阵。

可微最优传输与 Sinkhorn 迭代

为保持整个流程可微，DOT-MoE 采用 Sinkhorn-Knopp 算法 在 log 域迭代求解最优传输计划。该算法交替缩放行和列以满足边缘分布约束（即每个神经元都被分配、每个专家满载），从而输出一个双随机矩阵作为软分配。通过调整熵正则化系数，可控制分配的离散程度。为了在反向传播中获得硬分配，使用 Straight-Through Estimator (STE)：前向传播取 argmax 得到离散的神经元-专家映射，反向传播将梯度直通至软分配 logits，使得整个分配过程参与梯度更新。

联合学习分配与路由

与许多两阶段方法不同，DOT-MoE 在同一个训练循环中端到端学习神经元分配和 token 路由。路由网络以 token 表示为输入，输出选择 top-k 专家的概率；分配模块则通过上述可微最优传输确定每个专家包含哪些原稠密层的神经元。两者通过重构损失（如蒸馏损失或下一 token 预测损失）和负载均衡损失联合优化。训练初期采用一个 对齐阶段，仅微调路由和分配 logits，冻结专家权重，以稳定初始化。

跟同类方法的差异点

传统 Dense-to-MoE 方法依赖启发式聚类（如基于权重的 k-means）或随机分割，属于静态后处理，无法根据下游任务反馈优化分配；DOT-MoE 首次将神经元分组表述为可微运输问题，使得分配与路由能在同一训练目标下全局优化，从而显著提升性能保留率。

实验

实验在多个Transformer架构和常见基准上评估 DOT-MoE。

实验设计

将预训练稠密模型的 FFN 层 分解为 MoE 专家，采用 可微最优传输 分配神经元，并端到端联合学习 token-to-expert 路由。基线包括结构化剪枝、启发式聚类与随机分割。消融研究探讨了 专家粒度、训练稀疏度、初始化 等因素对重建质量与推理速度的影响，并在 32B 参数规模验证扩展性。

关键发现

• DOT-MoE 在减少 50% 激活参数 时，仍保留原始模型约 90% 的性能，大幅领先所有基线。
• 可微分配与端到端训练使 专家利用率 更均衡，专门化 更明显，克服了随机或启发式方法的负载不均问题。
• 适当增大专家粒度与训练稀疏度可进一步提升效率，且方法在长序列和大模型上表现稳健。

与基线对比

与静态聚类或随机分割相比，DOT-MoE 将分配形式化为 平衡传输问题，通过可微 Sinkhorn-Knopp 迭代 严格满足专家容量约束，避免局部最优。相较于结构化剪枝，MoE 化用更少的激活参数（50%）实现相近的模型质量，在 参数效率 与 推理速度 间取得更优平衡，为稠密模型稀疏化提供了更灵活、更高效的路线。

行业影响

DOT-MoE 将预训练密集模型转换为稀疏 MoE，在保持 90% 性能的同时减少 50% 活跃参数，为 LLM 推理降本增效提供了高性价比路径，尤其适合对延迟和成本敏感的实时服务场景。

落地场景

• 云端 LLM API 服务：将已有大模型（如 GPT 系列开源复现）转换为 MoE，降低每 token 推理成本，使 API 定价更具竞争力，同时维持服务质量。
• 边缘侧与私有化部署：在资源受限的端侧设备或企业私有环境中，压缩模型规模但保留核心能力，支撑智能客服、文档摘要等应用。
• 高吞吐批处理任务：如电商平台的商品描述生成、内容审核、代码补全等，利用 MoE 稀疏激活特性显著提升吞吐。

商业价值

• 推理成本直接降低：活跃参数减半意味着 GPU 内存与计算量下降，单卡可承载更高并发，或改用更经济的硬件，节省约 50% 的推理开支。
• 收入增长：更低的延迟与更高的吞吐可改善用户体验，提升付费 API 调用量；同时，对已有密集模型进行 MoEfication 无需从头预训练，节省数百万美元的训练成本与时间。
• 模型资产复用：企业可对内部定制模型进行转换，延长模型生命周期，避免频繁重新训练带来的碳足迹与运维负担。

与现有产品 / 工作流的集成

• 建模阶段：在标准 fine-tuning 之后增加一个 DOT-MoE 转换步骤（alignment phase），使用与原始训练相同的数据进行少量迭代，即可产出 MoE 版本。
• 推理部署：转换后的 MoE 模型可直接适配主流 MoE 推理引擎（如 vLLM、DeepSpeed-MoE、FasterTransformer），无需额外定制。
• MLOps 流程：作为模型优化阶段的可插拔组件，与模型量化、蒸馏等一同纳入 CI/CD 管线，自动生成稀疏化部署候选。

具体落地用例

1. 电商平台智能客服：某全球电商平台部署 LLM 处理海量用户咨询，高峰时推理压力大。使用 DOT-MoE 将客服模型转换为 MoE，吞吐提升 2×，延迟降低 40%，在促销季保障服务稳定性的同时，节省 45% 的 GPU 资源开销。
2. 金融研报摘要生成：一家资产管理公司需对大量研究报告进行自动摘要。原本依赖昂贵的高性能实例，DOT-MoE 转换后可将模型部署到中端 GPU 服务器，在几乎不损失摘要质量的前提下，推理成本下降 48%，满足合规要求下的私有化部署。

局限

• **训练开销与微调依赖**：DOT-MoE 在求解最优传输分配及端到端学习路由时引入额外计算量，尤其在大型模型上训练开销不容忽视；转换后的模型仍需下游任务微调才能恢复性能，无法直接零样本部署。
• **分配与路由的鸿沟**：论文附录 A 明确指出，即便通过最优传输获得最优的静态神经元分组，当路由策略与分配不是联合优化时仍会出现性能退化，这一 gap 在动态路由场景下可能进一步放大，影响实际推理精度。
• **实验覆盖与泛化性**：目前实验主要基于中等规模模型（最高 32B 参数）和常见基准，未验证在超大规模 MoE（如 100B+ 且专家数 > 64）或不同模态架构下的稳定性；另外对训练稀疏度的敏感性和长序列鲁棒性依赖特定温度退火策略，实际调参成本较高。