AI Research 技术调研

Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning

23分钟阅读

一句话定位

MoDE（Mixture-of-Denoising Experts）把 Mixture-of-Experts 引入 Diffusion Policy 的去噪 Transformer，但路由只看当前噪声等级 σ_t、不看 token 内容——这让每个噪声等级对应的专家选择在推理前就能确定性地预计算，把选中的专家 MLP 提前融合成一个复合 MLP 并行跑，从而在扩大模型容量（460M/685M 总参数）的同时把推理 FLOPs 砍掉九成、比稠密 Transformer 基线少 40% 激活参数；在 CALVIN 与 LIBERO 共 134 个任务上刷新 SOTA（CALVIN ABC→D 预训练后 4.01、LIBERO-90 预训练后 0.95），平均比第二名基线高 57.5%（ICLR 2025）。

背景与定位

扩散策略（Diffusion Policy）系谱里，Diffusion Policy（Chi et al. 2023）证明条件 DDPM 能建模模仿学习中的多模态、非连续动作分布，Octo、BESO（同一作者组前作）、3D-Diffusion-Actor 等把骨干换成 Transformer 以获得更好的扩展性。但这条路线的代价是计算量随模型变大而线性增长：标准架构动辄数亿参数，且每次决策要跑数十步去噪，图像/文本编码器进一步推高算力需求，限制了在算力受限的移动机器人等场景的实时部署。

MoE 是另一条已在语言模型里验证过的”扩容不扩算力”路线——Switch Transformers（Switch Transformers，Fedus et al. 2022）把 MoE 引入 Transformer 的 FFN 层，论文原文点名 Mixtral（Jiang et al. 2024）与 GLaM（Du et al. 2022）延续这一范式；标准做法是路由器依据 token 内容 动态选专家（content-based routing），需要运行时逐 token 派发，并伴随专家坍塌（expert collapse）、路由器坍塌等已知优化难题。MoDE 的路由器初始化方法本身则借鉴自 OLMoE（Muennighoff et al. 2024，见”模型架构”节）。机器人领域此前唯一尝试 MoE 的扩散策略是并行工作 Sparse-DP（Wang et al. 2024），但它用 MoE 学习面向不同任务的专家（task-specialist experts），而非针对去噪过程本身做效率优化。

MoDE 的立意来自一个不同的观察：扩散去噪本身是多任务的（Hang et al. 2024 指出去噪不同阶段之间迁移很少），此前已有工作（Denoising Task Routing、Switch Diffusion Transformer 等,均来自图像扩散领域）按去噪阶段切换参数或调制表征。MoDE 把这一思路系统化为噪声条件路由：路由器只依据 φ(σ_t)（噪声等级的编码），不依赖 token 内容。这个设计选择本身带来两个收益：(1) 同一噪声等级下所有 token 的专家选择路径是确定性的，可以在推理前一次性预计算并融合专家 MLP，省去逐 token 动态派发的开销；(2) 消融显示这种”只看噪声、不看 token”的路由在性能上与传统 token-content 路由相当甚至略优（5 个基准平均 0.851 vs 0.845）。论文的落点是”扩散策略如何在扩容的同时保持工程可部署性”，为后续讨论扩散策略计算效率的工作提供了一个可复现的强基线。作者来自 KIT 直觉机器人实验室（Moritz Reuss、Rudolf Lioutikov）与 MIT CSAIL（Jyothish Pari、Pulkit Agrawal），同实验室后续工作 MDT（Multimodal Diffusion Transformer）在项目页被列为关联项目。

模型架构

总体形式：策略 π_θ(ā|s̄,g) 以状态历史 s̄（长度 h=1）、语言目标 g 为条件，用连续时间扩散模型（EDM 框架，Karras et al. 2022）对未来动作序列 ā（chunk 长度 j=10）做分数匹配去噪，推理用 DDIM 求解器、10 步去噪（主实验默认；HF 发布的预训练权重卡标注”Sampling Steps: 5（性能最优配置）“，与论文正文的 10 步略有出入，两个数字均来自各自一手来源，未强行统一）。

