← 返回 2026-06-30

聚焦关键：基于显著性的扩散MoE精准路由 Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE

Haoyou Deng, Keyu Yan, Chaojie Mao, Xiang Wang, Yu Liu, Changxin Gao, Nong Sang 📅 2026-06-25 👍 6 2026-07-13 08:37
扩散模型 混合专家架构 计算资源分配
arXiv HuggingFace PDF

利用清洁潜特征指导路由，解决扩散MoE中的资源分配问题

前置知识

Mixture-of-Experts (MoE)

MoE是一种稀疏激活架构，通过路由器动态选择专家网络处理不同输入。每个token只被路由到top-k个专家，专家输出加权求和得到最终结果。这使得模型可以在不显著增加计算开销的情况下大幅扩展容量。在视觉生成中，MoE可以根据不同区域的重要性分配计算资源。

本文的核心创新就是对MoE的路由机制进行改进，理解MoE的基本原理和工作流程是掌握本文技术路线的基础。

扩散模型去噪过程

扩散模型通过逐步添加噪声将数据转换为纯高斯噪声，然后训练神经网络学习反向去噪过程。在Rectified Flow框架下，模型直接回归速度场$v_\theta(x_t, t)$，目标为$v = x_1 - x_0$。生成时通过欧拉积分逐步还原干净图像。关键特征是中间步骤的潜特征$x_t$包含不同强度的噪声。

本文要解决的核心问题就是扩散过程中噪声对路由决策的干扰，理解去噪过程的时间步特性和噪声特性对论文动机和方法都很重要。

显著性(Saliency)

在图像生成中，显著性指图像中包含关键信息的区域，如主要物体结构、复杂纹理、锐利边缘等。这些区域对生成质量影响更大，需要更多计算资源。本文使用拉普拉斯算子估计显著性，通过计算干净图像的二阶导数识别高频成分。

显著性感知是本文的核心思想，理解如何量化和利用显著性对于掌握路由改进机制至关重要。

研究动机

现有扩散MoE框架存在严重的路由分配问题：路由器无法准确向显著token分配更多计算资源。论文通过拉普拉斯算子提取每个token的纹理信息作为显著性表示，分析了不同显著性级别token的专家分配分布。结果如图1(a)所示，尽管现有方法如DiffMoE声称实现显著性感知分配，但实际路由结果在很大程度上对显著性不敏感，不同显著性级别的token之间专家分配几乎没有变化。这个问题可以归因于噪声路由——路由器在整个多步去噪过程中一直基于噪声污染的潜特征$x_t$进行条件化。这种普遍存在的噪声，特别是在早期高噪声时间步，掩盖了关键的结构和纹理细节，最终损害了路由器有效区分显著区域的能力，导致计算资源分配不准确。

本文的目标是本文提出SharpMoE，一个后训练框架，整合了显著性感知的精确路由机制，以促进扩散MoE的清洁路由。SharpMoE的核心思想是利用来自前一个去噪时间步的清洁潜特征$\hat{x}_0^{t_{k-1}}$作为路由器的输入，为当前时间步$t_k$提供稳定的无噪声指导信号。这种设计带来两个优势：首先是显著性感知，预测的清洁特征明确捕捉和突出图像的显著区域，即使在早期高噪声阶段也能提供稳健且定义清晰的结构指导；其次是时间稳定性，潜特征$\hat{x}_0^{t_{k-1}}$的无噪声特性确保了在整个去噪轨迹上为路由提供清洁而鲁棒的输入，有效减轻了早期方法中噪声污染输入导致的资源分配不准确问题。

与已有工作不同的是，本文的独特切入角度是重新审视扩散MoE中路由器的输入表示问题，而不是简单地改进路由架构或损失函数。现有方法如DiffMoE采用批级池化方法，EC-DiT采用专家选择策略，但都忽视了最根本的问题：在扩散过程中，用于路由决策的输入$x_t$本身就被噪声严重污染。本文第一个明确指出并系统分析这个噪声路由问题，提出一个简单但有效的解决方案——使用清洁潜特征替代噪声潜特征。这个视角转换非常关键，因为它揭示了为什么之前显著感知路由方法无法达到预期效果，也为后续改进指明了正确方向。

核心方法

SharpMoE的方法整体思路非常直观：既然噪声污染的潜特征无法准确表示图像显著性，那就用清洁的潜特征来指导路由。具体技术路线包括三个关键组件：首先，在时间步$t_k$时，不再依赖噪声潜特征$x_{t_k}$进行路由决策，而是提出显著性感知路由机制，整合来自前一个时间步$t_{k-1}$的清洁预测$\hat{x}_0^{t_{k-1}}$作为路由输入；其次，引入$\hat{x}_0^{t_{k-1}}$创造了时间步之间的递归依赖，为此开发递归全轨迹训练方案；最后，为了调节累积专家分配，引入轨迹路由损失，对齐所有时间步上的总计算资源与图像的显著性分布。

