Diffusion Alignment as Variational Expectation-Maximization

2025/10/01 by Jae-Woo Lee, Lee, Jaewoo, Min-Su Kim +14
Computer Science · #Domain Adaptation and Few-Shot Learning #Evolutionary Algorithms and Applications #FOS: Computer and information sciences #Generative Adversarial Networks and Image Synthesis #Machine Learning (cs.LG)

paper · pdf · doi:10.48550/arxiv.2510.00502

openalex publication_date 2025/10/01 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Diffusion alignment aims to optimize diffusion models for the downstream objective. While existing methods based on reinforcement learning or direct backpropagation achieve considerable success in maximizing rewards, they often suffer from reward over-optimization and mode collapse. We introduce Diffusion Alignment as Variational Expectation-Maximization (DAV), a framework that formulates diffusion alignment as an iterative process alternating between two complementary phases: the E-step and the M-step. In the E-step, we employ test-time search to generate diverse and reward-aligned samples. In the M-step, we refine the diffusion model using samples discovered by the E-step. We demonstrate that DAV can optimize reward while preserving diversity for both continuous and discrete tasks: text-to-image synthesis and DNA sequence design. Our code is available at https://github.com/Jaewoopudding/dav.

Citations

• Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
• Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design
• Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
• Psi-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models
• Inference-time Scaling of Diffusion Models through Classical Search
• DanceGRPO: Unleashing GRPO on Visual Generation
• Flow-GRPO: Training Flow Matching Models via Online RL
• Inference-Time Scaling for Flow Models via Stochastic Generation and Rollover Budget Forcing
• Offline Model-Based Optimization: Comprehensive Review
• Dynamic Search for Inference-Time Alignment in Diffusion Models
• Split Gibbs Discrete Diffusion Posterior Sampling
• Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps
• A General Framework for Inference-time Scaling and Steering of Diffusion Models
• Test-time Alignment of Diffusion Models without Reward Over-optimization
• Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets
• Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design
• Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control
• Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding
• Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review
• Simple and Effective Masked Diffusion Language Models
• ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search
• Unlocking Guidance for Discrete State-Space Diffusion and Flow Models
• Amortizing intractable inference in diffusion models for vision, language, and control
• Gradient Guidance for Diffusion Models: An Optimization Perspective
• Aligning Diffusion Models by Optimizing Human Utility
• Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control
• Confronting Reward Overoptimization for Diffusion Models: A Perspective of Inductive and Primacy Biases
• BioXP-0.5B: Explainable Medical-AI via RL-GRPO
• Beyond Human Data: Scaling Self-Training for Problem-Solving with Language Models
• Diffusion Model Alignment Using Direct Preference Optimization
• Aligning Text-to-Image Diffusion Models with Reward Backpropagation
• Directly Fine-Tuning Diffusion Models on Differentiable Rewards
• Reinforced Self-Training (ReST) for Language Modeling
• Bootstrapped Training of Score-Conditioned Generator for Offline Design of Biological Sequences
• Direct Preference Optimization: Your Language Model is Secretly a Reward Model
• DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models
• Training Diffusion Models with Reinforcement Learning
• ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation
• FreeDoM: Training-Free Energy-Guided Conditional Diffusion Model
• Diffusion Policy: Visuomotor Policy Learning via Action Diffusion
• Diffusion policy: Visuomotor policy learning via action diffusion
• Consistency Models
• Universal Guidance for Diffusion Models
• Optimizing DDPM Sampling with Shortcut Fine-Tuning
• Diffusion Posterior Sampling for General Noisy Inverse Problems
• Classifier-Free Diffusion Guidance
• RL with KL penalties is better viewed as Bayesian inference
• Video Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• High-Resolution Image Synthesis with Latent Diffusion Models
• Greedification Operators for Policy Optimization: Investigating Forward and Reverse KL Divergences
• Conservative Objective Models for Effective Offline Model-Based Optimization
• Structured Denoising Diffusion Models in Discrete State-Spaces
• Diffusion Models Beat GANs on Image Synthesis
• Effective gene expression prediction from sequence by integrating long-range interactions
• Learning Transferable Visual Models From Natural Language Supervision
• Oops I Took A Gradient: Scalable Sampling for Discrete Distributions
• Score-Based Generative Modeling through Stochastic Differential Equations
• Denoising Diffusion Probabilistic Models
• AWAC: Accelerating Online Reinforcement Learning with Offline Datasets
• Behavior Regularized Offline Reinforcement Learning
• Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction
• Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review
• Maximum a Posteriori Policy Optimisation
• The Unreasonable Effectiveness of Deep Features as a Perceptual Metric
• Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm
• Deep Unsupervised Learning using Nonequilibrium Thermodynamics
• Tweedie’s Formula and Selection Bias
• Levenshtein Distance Technique in Dictionary Lookup Methods: An Improved Approach

Related

• Dynamic Search for Inference-Time Alignment in Diffusion Models
• Discrete Diffusion Trajectory Alignment via Stepwise Decomposition
• Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
• Aligning Few-Step Diffusion Models with Dense Reward Difference Learning
• Test-time Alignment of Diffusion Models without Reward Over-optimization