Explicit Critic Guidance for Aligning Diffusion Models
This work addresses the problem of fine-grained credit assignment and stable value-based optimization in aligning diffusion models with non-differentiable objectives, which is important for improving generation quality in text-to-image and similar tasks.
The paper proposes a state-aligned latent actor-critic framework for aligning diffusion models, where the diffusion model acts as its own value function, enabling stable PPO training and inference-time steering. The method outperforms prior RL and actor-critic baselines on single- and multi-reward benchmarks across UNet and DiT backbones.
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.