🤖 machine learning

Explicit Critic Guidance for Aligning Diffusion Models

This paper proposes a state-aligned latent actor-critic framework that enables diffusion models to serve as their own timestep-conditioned value functions, thereby facilitating stable trajectory-level PPO training, multi-reward optimization, and effective inference-time steering to outperform existing alignment methods.

Original authors: Zhengyang Liang, Qihang Zhang, Ceyuan Yang

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhengyang Liang, Qihang Zhang, Ceyuan Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a talented artist (the Diffusion Model) to paint a picture based on a specific description, like "a cat wearing a hat on Mars."

The artist doesn't paint the whole picture in one go. Instead, they start with a canvas full of static noise and slowly, step-by-step, refine it until the cat appears. This is the "denoising trajectory."

The problem the paper tackles is: How do you teach the artist to do this better?

Usually, you only look at the final painting to see if it's good. If the cat looks wrong, you tell the artist, "Bad job." But the artist doesn't know which step went wrong. Was it step 10? Step 30? Did they mess up the hat or the background? This is called the credit assignment problem.

Here is how the authors' new method, State-Aligned Latent Actor-Critic, solves this, using simple analogies:

1. The "Internal Coach" (The Latent Critic)

In older methods, the "coach" (the critic) would wait until the artist finished the painting, look at the final image, and then try to guess what went wrong earlier. This is like a coach watching a soccer game only at the final whistle and then trying to tell the player what to do differently during the first half. It's inaccurate and slow.