SharpMoE的核心创新点是显著性感知的精确路由机制，这与已有方法有本质区别。现有方法如DiffMoE和EC-DiT都直接基于当前噪声潜特征$x_{t_k}$进行路由决策，这在扩散生成过程中被剩余噪声污染。SharpMoE则使用清洁潜特征$\hat{x}_0^{t_{k-1}}$，这是模型在前一个时间步对干净图像的预测。关键洞察是：潜特征空间本身就被训练为编码视觉信息到语义密集的表示中，天生强调结构重要区域；而$\hat{x}_0^{t_{k-1}}$提供图像语义骨架的稳定无噪声估计，即使早期阶段局部纹理尚未出现。

方法步骤详情

SharpMoE采用双路由器机制：预训练路由器$R_{pre}$处理当前生成状态$x_{t_k}$，显著性感知路由器$R_{sal}$处理前一步的预测清洁token$\hat{x}_0^{t_{k-1}}$。最终路由得分通过融合两个路由器的输出得到：$S = R_{pre}(x_{t_k}) + R_{sal}(\hat{x}_0^{t_{k-1}})$。递归全轨迹训练模拟短程生成展开，采样$T$个连续时间步$\{t_k\}_{k=1}^T$，在每个时间步处理当前噪声潜特征和前一步的预测清洁潜特征来估计速度场$v_k$，然后推导后续状态：$x_{t_{k+1}} = x_{t_k} + (t_{k+1}-t_k) \cdot v_k$，$\hat{x}_0^{t_k} = x_{t_k} - t_k \cdot v_k$。轨迹路由损失通过最小化归一化分配和显著性的KL散度对齐累积专家分配与显著性分布。

技术新颖性

SharpMoE的技术新颖性体现在多个方面。首先是首次系统分析了扩散MoE中的噪声路由问题，揭示了为什么现有显著性感知方法失效的根本原因。其次是递归全轨迹训练方案，解决了标准单步去噪训练与使用前一步清洁预测之间的兼容性问题。第三是轨迹路由损失，利用全轨迹训练提供的全局视角调节整个去噪序列上的计算资源分配。最后是插拔式后训练设计，SharpMoE只需要少量训练步骤就能显著提升预训练模型的性能，这种实用性在扩散模型快速发展的背景下尤为重要。

Overview of SharpMoE architecture.
Fig. 2: Overview of SharpMoE architecture.

实验结果

实验结果表明SharpMoE在多个方面都取得了显著成果。在ImageNet 256×256上，经过500K预训练步骤后，SharpMoE在100K后训练步骤内一致优于所有竞争扩散MoE方法，包括TC-DiT、EC-DiT和DiffMoE，在所有评估指标、模型规模和CFG尺度上都表现更好。最强结果是在DiffMoE-L骨干上使用cfg=1.5时达到FID 3.10和IS 228.88。消融研究显示，整合显著性感知路由机制后，在cfg=1.5时FID从8.03降至6.95；进一步加入轨迹路由损失后，FID进一步降至6.66。对预训练阶段的适应性评估显示，无论从400K步还是700K步预训练检查点初始化，SharpMoE都能实现一致的改进。

Quantitative comparison on ImageNet (256×256).
Table 1: Quantitative comparison on ImageNet (256×256).
Ablation study of each component on ImageNet (256×256).
Table 2: Ablation study of each component on ImageNet (256×256).
Samples generated by SharpMoE after 100K post-training steps.
Fig. 3: Samples generated by SharpMoE after 100K post-training steps.
Ablation studies on our critical designs.
Fig. 4: Ablation studies on our critical designs.
Aggregated distribution of saliency levels and the number of assigned experts averaged across multiple generated images.
Fig. 5: Aggregated distribution of saliency levels and the number of assigned experts averaged across multiple generated images.
查看结构化数据
任务指标本文基线提升
ImageNet 256×256图像生成 FID (cfg=1.5) 3.10 DiffMoE-L: 3.86 19.7%提升
ImageNet 256×256图像生成 IS (cfg=1.5) 228.88 DiffMoE-L: 203.00 12.8%提升
ImageNet 256×256图像生成 FID (cfg=1.5, EC-DiT-L) 3.27 EC-DiT-L: 4.09 20.0%提升
ImageNet 256×256图像生成 FID (cfg=1.5, DiffMoE-B) 6.66 DiffMoE-B: 8.03 17.1%提升

