[2510.00502] Diffusion Alignment as Variational Expectation-Maximization
Summary
The paper introduces Diffusion Alignment as Variational Expectation-Maximization (DAV), a novel framework that optimizes diffusion models to enhance reward alignment while maintaining diversity in generated samples.
Why It Matters
This research addresses key challenges in machine learning, particularly in optimizing diffusion models for tasks like text-to-image synthesis and DNA sequence design. By improving reward alignment and diversity, it has implications for various applications in generative AI and data science.
Key Takeaways
- DAV optimizes diffusion models through an iterative E-step and M-step process.
- The E-step generates diverse, reward-aligned samples via test-time search.
- The M-step refines the diffusion model based on samples from the E-step.
- DAV effectively addresses issues of reward over-optimization and mode collapse.
- Applications include text-to-image synthesis and DNA sequence design.
Computer Science > Machine Learning arXiv:2510.00502 (cs) [Submitted on 1 Oct 2025 (v1), last revised 23 Feb 2026 (this version, v2)] Title:Diffusion Alignment as Variational Expectation-Maximization Authors:Jaewoo Lee, Minsu Kim, Sanghyeok Choi, Inhyuck Song, Sujin Yun, Hyeongyu Kang, Woocheol Shin, Taeyoung Yun, Kiyoung Om, Jinkyoo Park View a PDF of the paper titled Diffusion Alignment as Variational Expectation-Maximization, by Jaewoo Lee and 9 other authors View PDF 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 this https URL. Comments: Subjects: Machine Learning (cs.LG) Cite as: arXiv:2510.00502 [cs.LG] (or arXiv:2510.00502v2 [cs.LG] for this versio...