Backbone：L 层 Transformer block，每层由一个噪声条件自注意力（noise-conditioned self-attention）层和一个MoE 层组成： f^i(X, φ(σ_t)) = MoE(SA(X̂) + X, φ(σ_t)) + X，其中 X̂ = φ(σ_t) + X（把噪声编码token加到所有输入token上再做自注意力，不引入额外参数、也不改变网络结构，是论文消融证明有效的一个”免费”设计）。

MoE 层与噪声条件路由（核心创新）：给定 N 个专家 E_i，MoE(X, φ(σ_t)) = Σ R(φ(σ_t)) · E_i(X)，路由函数 R(φ(σ_t)) = topk(softmax(φ(σ_t))·W_R), k) —— 路由只依赖噪声等级编码 φ(σ_t)，与 token 内容 X 无关（这是与传统 LLM MoE 路由最本质的区别）。topk 用不放回多项式采样实现，为保持采样过程的梯度流，用路由概率缩放专家输出并重新归一化选中概率。路由器按 OLMoE（Muennighoff et al. 2024）的方法用截断正态分布（std=0.02）初始化。为防专家坍塌加入负载均衡损失（Fedus et al. 2022 形式），权重 γ=0.01（默认，见消融）。默认 topk=2、N=4 专家（这是 8 层 / 4 专家 / 1024 维配置下的默认组合）。每个专家是一个 Swish-GLU 激活的简单 MLP（项目页描述）。

Router/Expert Caching（3.3.1 节，效率的关键机制）：因为路由只看噪声等级，同一噪声等级下的路由路径是确定性且可提前计算的——推理前枚举每个噪声等级对应的选中专家，把这些专家的 MLP 权重融合成一个复合 MLP,原本需要逐专家循环的计算变成并行单次前向，省去动态派发专家的运行时开销和内存访问开销。论文报告该缓存策略比标准 MoE rollout 减少超过 80% 的 FLOPs 开销，推理提速 2 倍。

条件编码：语言用冻结 CLIP 语言编码器（ViT-B/32）生成目标向量；图像用 FiLM 条件 ResNet（CALVIN ABCD / LIBERO-10 / LIBERO-90 用 ResNet-18，CALVIN ABC 与预训练大模型变体用 ResNet-50，见 Table 3）。输入 token 共 14 个（1 噪声 + 1 目标 + 2 图像 + 10 个带噪动作 token，按论文 A.4 节口径）。

关键配置数字（Table 3，附录超参表）：

  • 主实验（CALVIN ABCD/ABC、LIBERO-10/90）：8 层 Transformer、4 个专家、8 个注意力头、嵌入维度 1024、动作 chunk 10、历史长度 1、attention dropout 0.3、residual/MLP dropout 0.1；参数量 460M（该配置下的模型规模）。
  • 预训练大变体（Pret-MoDE）：12 层、4 个专家、嵌入维度同为 1024；参数量 685M（Table 3 口径）。
  • 效率对比表（Table 8，A.4 节）单独给出的口径为：MoDE(ours) 激活参数 436M / 总参数 740M——与 Table 3 的 685M 存在不完全一致（两表分属不同附录小节，本页按各自表格如实呈现，未强行折算统一）。
  • σ_max=80、σ_min=0.001，另有超参表列出的 σ_t=0.5（Table 3 原始记法，正文未展开解释具体含义）、时间步指数（exponential）调度、DDIM 采样、EMA 开启。

数据

