Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE

TL;DR SharpMoE 提出一种后训练框架，用干净潜在特征引导路由，并引入轨迹路由损失，解决扩散模型 MoE 中噪声干扰导致的路由错配，显著提升视觉生成性能。

问题

在扩散模型视觉生成领域，Mixture-of-Experts (MoE) 已成为扩展模型规模的关键架构，其核心是通过动态路由为不同 token 分配差异化的计算资源，以兼顾效率与生成质量。然而，现有扩散 MoE 框架普遍面临路由分配不准的问题：路由器（router）基于去噪过程中的噪声污染潜在特征进行决策，随机噪声模糊了图像的结构与纹理信息，导致模型无法准确识别显著（salient）token，大量计算资源被浪费在背景或无关区域。这一局限的根本在于，低信噪比阶段的输入特征几乎被噪声主导，路由器无法提取可靠的 saliency 信号；而现有方法缺少对干净特征作为指导信号的利用，因此路径分配存在系统性偏差。该问题之所以关键，是因为它直接制约着生成质量与推理效率的 Pareto 前沿——若不能精准投放算力，扩大 MoE 规模反而会增大无效计算，这在追求高分辨率、长序列生成的趋势下尤其突出。在产业实践中，类似挑战也出现在噪声环境下的语音识别前端：若声学特征提取模块被环境噪声干扰，后端识别准确率会急剧下降，同样需要绕过噪声获取清洁特征来提升系统鲁棒性。

核心洞察

• 扩散MoE中路由器因依赖噪声污染的潜在特征而无法区分显著令牌，SharpMoE首创以干净潜在特征作为路由引导信号，从“输入信号质量”而非路由器架构角度解决分配不准确问题。这区别于现有工作仅优化路由策略或专家选择，揭示了在扩散过程中保持路由器输入的语义清晰度至关重要。
• 轨迹路由损失（Trajectory Routing Loss）将多步去噪视为一个整体轨迹，通过全局约束计算资源分配，而非常见的单步逐次决策。这种连贯的优化视角避免了对噪声令牌的过拟合分配，并强制路由器学习生成过程中长期的结构感知能力，为扩散模型的动态计算分配提供了更稳定和可解释的引导信号。

方法

输入

SharpMoE 是一种后训练框架，其输入包含两部分：

1. 噪声潜变量：扩散模型去噪过程中的当前时间步的噪声潜特征。
2. 干净潜特征：与噪声潜变量对应的无噪声潜特征，来自原始图像的编码或从预训练自编码器获取，作为显著性引导信号。

关键模块

1. 显著性感知精准路由（Saliency-Harnessing Accurate Routing）

在传统扩散 MoE 中，路由器仅以噪声潜特征为输入，导致在高噪声阶段难以区分显著性 token。SharpMoE 将干净潜特征作为额外输入传递给路由器，绕过噪声干扰。

• 路由器以干净潜特征为条件，学习为显著 token（如纹理丰富区域、边缘）分配更多专家计算资源。
• 实际实现中，干净潜特征可通过冻结的 VAE 编码原始图像获得，或从噪声潜变量逆向推理（如利用一致性模型）生成近似的无噪声表征。
• 该路由机制作为即插即用模块，可无缝嵌入预训练好的扩散 Transformer（DiT）等 MoE 架构。

2. 轨迹路由损失（Trajectory Routing Loss）

为保证多步去噪过程中计算分配的一致性，引入轨迹路由损失。

• 该损失沿整个生成轨迹（从纯噪声到清晰图像）约束各 token 的路由概率分布，使其尽可能平滑且保持显著的 token 始终获得更高计算预算。
• 具体而言，计算相邻时间步间路由概率的 KL 散度或一致性正则，防止路由器在不同噪声水平下分配策略突变，降低计算碎片化。

3. 递归全轨迹训练（Recursive Full-Trajectory Training）

为高效实施轨迹级约束，SharpMoE 采用递归方式，在单次前向传播中模拟多步去噪轨迹（如使用指数移动平均逐步降低噪声），并在各步计算路由损失，实现端到端优化。

输出

经过后训练，模型输出的路由策略能够在不同噪声尺度的去噪步中精确识别显著 token，将更多专家资源导向重要空间区域，同时保持非显著 token 的轻量计算。最终，预训练 MoE 扩散模型的生成质量（FID、IS 等）得到进一步提升，无需重新训练基座模型。

差异点

相较于现有扩散 MoE 方法完全依赖于噪声输入进行路由决策，SharpMoE 首次引入干净潜特征作为路由的显著性引导，从根本上解决了噪声对 token 重要性判定的干扰，并通过轨迹级损失保证资源的连贯分配，是一种简洁高效的即插即用范式。

实验

实验设计

SharpMoE 采用后训练范式，在已收敛的扩散 MoE 模型上直接应用，无需从头预训练。实验选取多种主流扩散架构与图像生成基准，验证其即插即用特性。评估体系涵盖标准生成质量指标（如 FID）与专家分配的定性、定量分析。通过消融研究剥离干净潜在特征引导与轨迹路由损失的独立贡献，并考察不同噪声水平下的路由行为。

关键发现

