Research map / Image generation / Fast diffusion sampling
Variational Distillation of Diffusion Policies into Mixture of Experts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Distills pre-trained diffusion policies (DDPM/BESO) into Gaussian mixture-of-experts policies via a decomposed variational (reverse-KL) upper bound with EM-style updates, evaluated on D3IL and other imitation tasks.
Problem
Diffusion policies for learning from demonstrations are slow at inference (many denoising steps) and have intractable likelihoods, while mixture-of-experts policies are fast and tractable but hard to train stably.
Approach
Derive an upper bound on the reverse-KL variational objective that decomposes into per-expert objectives driven by the teacher's score function, with a repulsion term from the auxiliary posterior and EM-style alternation; the gating is trained with a cross-entropy-style objective, and score is sampled over interval time steps.
Key results
- On D3IL, VDD-VE reaches 0.95 success on Avoiding and 0.85 on Pushing, outperforming CD and 1-step teachers in 6 of 7 tasks (except Aligning, where it has higher entropy).
- VDD-VE attains task entropy 0.69 on Pushing versus 0.53 for CD and 0.50 for VE-1 (ablation of entropy across seven D3IL tasks).
- VDD-VP achieves 3.24 on Relay Kitchen, below CD-VE's 3.87, and is faster: distilled MoE samples in a single forward pass, with the 1-step BESO remaining strong on Kitchen and Block Push.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Distill BESO into VDD-VE instead of using BESO directly | VDD-VE · D3IL tasks | success rate | Mixed |
| Distill the diffusion policy into VDD rather than train an MoE from scratch with EM-GPT or IMC-GPT | VDD · D3IL tasks | success rate | Helped |
| Use VDD rather than DDPM for image-based stacking inference | VDD · Image-based stacking | inference time (ms): 6.49 -> 4.91 | Helped |
| Use one expert instead of multiple experts (ablation) | VDD · Avoiding | task entropy | Hurt |
| Increase the number of experts from one to multiple experts (ablation) | VDD · Avoiding | task entropy | Helped |
| Train a parameterized gating distribution instead of using uniform gating (ablation) | VDD · D3IL tasks | success rate | Helped |
| Train a parameterized gating distribution instead of using uniform gating (ablation) | VDD · D3IL tasks | task entropy | Hurt |
| Use interval time-step sampling rather than minimum, maximum, or uniform time-step sampling (ablation) | VDD · Avoiding | task entropy | Helped |
| Use interval time-step sampling instead of other time-step selection variants (ablation) | VDD · Avoiding | success rate | No clear effect |
Surprising findings
- One-step continuous diffusion (BESO-1) already performs well, often outperforming distillation results on Relay Kitchen and Block Push.
- Training a gating distribution improves success rates but slightly reduces task entropy in two of three tasks.
Limitations
- Not straightforwardly applicable to very high-dimensional data such as images without latent-space extensions
- Number of experts must be pre-defined by the user
- Performance is upper-bounded by the teacher diffusion model
Experimental setup
Datasets: D3IL Relay Kitchen XArm Block Push
Benchmarks: D3IL Relay Kitchen XArm Block Push
Metrics: task success rate environment return task entropy inference time (ms) number of function evaluations
Baselines: DDPM BESO DDPM-1 BESO-1 Consistency Distillation Consistency Trajectory Model EM-GPT IMC-GPT
Builds on: EM algorithm variational inference consistency distillation consistency trajectory models BESO
Closest papers by meaning
- DiffEnc: Variational Diffusion with a Learned Encoder (ICLR 2024 · similarity 0.86)
- Distilling ODE Solvers of Diffusion Models into Smaller Steps (CVPR 2024 · similarity 0.85)
- Denoising Diffusion Variational Inference: Diffusion Models as Expressive Variational Posteriors (AAAI 2025 · similarity 0.85)
- Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling (NeurIPS 2024 · similarity 0.85)
- Diffusion Model Patching via Mixture-of-Prompts (AAAI 2025 · similarity 0.85)
- Data-free Distillation of Diffusion Models with Bootstrapping (ICML 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.