Variational Distillation of Diffusion Policies into Mixture of Experts

Summary

This paper introduces a theoretical method for extracting an MoE policy from a pretrained diffusion policy. The advantage of MoE policy is that it allows faster sampling speed compared with the diffusion policy more stable training and better performance compared with methods that train MoE from scratch.

Strengths

The proposed method solves the policy diversity issues in RL while maintaining high training/inference efficiency. The experiments show clear improvement over previous MoE-based imitation learning methods. The theories are well-proved and solid.

Weaknesses

1. One main advantage of MoE policy against deterministic policy is its diversity. However, this paper seems to only focus on evaluating the final performance of algorithms and does not compare the diversity of different algorithms. Regarding on performance solely, it seems to me Gaussian policy + reverse KL objective is good enough, in what cases do we need more diverse policies? This needs to be justified. 2. Variantional distillation is a known technique in the 3D field. Although its application in RL is new and meaningful, this still indicates limited novelty. Could emphasize more on the new technique introduced in this paper when applying variational distillation for policy extraction.

Questions

1. What is the mixture number n in your experiments? Can you provide any experiment results ablating n? This could be very informative. 2. I highly suggest the authors provide the source code during rebuttal. Reproducibility is crucial for this kind of paper. (Even if it is just unsorted/raw code)

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Above

Summary

This study introduces Variational Diffusion Distillation (VDD), a novel method that distills pre-trained diffusion models into Mixture of Experts (MoE) frameworks. VDD addresses diffusion models' drawbacks of intractable likelihoods and long inference times, while leveraging their ability to represent complex distributions. By leveraging a decompositional variational objective, VDD trains MoEs efficiently, enabling real-time applications. VDD outperforms existing distillation and MoE training methods in several complex behavior learning tasks.

Strengths

1) The technical approach appears to be sound, demonstrating a rigorous and well-grounded methodology. 2) The visualizations presented in the paper are meticulously crafted, showcasing a high level of attention to detail and aesthetics. 3) The experimental section is quite impressive, offering a substantial evaluation of the proposed method.

Weaknesses

1) The challenges posed, although relevant to diffusion models and their potential difficulties in handling certain tasks or long inference times, lack sufficient evidence to conclusively state that these issues persist across all novel variants of diffusion models. 2) The use of identical phrasing in the abstract and introduction could potentially limit the reader's engagement, as it fails to provide a nuanced progression of ideas. 3) A missing period on line 218 detracts slightly from the overall readability and professionalism of the manuscript. 4) It appears that the experiments for the DDPM component in Table 1(a) are incomplete for the kitchen and block push datasets, which limits the comprehensiveness of the evaluation. 5) While the proposed method achieves commendable results, it does not consistently outperform all metrics, necessitating a deeper analysis to explain the observed variations and identify potential avenues for improvement. 6) The inconsistency in Table 3 left, where 2.16 is highlighted as better than 2.15 is confusing and suggests that the experiment may not have been fully completed. 7) The comparison to a limited number of methods may limit the ability to comprehensively assess the strengths and weaknesses of the proposed approach. 8) The lack of illustrative visualizations, such as good and bad case studies, in the main text hinders the reader's ability to fully grasp the performance and limitations of the method in real-world scenarios.

Questions

see weaknesses

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

see weaknesses

Summary

This paper presents a variational inference method for distilling denoising diffusion policies into Mixture-of-Experts (MoE) policies. The primary motivation is to combine the strengths of both worlds - the ability to learn complex, multi-modal distributions of diffusion models - and the efficiency of MoEs offering faster inference and tractable likelihoods. Importantly, the decomposed upper bound of the variational objective enables separate, and thus more robust training for different experts. The proposed method is empirically evaluated on 9 behavior learning tasks. The author demonstrate its ability to distill complex distributions, outperforming existing distillation methods.

Strengths

1. clear motivation of combining the strengths of both diffusion models and MoE. 2. combining decomposing objective with EM is shown to be effective to escape certain challenges of training MoE. 3. the paper is overall well-organized, with clear explanations of the method and detailed experimental results.

Weaknesses