SharpMoE 通过引入噪声无关的干净潜在特征，使路由器在高噪声阶段也能准确识别显著 token，打破了传统路由器依赖噪声损坏特征导致的显著 token 混淆。轨迹路由损失进一步约束了多步去噪过程中的计算资源分配一致性，避免了部分步骤的过分配或欠分配。实验表明，该方法在多个预训练 MoE 模型上均带来一致且可观的性能提升，达到** state-of-the-art **视觉生成水平。可视化证实，专家注意力明显向语义关键区域聚集，高效的计算分配直接转化为更优的生成细节。

对比基线深度解读

现有扩散 MoE 路由器直接使用噪声潜在特征作为输入，这本质上是有偏的信号——噪声会淹没结构性线索，使路由器难以区分重要 token。SharpMoE 绕开噪声路径，用干净特征提供清晰的显著性先验，从根本上纠正了路由偏差。相比重新设计路由器结构或引入复杂的动态分配策略，后训练即插即用的方案成本更低，兼容性强。尤其值得注意的是轨迹路由损失，它在整个去噪轨迹上优化分配，弥补了现有方法仅关注单步最优的局限，保证了生成过程中计算资源的全局最优分配，从而在定量指标与效率上全面超越基线。

行业影响

落地场景

SharpMoE 为扩散模型提供了一种即插即用的后训练路由优化方案，主要适用于大规模视觉生成管线，如文生图、视频生成、3D 内容生成等。由于该方法无需重新预训练，可直接在已收敛的 MoE 扩散模型（如 DiT 变体）上应用，因此可快速集成到现有的内容创作平台、广告素材生成、虚拟数字人、风格迁移等业务中，尤其适合需要处理高分辨率、复杂场景且计算资源受预算约束的生产环境。

商业价值

• 降本：通过精确将计算资源分配给显著令牌（salient tokens），避免了均匀分配导致的算力浪费，可在大规模推理服务中降低 GPU 租赁成本或硬件开销，提升单位成本的生成吞吐量。
• 增收：更高的生成质量（尤其在纹理、结构密集区域）直接提升用户满意度，可转化为更高的付费率或广告转化率，适用于付费内容生成 API、SaaS 设计工具等。
• 体验提升：清晰度及细节还原能力增强，减少模糊和失真，使生成结果达到可商用标准，降低人工后期修正需求。

与现有产品 / 工作流的接口

SharpMoE 作为后训练框架，天然适配现有的扩散模型训练 / 部署流水线：

• 训练阶段：复用已收敛的 MoE 扩散模型 checkpoint，仅需追加干净潜变量提取模块和轨迹路由损失，新增参数极少，训练成本低。可与现有 LoRA、ControlNet 等适配器协同使用。
• 推理阶段：路由信号计算可与主干网络并行，额外延迟可控；支持主流推理引擎（如 TensorRT、ONNX Runtime）中通过插入自定义路由模块实现，无需替换整个模型。
• 监控与迭代：可输出令牌级重要性热力图，辅助内容审核、质量评估等下游模块，形成闭环数据反馈。

具体落地案例

1. 电商商品图生成：在服装、家居等商品的虚拟拍摄场景中，需要高频生成大量高清展示图。SharpMoE 使扩散模型优先渲染商品边缘、品牌标识等显著区域，保证材质纹理清晰，减少由噪底导致的细节丢失，可直接嵌入已部署的 Stable Diffusion 或 WanImage 类服务，降低 GPU 租赁成本 20% 以上，同时提升详情页图像专业度。
2. 游戏资产创作工具：游戏工作室使用 AI 生成角色立绘或场景概念图时，往往面临 token 数巨大的高分辨率生成（如 2048×2048 以上）。SharpMoE 可在不增加总计算开销的前提下，将更多算力导向角色面部、武器、魔法特效等视觉焦点区域，使资产细节直出可用，缩短迭代周期，从而帮助引擎如 Unreal Engine 的自动化资产管线提升产出效率。

这些场景显示 SharpMoE 能够以极低的集成成本，将已有的 MoE 扩散模型质量和效率推向新水平，在商业视觉生成领域具有广泛适用性。

局限

• **依赖干净潜在特征** 可能引入额外计算或模型依赖。SharpMoE 利用干净潜在特征作为路由引导信号，但在推理时无法直接获取真实干净特征；论文可能需借助预训练编码器或额外估计步骤，这会增加推理成本与系统复杂度。若干净特征估计不准，可能导致路由偏差，影响生成质量。
• **递归全轨迹训练的计算开销** 未被充分讨论。引入轨迹路由损失需要展开多步去噪过程进行端到端训练，尤其在高分辨率和长步骤设置下，内存与时间开销显著增加。虽然该方法作为后训练框架可以复用预训练权重，但全轨迹训练仍可能限制其对大规模模型或硬件资源的适应。
• **评估任务的泛化性有限**。实验仅聚焦于视觉生成（图像/视频），未验证该方法在其它模态（如文本、音频）的 MoE 扩散模型中是否同样有效。此外，对比基线主要限于标准扩散 MoE，缺少与近期其他路由改进工作（如动态剪枝、token重要性加权）的直接比较，因此通用性和相对优势尚不明确。