局限与改进

论文存在一些局限性。首先是初始时间步的处理问题，在推理时最初始时间步$t=1$没有可用的前一步$\hat{x}_0$，作者使用噪声潜特征$x_1$作为代理，设置第一个时间步$t_1 = 0.999$作为近似。然而，这仍然是一个近似，可能在理论上限制早期路由的准确性。其次是计算开销，递归全轨迹训练需要在每个训练迭代中模拟多步生成展开，比标准单步训练计算成本更高。第三是对TC-DiT的适配问题，由于TC-DiT均匀分配计算资源给每个token，轨迹路由损失变得不适用，只能使用Flow Matching损失进行训练。最后是方法依赖预训练模型，SharpMoE的设计理念是增强现有预训练扩散MoE，这意味着其效果在一定程度上依赖于预训练模型的质量和收敛状态。

独立分析的弱点

从独立分析的角度看，SharpMoE存在几个可以改进的弱点。首先是显著性估计方法的简单性，论文使用拉普拉斯算子估计显著性，虽然能有效隔离高结构密度区域，但这种方法主要基于纹理复杂度，可能无法完全捕捉人类感知的显著性，特别是对于语义重要性而非纹理复杂性的区域。未来可以考虑引入更先进的显著性估计方法。其次是双路由器融合机制的简单性，当前方案只是将两个路由器的输出相加，没有学习更复杂的融合策略。可以考虑引入可学习的融合权重或注意力机制，让模型能够根据去噪阶段动态调整两个路由器的贡献。第三是评估的局限性，论文主要在ImageNet上进行评估，但更复杂的真实世界图像生成任务可能需要不同的显著性定义和路由策略。

未来方向

基于论文成果和当前研究趋势，有几个有前景的未来研究方向。首先是扩展到更多模态，虽然本文专注于图像生成，但显著性感知路由的思想可能适用于视频生成、3D生成等其他模态。视频生成中的时间显著性和3D生成中的几何显著性都值得探索。其次是与其他先进架构的结合，如将SharpMoE的思想扩展到流匹配、一致性模型等新兴扩散变体，或者结合稀疏注意力机制实现更全面的计算效率优化。第三是动态显著性感知，当前方法基于静态的清洁预测进行路由，但可以考虑在生成过程中动态调整显著性估计，或者引入强化学习框架让模型学习更智能的路由策略。最后是理论分析，虽然论文提供了实证验证，但进一步的理论分析将有助于理解为什么清洁潜特征能够更好地表示显著性，以及这种表示的泛化界限。

复现评估

论文的可复现性评估需要从多个角度考虑。首先，作者提到架构细节在补充材料的A节中提供，但论文没有明确提到代码是否开源，这对完整复现是一个潜在障碍。其次，实验设置相对标准，使用ImageNet数据集和256×256分辨率，训练采用Rectified Flow范式，使用AdamW优化器和学习率$1 \times 10^{-4}$，批量大小256。对于SharpMoE，使用全轨迹训练方案，T=10采样步骤，使用总损失$L_{total}$优化所有网络参数。第三，计算资源方面，实验涉及多个模型规模，其中最大的DiffMoE-L有约20亿总参数，这需要相当大的GPU资源才能复现。然而，SharpMoE作为后训练方法，只需100K训练步骤，这比完整的500K预训练要经济得多。最后，评估指标使用FID和IS，但论文没有提供统计显著性测试或多次运行的方差信息，这使得性能提升的统计稳健性无法完全评估。

论文图表

(Left) Visualization of generated samples and per-channel router inputs. (Right) Distribution of saliency level and the number of assigned experts.
Fig. 1: (Left) Visualization of generated samples and per-channel router inputs. (Right) Distribution of saliency level and the number of assigned experts.

图1展示了两个关键对比：左侧是生成样本和每通道路由器输入的可视化，右侧是显著性级别和分配专家数量的分布。子图(a)显示现有方法(如DiffMoE)由于使用噪声潜特征进行路由，难以区分显著token，导致资源分配失败，呈现显著性不敏感的特征。子图(b)显示SharpMoE使用清洁潜特征进行路由，获得更好的显著性感知，实现与token显著性高度相关的计算分配，形成了明确的单调正相关关系。

这张图对理解论文至关重要，因为它直接可视化了论文要解决的核心问题——噪声路由导致的显著性不敏感分配，以及SharpMoE的解决方案——清洁路由带来的显著性相关分配。通过对比左右两侧的专家分配分布，读者可以直观理解论文的动机和方法的直观效果。