MoDE 本体是策略架构论文，各基准均复用公开示范数据集，不采集新数据；另有一次面向泛化能力的多机器人数据预训练：

  • CALVIN ABCD→D / ABC→D：4 个环境（A/B/C/D，桌面纹理与物体位置不同）的人类遥操作 play 数据，22,966 条交互序列、每条 64 时间步、34 个多样任务，标准协议下测试 1000 条指令链（每链 5 个连续子任务）。ABC→D 设定只用 A/B/C 训练、在未见环境 D 上做零样本评测。
  • LIBERO-10（LONG，长时程 10 任务）/ LIBERO-90（90 个短时程任务）：各任务 50 条示范，Franka Emika Panda 末端控制器,静态相机+腕部相机双路图像输入。
  • SIMPLER（real2sim 泛化评测）：评测集本身来自 BridgeV2 与 Google Fractal（RT-1 训练数据）的 real2sim 复刻,MoDE 不用它训练,只用于零样本对比 Octo（800K 条轨迹训练）与 OpenVLA（1M 条轨迹训练）。
  • 预训练混合数据（OXE 子集,Table 4）：196K 条轨迹,来自 6 个数据源,加权采样比例——BC-Z 0.258768、BRIDGE 0.188043、DOBB-E 0.245176、Google Fractal 0.162878、CMU Play-Fusion 0.101486、LIBERO-10 0.043649（合计 1.0）。覆盖 Google 机器人、Franka Panda、Hello-Stretch 等多种机器人平台。用 400K 样本的大 shuffle buffer 保证混合均匀,每个数据集单独做归一化以应对不同机器人平台的尺度差异。
  • 训练/微调不引入额外的动作标注或过滤流程——直接使用各基准/数据集的原始动作标签（CALVIN 用 Delta 末端执行器空间+离散夹爪,LIBERO 用末端控制器）。

训练方法

目标函数：EDM 框架下的分数匹配（score matching）损失 L_SM = E[α(σ_t)·‖D_θ(ā+ε,s̄,g,σ_t) − ā‖²]，即让网络在给定噪声等级下直接回归干净动作序列。

训练/微调两阶段：

  • 主实验（不预训练）：CALVIN ABCD 训练 30k 步、CALVIN ABC 25k 步（Table 3 超参表与附录 A.2.2 实验细节一致）；LIBERO-10 按 Table 3 为 15k 步、附录 A.2.3 描述为”50 epochs”（两处未给出可直接互算的 batch/数据量对应关系）；LIBERO-90 附录 A.2.3 正文明确写”50k 步”，与 Table 3 超参表标注的 30k 步不一致（本页如实并列两处口径，未强行统一）；batch size 512；AdamW，betas=[0.9,0.95]，lr 1e-4，weight decay 0.05。
  • 预训练（Generalist Policy Pre-Training）：在上述 OXE 子集上跑 300k 步，batch size 1024，lr 1e-4，weight decay 0.1。
  • 微调（从预训练权重出发）：冻结每层预训练好的路由器、去掉负载均衡损失、只微调其余模型组件；LIBERO 微调 10k 步、CALVIN 微调 15k 步，batch size 64/GPU × 4 GPU。

路由与专家超参选择（消融确定的默认值，均见下文”评测”消融部分）：topk=2、γ_LB=0.01、4 个专家、噪声条件路由（而非 token-content 路由）、噪声输入 token + 噪声条件自注意力同时使用。

