Variational Distillation of Diffusion Policies into Mixture of Experts

We would like to thank the reviewer for taking the time to review our work and the many helpful comments and suggestions. We hope the following replies sufficiently address the raised questions and concerns. We will update the paper accordingly.

---

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.

We thank the reviewer for sharing their concern. If we understand correctly, when the reviewer mentions "final performance" they are referring to the success rate. In our evaluation, we additionally assess the diversity of policies using task entropy (TE).

TE was originally introduced by [1] as a scalar metric ranging from 0 to 1, where 0 indicates that the model has only learned one way to solve the task, and 1 indicates that the model has learned all skills present in the demonstration data.

We apologize for any confusion caused by our description of TE as a measure of versatility rather than diversity. It is indeed intended to quantify the model's ability to exhibit diverse behaviors. We will make the necessary changes to the experiment section to further enhance clarity.

---

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.

We thank the reviewer for their suggestion to enhance the motivation for policies capable of learning diverse behavior. In response, we have incorporated additional points and references into the introduction of our paper.

Additionally, we summarized these points below:

1. Improving Generalization: If the learned policy overfits a specific set of demonstrated behaviors, it may not generalize well to new situations. By exposing the model to diverse behaviors, the risk of overfitting is reduced, and the learned policy is more likely to capture the underlying principles of the task rather than memorizing specific trajectories [1, 2, 3].
2. Enhancing Skill Transfer: Learning diverse behaviors facilitates better skill transfer across different but related tasks. If the agent can imitate a wide range of behaviors, it is more likely to possess a set of skills that can be applied to various tasks, making it a more versatile and capable learner [1, 3].
3. Behavior in complex multi-agent settings: Following a multimodal policy through a diverse set of skills, has been shown to be advantageous in multi-agent settings, such as table tennis [4].

---

Variational 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.

We assume the reviewer is referring to [6] when mentioning that variational distillation is a known technique in the 3D field. We will add further details to the 'Related Works' section, highlighting the distinctions between their work and ours to emphasize the novelty introduced in our approach.

In essence, [6] aims to generate 3D scenes using a pre-trained 2D teacher model for distillation. Their approach fundamentally differs from ours, as they focus on distilling knowledge from one diffusion model into another diffusion model. In contrast, our goal is to distill a diffusion model into another family of generative models, specifically Mixture of Experts, to achieve favorable properties such as faster inference and tractable likelihoods.

---

What is the mixture number n in your experiments?

To determine the value of n, we conducted a grid search and found that using n=8 yields consistently good performance across all D3IL tasks, and n=4 yields good results for Franka kitchen and block push. Additional details about the hyperparameters can be found in Appendix E.2, Table 4.

We will also include more detailed information about the mixture number in the main section of the paper.

---

Can you provide any experiment results ablating n? This could be very informative.

We agree that such an ablation is very informative and in fact already included such this ablation in the original submission, see Figure 3a, with accompanying text in Section 5.5. The main takeaway is that increasing the mixture number improves task entropy and therefore the skill diversity of the learned policy. We will emphasize this ablation stronger in the next iteration of the paper.

---

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)

We apologize for not including the source code in our initial submission. We have now provided a link to an anonymous GitHub repository containing the code for all our experiments. Additionally, we have included an IPython notebook to test our method on a small toy task, which essentially reproduces Figure 2. We have sent these links to the AC in an Official Comment, as the rebuttal guidelines prevent us from directly sharing these links in our response to the reviewers.

---

We would like to thank the reviewer again and welcome the opportunity to address any additional concerns or questions that the reviewers may have.

---

[1] Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations, ICLR ‘2024

[2] Neural Probabilistic Motor Primitives for Humanoid Control, ICLR ‘19

[3] InfoGAIL: Interpretable Imitation Learning from Visual Demonstrations, NeurIPS ‘17

[4] One Solution is Not All You Need: Few-Shot Extrapolation via Structured MaxEnt RL, NeurIPS ‘20

[5] Specializing Versatile Skill Libraries using Local Mixture of Experts, CoRL ‘21

[6] High-fidelity and diverse text-to-3d generation with variational score distillation, NeurIPS ‘23

We would like to thank the reviewer for taking the time to review our work and the many helpful comments and suggestions. We hope the following replies sufficiently address the raised questions and concerns. We will update the paper accordingly.

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.

To the best of our knowledge, there are currently no variants of diffusion models that enable precise one-step generation with tractable likelihoods. Therefore, the study of diffusion distillation remains an ongoing area of investigation. Recent studies include [1, 2]. However, if the reviewer is aware of any novel diffusion model variants that address the challenges outlined in our work, we would be very happy to include them as references/baselines in our work.

