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

Why optimize diffusion model efficiency when perception models seem more computationally demanding? Can you include FLOPs and inference speed comparison?

Thank you for this important question. Perception models like ResNet-50 are considerably smaller (27M parameters vs 760M for MoDE’s Diffusion Transformer), our comprehensive benchmarks show why optimizing diffusion models is crucial:

Per-action computation:

• ResNet-50: Single pass (2ms)
• MoDE: 10 denoising steps (~12 ms each)
• Total MoDE time per action: 12.2 ms with caching, would be around 30ms without

Our efficient design:

• Reduces FLOPs by 80% through expert caching and increases the inference speed form 30ms to 12.4 ms for batch size of 1 Achieves fastest inference among transformers (12.2ms vs 16.2ms for DP-T) while DP-T has only 30% of the total number of parameters of MoDE
• Uses only 1.53 GFLOPs per action vs 2.62 for standard transformers thanks to our efficient expert caching and sparse MoE Design

With diffusion models growing beyond 1B parameters [1], these efficiency gains become even more critical for real-world robotics. We included a new Table 6 in the appendix for detailed latency and FLOP comparisons across all policies on CALVIN. On average MoDE has fast inference with very low FLOPS demonstrating its efficiency.

[1]: Liu, Songming, et al. "RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation." arXiv preprint arXiv:2410.07864 (2024).

MoE may bring additional parameters. How do the overall parameters (instead of active parameters) in Table 1.

We thank the reviewer for their suggestion and added a new Table 6 in the Appendix to show the absolute number of parameters for each baseline in the CALVIN experiment. The table:

Caption: Comparison of total and active number of parameters of methods used in the CALVIN benchmark. Additional overview of average FLOPS required by the different methods together with their average performance on the ABC benchmark. SF-Ratio compares average rollout length with GFLOPS.

Some typos. 285, wrong usage for quotes. You should use `` instead.

We thank the reviewer for pointing out that mistake, we fixed it for the final version.

Table 1 shows that MoDe works better than other methods without pretraining, which is good. However, it can be better to report MoDe with pretraining for a fail comparison with GR-1.

We have now conducted fair comparisons with pretrained models by training MoDE on 200k robot trajectories from Open-X-Embodiment:

Performance gains:

• CALVIN ABC: 3.98 average length (vs GR-1's 3.06)
• LIBERO-90/10: 0.95/0.94 success rate
• Achieved with just 300k pretraining steps + 15k finetuning steps

Training efficiency:

MoDE achieves superior performance with significantly lower computational resources and training time. We will release all pretrained weights of MoDE to support future research.

Continued in next answer

Since the discussion period is ending soon, we would greatly appreciate it if you could let us know whether our response addressed all of your remaining concerns. If you have any remaining questions or points requiring clarification, we are ready to provide additional information. Thanks a lot!

Dear Authors, Thank you for your response. However, I still have some remaining questions.

1. Efficiency The first question is still about the computation costs of diffusion policy networks in the whole system of embodied AI. I have searched some papers in this domain, and some of the policy networks have much fewer parameters and computation costs than the numbers listed by you. I would like to cite some words from a technical report as follows. "For instance, the official implementation of DD utilizes around 60M parameters[1], while Diffuser uses 4M[30], and IDQL[23] has approximately only 1.6M parameters." You can find them in the D.1 of this paper: https://arxiv.org/pdf/2406.09509. Besides, I have also known that there are some papers using flow matching and consistency models to accelerate the diffusion models, which reduce the sampling timesteps to only 1 or very few steps, just as how people accelerate the diffusion models for image generation. I believe these models have much better efficiency, and I hope authors can introduce them and compare them.

Authors have given the computation costs of a ResNet forward passing to prove that the costs of diffusion policy networks are significant. However, what are the detailed configure of the ResNet (depth, width, input resolutions). What's the whole computation of an embodied AI system. Is there only a ResNet and a diffusion model? What are the costs of other modules.

If this question can not be answered properly, the advantage of efficiency of this paper is still a problem.

2. MOE Routing I'm confused about the noise-level conditioned router. Is this router conditioned on both x_t and the noise level which is decided by t? If it is still influenced by the t, how can you pre-compute the routing strategy in each timestep before getting x_t. If it is only conditioned on the noise-level (t), can it still be considered a MOE? For me, it is more similar to the diffusion models which have different parameters in different timesteps. In word words, you are training multiple diffusion models for different timesteps.

Besides, it is confusing for me why the model still needs the noise-level when x_t has already been given. Since x_t has different distributions in different timesteps, giving x_t to the diffusion model has also given information from t. The improvements from noise-level in the ablation study is also not significant, since only 2 of 5 show improvements.

I believe my understanding is not accurate, please correct me. Thanks.

The detailed computational costs and configuration of all components in your full embodied AI system?

Thank you for this important question about computational costs. Let us provide a detailed breakdown of our system's architecture and computational requirements:

Vision Processing Pipeline:

ResNet Configuration:

• Pretrained MoDE: ResNet-50 (25.6M parameters)
• Resolution: 224x224x3 input
• FLOPs: 8.27 GFLOPs for both images
• Standard MoDE: ResNet-18 (11.7M parameters)
• Same resolution: 224x224x3 input
• FLOPs: 3.62 GFLOPs for both images

Language Processing:

CLIP ViT-B/32 for text encoding

• Parameters: 151M
• Compute cost effectively zero during deployment:
• All language commands are predefined
• Embeddings pre-computed and cached
• Both CALVIN and LIBERO provide cached embeddings in their dataset

Policy Network (MoDE):

Core MoDE MoE architecture:

• Pretrained MoDE: 436M active parameters (685M total)
• Standard MoDE: 277M active parameters (460M total)
• FLOPs per Denoising Step: 0.7 GFLOPs with router caching

Total System Costs (per action):

Pretrained MoDE variant:

• Total Parameters: ~740M (excluding CLIP)
• Inference Time: 12.2ms on NVIDIA A6000

This analysis shows our efficient MoDE architecture helps minimize the policy network's overhead.

Is noise-level routing truly a Mixture of Experts, or just separate models for different timesteps?

Routing Mechanism: The router in MoDE is exclusively conditioned on the noise level σ_t, not on the noisy input x_t. This design choice is motivated by recent findings showing that different denoising phases require specialized processing [4]. While traditional MoE architectures route based on input content, MoDE uses the natural structure of the diffusion process itself.

Why This Is Still an MoE: Our architecture differs fundamentally from training separate models for different timesteps in several key ways:

a) Shared Core Components:

