Variational Distillation of Diffusion Policies into Mixture of Experts
Authors: Hongyi Zhou
Denis Blessing
Ge Li
Onur Celik
Xiaogang Jia
Gerhard Neumann
Rudolf Lioutikov
Venue: NeurIPS 2024 poster (2024)
Keywords: Diverse Behavior Learning
Model Distillation
Diffusion Models
Variational Inference
Primary area: robotics
TL;DR: This work introduces Variational Diffusion Distillation (VDD), a novel method for distilling denoising diffusion policies into a Mixture of Experts (MoE).

Abstract: This work introduces Variational Diffusion Distillation (VDD), a novel method that distills denoising diffusion policies into Mixtures of Experts (MoE) through variational inference. Diffusion Models are the current state-of-the-art in generative modeling due to their exceptional ability to accurately learn and represent complex, multi-modal distributions. This ability allows Diffusion Models to replicate the inherent diversity in human behavior, making them the preferred models in behavior learning such as Learning from Human Demonstrations (LfD).
However, diffusion models come with some drawbacks, including the intractability of likelihoods and long inference times due to their iterative sampling process. The inference times, in particular, pose a significant challenge to real-time applications such as robot control.
In contrast, MoEs effectively address the aforementioned issues while retaining the ability to represent complex distributions but are notoriously difficult to train.
VDD is the first method that distills pre-trained diffusion models into MoE models, and hence, combines the expressiveness of Diffusion Models with the benefits of Mixture Models.
Specifically, VDD leverages a decompositional upper bound of the variational objective that allows the training of each expert separately, resulting in a robust optimization scheme for MoEs.
VDD demonstrates across nine complex behavior learning tasks, that it is able to: i) accurately distill complex distributions learned by the diffusion model, ii) outperform existing state-of-the-art distillation methods, and iii) surpass conventional methods for training MoE. The code and videos are available at https://intuitive-robots.github.io/vdd-website.

Variational Distillation of Diffusion Policies into Mixture of Experts
Hongyi Zhou +
Denis Blessing Ge Li Gerhard Neumann 18
Onur Celik Rudolf Lioutikov
Xiaogang Jiaft
Intuitive Robots Lab, Karlsruhe Institute of Technology
t Autonomous Learning Robots, Karlsruhe Institute of Technology
8 FZI Research Center for Information Technology
Abstract
This work introduces Variational Diffusion Distillation (VDD), a novel method that distills denoising diffusion policies into Mixtures of Experts (MoE) through variational inference. Diffusion Models are the current state-of-the-an in generative modeling due their exceptional ability to accurately learn and represent complex, multi-modal distributions. This ability allows Diffusion Models to replicate the inherent diversity in human behavior, making them the preferred models in behavior learning such as Learning from Human Demonstrations (LfD). However, diffusion models come with some drawbacks, including the intractability of likelihoods and long inference times due to their iterative sampling process. The inference times, in particular, pose a significant challenge to real-time applications such as robot control. In contrast, MoEs effectively address the aforementioned issues while retaining the ability to represent complex distributions but are notoriously difficult train. VDD is the first method that distills pre- pre-trained diffusion models into MoE models, and hence, combines the expressiveness of Diffusion Models with the benefits of Mixture Models. Specifically. VDD leverages decompositional upper bound of the variational objective that allows the training of each expert separately, resulting in robust optimization scheme for MoEs. VDD demonstrates across nine complex behavior learning tasks, that it able to: i) accurately distill complex distributions learned by the diffusion model, ii) outperform existing state-of-thh--rt distillation methods, and iii) surpass conventional methods for training MoE. The code and videos available https://nttuiiiv--roosss wwt.::ooww.giiiibbboooodd-eeiiie
Introduction
Diffusion models 1 have gained increasing attention with their great success in various domains such as realistic image generation (5-8). More recently, diffusion models have shown promise in Learning from Human Demonstrations (LiDs) [Н13]. A particularly challenging aspect of LID is the high variance and multi-modal data distribution resulting from the inherent diversity in human behavior [14]. Due to the ability to generalize and represent complex, multi-modal distributions, diffusion models are particularly suitable class of policy representations for LiD. However, diffusion models suffer from several drawbacks such as long inference time and intractable likelihood calculation. Many diffusion steps are required for nigh-quality samples leading to a long inference time, limiting the use in real-time applications such as robot control, where decisions are needed high frequency. Moreover, important statistical properties such as exact likelihoods are not easily obtained for diffusion models, which poses significant challenge for conducting post hoc
"Corespondence to hongyi www.wkit.ttcccoo edu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