Infra（训练 / 推理工程）

  • 训练硬件：主实验（CALVIN ABC/ABCD/LIBERO-10/90）均在 4× NVIDIA A6000（40GB） 集群节点上训练；CALVIN ABC 训练耗时 6.5 小时，CALVIN ABCD 7.5 小时，LIBERO-10 2 小时，LIBERO-90 12 小时（均含训练结束后完整 rollout 评测耗时）。预训练用 6× NVIDIA A6000（40GB），300k 步耗时约 3 天。微调阶段用 4 GPU、batch size 64/GPU。State-based 附加实验（Block Push / Relay Kitchen）在本地 RTX 3070 单卡上训练，各约 2–3 小时。
  • 精度：论文未披露是否使用混合精度训练（未披露）。
  • 推理测量环境：单张 NVIDIA A6000（40GB），对每个方法测 100 次去掉离群值取平均。
  • 推理 FLOPs / 延迟（Table 8, A.4 节，预测单个动作的平均口径）：
    • MoDE（预训练大变体，12 层/4 专家/ResNet-50）：论文给出的口径——Transformer 单次前向 0.7 GFLOPS（有专家缓存）/ 5 GFLOPS（无缓存，论文表述为”缓存降低超过 90% 计算成本”，与 0.7/5 两个数字字面换算出的约 86% 略有出入，原文如此，未代为改写）；ResNet-50 图像编码 8.27 GFLOPS；最终论文报告的单动作平均值为 1.53 GFLOPS、12.2 ms。
    • 对照基线：Diff-P-CNN 1.28 GFLOPS/11.7ms、Diff-P-T 2.16 GFLOPS/16.2ms、GR-1 27.5 GFLOPS/12.6ms、RoboFlamingo 690 GFLOPS/65ms、SuSIE 60 GFLOPS/199ms（细节见”评测”节表格）。MoDE（436M 激活/740M 总参数）与 GR-1（130M）几乎同速（12.2ms vs 12.6ms）——论文原文称 MoDE 参数量”6 倍以上”于 GR-1（未注明是激活还是总参数口径；740M/130M≈5.7×，436M/130M≈3.4×，均未精确等于 6×，原文表述如实转引），同时 FLOPs 只有 GR-1 的约 1/18（1.53 vs 27.5，简单除法）。
    • Router/Expert 缓存对 MoE 层本身的效果（3.3.1 节独立表述）：比标准 MoE rollout 减少 >80% FLOPs、推理提速 2 倍。
  • 稠密 Transformer 基线对比（Figure 5,相近参数量下）：batch size 1 时 MoDE 推理 12ms vs 稠密基线 15ms（快 20%）；batch size 512 时 MoDE 361 GFLOPS vs 稠密基线 5,772 GFLOPS（少 16 倍），推理 64ms vs 104ms（快约 40%）。
  • 未披露：边缘硬件（机载计算单元）部署数据；训练是否用分布式数据并行/张量并行等具体并行策略细节。

评测 benchmark

所有数字取自论文正文与附录表格（arXiv HTML 全文），基线取自各自论文报告的标准协议结果。

CALVIN ABCD→D（Table 1 上半，1000 条指令链，3 seed 均值,数字为完成 1/2/3/4/5 个连续子任务的成功率与平均完成长度 Avg.Len）：

方法激活参数(M)预训练12345Avg.Len
Diff-P-CNN321×86.3%72.7%60.1%51.2%41.7%3.16±0.06
Diff-P-T194×78.3%53.9%33.8%20.4%11.3%1.98±0.09
RoboFlamingo1000✓96.4%89.6%82.4%74.0%66.0%4.09±0.00
GR-1130✓94.9%89.6%84.4%78.9%73.1%4.21±0.00
MoDE277×96.6%90.6%86.6%80.9%75.5%4.30±0.02
MoDE436✓97.1%92.5%87.9%83.5%77.9%4.39±0.04

CALVIN ABC→D（Table 1 下半,零样本泛化：训练于 A/B/C,测试于未见环境 D）：

方法激活参数(M)预训练12345Avg.Len
Diff-P-CNN321×63.5%35.3%19.4%10.7%6.4%1.35±0.05
Diff-P-T194×62.2%30.9%13.2%5.0%1.6%1.13±0.02
RoboFlamingo1000✓82.4%61.9%46.6%33.1%23.5%2.47±0.00
SuSIE860+✓87.0%69.0%49.0%38.0%26.0%2.69±0.00
GR-1130✓85.4%71.2%59.6%49.7%40.1%3.06±0.00
MoDE307×91.5%79.2%67.3%55.8%45.3%3.39±0.03
MoDE436✓96.2%88.9%81.1%71.8%63.5%4.01±0.04（新 SOTA）

LIBERO-10（LONG）与 LIBERO-90（附录 A.2”平均性能提升”表)：

