MMODELYST
Papers/Explicit Critic Guidance for Aligning Diffusion Models
PAP

Explicit Critic Guidance for Aligning Diffusion Models

May 26, 2026

arXiv
Abstract

Online reinforcement learning is becoming increasingly important for aligning diffusion models with non-differentiable objectives. However, existing methods still face limitations in assigning fine-grained credit along denoising trajectories and in realizing stable value-based optimization. We propose a state-aligned latent actor-critic framework for diffusion post-training, in which the diffusion model serves as its own timestep-conditioned value function and predicts values directly on noisy latent states. This enables trajectory-level PPO training, supports stable actor-critic optimization with simple conditioning and value pretraining strategies, and naturally allows the learned critic to be reused for inference-time steering. We further extend the framework to multi-reward optimization, where joint training with complementary rewards helps alleviate reward hacking. Across both UNet- and DiT-based backbones, our method consistently outperforms prior group-relative RL and actor-critic baselines on single-reward and multi-reward benchmarks, while test-time steering provides additional gains in generation quality.

Select text to highlight · click a highlight to remove · saved in this browser only
Authors
Zhengyang Liang, Qihang Zhang, Ceyuan Yang
Your notes (browser-local)
saved
arXiv:2605.27736