English

Explicit Critic Guidance for Aligning Diffusion Models

Machine Learning 2026-05-28 v1 Computer Vision and Pattern Recognition

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.

Keywords

Cite

@article{arxiv.2605.27736,
  title  = {Explicit Critic Guidance for Aligning Diffusion Models},
  author = {Zhengyang Liang and Qihang Zhang and Ceyuan Yang},
  journal= {arXiv preprint arXiv:2605.27736},
  year   = {2026}
}