• Attention layers process all inputs regardless of noise level
• Layer normalization and embedding modules are unified
• Common backbone learns generalizable representations

b) Dynamic Expert Selection:

• Each noise level uses top-k=2 experts, allowing interpolation
• Experts specialize in noise regions rather than specific timesteps
• Router learns smooth transitions between noise regimes

Empirical Evidence:

Our ablation studies (Figure 6 & 12 in paper) demonstrate that:

• Experts develop complementary specializations
• Smooth transitions occur between noise levels
• Performance improves over both single-expert and fully separated models

Computational Benefits:

This design offers unique advantages:

• 80% reduction in FLOPs through router caching
• 40% of parameters not active during inference

Maintains robust performance despite sparsity

Comparison to Separate Models:

Training separate models for each timestep would:

• Require N times more parameters (N = number of denoising steps)
• Struggle to learn robust vision/language representations
• Miss opportunities for parameter sharing
• Need significantly more training compute and memory

MoDE's architecture thus combines the efficiency of MoE with the structure of diffusion processes, enabling more effective scaling than either approach alone.

Continue in next response

The methods section of the paper is too brief, and much of the content is drawn from existing research, with the authors not highlighting their contributions effectively.

Thank you for this valuable feedback. We agree that we should have highlighted our contributions stronger. We revised the methods section accordingly to talk about our noise-only routing in more detail and also introduced two new sections to talk about expert caching used in our framework which we have omitted before and shortly discuss our new pretraining of MoDE.

In the experimental section, the authors do not provide a comparison of FLOPs and inference times with other baseline methods, which may lead to unfair comparisons.

We thank the reviewer for raising this point. We have added a section to provide an estimate of FLOPS for all methods in the appendix and summarize our findings in Table 6. We explain the average GFLOPs for the tested baselines on CALVIN ABC and introduce the SF-Ration that represents the average rollout length per GFLOPs. Thanks to our efficient design, MoDE has by far the highest SF-Ration and a very low GFLOPs count overall. MoDE is also the 2nd fastest tested policy although it has more than 700M parameters. This provides more evidence of the computational efficiency and strong performance of MoDE.

Caption: Comparison of total and active number of parameters of methods used in the CALVIN benchmark. Additional overview of average FLOPS required by the different methods together with their average performance on the ABC benchmark. SF-Ratio compares average rollout length with GFLOPS.

In Figure 1, the authors propose a new router strategy, but the figure does not clearly illustrate how this strategy selects different experts, failing to convey the core idea of the paper. Furthermore, the framework diagram is overly simplified and lacks sufficient depth.

We appreciate the reviewer’s comment regarding Figure 1 and agree that it did not adequately highlight our contributions. To address this, we have revised the figure to better showcase the core of our proposed routing strategy. The updated figure now consists of two parts: the first illustrates the general architecture, while the second focuses specifically on our novel noise-conditioned routing mechanism. We hope this updated figure addresses the initial concerns effectively.

Can you elaborate on the specific innovations introduced by your Mixture-of-Experts Diffusion Policy that differentiate it from existing methods in the literature?

• a novel noise-conditioned routing mechanism that dynamically partitions experts based on noise levels, eliminating the need for predefined task clustering or fixed channel partitioning seen in prior works [1]. As shown in Figure 5, MoDE learns to allocate experts to handle different noise scales, enabling greater flexibility and task alignment.
• Fast inference thanks to our proposed router caching and MLP Fusion. These changes reduce the FLOPS by over 80% and double the inference speed.
• MoDE achieves superior performance with fewer parameters and steps compared to prior Diffusion Policies, offering an efficient and generalizable solution with noise-conditioned expert utilization

