We present a maximum entropy inverse reinforcement learning (IRL) approach for improving the sample quality of diffusion generative models, especially when the number of generation time steps is small. As IRL trains a policy using the reward function learned from expert demonstrations, we train (or fine-tune) a diffusion model using the log probability density estimated from training data. Since we employ an energy-based model (EBM) to represent the log density, our approach boils down to the joint training of a diffusion model and an EBM. Our IRL formulation, named Generalized Contrastive Divergence (GCD), is a minimax problem that reaches equilibrium when both models converge to the data distribution. The entropy maximization plays a key role in GCD, facilitating the exploration of the diffusion model and ensuring the convergence of the EBM. We also propose Diffusion by Dynamic Programming (DiDP), a novel RL approach for diffusion models, as a subroutine in GCD. DiDP makes the diffusion model update in GCD efficient by transforming the original problem into an optimal control formulation where value functions replace back-propagation in time. Our empirical studies show that diffusion models fine-tuned using GCD and DiDP can generate high-quality samples with fewer steps. Additionally, GCD-DiDP enables the training of an EBM without MCMC, stabilizing EBM training dynamics and enhancing outlier detection performance.

유튜브 스트리밍 예정입니다.