Research map / Reinforcement learning / Imitation learning
Maximum Entropy Inverse Reinforcement Learning of Diffusion Models with Energy-Based Models
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.
placeholder
Problem
Diffusion models produce degraded samples when the number of generation steps is small, due to the train-generation distribution gap; training an EBM typically requires MCMC, which is costly and hyperparameter-sensitive.
Approach
Treat diffusion sampling as a maximum-entropy RL problem with the EBM's energy as the reward, jointly train the diffusion model and EBM as a generalized contrastive divergence minimax, and solve the diffusion update via dynamic programming with value functions (no back-propagation through time) using conditional-entropy running costs.
Key results
- On MVTec-AD, DxMI reaches 97.0 mean detection AUROC and 97.1 localization AUROC, above UniAD (96.5/96.8) and MPDR (96.0/96.7).
- On LSUN Bedroom 256x256, a 4-step DxMI model reaches FID 5.93, versus 5.22 for the 2-step consistency model and 2.44 for 79-step EDM.
- DxMI outperforms SFT-PG, the policy-gradient IRL baseline, on CIFAR-10 and ImageNet 64 with identical backbones and starting checkpoints; the best CIFAR-10 T=10 run reaches FID 3.19 versus 4.32 for SFT-PG.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Set entropy regularization τ to 0 (ablation) | DxMI · MVTec-AD mean | Detection AUC: 97.0 -> 67.9 | Hurt |
| Set entropy regularization τ to 0 (ablation) | DxMI · MVTec-AD mean | Localization AUC: 97.1 -> 84.6 | Hurt |
| Set sampler-update temperature parameters τ₁ and τ₂ to zero (ablation) | DxMI image generation · Image generation experiments | FID | Hurt |
| Replace sigmoid time cost with linear time cost (ablation) | DxMI image generation · Image generation experiments | FID | Hurt |
| Add value-gradient guidance during sampling (ablation) | DxMI · Image generation experiments | Sample quality | Helped |
| Pretrain the sampler as DDPM (ablation) | DxMI · 8 Gaussians | SW | Helped |
| Use τ=0.1 entropy regularization instead of τ=0 (ablation) | DxMI · 8 Gaussians | SW | Helped |
| Remove entropy regularization by setting τ=0 (ablation) | DxMI · 8 Gaussians | Energy density estimation | Hurt |
| Increase entropy regularization to τ=1 (ablation) | DxMI · 8 Gaussians | SW | Hurt |
| Use DxMI rather than SFT-PG | DDPM, CIFAR-10 T=10 · CIFAR-10, T=10 | FID: 4.32 -> 3.19 | Helped |
| Use DxMI with T=4 instead of the one-step Consistency Model | ImageNet 64×64 generation · ImageNet 64×64 | FID: 6.20 -> 3.21 | Helped |
Surprising findings
- Setting the entropy temperature to zero (tau=0) collapses anomaly detection AUROC from 97.0 to 67.9, showing entropy maximization is essential for the EBM.
Limitations
- Introduces several hyperparameters (tau1, tau2, gamma, time cost) that need tuning.
- Not directly applicable to single-step generators; does not allow changing the number of steps T at test time.
- No formal theoretical convergence analysis for DxMI.
Experimental setup
Datasets: CIFAR-10 ImageNet 64x64 LSUN Bedroom MVTec-AD 8 Gaussians
Benchmarks: CIFAR-10 ImageNet 64x64 LSUN Bedroom MVTec-AD 8 Gaussians
Metrics: FID precision recall sliced Wasserstein distance AUC AUROC
Baselines: SFT-PG EDM Consistency Model StyleGAN2 DRAEM MPDR UniAD
Builds on: contrastive divergence maximum entropy inverse reinforcement learning SFT-PG DDPM EDM
Closest papers by meaning
- DIME: Diffusion-Based Maximum Entropy Reinforcement Learning (ICML 2025 · similarity 0.92)
- Scalable Maximum Entropy Reinforcement Learning for Diffusion Policies via Adjoint Matching (NeurIPS 2026 · similarity 0.90)
- Learning to Sample From Diffusion Models via Inverse Reinforcement Learning (NeurIPS 2026 · similarity 0.90)
- Maximum Entropy Reinforcement Learning with Diffusion Policy (ICML 2025 · similarity 0.89)
- Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages (NeurIPS 2026 · similarity 0.89)
- Diffusion Spectral Representation for Reinforcement Learning (NeurIPS 2024 · similarity 0.89)
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.