基准MoDEDiff-P-TDiff-P-CNN基线均值提升幅度
CALVIN ABC→D（归一化）0.6780.2260.2700.248+151.1%
CALVIN ABCD→D（归一化）0.8600.3960.6320.514+36.1%
LIBERO-900.9100.6900.7800.735+16.7%
LIBERO-100.9200.5100.7300.620+26.0%
4 基准平均提升（相对第二名）+57.5%

正文另指出：MoDE 是首个在 LIBERO-10 上突破 90% 成功率的策略；QueST（Mete et al. 2024,离散动作 VQ 表征）在 LIBERO-90 上是第二名，DP-CNN 在 LONG 设定上是第二名（论文只给出排名描述，具体数值来自 Figure 3(b) 柱状图，未在正文表格中单独列出，故此处不重复引用其精确读数）。预训练后的 MoDE 在两个 LIBERO 基准上均更高，摘要给出的头条数字为 LIBERO-90 = 0.95（预训练变体）。

SIMPLER real2sim 泛化评测（Table 6,对比 Octo-Base 与 OpenVLA-7.7B，MoDE 为预训练变体，零样本）：

指标OpenVLAOcto-BaseMoDE
平均成功率23.70%17.75%26.30%
平均排名1.952.11.65

Table 6 逐任务列出 20 行成绩（同名任务各出现两次，对应论文未在正文文字中展开说明的两套子评测口径），区间从 0%（stack green block on yellow block，三个模型全部失败）到单项最高 76.85%（drawer close 其中一个子评测口径下的 MoDe 分数）不等；完整逐任务数字见已归档的项目主页/arXiv 附录 Table 6。

计算效率对比（Table 8,A.4 节,预测单个动作的平均口径,详见 Infra 一节）：MoDE 436M 激活/740M 总参数,1.53 GFLOPS,12.2ms；对照 Diff-P-CNN 1.28GFLOPS/11.7ms、Diff-P-T 2.16GFLOPS/16.2ms、GR-1 27.5GFLOPS/12.6ms、RoboFlamingo 690GFLOPS/65ms、SuSIE 60GFLOPS/199ms。

消融研究（LIBERO-10,Table 2）：

  • 噪声注入：完整 MoDE（噪声 token + 噪声条件自注意力）0.92；去掉噪声 token 0.90；只用噪声 token（无噪声条件自注意力）0.85；FiLM 噪声条件（图像扩散常用做法）0.81——两种机制互补，FiLM 明显更差。
  • MoE 设计：topk=1（默认 topk=2）0.91；共享专家（shared expert，始终用同一专家）0.90；γ_LB=0.1 → 0.90；γ_LB=0.001 → 0.86（默认 γ_LB=0.01 对应 0.92）。
  • 嵌入维度：256维 0.86、512维 0.87、1024维（默认）0.92。

路由策略对比（Table 7,5 个基准：Block Push / Relay Kitchen / CALVIN ABC / CALVIN ABCD / LIBERO-10,3 seed 均值）：

路由方式Block PushRelay KitchenCAL ABCCAL ABCDL-105 基准平均
Dense Transformer（无 MoE）0.96±0.023.73±0.122.83±0.194.13±0.110.91±0.020.839±0.144
Token-content 路由0.97±0.013.85±0.032.67±0.044.29±0.080.90±0.010.845±0.161
噪声条件路由（MoDE 默认）0.97±0.013.79±0.042.79±0.164.30±0.020.92±0.020.851±0.151

噪声条件路由与 token-content 路由性能接近（0.851 vs 0.845），但前者可预计算专家路径、支持缓存加速，故为默认选择。

专家数量扩展（CALVIN ABCD/ABC，2/4/6/8 专家 + Dense-small/Dense-large 对照）：4 专家在两个设定上均为最优；超过 4 个专家性能反而下降（推测因表征学习被路由复杂度稀释）；Dense-small（参数量与 MoDE 相当但稠密）持续弱于 MoE 版本，凸显 MoE 架构的参数利用效率。