The Paper's Solution:
They turn the artist into their own coach. They give the artist a special "internal eye" that can see the painting while it is still being made (while it's still noisy and blurry).

The Magic: Instead of waiting for the final image, this internal coach looks at the messy, noisy canvas at every single step and says, "If you keep going this way, you'll get a good score," or "Stop, that path leads to a bad result."
Why it's better: Because the coach is looking at the actual messy steps the artist is taking (not a cleaned-up version), the advice is much more accurate. It's like having a GPS that updates your route every second based on your current location, rather than guessing where you are based on a map from yesterday.

2. The "Rehearsal" (Value Pretraining)

Teaching a new coach to give advice is hard. If you start training the coach and the artist at the exact same time from scratch, they might get confused and argue with each other, making the training unstable.

The Paper's Solution:
They give the coach a short "rehearsal" first. Before the real training starts, they let the coach practice just watching the artist and guessing the final score, without actually changing the artist's behavior yet.

The Result: By the time the real training begins, the coach is already pretty good at predicting outcomes. This makes the whole learning process much smoother and faster, preventing the "argument" between the artist and the coach.

3. The "Multi-Task" Challenge (Multi-Reward Optimization)

Sometimes, you want the painting to be good at many things at once: it needs to look like the prompt, look beautiful to humans, and have the right number of objects.

The Problem: If you only focus on one thing (e.g., "make sure there are exactly 4 chairs"), the artist might get lazy and draw 4 chairs but make the background look terrible or the colors look weird. This is called "reward hacking"—finding a shortcut to win the game without actually playing well.
The Paper's Solution: They give the artist a team of coaches, each specializing in a different skill (one for beauty, one for object count, one for text). These coaches share the same "brain" (the underlying model) but have their own specific scoreboards.
The Result: The artist learns to balance everything. They can't just hack the "chair count" reward because the "beauty" coach will penalize them for making a ugly picture. It forces a more well-rounded performance.

4. The "Steering Wheel" (Inference-Time Steering)

Once the artist is trained, you don't just have to accept whatever they paint. You can use the coach's "internal eye" to steer the painting while it's being made, even after training is over.

How it works: Imagine the artist is painting a path. The coach can gently nudge the brush toward areas that look promising.
The Benefit: This allows you to get even better results without retraining the whole model. It's like having a GPS that suggests a better route while you are driving, improving your trip without needing to learn a new map.

Summary of Results

The authors tested this on two different types of "artists" (one based on older technology called UNet, and one based on newer technology called DiT).

Better Scores: Their method consistently produced better images than previous methods, especially for tricky tasks like counting objects or reading text in images.
More Stable: The training didn't crash or get confused as easily as other methods.
Efficient: It was faster to train and required less computing power because the coach didn't have to wait for the final image to give advice.

In short, the paper teaches AI artists to be their own best critics, giving them instant feedback during the creative process, rehearsing before the big game, and using a team of coaches to ensure they don't cut corners.

Technical Summary: Explicit Critic Guidance for Aligning Diffusion Models

Problem Statement

Post-training diffusion models to align with non-differentiable objectives (e.g., human preference, compositional correctness, text rendering) is increasingly critical, as pretraining alone often yields suboptimal results on user-centric properties. A central challenge in this domain is credit assignment. Diffusion models generate samples via multi-step denoising trajectories, yet rewards are typically defined only on the final decoded image. Existing solutions face significant limitations:

Group-Relative Methods (e.g., GRPO): These bypass explicit critics by comparing samples within a group. While they avoid training a value function, the resulting reward signal is coarse and trajectory-level, lacking fine-grained credit assignment for intermediate denoising decisions.
Pixel-Space Critics: Existing explicit actor-critic methods often define critics in pixel space, evaluating reconstructed "clean" images from intermediate noisy latents. This creates a state mismatch: the policy operates on noisy latent states, while the critic evaluates proxy states that are often out-of-distribution for standard vision encoders (e.g., CLIP, BLIP). This mismatch leads to brittle value estimation, instability in optimization, and substantial decoding overhead.
Reward Hacking: Optimizing for specific verifiable rewards often leads to shortcut solutions (e.g., simplified backgrounds, repetitive patterns) that maximize the target metric while degrading unmeasured aspects of image quality.

Methodology

The authors propose a State-Aligned Latent Actor-Critic Framework that addresses these issues by integrating the critic directly into the diffusion generation process.

1. State-Aligned Latent Critic

Instead of operating in pixel space, the framework learns a value function $V_\phi(z_t, t, y)$ directly on the noisy latent states $z_t$ visited by the policy.

Architecture: The critic is initialized from the pretrained diffusion backbone (UNet or DiT) and modified with lightweight architectural changes to predict a scalar value.
- UNet: The feature map is patchified, and a learnable summary token aggregates information via an attention head and MLP.
- DiT: Lightweight critic branches are inserted at intermediate layers (e.g., layers 7, 15, 23) using learnable query tokens to aggregate visual information.
State Alignment: By operating in the same latent space and sharing the generative prior, the critic is naturally aware of the timestep and noise level, avoiding the distribution shift issues of pixel-space critics.

2. Stabilizing Training Techniques

To ensure stable actor-critic optimization, two specific techniques are introduced:

Timestep-Aware Conditioning: The value head is explicitly conditioned on the timestep $t$ (via AdaLN-style modulation), providing direct access to the noise level to improve value prediction across the trajectory.
Value Pretraining: Before joint actor-critic optimization, a short warm-up stage is performed where only the critic is updated using regression loss against Monte Carlo return targets. This yields better initial value estimates, stabilizing subsequent PPO updates.

3. Inference-Time Steering

Because the critic is defined on the same latent states as the policy, it can be reused at inference time as a guidance module. The gradient of the value function, $\nabla_{z_t} V_\phi$ , points toward latent regions predicted to yield higher terminal rewards. This gradient is injected into the sampler update (modifying noise prediction for DDIM or velocity for flow-matching), effectively steering the denoising trajectory without requiring additional training.

4. Multi-Reward Optimization

To mitigate reward hacking, the framework is extended to multi-reward optimization. A shared latent backbone is used with separate, lightweight value heads for each reward type. This design allows heterogeneous rewards (e.g., compositional vs. aesthetic) to regularize one another while sharing expensive feature extraction, preventing the collapse into shortcut solutions common in single-reward training.

Key Contributions

State-Aligned Latent Critic: A framework converting a pretrained diffusion model into its own timestep-conditioned value network, enabling value prediction on noisy latents and inference-time steering.
Stabilization Strategies: Identification of timestep-aware conditioning and value pretraining as essential techniques for stabilizing explicit actor-critic diffusion training.
Multi-Reward Extension: A mechanism using separate value heads on a shared backbone to jointly optimize complementary rewards, mitigating proxy overoptimization and reward hacking.
Empirical Superiority: Consistent outperformance over group-relative (GRPO) and prior actor-critic baselines on both single-reward and multi-reward benchmarks across UNet and DiT backbones.

Experimental Results

The method was evaluated on Stable Diffusion 1.5 (UNet) and Stable Diffusion 3.5M (DiT) across diverse reward types:

Single-Reward Benchmarks: The method achieved state-of-the-art results on CLIP, HPSv2.1, PickScore, GenEval, and OCR metrics. Notably, it showed significant gains on verifiable rewards (GenEval, OCR), where sparse and discrete signals benefit most from per-timestep value estimates. It also demonstrated stability where DDPO failed to converge on SD3.5M.
Multi-Reward Benchmarks: In joint optimization tasks (e.g., CLIP + HPS + GenEval), the method outperformed GRPO, particularly on GenEval, while maintaining comparable performance on preference metrics.
Inference-Time Steering: Applying the learned critic as a guidance signal during sampling further improved generation quality. On SD1.5, critic guidance alone surpassed Best-of-N (BoN) sampling; on SD3.5M, combining guidance with BoN yielded the best results.
Efficiency: The latent critic was significantly faster than pixel-space critics (CLIP/BLIP/HPS) in both value computation (10s vs. 28–35s) and training steps (20s vs. 52–86s).
Ablations:
- Removing value pretraining degraded performance, especially on sparse rewards.
- Separate value heads for multi-reward tasks outperformed shared-head designs.
- Training without Classifier-Free Guidance (CFG) yielded the most stable optimization and highest final rewards, avoiding the instability and artifacts associated with high CFG scales.

Significance and Claims

The paper claims to establish a practical recipe for applying reinforcement learning to modern diffusion models. Its significance lies in:

Resolving the State Mismatch: By aligning the critic with the policy's latent state space, it eliminates the brittleness and overhead of pixel-space evaluation.
Enabling Stable Actor-Critic Training: The combination of latent alignment, timestep conditioning, and value pretraining makes explicit actor-critic training viable and stable for diffusion models, offering a more sample-efficient alternative to group-relative methods.
Dual Utility: The learned critic serves a dual purpose: optimizing the policy during training and steering the generation process at inference time, providing performance gains beyond the training phase.
Mitigating Reward Hacking: The multi-reward extension offers a practical mechanism to balance competing objectives, reducing the prevalence of shortcut solutions in diffusion RL.

The authors acknowledge limitations, noting the current focus on image diffusion (not video), potential efficiency degradation in multi-reward settings compared to single-reward training, and discretization errors inherent in SDE sampling.