---

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

We thank the reviewer and will do our best to remove respective phrasing and make the paper more engaging for the reader.

---

[…] missing period on line 218 […] The inconsistency in Table 3 left, where 2.16 is highlighted as better than 2.15 is confusing […]

We thank the reviewer for carefully reading the paper and identifying the formatting issues. We have addressed and corrected them.

---

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.

We apologize for not including these results in the initial submission due to time constraints. We added the missing results to the Table R1 in PDF that accompanies the rebuttal.

---

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.

We thank the reviewer for sharing their concern. However, we want to emphasize that the primary goal of our paper is not to outperform all metrics, as the performance of our approach is limited by the diffusion ‘teacher’ model.

Nevertheless, we agree with the reviewer that the submitted version of the paper did not sufficiently discuss future directions to enhance distillation capabilities. We will add this section to the manuscript and provide a summary below.

Future Work. A promising avenue for further research is to utilize the features of the diffusion 'teacher' model to reduce training time and enhance performance. This can be achieved by leveraging the diffusion model as a backbone and fine-tuning an MoE head to predict the means and covariance matrices of the experts. The time-dependence of the diffusion model can be directly employed to train the MoE on multiple noise levels, effectively eliminating the need for the time-step selection scheme introduced in Section 4.3.

---

The comparison to a limited number of methods may limit the ability to comprehensively assess the strengths and weaknesses of the proposed approach.

To the best of our knowledge, we have considered the most recent and prominent baselines for training MoE models. However, we acknowledge that novel methods for distilling diffusion models, such as Consistency Trajectory Model (CTM) [3], were not included in our evaluation as they were not published at the time of submission. In order to better evaluate the strength of VDD, we additionally add CTM as a new distillation baseline. The evaluation results are presented in the PDF that accompanies the rebuttal.

---

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.

We thank the reviewer for bringing the need for illustrative visualizations to our attention. In response, we included additional visualizations (Figure R1) in the PDF that accompanies the rebuttal.

Figure R1 highlights the most likely experts according to the gating probability at a given state. We can see that using a single component can achieve a perfect success rate at the cost of losing behavior diversity. Contrary, using many experts potentially results in a slightly lower success rate but increased behavior diversity. These qualitative results are consistent with the quantitative results from the ablation study presented in Figure 3a of the paper.

---

We express our gratitude to the reviewer for their valuable comments and suggestions. We are pleased to address any additional questions or concerns that may arise.

---

[1] Song, Yang, and Prafulla Dhariwal. "Improved techniques for training consistency models.” ICLR 2024

[2] Xie, S., et al. “EM Distillation for One-step Diffusion Models”, Preprint

[3] Kim, Dongjun, et al. "Consistency trajectory models: Learning probability flow ode trajectory of diffusion." ICLR 2024.

Dear Reviewer,

Please reply to the rebuttal.

AC.

We would like to thank the reviewer for taking the time to review our work and the many helpful comments and suggestions. We hope the following replies sufficiently address the raised questions and concerns. We will update the paper accordingly.

---

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 […] .

We thank the reviewer for bringing this issue to our attention. We apologize for not reporting the training times for a varying number of experts in the initial draft of the paper. We added a table showing the training costs using $n \in \{1,2,4,8,16\}$ components to the PDF that accompanies the rebuttal. Figure R2 shows that both, training time and the number of parameters increase sub-linearly with $n$.

---

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.

Since the approach leverages variational inference, and hence the inherit mode seeking behavior of the reverse KL, VDD is relatively robust to the number of components. In fact, the results reported in the paper use $n=8$ for all D3IL tasks and $n=4$ for both kitchen and block push. Moreover, using a large number of components introduces only slight variations in performance criteria, as the gating network can deactivate those that are not needed to solve a task, as shown by the ablation study in Figure 3a in the paper and Figure R1 in the PDF that accompanies the rebuttal.

---

Is it possible to take task similarity into account and perform kind of meta-training like policy distillation?

We thank the reviewer for the interesting suggestion. We believe that using a shared feature embedding across tasks could potentially speed up training time and is a promising avenue for future research. Currently, we are exploring how to efficiently leverage the features learned by the teacher diffusion model to improve training times.

---

The authors do not explicitly discuss how well VDD can capture long-range temporal dependency that might be presented in the original diffusion model […]

VDD shares the same capacity for modeling temporal dependencies as the teacher diffusion model, as it employs the same transformer backbone. We argue that VDD's comparable success rates to the diffusion model demonstrate its ability to capture the long-range temporal dependencies inherent in the original diffusion model.

---

We hope that our responses have addressed your concerns. We would be more than happy to address any remaining questions/concerns.