1. While EM is arguably be able to handle some limitations, it may also suffer from more longer convergence time during training. Moreover, while the authors discuss inference time improvements, this paper does not analyze training cost. As the training involves optimizing multiple experts and a gating - when the “optimal” number of experts is difficult to be pre-determined, it could potentially be computationally intensive. Is it possible to take task similarity into account and perform kind of meta-training like policy distillation? 2. The authors do not explicitly discuss how well VDD can capture long-range temporal dependency that might be presented in the original diffusion model, provided that it improves the inference speed. However, this can be important for sequential decision-making tasks.

Questions

Please see Weaknesses section

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors do have included discussion about the limitation and shed light on how they might be approaching in future studies.

Summary

This paper presents Variational Diffusion Distillation (VDD), a method that distills denoising diffusion policies into Mixtures of Experts (MoE) using variational inference. Diffusion Models excel in learning complex distributions for behavior learning but have drawbacks like slow inference times. MoEs address these issues but are hard to train. In VDD, each expert can be trained separately under the corresponding guidance from the diffusion teacher. VDD demonstrates convincing performance in distilling complex distributions and outperforming existing methods across nine behavior learning tasks.

Strengths

The paper addresses the problem of distilling an expressive but slow diffusion model into a faster Mixture of Experts (MoE) model. MoE is an ideal target for distillation due to its intrinsic structure and rapid computational process. Though closely related to Variational Score Distillation, the proposed extension to MoE generators is novel. The claims appear valid, and I found no issues in the derivations. Overall, the paper is well-written and clear.

Weaknesses

- I like the idea of updating the gating of MoE, which has been proven to be crucial for the success of VDD in an ablation study. However, it appears a bit unnatural to me that the experts are learned with reverse KL with the gating is learned with forward KL, even though the minimizer for forward and reverse KLs should align. Authors may find it interesting to visit another perspective from a probably concurrent work (Xie et al., 2024), in which a new distillation framework is derived from forward KL. - A very relevant work on distilling diffusion models, Luo et al. 2023, is missing. - I observed some performance gaps between VDD-DDPM and VDD-BESO across various tasks even though the teachers perform similarly. I wonder if VDD is sensitive to the specific forward process used in the diffusion model. - In my humble opinion, the author could try some more complex tasks to demonstrate the capabilities of VDD. For instance, implementing a fast and expressive model on a dexterous hand would benefit from an expressive policy like diffusion while also requiring speed. Prior work exists in directly learning an MoE image generation model with the EM algorithm from data (Yu et al., 2019). Given the power of diffusion teachers, I would expect a distillation method to perform better than directly learning from data. Xie et al. 2024, EM Distillation for One-step Diffusion Models Luo et al. 2023, Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models Yu et al. 2021, Unsupervised Foreground Extraction via Deep Region Competition

Questions

- It appears that in simpler tasks like kitchen and block push, the distilled model underperforms compared to the original model. However, in more complex tasks, many distilled models seem to outperform the original. I am curious if the authors have some insights on this.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors provided the limitation statement, mainly concerning the expressity of MoE.

Summary

This paper studies the knowledge distillation problem in diffusion models by distilling denoising diffusion policies into Mixtures of Experts (MoE) using variational inference. The goal is to combine the advantages of diffusion models, with the fast inference capabilities of MoE models. The authors construct an upper bound of a KL divergence-style loss function between the diffusion policy and the MoE policy and implement EM steps for variational inference to achieve the distillation. The proposed method, VDD (Variational Distillation of Diffusion), demonstrates superior performance compared to existing distillation and MoE methods.

Strengths

1. The paper studies an interesting and practical problem in the distillation of diffusion models, achieving 1-step inference with performance comparable to the original diffusion policy. This type of fast inference can be highly beneficial for real-world control problems utilizing diffusion policies. 2. The proposed method, VDD, introduces a robust approach to distilling diffusion policies to MoE models using VI. 3. Overall, the paper is well-structured, easy to understand, and provides a clear presentation of the proposed method.

Weaknesses