prediction (a)s)
observation
VDD Experts
distillation using score function
diu
[a 1)
(0 Va, log (a. + de (())aw,
Bo
BT
Figure 1: VDD distills diffusion policy into MoE. LfD is challenging due the multimodality of human behaviour. For example, tele-operated demonstrations of an avoiding task often contain multiple solutions [3] Lower: A diffusion policy predict high quality actions but relies iterative sampling process from noise to data, shown as the red arrows. Upper: VDD uses the score function distill a diffusion policy into an MoE, unifying the advantages of both approaches.
optimization such as fine-tuning through well-established reinforcement learning (RL) approaches like policy gradients or maximum entropy RL objectives.
well-studied approach that effectively addresses these issues are Mixture Experts (MoE). During inference, the MoE first selects expert that is subsequently queried for forward pass. This hierarchical structure provides a fast and simple sampling procedure, tractable likelihood computation, and the ability to represent multimodal distributions. These properties make them well-suited policy representation for complex, multimodal behavior. However, training Mixture of Experts (MoEs) often difficult and unstable [15]. The commonly used maximum likelihood objective can lead to undesired behavior due modes. Yet, this limitation has been alleviated by recent methods that use alternative objectives, such as reverse KL-divergence, which do not exhibit mode-averaging behavior [14 [6].
averaging where the model fails to accurately represent certain
To obtain the benefits of both models, i.e., learning highly accurate generative models with diffusion and obtaining simple, tractable models using a mixture of experts, this work introduces Variational Diffusion Distillation (VDD), a novel method that distills diffusion models to MoEs. Starting from the variational inference objective [77 I8]
we derive lower bound that decomposes the objective into separate per-expert objectives, resulting in robust optimization scheme. Each per-expert objective elegantly leverages the gradient of the pre-trained score function such that the MoE benefits from the diffusion model's properties. The resulting MoE policy performs and covers the same modes, while being interpretable, faster during inference, and has tractable likelihood. This final policy is readily available to the user for post hoc analysis or fast fine-tuning for more specific situations. high-level architecture of the VDD model and its relation to the diffusion policy is shown in Fig. D VDD is thoroughly evaluated that demonstrate the afforementioned properties. As an additional insight, this paper observed that one-step continuous diffusion models already perform well, finding not discussed in prior work.
par with the diffusion model
nine complex behavior-leamin
tasks
In summary, this work presents VDD, novel method to distill diffusion models to MoEs, by propos- ing variational objective leading to individual and robust expert updates, effectively leveraging pre-trained diffusion model. The thorough experimental evaluation on nine sophisticated behavior learning tasks show that VDD i) accurately distills complex distributions, ii) oulperforms existing SOTA distillation methods and iii surpasses conventional MoE training methods.
2 Related Works
Diffusion Models for Behavior Learning. Diffusion models have been used in acquiring complex behaviors for solving sophisticated tasks in various learning frameworks. Most of these works train diffusion policies using offline reinforcement learning [19 24, or imitation learning ( 12
25). In contrast, VDD distills diffusion models into
MoE policy to overcome diffusion-

---

policy drawbacks such as long inference times or intractable likelihoods instead of optimizing policies directly from the data.
Mixture of Experts (MoE) for Behavior I Learning. MoF. models are well-studied, provide tractable likelihoods, and can represent multi-modality which makes them a popular choice in many domains such as in imitation learning [14] 26 16 [3), reinforcement learning (32H38) and motion gen-
eration [39 to obtain complex behaviors. Although VDD also uses an MoE model, the behaviors are distilled from a pre-trained model using a variational objective and are not trained from scratch,
The empirical evaluation demonstrates that VDD's s stable training procedure results in improved performance compared to common MoE learning techniques.
Knowledge Distillation from Diffusion Models Knowledge distillation from diffusion models has been researched in various research areas. For instance, in text-to-3D modeling, training NeRF- based text-to-3ID model without any 3D data by mapping the 3D scene to 2D image and leveraging text-to-2D diffusion is proposed 77] The work proposes minimizing the Score Distillation Sampling (SDS) loss that is inspired by probability density distillation [40] and incentivizes the 3D-model to be updated into higher density regions as indicated by the score function of the diffusion model. To overcome drawbacks such as over-smoothing and low-diversity problems when using the SDS loss, variational score distillation (VSD) treats the 3D scene as a random variable and optimizes a distribution over these scenes such that the projected 2D image aligns with the 2D diffusion model [8] In a similar context, the work in [41] proposes distilling a trained diffusion model into another diffusion model while progressively reducing the number of steps. However, even though the number of diffusion steps is drastically reduced, a complete distillation, i.e. one-step inference as for VDD not provided. Additionally, the resulting model suffers from the same drawbacks of diffusion models such as intractable likelihoods. In contrast, in consistency distillation (CD), diffusion models are distilled to consistency models (CM) 452 such that data generation is possible in one step from noise to data. However, one-step data generation typically results in lower sample quality, requiring a trade-off between iterative and single-step generation based on the desired outcome, As CMs. VDD performs one-step data generation but distills the pre-trained diffusion model to an MoE which has a tractable likelihood and is efficient in inference time. The experimental evaluations show the advantages of VDD over CMs. Diff- Instruct [46] proposes two-step framework for distilling diffusion models into implicit generative models, whereas VDD considers an explicit generative model where the model's density can be directly evaluated. In addition, Diff- requires training an auxiliary diffusion model, while VDD only optimizes a single model. Score Regularized Policy Optimization (SRPO) 47 also leverages diffusion behavior policy to regularize the offline RL-based objective. However, in contrast to SRPO, VDD learns an MoE policy instead of a uni-modal Gaussian policy and explicitly distills diffusion model instead of using it as guidance during optimization. Furthermore, VDD trains Možs policies in imitation learning instead of reward-labele offline RL. A concurrent work, EM- EMD-istillatton (EMD) introduces an EM-style distillation objective derived from the mode-covering forward KL divergence. In contrast, VDD proposes an EM-style objective based on the mode-seeking reverse KL but encourages mode-covering behavior by having multiple experts.
data as in
3 Preliminaries
Here, we introduce the notation and foundation for Denoising Diffusion and Mixture of Experts policies. Throughout this work, we assume access to samples from a behavior policy x² and the corresponding state distribution μu, that is a ~ TT" (-Is) and p(-), respectively.
Denoising Diffusion Policies. Denoising diffusion policies employ a diffusion process to smoothly convert data into noise. For given state s', a diffusion process is modeled as stochastic differential equation (SDE) 3
da, = f(a,, + g(t)dw, ap ~ x" (-Is'), s' ~ ((-)
(1)
with drift f, diffusion coefficient g(t and Wiener process w, E Rdd The solution of the SDE is a diffusion process (at) with marginal distributions #f such that का Rs (N I) and mo m².
49
50 showed that the time-reversal of Eq. D again an SDE given by
da, = [f(ar,)) - v1) log मं (ar]s')] di + g(!)dw, ar ~ (o
(2)
Simulating the SDE generates samples from n= (-Is") starting from pure noise. For most distributions n², however, we do not have access to the scores (V a. log म. (a (s) The goal of diffusion-
3

---

(a) Action distribution
(b) Initialize VDD
Training
(d) Training
Training
(i) Convergence
Illustration of training VDD using the score function for fixed state in a 2D toy task. (a) The probability density of the distribution is depicted by the color map. The score function is shown by the gradient field, visualized as white arrows. From (b) to (f), we initialize and train VDD until convergence. We initialize 8 components, each represented by an orange circle. These components are driven by the score function to match the data distribution and avoid overlapping modes by utilizing the learning objective in Eq. 11 Eventually, they align with all data modes.
Figure 2:
based modeling is therefore to approximate the intractable scores using a parameterized score function, i.e., fo 8, & log (a|s). To that end, several techniques have been proposed 52), allowing for sample generation by approximately simulating Eq. 2 The most frequently employed SDEs in behavior learning are variance preserving (VP) 12 and variance exploding (VE) [9]. For further details on diffusion-based generative modeling, we refer the reader to S 4. While we only consider VE and VP this work, VDD can be applied to any score-base method.
Gaussian Mixtures of Expert Policies. Mixtures of expert policies are conditional discrete latent variable models. Denoting the latent variable 2, the marginal likelihood can be decomposed as
(ais) Σ
where (z(s) and q"* (als, =) are referred to as gating and experts respectively. & and denote the gating and expert parameters and thus o { {v=}2. The gating is responsible for soft coft-partitioning the state space into sub-regions where the corresponding experts approximate the target density. To sample actions, that is, a' (-Is') for some state s', we first sample component index from the gating, i.e., z' ~ (-Is') The component index selects the respective expert to obtain a' ~ q"* (-Is', ='). For Gaussian MoE the experts are chosen as g"* (a|s, z) = N (alu" " (s), E (s)), where n²* " could be neural networks parameterized by Vz. From the properties of Gaussian distributions, it directly follows that this model class admits tractable likelihoods and fast sampling routines. Furthermore, given enough components. Gaussian MoEs are universal approximators of densities [53], which makes them a good representation for distillation of diffusion policies.
4 Variational Distillation of Denoising Diffusion Policies
In this section, we outline the mathematical formulations of the VDD model. Detailed descriptions of the model architecture and algorithms can be found in Appendix [ We aim to distill a given diffusion policy π (l))) by using a different policy representation ge with parameters Φb. This useful, e.g.. (a(s) has favorable properties such as likelihood tractability or admits fast inference schemes. Assuming that we can evaluate श point-wise, a common approach is to leverage variational inference (VI) to frame this task as an optimization problem by minimizing the reverse Kullback-Leibler (KL)
54 divergence between g and RT, that is.
min (KL(g (als') (T((ss) = min Egg(al>)) [log (als') - log π (als')]
(4)
for specific state s'. To obtain scalable optimization scheme, we combine amortized and stochastic VI [55]. The former allows for learning a conditional model (als) instead of learning separate 9d for each state, while the latter allows for leveraging mini-batch computations, that is.
min J (() = min En(s) DKL (als) so min
M N
Σ (allo (5) 8:rip
with batch size M < N. Thus, (()) can be minimized using gradient-basee optimization techniques with a gradient estimator such as reinforce [56] 57) or the reparameterization trick (58). Note that, while the states are sampled from the given data set of the behavior policy, the actions needed to

---

evaluate the KL are generated using our estimated model (als:). Yet, there are two difficulties to directly apply this scheme in distilling a diffusion model into an MoE: i) we are not able to evaluate the likelihood x (alss) of a diffusion model, ii) training of MoE models is notoriously difficult [15 We will address these two issues in Section 4.1 and Section respectively.
4.1 Scalable Variational Inference for Denoising Diffusion Policy Distillation
Although we cannot directly evaluate the likelihood of the diffusion policy T r[als), we have access to its score functions Va, log Te (a (s) fo (a,, 8, 1), where E [o, T is the diffusion time step. In prac- tice, we would like to evaluate the score in the limit of + 0 as V a log n (als) R lim;- Se a 8, (). Yet, this might lead to an unstable optimization [59 this score often not estimated well through- out the action space, and, hence other diffusion time-step selection processes are needed [47]. For now, we will omit the diffusion time-step for the sake of simplicity and refer to Section 4.3 for detailed discussion about time-step selection. Moreover, we will, for now, assume that the parametrization of qº amendable to the reparameterization trick [58]. In this case, we can express (-[s) using transformation h (e, s) with auxiliary variable p(-) such that (e s). We then express the gradient of J w.r.t. Φ as
M
V JΦ)
Σ
Bp(e)
- V log x (h 8;)ls:]
(6)
N
Using the chain rule for derivatives, it is straightforward to see that
V log (als,) (Va log T(als:)) V (c.s) s, (e, As Va log π (a)s:) (s)) can be replaced by the given score of the pre-trained diffusion policy, we can directly use of VI for optimizing J without evaluating the likelihoods æ.
(7)
4.2 Variational Inference via Mixture of Experts
To distill multimodal distributions learned by diffusion models, we require a more complex family of distributions than conditional diagonal Gaussian distributions, which are commonly used in amortized We will therefore use Gaussian mixture of experts. To that end, we construct an upper bound of which is decomposable into single objectives per expert, allowing for reparameterizing each expert individually and therefore avoiding the need for techniques that perform repurameterization for the entire MoE [60] 61. The upper bound (() g) can be obtained by making use of the chain rule for
KL divergences 662 63 15), i.e.,
((,) U( ② (o DKL (z/a, ()
(8)
where q is an arbitrary auxiliary distribution and upper bound U (6)
E B(((s) (x)s) (alss)) [log q²* (als, 2) log x (als) - log q(= a s)) + log of (2|s)]].
(9)
(g()..
with (p 9) being the objective function for single expert and state For further details Appendix B Since the expected KL term on the right side of Eq. is always positive, it directly follows that U is an upper bound on J for any q. This gives rise to an optimization scheme similar to the expectation-maximizattte maximization algorithm [64], where we alternate between minimization (M-Step) and tightening of the upper bound U (E-Step), that is,
min ( q)
and
min B DKL (2la, {)
(10)
respectively. Please note that is fixed during the M-Step and o during the E-Step. In what follows, we identify the M-Step as a hierarchical VI problem and elaborate on the E-Step.
M-Step for Updating the Experts. The decomposition in Eq. g allows for optimizing each expert separately. The optimization objective for a specific expert and state s, that is, Ug ÿ), is
min U2(V. 9) = min Ege (a x) Ig (als, 2) log ा (als) log ((2|a,
(11)
Note that this objective corresponds to the standard reverse KI. objective from variational inference (c.f. Eq. 4, with an additional term log ((z|a, la s), which acts as repulsion force, keeping the
5

---

individual components from concentrating on the same mode. Assuming that q is differentiable and following the logic in Section[4.1] it is apparent that the single component objective Ta be optimized only by having access to scores V, log (als). Moreover, we can again leverage amortized M N Es. U (V:. 9)-
stochastic VI, that is miny, Bys)) [U] (V:. Q)] z miny,
M-Step for Updating the Gating. The M-Step for the gating parameters, i.e., minimizing Ü 7) with respect to { C Φ given by
min U(6 0) max (e (x]s) U3( {
(12)
Using gradient estimators such as reinforce [57], requires evaluating (12, q) which is not possible as we do not have access to log π (als). This motivates the need for a different optimization scheme. We note that of is a categorical distribution and can approximate any distribution over 2. does therefore not suffer from problems associated with the forward KL divergence such mode averaging due to limited complexity of qs. We thus propose using the following objective for optimizing &, i.e.,
min Eu(s) DKL (I (1 (als)) = max (x (alls) (q(1),s)))))) [log [ (=[s)] {
(13)
resulting in cross-entropy loss that does require evaluating the intractable log raals). Moreover, we note that using the forward KI. does not change the minimizer as
5 = arg min Eacc)) DKL. (m = arg min Epts) DKI. (alss)
(14)
and, therefore does not affect the convergence guarantees of our method. Please note that the forward KL requires samples from the teacher model æ. In practice, if the dataset used for training the diffusion model is available, it can be used as proxy and prevent costly data regeneration.
E-Step: Tighening the Lower Bound. Using the properties of the KI. divergence, it can easily be seen that the global minimizer of the E-Step, i.e., the optimization objective defined in Eq. 10 can be found by leveraging Bayes' rule, i.e.,
(als, 2)4 COM (2|s)
7()aa, g0dh (zla,s) = E. gdd (a/s, s, 2 2)9 (2|s)
(15)
The superscript 'old' refers the previous iteration. Numerically, Gold easily be obtained by using a stop-gradient operation which is crucial as ② is fixed during the subsequent M-step which requires zero after this update, the upper
blocking the gradients of q with respect to Φb As the KL bound is tight after every E-step ensuring Ey(e) DKL (q (a|s)][z* (ais)) = U U(d g). Hence, VDD has similar convergence guarantees to EM, i.e., every update step improves the original objective.
4.3 Choosing the Diffusion-Timester
In demoising diffusion models, the score function is usually characterized as a time-dependent function V a, log (a(ls)) = fo(au, [a,, s, t), where t is the diffusion timestep. Yet, the formulation in Section 4.1 only leverages the pretrained diffusion model at time + O. However, [47]showed that using an ensemble of scores from multiple diffusion time steps significantly improves performances. We, therefore, replace Eq. [1 with a surrogate objective that utilizes scores at different time steps, that is,
min (V: ! min (al=s)))).)) Eppc) [q" (a|s, 2) - log (a (s, t) log (= s)] Pz
(16)
Ps
with p(t) being distribution on [0, T}. Furthermore, we provide an ablation study for different time step selection schemes and empirically confirm the findings from 47
5 Experiments
We conducted imitation learning experiments by distilling two types of diffusion models: variance preserving (VP) 12 and variance exploding (VE) 65 We selected DDPM as the representative for VP and BESO as the representative for VE. We adopt the choices of samplers and the number of denoising steps in ( and [3]] Additional evaluation of teacher models with different numbers of denoising steps can be found in Appendix F In the experiments, VP-1 and VE-1 denote the

---

VP (DDPM)
VE (BRSO)
VP-1
YE-1
CD-VR
CTA-YF
VDD-VPlOuRSI
Kitchen
4.06
0.22
3.87
8.89 0.11
0.112
3.85 + 0.10
0.96
0.99
0.94
0.89
0.89
0.93 0.03
0.91 0.03
Avoiding
0.84
0.96
0.99
0.84
0.82 = 0.05
0.93 0.02
0.92 = 0.02
0.95 + 0.01
Aligning
0.90
0.93
0.08
0.81 0.11
0.07
Pushing
0.74
0.78
0.00
0.70
0.66 1: 0.05
0.80 0.07
0.61 1: 0.04
0.85 + 0.02
Stucking-1
0.75
0.11
0.02
Stacking-2
0.68
0.70
0.00
0.63
0.46
0.11
0.30 2 0.00
0.60 + 0.07
0.07
0.57 = 0.06
Sorting (Iusage)
0.80 0.04
0.04
Stacking (Image)
0.58
0.86
0.00
0.58
0.63 4: 0.01
0.59 0.10
0.78 +: 0.02
0.60 +: 0.04
(a) Task Success Rate (or Environment Return for Kitchen)
(DDPM)
(BESO))
VI-I
VE-1
CD-VE
CTM-VE
VDD-VFourr)
VDD-VElours)
Avoiding
0.60
0.87
0.25
0.76
0.002
0.79 + 0.04
0.37 I 0.01
0.72 0.12
Aligning
0.90
0.34
0.32 0.14
0.31 0.28
- 0.09
0.40 0.04
Pushing
0.74
0.76
0.00
0.50
0.53 1033:000 0.07
0.54 2 0.08
0.66 4 0.05
0.69 1 0.08
Stacklug-1
0.24
0.26
0.19 - 0.05
0.16
Stucking-2
0.12
0.13
0.00
0.07
0.03 : 0.05
0.09 1 0.06
0.07 2: 0.04
0.13 0.08
Sorting (Image)
0.16
0.19
0.09
0.14
0.14
0.06
0.08 0.05
0.12 : 0.03
0.22 - 0.03
Stacking (Insage)
0.15
0.00
0.30
0.06
0.01
0.04 0.04
+ 0.02
0.11 0.08
(b) Task Entropy
Table Comparison of distillation performance, (a) VDD achieves on-par performance with Consis- tency Distillation (CD) (b) VDD is able to possess versatile skills (indicated by high task entropy) while keeping high success rate. The best results for distillation are bolded, and the highest values except origin models are underlined. In most tasks VDD achieves both high success rate and entropy. Note: to better compare the distillation performance, we report the performance of origin diffusion model, therefore only seed 0 results of diffusion models are presented here.
results when performing only one denoising step of the respective diffusion models during inference. VDD-VP and VDD-VE denote the results of distilled VDD Additionally, we consider the SoTA Consistency Distillation (CD) 42 and Consistency Trajectory Model (CTM) 44 as baselines for comparing VDD's performance distillation. For CD and CTM, we distill from the VE following the original works. For CTM we adapt the implementation and design choices from Consistency Policy 45). which are specialized for behavior learning. We compare VDD against MoE learning baselines, namely the widely-used of the maximum likelihood based objective and the recently introduced So'TA method Information Maximizing Curriculum (IMC) as representative of the reverse KL-based objective. To make them stronger baselines, we extend them with the architecture described in Figure and name the extended methods as EM-GPT and IMC-GPT, respectively. For fair comparison, we used the same diffusion models as the origin model for all distillation methods that we have trained on seed 0. For statistically significant comparison, all methods have been run on random seeds, and the mean and the standard deviation are reported throughout the evaluation. Detailed descriptions regarding the baselines implementation and hyperparameters selection can be found in Appendix D and E
(EM) [66] approach representative
The evaluations are structured as follows. Firstly, we demonstrate that VDD is able to achieve competitive performance with the SoTA diffusion distillation method and the original diffusion models on two established datasets. Next, we proceed to a recently proposed challenging benchmark with human demonstrations, where VDD outperforms existing diffusion distillation and SoTA MoE learning approaches. We then highlight the faster inference time of VDD, Following this, a series of ablation studies reflect the importance of VDD's s essential algorithmic properties. Finally, we provide a visualization to offer deeper insights into our method.
5.1 Competitive Distillation Performance in Imitation Learning Datasets
We first demonstrate the effectiveness of VDD using two widely recognized imitation learning datasets: Relay Kitchen [67] and XArm Block Push 68]. A detailed description of these environments is provided in Appendix C To ensure a fair comparison, we follow the same evaluation process as outlined in [9]. The environment rewards for Relay Kitchen and the success rate for XArm Block Push are presented in Table Ia with mean and standard deviation resulting from 100 environment rollouts. The results indicate that VDD achieves performance comparable to CD both tasks, with slightly better outcomes in the block push dataset. An additional interesting finding is that BESO, with only one denoising step (VE-1), already proves to strong baseline these tasks, as the original models outperformed the distillation results in both cases. We attribute this interesting
7

---

Environments
EM-GPT
INC-GPT
VDD-VP
VDD-VE
EM-GPT
IMC-GPT
VDD-VP
VDD-VE
Avoiding
0.65.10..0
0.75 = 0.08
0.92 a 0.02
0.95 2: 0.01
0.17 0.13
0.82 0.05
0.37 4 0.01
0.73 A 0.09
Aligning
0.78 4: 0.04
0.83 =.00=0.02
0.70 0.07
0.8G +: 0.04
0.38 0.111
0.27 :: 0.00
0.25 + 0.09
0.40 :: 0.04
Pushing
0.36
0.76
0.61 0.04
+
0.14 0.10
+ 0.03
0.66 0.05
0.69 0.08
Stucking
0.58
0.54 0.05
0.81 0:08
0.83 1: 0.09
0.43 0.08
0.37 0.04
0.19 + 0.05
0.16 0.03
Stacking-2
0.34 0.07
0.29 = 0.07
0,60 +: 0.07
0.57 + 0.06
0.27 +: 0.03
0.17 + 0.07
0.07 + 0.04
0.13 + 0.06
Sorting (image)
0.60 + 0.02
0.74 = 0.04
0.80 +: 0.04
0.76 + 0.03
0.13 + 0.03
0.10 0.03
0.12 + 0.03
0.22 + 0.03
Stacking (image)
0.04
0.30
0.60 +: 0.04
0.00
0.08
0.11
Relay Kitchen
3.62 +
3.67 0.05
3.24 + 0.12
3.85
Block Push
0.88 6.8-++0.04 0.04
0.89 0.04
0.98 + 0.03
0.91 + 0.03
Table 2: Comparison between VDD and SoTA MoE approaches, with left: success rate and right: entropy. VDD consistently outperforms EM and IMC in terms of task success. For behavior versatility, VDD outperforms in 4 out of 7 D3IL tasks.
NFE
8
32
64
NFE
16
VE (BESO)
4.03
10.38
32.14
60.64
VE (BESO)
9.69
12.72
16.39
23.68
VP (DDPM)
2.15
8.25
29.47
55.62
VP (DDPM)
6.49
9:33
12.79
19.90
VDD (Ours)
2.16
VDD (Ours)
4.91
Table 3: Inference time in state-based pushing (left) and image-based stacking (right). The gray shaded area indicates the default setting for diffusion models.
observation to the possibility that the Relay Kitchen and the XArm Block Push tasks are comparably easy to solve and do not provide diverse, multi-modal data distributions. We therefore additionally evaluate the methods on more recently published dataset (D3IL) [33 which is explicitly generated for complex robot imitation learning tasks and provides task entropy measurements.
5.2 Replicating Diffusion Performance while Possessing Versatile Behavior
The D3IL benchmark provides human demonstrations for several challenging robot manipulation tasks, focusing on evaluating methods in terms of both success rate and versatility. i.c. the different behaviors that solve the same task. This benchmark includes a versatility measure for each task, referred to as task entropy. Task entropy is a scalar value ranging from 0 to 1, where O indicates the model has only learned one way to solve the task, and indicates the model has covered all the skills demonstrated by humans. Detailed descriptions of the environments and the calculation of task entropy are provided in the Appendix C The task success rate of the distilled polices is presented in Table la The results show VDD outperforms consistency distillation and 1 1-step variants of origin models in out of tasks except the Aligning. However, in the Aligning task VDD achieves higher task entropy, indicating more diverse learned behaviors. The task entropy is presented in Table Tb The results demonstrate that VDD achieves higher task entropy compared both consistency models (CD, CTM) and the 1-step diffusion models (VP-1, VE-I) in 4 out of 7 tasks, which shows that our method replicates high-quality versatile behaviors from diffusion policies.
5.3 Comparison with MoE learning from scratch
The previous evaluation demonstrated that VDD can effectively distill diffusion models into MoEs, preserving the performance and behavioral versatility. In this section, we discuss the necessity of using VDD instead of directly learning MoEs from scratch by comparing VDD against EM-GPT and IMC-GPT. Both methods train MoE models from scratch but differ in their objectives. While EM is based on the well-know maximum likelihood objective, IMC is based on a reverse KL objective. show that VDD consistently outperforms both, EM-GPT and IMC-GPT across
The results in Table 2 a majority of all tasks in terms of both success rate and task entropy. We attribute the performance boost leveraging the generalization ability of diffusion models and the stable updates provided by our decomposed lower bound ETI
5.4 Fast Inference with distilled MoE
Inference with MoE models do not require an iterative denoising process and are therefore faster in sampling. We evaluate the inference time of VDD against DDPM and BESO on the state-based pushing task and the image-based stacking task and report the average results from 200 predictions in Table 3 In addition to the absolute inference time in milliseconds, we report the number of function
8

---

Success
entropy
wro guting
with gating
t wio guzing
with gating
success extrppy
2
Avoiinng
Algning
Stecking
Aveidng
Aligaing
Staskng
Number of Experts
mil
100
(0) Avoiding: number of experts
(b) Success rate, w/o gating
(c) Entropy, w/o gating
(d) Choices of noise level
Figure 3: Ablation studies for key design choices used in VDD. (a) Using only expert leads to higher success rate but is unable to solve the task in diverse manners. Sufficiently more experts can trade off task success and action diversities. Leamnng the gating distribution improves the success rates in three D3IL tasks. (c) A Uniform gating leads to higher task entropy in three out of two tasks. (d) Sampling the score from multiple noise levels leads better distillation performance
evaluations (NFE) in Table 3 for better comparability. The results show that VDD is significantly faster than the original diffusion models in both cases, even when the diffusion model takes only one denoising step. For a fair comparison, all methods used an identical number of transformer layers. The predictions were conducted using the same system (RTX 3070 GPU, Intel i7- CPU).
5.5 Ablation Studies
We assess the importance of VDD's key properties on different environments by reporting the task performance and task entropy averaged over four different seeds.
Number of experts matters for task entropy. We start by varying the number of experts the MoE model while freezing all other hyperparameters on the avoiding task. Figure Saowwwss the average task success rate and task entropy of MoE models trained with VDD. The success rate almost constantly high for all numbers of experts, except for the single expert (i.e. a Gaussian policy) case which shows slightly higher success rate. However, the single expert can only cover single mode of the multi-modal behavior space and hence achieves a task entropy of O. With an increasing number of experts, the task entropy increases and eventually converges after small drop.
Training a gating distribution boosts performance. Figure 36 shows the success rates when training parameterized gating network (=|s) (red) and when fixing the probability of choosing expert 2 to (=) = 1/N, where N is the number of experts (blue). While training a gating distribution increases the success rate over three different tasks, the task entropy (see Figure 3cb slightly decreases in two out of three tasks. This observation makes sense as the MoE with a trained gating distribution leads to an input-dependent specialization of each expert, while the experts with fixed gating are forced to solve the task in every possible input.
Interval time step sampling increases task entropy. Here, explore different time step distribu- tions p()), as introduced in Eq. We consider several methods in Fig. 3d using the minimum time step, i.e., p(t) lim (t), where denotes Dirac delta distribution, the maximum time step p(t) = 5(T), a uniform distribution on [0, and on sub-intervals (4a, 4 C [0. 7] with interval bounds to 1 being hyperparameters. While success rates
comparable across the variants, inter- val sampling yielded the highest task entropy with very high success rates. Thus, interval time-step sampling is adopted our default setting. The results were obtained from the avoiding task.
5.6 Visualization of the per-Expert Behavior
We provide additional visualizations on the Avoiding task from the D3IL task suite aiming to provide further intuition on how VDD leverages the individual experts. according to the likelihood of the gating distribution at a given state, offering several key insights. First, VDD effectively distills experts with distinct behaviors, e.g., 21 typically moves downward, 22 tends to move upward, while 2y and tend to generate horizontal movements. Second, the gating mechanism effectively deactivates redundant experts (=6, 27; #8) most states, demonstrating that a larger number of components can be used without harming performance, as the gating mechanism deactivates redundant experts. Lastly, using single component (Z 1) can achieve a perfect success rate at the cost of losing behavior diversity. On the contrary, using many experts potentially results in
the expert selection

---

slightly lower success rate but increased behavior diversity. These qualitative results are consistent with the quantitative results from the ablation study presented in Figure
z-1
Z 2
Z Z=
Z 8
( Figure 4: Trajectory visualization for VDD with different number of components 2 E (1, 2, 8] on the Avoiding task (left). Different colors indicate components with highest likelihood according to the learned gating network of (z|s) at a state s. For each step we select the action by first sampling an expert from the categorical gating distribution and then take the mean of the expert prediction. We decompose the case Z = 8 and visualize the individual experts 2 (bottom row). Diverse behavior emerges as multiple actions are likely given the same state. For example, moving to the bottom right (21) ) and top right (z2). An extreme case of losing diversity is seen with Z = 1, where the policy is unable to capture the diverse behavior of the diffusion teacher, leading to deterministic trajectories.
6 Conclusion
This work introduced Variational Diffusion Distillation (VDD), a novel method that distills a diffusion model to an MoE. VDD enables the MoE to benefit from the diffusion model's properties like generalization and complex, multi-modal data representation, while circumventing its shortcomings like long inference time and intractable likelihood calculation. Based on the variational objective, VDD derives lower bound that enables optimizing each expert individually. The lower-bound leads to stable optimization and elegantly leverages the gradient of the pre-trained score function such that the overall MoE model effectively benefits from the diffusion model's properties. The evaluations on nine sophisticated behavior learning tasks show that VDD achieves on-par or better distillation performance compared to SOTA methods while retaining the capability of learning versatile skills. The ablation the number of experts reveals that single expert is already performing well, but can not solve the tasks in a versatile manner. Additionally, the results show that training the gating distribution greatly boosts the performance of VDD, but reduces the task entropy.
Limitations. VDD is not straightforwardly applicable to generating very high-dimensional data like images due to the MoE's contextual mean and covariance prediction. Scaling VDD to images requires further extensions like prediction in a latent space. Additionally, the number of experts needs to be pre-defined by the user. However, redundantly high number of experts could increase VDD's s training time and potentially decrease the usage in post hoc fine-tuning using reinforcement learning. Similar to other distillation methods, the performance of VDD is bounded by the origin model.
Future Work. 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 iine-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.
Broader Impact. Improving and enhancing imitation learning algorithms could make real-world applications like robotics more accessible, with both positive and negative impacts. We acknowledge that it falls on sovereign governments' responsibility to identify these potential negative impacts.
Acknowledgement
We thank Moritz Reuss for the valuable discussions and technical support. H.Z. and R.I., acknowl- edges funding by the German Research Foundation (DFG) 448648559. D.B. is supported by
10

---

funding from the pilot program Core Informatics of the Helmholtz Association (HGF). G.L. is supported in part by the Helmholtz Association of German Research Centers. G.N. was supported in part by Carl Zeiss Foundation through the Project JuBot (Jung Bleihen mit Robotern). The authors also acknowledge support by the state of Baden- Baden- Württemberg through KoreKa supercomputer funded by the Ministry of Science, Research and the Arts Maden-Württemberr and by the German Federal Ministry of Education and Research.
References
[1] Jascha Bohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages
PMLR, 2015.
[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688, 2020.
[3] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
[4] Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of disfusion-based generative models. Advances in Neural Information Processing Systems, 35: 2022.
[5] Alexander Quinn Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob Megrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image genera- tion and editing with text-guided diffusion models. In International Conference on Machine Learning, pages 16784-16804, PMLR, 2022.
[6] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents.
[7] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. In The Eleventh International Conference on Learning Representations, 2022.
[8] Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Pro- lifiedreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[9] Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal conditioned imitation learning using score-based diffusion policies. In Robotics: Science and Systems, 2023.
[10] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
[11] Paul Maria Scheikl, Nicolas Schreiber, Christoph Haas, Niklas Freymuth, Gerhard Neumann, Rudolf Lioutikov, and Franziska Mathis-Ullrich. Movement primitive diffusion: Learning gentle robotic manipulation of deformable objects. IEEE Robotics and Automation Letters, 2024.
[12] Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, et al. Imitating human behaviour with diffusion models. In The Eleventh International Conference on Learning Representations (ICLR 2023), 2023.
[13] Xiaogang Jia, Denis Blessing, Xinkai Jiang, Moritz Reuss, Atalay Donat, Rudolf Lioutikov, and Gerhard Neumann. Towards diverse behaviors: benchmark for imitation learning with human demonstrations. arXiv preprint wrXiv:2402.44600 2024.
[14] Denis Blessing, Onur Celik, Xiaogang Jia, Moritz Reuss, Maximilian Li, Rudolf Lioutikov, and Gerhard Neumann. Information maximizing curriculum: curriculum-based approach for learning versatile skills, Advances in Neural Information Processing Systems, 36, 2024.

---

[15] Oleg Arenz, Philipp Dahlinger, Zihan Ye, Michael Volpp, and Gerhard Neumann. A unified perspective on natural gradient variational inference with gaussian mixture models. arXiv preprini arXiv:2209. 1533, 2022.
[16] Philipp Becker, Oleg Arenz, and Gerhard Neumann. Expected information maximization: Using the i-projection for mixture density estimation. In International Conference on Learning Representations, 2019.
[17] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):55-887
[18) Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Beyond elbos: large-scale evaluation of variational methods for sampling. arXiv preprini arXiv:2406.07423, 2024.
[19] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pages 9902-9915. PMLR, 2022.
[20] Zhendong Wang. Jonathan J Hunt, and Mingyuan Zhou Diffussion policies as an expres- sive policy class for offline reinforcement learning. International Conference on Learning Representations, 2023.
[21] Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative hehavior modeling. In The Eleventh International Conference on Learning Representations, 2023.
[22] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
[23] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, 36, 2023.
[24] Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. Advances in Neural Information Processing Systems, 36, 2023.
[25] Ajay Sridhar, Dhruv Shah, Catherine Glossop, and Sergey Levine. Nomad: Goal masked diffusion policies for navigation and exploration. arXiv preprint 10.00223 2023.
[26] Maximilian Xiling Li, Onur Celik, Philipp Becker, Denis Blessing, Rudolf Lioutikov, and Gerhard Neumann. Curriculum-base imitation of versatile skills. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 2951-2957. IEEE, 2023.
[27] Niklas Freymuth, Nicolas Schreiber, Aleksandar Taranovic, Philipp Becker, and Gerhard Neumann. Inferring versatile behavior from demonstrations by matching geometric descriptors. In 6th Annual Conference on Robot Learning, 2022.
[28] Katharina Mülling. Jens Kober, Oliver Kroemer, and Jan Peters. Learning to select and generalize striking movements in robot table tennis. The International Journal of Robotics Research, 3263)263-- 2013.
[29] Marco Ewerton, Gerhard Neumann, Rudolf Lioutikov, Heni Ben Amor, Jan Peters, and Guil- herme Maeda. Learning multiple collaborative tasks with a mixture of interaction primitives. In 2015 IEEE International Conference on Robotics and Automation (ICRA), pages 1535-1542. IEEE, 2015.
[30] You Zhou, Jianfeng Gao, and Tamim Asfour. Movement primitive learning and generalization: Using mixture density networks. IEEE Robotics & Automation Magazine, 27(2):22 32, 2020.
[31] Vignesh Prasad, Alap Kshirsagar, Dorothea Koert Ruth Stock-Homburg. Jan Peters, and Georgia Chalvatzaki. Moveint: Mixture of variational experts for learning human-robot interactions from demonstrations. IEEE Robotics and Automation Letters, 2024.
12

---

[32] Riad Akrour, Davide Tateo, and Jan Peters. Continuous action reinforcement learning from a mixture of interpretable experts. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44:6795-6806, 2020. URL htl.bppppe//p....eenntttccoooooo waapoolaarrrrrooooopeee C 219558472
semanticscholarr
[33] Xue Bin Peng, Michael Chang, Grace Zhang, Pieter Abbeel, and Sergey Levine, Mcp: Learning composable hierarchical control with multiplicative compositional policies. Advances in Neural Information Processing Systems, 32, 2019.
[34] Jie Ren, Yewen Li, Zihan Ding, Wei Pan, and Hao Dong. Probabilistic mixture-of-experts for efficient deep reinforcement learning. arXiv preprint arXiv:2104.09112, 2021.
[35] Onur Celik, Dongzhuoran Zhou, Ge Li, Philipp Becker, and Gerhard Neumann. Specializing versatile skill libraries using local mixture of experts. In Conference on Robot Learning, pages 1433. PMLR, 2022.
[36] Onur Celik, Aleksandar Taranovic, and Gerhard Neumann. Acquiring diverse skills using curriculum reinforcement learning with mixture of experts. arXiv preprint arXiv:2403.06966 2024.
[37] Ahmed Hendawy. Jan Peters, and Carlo D'Eramo. Multi-task reinforcement learning with mixture of orthogonal experts. In The Twelfth International Conference on Learning Represen- tations, 2023.
[38] Samuele Tosatto, Georgia Chalvatzaki, and Jan Peters. Contextual latent- movemmtss off- off-policy optimization for robotic manipulation skills. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 10815-10821. IEEE, 2021.
[39] Kay Hansel, Julen Urain, Jan Peters, and Georgia Chalvatzaki. Hierarchical policy blending as inference for reactive robot control. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 10181-10188. IEEE, 2023.
[40] Aaron Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George Driessche, Edward Lockhart, Luis Cobo, Florian Stimberg, et al. Parallel wavenet: Fast high-fidelity speech synthesis. In International conference on machine learning, pages 3918-3926. PMLR, 2018.
[41] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In International Conference on Learning Representations, 2021.
[42] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Proceedings of the 40th International Conference on Machine Learning, pages 3221 1-32252, 2023.
[43] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In The Twelfth International Conference on Learning Representations, 2024. URL https
/openreview net/forummid ppporrrriew..et//ouuuut-e NN22999RDND
[44] Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. In The Twelfth International Conference on Learning Representations, 2023.
[45] Aaditya Prasad, Kevin Lin, Jimmy Wu, Linqi Zhou, and Jeannette Bohg. Consistency policy: Accelerated visuomotor policies via consistency distillation. Robotics: Science and Systems, 2024.
[46] Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang. Diff- instruct: universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
[47] Huayu Chen, Cheng Lu, Zhengyi Wang, Hang Su, and Jun Zhu. Score regularized policy optimization through diffusion behavior. The Twelfih International Conference on Learning Representations, 2024. URL /
13

---

[48] Sirui Xie, Zhisheng Xiao, Diederik P Kingma, Tingbo Hou, Ying Nian Wu, Kevin Patrick Murphy, Tim Salimans, Ben Poole, and Ruiqi Gao. Em distillation for one-step diffusion models. arXiv preprint arXiv:2405. 2024.
[49] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-322 1982.
[50] Uirich G Haussmann and Erienne Pardoux. Time reversal of diffusions. The Annals of Probability, pages 188-1205, 1986.
[51] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in Artificial Intelligence, pages 574-584, PMLR, 2020.
[52] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. Advances in neural information processing systems, 33:12438-12444 2020.
[53] Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, Cambridge, MA, USA, 2016. http::/www. http:://ww..eeppeeenntttit http///wwww..p.ee..ang.oooo.ooo
[54] Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. Advances in Neural Information Processing Systems, 34:2796-2809, 2021.
[55] Matthew D Hoffman, David M Blei, Chong Wang. and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 2013.
[56) Peter W Glynn. Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10):75-84 1990.
[57] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-256, 1992.
[58] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint 13Xiv:1312.611 14, 2013.
[59] Valentin De Bortoli, Michael Hutchinson, Peter Wirnsherger, and Arnaud Doucet. Target score matching. arXiv preprini arXiv:2402.08666 2024.
[60] Alex Graves. Stochastic backpropagation through mixture density distributions. arXiv preprint 1rXiv:1607.05699 2016.
[61] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv: 1.00001 144, 2016.
[62] Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
[63] Oleg Arenz, Gerhard Neumann, and Mingjun Zhong. Efficient gradient-free free variational inference using policy search. In International conference on machine learning, pages 234-243. PMLR, 2018.
[64] Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series B (methodological), 39 (1):1-22, 1977.
[65] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[66] Todd K Moon. The expectation-maxiiizaaioo algorithm. IEEE Signal processing magazine, 13 (6):47-60, 1996.
[67] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizan tasks via imitation and reinforcement learning. arXiv preprint 1910v:1100.11955 2019.
[68] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 158-168. PMLR, 2022.
14

---

A VDD Architecture and Algorithm Box
Algurithm VDD training
Require: observations actions A. pretrained denoising sevre function (s
Soft- max
f
Initialize MoB:
(als) (ala))
Gating Net
Mean Net
Cov Net
1,2.....M
Number of
Experts Z.
M-Step: Update Experts:
4: updatt (alaa, :) with Eq. 16
Transformer
5: E-Step: Update Gating:
Positional Encoding
compute gating targets
State Encoder
update of (=(s) using cross-entropy loss
8: end for
9: Output learned Moli yf (a)s)
Figure 5: VDD's architecture
Figure 6: Task Envs.
The architecture of VDD is depicted in Fig. 5 and the algorithm box is given in Algorithm D Align with the SoTA diffusion policies' architecture, such as [2] we leverage transformer architecture for VDD to encode a short observation history, seen are first encoded using a state encoder: a linear layer for state-based tasks or pre-trained ResNet- 18 for image-based based tasks. Positional encodings are added to the encoded sequence, which is then fed into a decoder-only transformer. The last output token predicts the parameters the MoE, including the mean and covariance of each expert, well the gating distribution for expert selection. This parameterization enables VDD to predict all experts and the gating distribution with single forward pass, enhancing inference time efficiency.
sequence of states {s The state sequences
B Derivations
B.1 Derivation of the Variational Decomposition (Eq. 8
Recall that the marginal likelihood of the Mixture of Experts is given
(a Σ of (als,,
(17)
where z denotes the latent variable, (2|s) and g" (als, 2) are referred to gating and experts respectively.
The expected reverse Kullback-Leibbe (KL) divergence between dº and is given
min E)(s) DKL (Kls( = min Ey(s) (as(a))) [log 8 (als) log rals))
(18)
= min J(o).
(19)
where
((6) u(s) / (als) [log q (als) log (aas)) dads.
(20)
We can write
[ ((s)) (=]s)) J g²r (als, 2) [log (a)s) - log न (a|s)] dads,
((
(21)
where we have used the definition of the marginal likelihood of the Mixture of Experts in Eq. With the identity
(als) (aals)g-"(als.......
(22)
(z(s,a)
15

---

we can further write
((Ø) 1 pls) E (2)ss) / qr (a|s, (fog((22)s + go (als, 2) log (2
log π (a(s)) dads.
(23)
We can now introduce the auxiliary distribution ğ(z (iz/a,p5 s) by adding and subtracting as
( = 1 ((s)) Σ of (z|s) / q"r (als, z) llog of (z|s) + log q" (als, 2) - log 9 (2)s,
- log π (als) + log (z]a,,s) - log ((z a. s) dads.
(24)
We can rearrange the terms such that
((0) = [ (ps) Σ (zls)) f qt (als,, 2) (logo (z|s) log 2) log (aass)
(uss) Σ (2|ss / q" (als, [log (()
log q(2) la s)] dads
log 8 (= s, a)] dads.
(25)
With can plug in the identity
(als,, (2lss) into the second sum and obtain
(26)
g((2|s)
(o) = [ p(s) Σ of (z|s) /. q" (als, [lo (z|s) + log q²= (als, 2) logm(a)s) log q(= aa s)] dads
+ 1 | (als) (z|s, [logū(z|a,s) log (=|s, a)] dads.
(27)
which is the expected negative KL as
(d) [m(S) p(s) Σ of (z|s) / q" (als, [log (x(s) log (als, 2) log ((ass) log @(=|a, s)] dads
- Eyy (als) (g (
(28)
note that
U( = / (u(s) Σ {f / (als, [log + log "ogg"*(asss =) - log (a
log (= a, s)] dads,
(29)
such that we arrive to the identical expression as in Eq. 8
( U ( q) (p(s) (9 ((ll((a)) (
(30)
B.2
Derivation of the Gating Update (Eq. 13
First, we note that
DEL (als)) = Er(als) [logs(a|s)] - En(als) [log g (ais))
(31)
we optimizing w.r.t d, write const. Ex(als) [log T (a|ss)
(au(T((ass)lo = Er(aas) (als) + const.
(32)
16

---

We can now introduce the latent variable
(kL( = - (alss)
Σ a(2la, (a, log logg (als)
+ const.
(33)
We use the identity in Eq. 22 to arrive at
DRL (TTass))
(als))
Σ ((zla,s) (log"" 2) + log log (2|s) (=(s,
(34)
const.
-aax(als)
Σ s) (log 4"' (als, 2) + logs log a))
+ log ((==a,,s)
log ((=|(((]))))) s)] + const.
(35)
--a(aa)s)
Σ 7(z|a, s) (als, + log (=|s) log ((z 8))
(KI s) (z|s, + const.
(36)
Since DRL. ( a)) 2 O we have the upper bound
U(o. ②) = (alsss))
Σ q(=|a, (log (als, + log {f(ss) log (z)a((/a,)))
(37)
We can now write
Ú ②) -Er(sss) Σ q(z|a, s) log (als, + ET(((s))
(38)
Hence optimizing () with respect to = & U (uz}= is equivalent to optimizing (aaaas))ll0a"aass) Specifically optimizing with respect to boils down to
min U( 7) = min (araas)) (KL s)]gf( = max (x Bylela.s) (log o (=|s)]-
(39)
Noting that the expectation concerning uu() does not affect the minimizer, concludes the derivation.
C Environments
Relay Kitchen: A multi-task kitchen environment with long-horizon manipulation tasks such as moving kettle, open door, and turn on/off lights. The dataset consists of 566 human-collected trajectories with sequences of 4 executed skills. We used the same experiment settings and the pre-trained diffusion models from [9]
XArm Block Push: We used the adapted goal-conditioned variant from The Block-Push Environment consists of an XARm robot that must push two blocks, red and green one, into a red and green squared target area. The dataset consists of 1000 demonstrations collected by deterministic controller with 4 possible goal configurations. The methods got 0.5 credit for every block pushed into one the targets with maximum score 1.0. We use the pretrained Beso model from 19..
D3IL is a simulation benchmark with diverse human demonstrations, which aims evaluate imitation learning models' ability to capture multi-modal behaviors. D3IL provides 7 simulation tasks consisting of 7DoF Franka Emika Panda robot and various objects, where each task has different solutions and the robot is required to acquire all behaviors. Except for success rate, D3IL proposes to use task behavior entropy to quantify the policy's 's capability of learning multi-modal distributions. Given the predefined behaviors 3 for each task, the task behavior entropy is defined
17

---

E xp)oo) (H(+((3s0))
Σ
Σ ( log (3)
(40)
somp(sa)
where So refers the initial state and Sp refers to the number of samples from the initial state distribution p(su). During the simulation, we rollout the policy multiple times for each So and use Monte Carlo estimation to compute the expectation of the behavior entropy.
Aligning
Pushing
Avoiding
Sorting
Stacking
Figure 7: Visualization of D3IL tasks. We further provide the figure of demonstrations for the Avoiding task, which indicates 24 solutions it.
In this paper, we evaluate our algorithm in Avoiding, Aligning, Pushing, and Stacking with state- based representations and Sorting and Stacking with image-based representations. The simulation environments can be found in Fig. 7 The Avoiding task requires the robot to reach the green line without colliding with any obstacles. The task contains 96 demonstrations, consisting of 24 solutions with 4 trajectories for each solution. The Aligning task requires the robot to push the box to match target position and orientation. The robot can either push box from inside or outside, thus resulting in 2 solutions. This task contains 1000 demonstrations, 500 for each solution with uniformly sampled initial states. The Pushing task requires the robot to push two blocks to the target areas. The robot can push the blocks in different orders and to different target areas, which gives the task 4 solutions. This task contains 2000 demonstrations, 500 for each solution with uniformly sampled initial states. The Sorting task requires the robot to sort red and blue blocks the corresponding box. The number solutions is determined by the sorting order. D3IL provides Sorting 2, 4, and 6 boxes, here only use the Sorting-4 task, which contains around 1054 demonstrations with 18 solutions. The Stacking task requires the robot to stack three blocks the target zone. Additionally, the blue block needs be stacked upright which makes it more challenging. This task contains 1095 demonstrations with 6 solutions.
D Baselines Implementation
a continuous time diffusion policy that uses a continuous stochastie-differenta equation to represent the denoising process. We implement this method from the D3IL benchmark, following the default which predicts one-step action conditional on past five-step observations.
BESO
DDPM
10 is a discrete diffusion policy. We do not directly use the code from DiffusionPolicy 10 which implements encoder-decoder transformer structure. For fair comparison, we evaluate this model from the D3IL benchmark which shares the same architecture as in BESO.
42 is designed to overcome the slow generation of diffusion models. Consistency models can directly map noise to data using one-step and few-step generation and they can be trained either through distilling pre-trained diffusion models or an independent generative model. Our implementation takes the main training part the model by integrating GPT-based diffusion policy as the backbone.
Consistency Distillation
18

---

Consistency Trajectory Models is an extension of the CD model originally used in image generation. augments the performance by integrating additional CTM and GAN loss terms in consideration. Later, it is used in robot policy prediction in 45 without taking the GAN loss. Our implementation of CTM is extended from our CD implementation, by modifying the loss computation.
IMC 14] a curriculum-based approach that uses curriculum to assign weights training data so that the policy can select samples to learn, aiming to address the mode-averaging problem in multimodal distributions. implement the model using the official IMC code with GPT structure.
EM 66 is based on the maximum likelihood objective and follows an iterative optimization scheme, where the algorithm switches between the M-step and the E-step in each iteration.
E
Hyper Parameters
E.1 Hyperparameter Selection
We executed a large-scale grid search to fine-tune key hyperparameters for each baseline method. For other hyperparameters, choose the value specified in their respective original papers. Below is list summarizing the key hyperparameters that we swept during the experiment phase.
BESO: None
DDPM: None
Consistency Distillation: μr: EMA decay rate, N: see Algorithm 2 in 42 All the other hyper- parameters reuse the ones from the diffusion policy (BESO), as CD requires to be initialized using a pre-trained diffusion model.
Consistency Trajectory Models:
Same as Consistency Distillation.
IMC-GPT: Eta
14] Number of components
EM-GPT: Number of components
VDD-DDPM: Number of components, 4mins Inxx
VDD-BESO: Number of components, Omius Fuax
19

---

E.2 Hyperparameter List
Methods Pessuctes
Grid Scadh
Asoiding
Aligning
Pushing
Stucking
SortingVision
Stacking-Vsinn
Kätche
Black Push
GFT (shared by all)
Nunber Layers
4
Number of Attention Heads
4
6
12
12
Embedding Dimeessio
72
72
12
120
120
240
192
Window Size
5
5
Oytimizer
Audem
Adam
Addam
Adam
Adam
Adam
Adam
Adam
Leorning Rate x10
DOPM
Numberr Time Stepe
64
16
16
RESO
Number of Sumpling Steps
16
64
16
16
16
0.1
0.01
0.1
0.3
0.3
0.3
00.3
0.1
IMC-GPT
Nunteer Component
(1.8. 200, 50) (,
8
20
20
Bla.6
EMAGPT
Number of Components
(4.8, 20,500
20
20
50
20
20
20
20
20
CD
N. Algurihm 2
+
(2. 20, 80, 120, 180)
2
EMA docay rase pl
(1.99, 0,999, 0.99009)
0,9999
0.9999
0,9999
0.9999
0.9999
0.9999
0.9999
0.9999
CTM
N. Algorihm in 42
(5. 20, 20, 40] (0.99, 0,990, 0,95000
20
20
20
20
20
20
10
10
BMA decay rate
0,9999
0.9999
0,9999
0.9999
0.9999
0.9999
0.9999
0.9999
YDI--DPP
Number of Components
8
[1,2,4]
2
(4.8, 10,
12
4
12
VDD-BESO
Number of Compunents
(4,8)
00.1,02,,044400000
0.2
0.8
0.2
0.4 0.4
0.2
0.1
0.2
0..1000,0088 101
0.5
1.0
0.5
1.0
0.8
0.5
0.1
0.6
Table 4: Hyperparameter algorithms proposed method and baselines. The 'Grid Serach' column indicates the values over which we performed a grid search. The values in the column which are marked with task names indicate which values were chosen for the reported results.
Evaluation of Teacher Diffusion Models with Different Denoising Steps
Unlike DDPM, which uses fixed of timesteps, BESO learns continuous-time representation of the scores. This continuous representation enables the use of various numerical integration schemes, which can impact the performance of the diffusion model. We conducted evaluation the BESO teacher we used with different denoising steps. The results are presented in Table
Environments
Euler Maru-16
Fuler Maru-32
Enler Maru-64
Fuler Maru-16
Enler Maru-32
Enler Maru-64
Avoiding
Aligning
0.67
0.32
Pushing
0.77
0.71
0.78
0.76
Stacking
0.91
0.87
0.88
0.30
0.32
0.32
Stacking-2
0.70
0.64
0.63
0.13
0.14
0.13
Sorting (image)
0.70
0.76
0.76
0.19
0.23
0.24
Stacking (imuge)
0.60
0.70
0.70
0.15
0.16
0.21
Table 5: BESO with varies denoising steps.
20

---

G
Compute Resources
We train and evaluate all the models based our private clusters. Each node contains 4 NVIDIA A100 and we use one GPU for each method. We report the average training time in Table 6
Training Time
state state-based based
image- based
EM-GPT IMC-GPT
2 3-3 h
6
3
6
VDD-DDPM
4
8
VDD-BESO
3 3-4h
8
Table 6: Training time for each method.
In addition, evaluate how the number of trainable parameters and training time scale with different numbers of components. The results are presented in Table 7
Num. of Experts
4
8
16
Parameters (x (x10*) Times/1k Iters (8)
144.23
144.45
144.89
145.76
147.50
64.48
70.19
74.62
106.98
166.54
Table 7: Adding more experts does not significantly increase the number of neural network parameters or the training time. We conducted the evaluation on the state-based avoiding task, using a machine with an RTX 3070 GPU and 17-13700 CPU. This result due the optimized network architecture of the VDD model, as shown in Figure 5 Adding more experts will only increase the number of output linear layers, i.e., the mean and covariance nets, while the transformer backbone which contains most of the parameters remains unchanged.
21

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA]-
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1 (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked While "[Yes] " is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providded a proper justification is given (e.g., "error bars expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
the checklist one of the factors in their evaluation.
reported because it would be too computationally
"[No] " or "[NA] is not grounds for rejection. While the questions are phrased acknowledge that the true answer often more nuanced, so please just use your best, judgment and write justification elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main contribution is novel method for distilling diffusion models into mixture of experts, which is outlined and described in the abstract and the introduction and the method section. Claims wrt to the performance of the distilled policies are verified in the experiment section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
22

---

Answer: [Yes]
Justification:
We discuss the limitations of this work in the conclusion section
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well sell-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor of transparency play an impor- lant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide a description of our learning objective in the method section, and a full derivation in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We describe how the baselines are implemented in Appendix C and corre- sponding hyperparameters to reproduce the experiment results in the appendix D
23

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will open source the codes in the near future once they are cleaned up and anomnymity is not concern anymore. All the experiments we conducted were using open-source datasets. the experiments section and appendix C we provide information to get access to the data.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c ec/ goblic/guidess http:///uuiis///eeuuubiiiiieee
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips cc/ wwwpcc/pplcc....c httpll/c/guidess/ddeeebbbiiiiiii
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
24

---

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide hyperparameter lists for each of the algorithms, how they were chosen and type of optimizer appendix E
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The experiments describe the number of trials and show the deviations in the result tables. There are no deviations for the origin models as they are the used as the base the distillation and hence a fixed seed
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
25

---

Answer: [Yes]
Justification:
The used compute resources are described in appendix G
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:/ppp...eeeepp..ppp/pppcacaa
Question:
Answer: [Yes]
Justification: Yes
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification:
We discussed the broader impact in the conclusion section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
26

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Yes, we do used pretrained models for diffusion model and open source code base for baselines, which is clearly stated in both experiment section and appendix.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwwpsswiteesseeeeeeeeee thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We plan to open source the code in the future
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
27

---

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "TfqcwHum9D", "reviewer": "Reviewer_14hC", "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.", "questions": "1. Why does the distilled model outperform the original diffusion policy in Sorting (Image) and Stacking (Image) tasks?\n\n2. Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models?\n\n3. Why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models?", "limitations": "The authors have discussed the limitations of the paper.\n\nPlease refer to the [weakness] section for detailed concerns and questions about the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n2. The proposed method, VDD, introduces a robust approach to distilling diffusion policies to  MoE models using VI. \n3. 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?\n\n2. 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.\n\n3. Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models?\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> […] the paper claims that \"VDD is not straightforwardly applicable to generating very high-dimensional data like images.\"\n> \n\nThe 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.\n\n---\n\n> Why does the distilled model outperform the original diffusion policy in Sorting (Image) and Stacking (Image) tasks?\n> \n\nWe 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.\n\nWe 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. \n\nWe add a discussion about this trade-off in the experiment section.\n\n---\n\n> Does the distilled MoE model have the same architecture as the EM-GPT and IMC-GPT models?\n> \n\nYes, to ensure fairness across all experiments, we use the same architecture for these methods.\n\n---\n\n> […] It would be beneficial to report this total time for a fair comparison between MoE methods.\n> \n\nWe 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).\n\n---\n\n> Why do we distill from diffusion models? How about distilling from large and robust MoE models into smaller MoE models?\n> \n\nDiffusion 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.\n\n---\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HtLQwj0iNA", "reviewer": "Reviewer_YkuG", "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.", "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.", "limitations": "The authors provided the limitation statement, mainly concerning the expressity of MoE.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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. \n\n- A very relevant work on distilling diffusion models, Luo et al. 2023, is missing. \n\n- 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.\n\n- 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. \n\nXie et al. 2024, EM Distillation for One-step Diffusion Models\n\nLuo et al. 2023, Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models\n\nYu et al. 2021, Unsupervised Foreground Extraction via Deep Region Competition", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n> 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.\n> \n\nWe 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.\n\n---\n\n> 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.\n> \n\nIndeed, 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.\n\n---\n\n> 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 […]\n> \n\nWe 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%.\n\nNevertheless, 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.\n\n---\n\n> 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.\n> \n\nWe 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.\n\nWe 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. \nWe add a discussion about this trade-off in the experiment section.\n\n---\n\nWe 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.\n\n---\n\n[1] Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations, ICLR ‘2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SVcgZiieOP", "reviewer": "Reviewer_MHge", "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.", "questions": "Please see Weaknesses section", "limitations": "The authors do have included discussion about the limitation and shed light on how they might be approaching in future studies.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. clear motivation of combining the strengths of both diffusion models and MoE.\n2. combining decomposing objective with EM is shown to be effective to escape certain challenges of training MoE.\n3. 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?\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n> 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 […] .\n> \n\nWe 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$.\n\n---\n\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.\n> \n\nSince 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.\n\n---\n\n> Is it possible to take task similarity into account and perform kind of meta-training like policy distillation?\n> \n\nWe 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.\n\n---\n\n> The authors do not explicitly discuss how well VDD can capture long-range temporal dependency that might be presented in the original diffusion model […]\n> \n\nVDD 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.\n\n---\n\nWe hope that our responses have addressed your concerns. We would be more than happy to address any remaining questions/concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BqBXwXmr8D", "reviewer": "Reviewer_8iR5", "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.", "questions": "see weaknesses", "limitations": "see weaknesses", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1) The technical approach appears to be sound, demonstrating a rigorous and well-grounded methodology.\n2) The visualizations presented in the paper are meticulously crafted, showcasing a high level of attention to detail and aesthetics.\n3) 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.\n2) 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.\n3) A missing period on line 218 detracts slightly from the overall readability and professionalism of the manuscript.\n4) 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.\n5) 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.\n6) 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.\n7) The comparison to a limited number of methods may limit the ability to comprehensively assess the strengths and weaknesses of the proposed approach.\n8) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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.\n> \n\nTo 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.\n\n---\n\n> 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\n> \n\nWe thank the reviewer and will do our best to remove respective phrasing and make the paper more engaging for the reader.\n\n---\n\n> […] missing period on line 218 […] The inconsistency in Table 3 left, where 2.16 is highlighted as better than 2.15 is confusing […]\n> \n\nWe thank the reviewer for carefully reading the paper and identifying the formatting issues. We have addressed and corrected them.\n\n---\n\n> 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.\n> \n\nWe 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. \n\n---\n\n> 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.\n> \n\nWe 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. \n\nNevertheless, 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.\n\n**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.\n\n---\n\n> The comparison to a limited number of methods may limit the ability to comprehensively assess the strengths and weaknesses of the proposed approach.\n> \n\nTo 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. \n\n---\n\n> 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.\n> \n\nWe 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. \n\nFigure 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.\n\n---\n\nWe 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.\n\n---\n\n[1] Song, Yang, and Prafulla Dhariwal. \"Improved techniques for training consistency models.” ICLR 2024\n\n[2] Xie, S., et al. “EM Distillation for One-step Diffusion Models”, Preprint\n\n[3] Kim, Dongjun, et al. \"Consistency trajectory models: Learning probability flow ode trajectory of diffusion.\" ICLR 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "24j3yEqTdY", "reviewer": "Reviewer_5fPT", "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.", "questions": "1. What is the mixture number n in your experiments?  Can you provide any experiment results ablating n? This could be very informative.\n2. 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)", "limitations": "Above", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n---\n\n> 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.\n> \n\nWe 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)*.\n\nTE 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. \n\nWe 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.\n\n---\n\n> 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.\n> \n\nWe 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.\n\nAdditionally, we summarized these points below:\n\n1. **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].\n2. **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].\n3. **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].\n\n---\n\n> 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.\n> \n\nWe 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.\n\nIn 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.\n\n---\n\n> What is the mixture number n in your experiments?\n> \n\nTo 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.\n\nWe will also include more detailed information about the mixture number in the main section of the paper.\n\n---\n\n> Can you provide any experiment results ablating n? This could be very informative.\n> \n\nWe 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.\n\n---\n\n> 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)\n> \n\nWe 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.\n\n---\n\nWe would like to thank the reviewer again and welcome the opportunity to address any additional concerns or questions that the reviewers may have.\n\n---\n[1] Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations, ICLR ‘2024\n\n[2] Neural Probabilistic Motor Primitives for Humanoid Control, ICLR ‘19\n\n[3] InfoGAIL: Interpretable Imitation Learning from Visual Demonstrations, NeurIPS ‘17\n\n[4] One Solution is Not All You Need: Few-Shot Extrapolation via Structured MaxEnt RL, NeurIPS ‘20\n\n[5] Specializing Versatile Skill Libraries using Local Mixture of Experts, CoRL ‘21\n\n[6] High-fidelity and diverse text-to-3d generation with variational score distillation, NeurIPS ‘23", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their valuable feedback.  We would like to reiterate our main action points in response to the reviewers’ suggestions and concerns. New results can be found in the PDF file that accompanies the rebuttal.

- **Open-Sourced Code-Base:** We provided a link to an anonymous Github repository containing the code for reproducing the results of the paper. Additionally, we provided a link to an IPython notebook for testing our method on a simple toy task. We have sent both links to the AC in an Official Comment, as the rebuttal guidelines prevent us from sharing them here.
- **Added Baselines:** In response to feedback, we have included a comparison of our method with Consistency trajectory models [1].
- **Completed Results-Table.** The initial submission was missing the results for DDPM on two tasks, *kitchen* and *block push*. We have now included these results.
- **Detailed Parameter and Training Time Information:** Additional specifics on parameter counts and training times across varying number of experts have been incorporated.
- **Enhanced Visualizations:** To offer deeper insights into our method, we have included supplementary visualizations. It highlights the most likely experts according to the gating probability at a given state.

---

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 during the upcoming discussion period.

---

[1] Kim D, Lai CH, Liao WH, Murata N, Takida Y, Uesaka T, He Y, Mitsufuji Y, Ermon S. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. ICLR 2024.

Meta-review: This paper introduces Variational Diffusion Distillation (VDD), a method for distilling pre-trained diffusion policies into Mixture of Experts (MoE) models, combining the expressiveness of diffusion models with the fast inference of MoEs. VDD uses variational inference with a decomposed objective to train MoE experts efficiently. The method is evaluated on 9 behavior learning tasks, showing strong performance.

Reviewers appreciated the paper's technical approach, methodology, and extensive evaluation. Concerns included limited comparison to recent baselines and need for more analysis and visualizations. The authors addressed these in their rebuttal by adding comparisons, providing missing results, offering deeper analysis, and including additional visualizations. They also released code and provided details on hyperparameters and training costs.

Given the paper's strengths, thorough evaluation, and the authors' careful addressing of concerns, I recommend accepting this paper. The method enables practical use of expressive diffusion policies in domains requiring fast inference, representing a contribution to the field.
Decision: Accept (poster)