We thank the reviewer for the very positive comments and are grateful for the valuable suggestions and comments. We hope the following replies sufficiently address the raised questions and concerns. We will update the paper accordingly.

---

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.

We thank the reviewer for pointing us towards these very interesting and related works. We will include these references in the main manuscript and further highlight both the commonalities and differences between these works and our approach.

---

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.

Indeed, in our experience BESO appears to be distilled easier. We credit this to the different forward SDE which in the case of BESO is a simple (scaled) Wiener Process. In contrast, the DDPM SDE contains an additional non-zero drift term, which makes the distillation harder.

---

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 […]

We thank the reviewer for the encouraging suggestion. However, we would like to emphasize that the tasks used in this work are taken from a very recent benchmark suite [1], allowing for the quantification of the behavior diversity of the learned policy. Maintaining this diversity through the distillation process is an integral part of our approach. Moreover, this task suite contains several complex manipulation tasks. These complexities are particularly evident in tasks such as Sorting (Image) and Stacking (Image), where the teacher model achieves success rates of less than 70%.

Nevertheless, we find the idea of incorporating the dexterous hand task very intriguing. In the future, we would like to further investigate this area and evaluate the performance of our model in this context.

---

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.

We assume the reviewer refers to the success rate as a performance criterion when mentioning ‘outperform’. We acknowledge that it may seem confusing at first that the 'student' outperforms the 'teacher' in terms of success rate. In these cases we observed that student usually exhibits lower task entropy, indicating less diverse behavior.

We hypothesize that there is a trade-off between mastering a single skill very well and learning multiple skills with slightly less accuracy. This hypothesis is further supported by the results of the ablation study presented in Figure 3a. Additionally, we have included further visual evidence supporting this hypothesis in the PDF file accompanying the rebuttal. Figure R1 shows that by reducing the number of experts the success rate increases, at the cost of losing skill diversity. The extreme case $Z=1$, has a success rate of 1, but 0 entropy. We add a discussion about this trade-off in the experiment section.

---

We would like to thank the reviewers again for assessing our work. We would be delighted to address any additional questions or concerns they may have.

---

[1] Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations, ICLR ‘2024

We would like to thank the reviewer for taking the time to review our work and the many helpful comments and suggestions. We hope the following replies sufficiently address the raised questions and concerns. We will update the paper accordingly.

[…] the paper claims that "VDD is not straightforwardly applicable to generating very high-dimensional data like images."

The statement in the limitation section refers to the dimensionality of the action space. However, in the image-based version of sorting and stacking, the images represent the state space (model input). The model output remains a low-dimensional control signal for the robot. VDD is not straightforwardly applicable to tasks that require high-dimensional output space such as image generation, because predicting the Cholesky decomposition of the Covariance matrix scales quadratically with the dimensionality of output space. We will clarify the statement in the limitation section.

---

Why does the distilled model outperform the original diffusion policy in Sorting (Image) and Stacking (Image) tasks?

We assume the reviewer refers to the success rate as a performance criterion when mentioning ‘outperform’. We acknowledge that it may seem confusing at first that the 'student' outperforms the 'teacher' in terms of success rate. In these cases we observed that student usually exhibits lower task entropy, indicating less diverse behavior.

We hypothesize that there is a trade-off between mastering a single skill very well and learning multiple skills with slightly less accuracy. This hypothesis is further supported by the results of the ablation study presented in Figure 3a. Additionally, we have included further visual evidence supporting this hypothesis in the PDF file accompanying the rebuttal. Figure R1 shows that by reducing the number of experts the success rate increases, at the cost of losing skill diversity. The extreme case $Z=1$, has a success rate of 1, but 0 entropy.

We add a discussion about this trade-off in the experiment section.

---

Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models?

Yes, to ensure fairness across all experiments, we use the same architecture for these methods.

---

[…] It would be beneficial to report this total time for a fair comparison between MoE methods.

We thank the reviewer for pointing this out. In the PDF file accompanying the rebuttal, we report the total training time as well as the separate training time of DDPM and VDD in Table R1(b).

---

Why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models?

Diffusion models have shown great performance in recent studies, often making them the preferred choice for addressing complex tasks. However, they suffer from slow inference due the iterative denoising and do not offer a tractable likelihood. The goal of our work is, therefore, to mitigate these inherent downsides of diffusion models by distilling them into MoEs through variational inference. Nevertheless, we find the idea of leveraging our approach to distill large MoE into smaller ones very intriguing and will look into it in future work.

---

We express our gratitude to the reviewer for their valuable comments and suggestions. We are pleased to address any additional questions or concerns that may arise.