1. In the limitations section, the paper claims that "VDD is not straightforwardly applicable to generating very high-dimensional data like images." However, the experimental results in Table 1(a) show that the distilled model VDD-DDPM outperforms the original DDPM diffusion policy in the Sorting (Image) and Stacking (Image) tasks, which contradicts the discussion. What is the reason for this? Why does the distilled model outperform the original diffusion policy? 2. Regarding training time, the time cost for distillation models (like VDD-DDPM) is higher than for MoE models. I assume this refers only to the knowledge distillation time. If we consider the original training time for diffusion policies, the total time cost may be significantly larger. It would be beneficial to report this total time for a fair comparison between MoE methods. 3. Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models? 4. While I agree that, according to this paper, the distillation of MoE achieves better performance than training from scratch, why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models? This could be especially relevant in low-dimensional state observation cases, where MoE models reportedly perform better.

Questions

1. Why does the distilled model outperform the original diffusion policy in Sorting (Image) and Stacking (Image) tasks? 2. Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models? 3. Why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have discussed the limitations of the paper. Please refer to the [weakness] section for detailed concerns and questions about the paper.

Response

I thank the authors for the very detailed responses, which have resolved most of my concerns. I raise the score from 5 to 6. However, I would like to share one of my very personal opinions: Most (if not all) current public RL benchmarks and their evaluation metrics actually lack a strong need for policy diversity, which hinders the necessity of diversified policies. I believe designing new tasks and more difficult benchmarks is of critical importance for designing these algorithms.

Response to the rebuttal

I appreciate the author's detailed responses, which have addressed most of my concerns. However, there are still some aspects that remain unclear to me. 1. Why does the student VDD model outperform the teacher diffusion model? The author suggests that this may be related to the trade-off between mastering a single skill very well and learning multiple skills with slightly less accuracy. However, when looking at the success rates in Table 1(a), the student model only outperforms in the Sorting (Image) and Stacking (Image) tasks, but not in others. Given a fixed teacher model and a fixed number of experts, it seems difficult to attribute this phenomenon to the acquisition of multiple skills or increased diversity. Could this be related to the use of image observations? Interestingly, the task entropy of VDD-BESO is comparable to BESO in the Sorting (Image) task. According to the paper, VDD-BESO is both more accurate and not less diverse than BESO. What could be the reason for this discrepancy? 2. How to choose the optimal number of experts in real-world control tasks? The results indicate that reducing the number of experts can increase the success rate, albeit at the cost of skill diversity. While this conclusion seems reasonable, it raises the question: how should we determine the optimal number of experts in a Mixture of Experts (MoE) model for a given real-world control task or set of tasks? In many safety-critical control problems, success rate is often the primary concern, leading to the possibility of opting for a smaller number of experts. Could you provide guidance on this decision, or better yet, demonstrate through experiments any potential drawbacks of this strategy?

Thank you for the insightful discussions and help to improve VDD!

Dear Reviewers, We sincerely thank all of the reviewers for the thoughtful feedback and engaging discussions throughout the rebuttal process. Your insights have been invaluable in helping us improve our work on VDD. Going into the reviewer discussion, we would like to summarize the main improvements that resulted from your feedback during the rebuttal: - **Related Work**: Expanded discussion on VDD's relation to recent diffusion distillation approaches to highlight our novelty. - **Baseline Comparisons**: Added evaluations against Consistency Trajectory Model (CTM) to show the strong distillation performance of VDD. - **Completed Results-Table:** Included results for DDPM and VDD-DDPM on kitchen and block push datasets. - **Enhanced Evaluation on Teacher Models:** Added evaluations on diffusion teacher model with different samplers and different integration steps to show that VDD is able to maintain comparable performance with the teacher model. - **Detailed Training Cost Information**: Provided clearer information on parameter counts and training times across varying number of experts. Reported the total training time of VDD including the training time of the diffusion teacher on kitchen and block push tasks. - **Enhanced Visualization**: Added new visualization to offer better insights in to our method. It highlights the most likely experts according to the gating probability at a given state. - **Open-Sourced Code-Base**: Provided a link to an anonymous Github repository containing the code for reproducing the results of the paper to the AC. Additionally, we provided a link to an IPython notebook for testing our method on a simple toy task. These changes have significantly strengthened our paper and all discussed points will be included in the final version. Thank you again for your time and valuable feedback and we hope you take these points into account for your final decision. Sincerely, The Authors