All of our changes combined lead to the strong performance of MoDE. We provide detailed ablation and insights into the Design for MoE-based Diffusion Policies to make them work in low-data settings.

[1] https://arxiv.org/abs/2310.07138

Response continues in the next message.

Since the discussion period is ending soon, we would greatly appreciate it if you could let us know whether our response addressed all of your remaining concerns. If you have any remaining questions or points requiring clarification, we are ready to provide additional information. Thanks a lot!

We thank the reviewer for the feedback on our work. We try to address your concerns below:

What makes MoDE novel beyond applying standard MoE techniques to diffusion policies?

We would like to emphasize the following contributions of MoDE:

• MoDE introduces the first sparse Mixture of Experts architecture for diffusion policies, reducing FLOPs by 80% while achieving state-of-the-art performance. Its novel noise-conditioned routing enables automatic expert specialization across noise levels, with pre-computed routing paths and fast inference with expert caching.

• MoDE demonstrates strong performance gains on challenging benchmarks - achieving 40% improvement over DP-T and 20% over DP-CNN, while being the first to exceed 90% success rate on LIBERO-10. and surpassing all baselines on CALVIN by 30%

• Our comprehensive empirical ablations provide crucial insights into making MoE work effectively for diffusion policies, where data constraints and real-time requirements pose unique challenges.

We believe that all these points make MoDE valuable for the robot learning community

How does MoDE generate actions during long-horizon tasks like CALVIN ABCD?

Action Generation Cycle:

• Takes current observation images and text goal
• Generates 10 actions using 10 denoising steps
• Executes all actions in simulation
• Gets new observation and repeats
• When a task is completed, MoDE receives the next goal in the sequence and starts the process again. Continues until all 5 task completion or failure (max 360 timesteps)

The improvements on the LIBERO-90 and CALVIN datasets are marginal.

We have updated our experimental results.
Below, we summarize MoDE against the 2nd best baseline in each tested benchmark:

Our method achieves significant improvements: 40% better performance than DP-T, 20% better than DP-CNN, and is the first to exceed 90% success rate on LIBERO-10. MoDE does outperform all specialist policies in the LIBERO setting, generalist policies in the SIMPLER benchmark and all baselines in CALVIN. Most notably, with our included pretraining, MoDE achieves a new state-of-the-art score of 3.98 on CALVIN ABC, surpassing the previous best baseline by over 30% with higher efficiency and less training time. These results demonstrate substantial improvements across all benchmarks, with an average improvement of 14.1% and gains of up to 30% in challenging zero-shot settings.

We thank the reviewer for their constructive feedback. Please don't hesitate to let us know if you have any other questions or comments.

I thank the reviewers for their efforts for the rebuttal. The rebuttal addressed most of my comments, so I raised my score to 6. I think as the paper is the first one to implement and verify the effectiveness of MoDE for imitation learning, it has some benefits to the relevant community. However, I didn't vote for a higher score because I believe the technical contribution of the paper is relatively limited.

How do load balancing weights affect expert utilization and performance?

We thank the reviewer for this suggestion. We have worked on a more in depth analysis of expert distribution, that we added to the appendix given the limited space for the main text. In our new Section A 5.1, we analyze how the load balancing loss affects expert distribution and how our newly added pretrained MoDE variant distributes experts across 10 denoising levels in detail. We tried to identify general patterns and analyze them with respect to hyperparameters.

Why does MoDE's performance peak at 4 experts and what happens with more experts?

Thank you for this important observation. We have conducted additional experiments with 6-8 experts and added a detailed discussion in the router ablations section of the Appendix. Our findings suggest that noise-only conditioning does not benefit from more than 4 experts. In addition, the router struggles with 6 experts and more to balance the load to the tokens effectively. There appears to be a general trade-off: experts must both learn noise-based latent representations from image and language goal tokens while specializing in specific noise ranges. We hypothesize that adding more experts beyond 4 does not improve this balance between representation learning and noise specialization.

A similar analysis can be done to study the impact of the top k routing parameter.

We thank the reviewer for his suggestion. We analyze the load balancing of topk=1 in the appendix. In addition, we also investigated the impact of having a shared expert, which is an expert that is always chosen. Overall, topk=2 seems crucial in more challenging settings to achieve the best performance, as the model is not able to distribute the loads effectively and the performance drops considerably in challenging settings like CALVIN ABC from 3.34 to 2.72. This indicates that topk=2 is required to achieve best performance.

Making MoE Work in Data-Limited Robot Learning: Key Findings and Impact

We thank the reviewer for their kind comment about our contributions. The reviewer's thoughtful review was very helpful and we hope our detailed ablation studies regarding expert utilization in diffusion policies further strengthens our understanding of how sparse MoE can be most effectively applied to data limited robotics settings.

We hope the new changes satisfy the initial concerns and are happy to answer any follow-up question.

Since the discussion period is ending soon, we would greatly appreciate it if you could let us know whether our response addressed all of your remaining concerns. If you have any remaining questions or points requiring clarification, we are ready to provide additional information. Thanks a lot!