topk 消融（专家利用可视化，附录 A.6.1）：topk=1 在 CALVIN ABC 上平均完成长度从 topk=2 的 3.34 降到 2.72，且专家利用容易坍塌到 4 个专家中的 2 个。

负载均衡权重 γ_LB 扫描（附录 A.6.1，LIBERO-10 专家分布可视化）：γ_LB=0.1 时专家利用接近均匀但性能降到 0.90；γ_LB=0.01（默认）在保持专家一定专精化的同时取得最佳性能；γ_LB=0.001/0.0001 时性能进一步降到 0.86/0.83，且 0.0001 出现多层专家坍塌。

专家利用可视化的定性发现（4.5.3 & A.6.1 节，预训练大模型）：不同专家在不同噪声区间形成明显分工，约 σ₈ 处出现利用模式的转折点；浅层（L0-L4）专精化明显，中层（L5-L7）利用更均衡，深层（L8-L11）重新表现出专精化；第一层（L0）尤其偏好在最后去噪步（低噪声）时激活特定专家。

State-based 附加实验（Table 9,Block Push 与 Franka Kitchen,4 seed）：MoDE 0.97±0.01（Block Push）/ 3.79±0.02（Relay Kitchen）,对照 BESO（同作者组前作稠密扩散 Transformer）0.96±0.02 / 3.73±0.05,VQ-BeT 0.87±0.02 / 3.78±0.04,C-BeT 0.87±0.07 / 3.09±0.12——MoDE 在两个 state-based 基准上同样优于稠密扩散基线。

创新点与影响

  • 噪声条件路由（noise-only routing）是核心创新：与传统 LLM MoE（依 token 内容路由）不同，MoDE 路由只依赖当前噪声等级，使得同一噪声等级下的专家选择路径完全确定、可在推理前预计算——这是论文能把标准 MoE rollout 的 FLOPs 开销削减 80%+、推理提速 2 倍的根本原因，而消融证明这一效率取舍几乎不损失性能（0.851 vs token-content 路由的 0.845）。
  • 噪声条件自注意力：把噪声编码 token 加到所有输入 token 上再做自注意力，不增加参数、不改变架构，消融显示与噪声输入 token 互补（同时使用 0.92 vs 单独使用噪声 token 0.85 / FiLM 条件 0.81）。
  • 首次把 MoE 用于扩散策略的去噪计算效率（而非像并行工作 Sparse-DP 那样用 MoE 学任务专精专家），论文自陈这是”首个把 MoE 架构与扩散结合用于行为生成”的工作。
  • 实证结果：134 个任务、4 个基准（CALVIN ABC/ABCD、LIBERO-10/90）全面 SOTA，平均比第二名基线高 57.5%；CALVIN ABC→D 零样本泛化取得 4.01 新 SOTA；SIMPLER 上以更少训练数据（196K vs OpenVLA 的 1M / Octo 的 800K 轨迹）取得更高平均成功率（26.30% vs 23.70% / 17.75%）。
  • 作者自陈局限：(1) MoDE 的标准差普遍高于基线，怀疑路由器初始化对优化过程影响显著,呼吁未来工作研究路由稳定化；(2) 部分实验中观察到专家坍塌（只有部分专家被实际利用），这是稀疏 MoE 训练的已知难题（Chi et al. 2022）,MoDE 虽用负载均衡损失缓解但未完全消除;(3) 结论部分提出未来想尝试 expert-choice routing（Zhou et al. 2022）等其他路由策略。
  • 影响：ICLR 2025 接收；PapersWithCode 上”CALVIN 零样本泛化""LIBERO-10""LIBERO-90”三个榜单的 SOTA 标记来源；官方开源代码与 4 套预训练权重（CALVIN ABC/ABCD/D 及 OXE 预训练大模型）；同实验室后续 MDT（Multimodal Diffusion Transformer）在项目页被列为关联后续工作。

原始链接

一手源存档（sources/）

关系图谱

反向链接