CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models

Shu Yu; Chaochao Lu

cs.CVSep 30, 2026

CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models

Authors: Shu Yu, Chaochao Lu

Organizations: Shanghai Artificial Intelligence Laboratory, Shanghai, China · Shanghai Innovation Institute, Shanghai, China · Fudan University, Shanghai, China

Abstract

Online reinforcement learning has been extended to flow matching for diffusion model (DM) image generation. However, this paradigm faces three limitations: (1) Window selection. Existing methods manually set the stochastic differential equation (SDE) sampling window, i.e., the denoising steps where exploration noise is injected. We instead determine it from each model's denoising trajectory. (2) Reward saturation. Current methods rely on scoring models trained on human annotations; we find that such scores are extremely high and nearly indistinguishable on the latest SOTA open-source DMs, making advantage estimation largely ineffective. (3) Sample inefficiency. A single scalar reward collapses different failure modes into almost identical scores, leaving minimal gradient guidance for targeted improvement. To address these issues, we propose CAST (Causal Advantage-Structured Training), an RL fine-tuning method for pretrained DMs, which (1) identifies the denoising step at which each model fixes the objects and their spatial arrangement in the image and uses that timing to set the SDE window, (2) decomposes each prompt via Causal Scene Graphs (CSG) into verifiable-atoms, i.e., minimal semantic units such as an object, count, attribute, or spatial relation that can each be checked independently, and rewards each atom separately, and (3) projects the signed atom-level advantages into pixel space through teacher-forced attention and uses them to spatially weight the SDE policy objective. We fine-tune two of the strongest open-source DMs, FLUX.2-dev and Qwen-Image-2512, with CAST, and evaluate them on GenEval 2, a compositional benchmark, and on Qwen-Image-Bench for overall quality. Within almost the same training budget, CAST's improvement over the base model on the most challenging GenEval 2 prompts is up to 3.07x that of Flow-GRPO, while overall generation quality also improves.

Figures & tables

Figure 1: Compositional failures in baselines and CAST’s improvement. GenEval 2 prompts: (a) “five striped monkeys in front of three pink cows to the left of four clocks.” (b) “a plastic horse under a glass raccoon.” GenEval 2 score reported below each image. Only CAST satisfies all verifiable-atoms in both prompts. More comparisons are provided in Appendix A .

Figure 2: PickScore saturates and weakly tracks structured correctness. (a) Score distributions pooled across the four DMs: PickScore concentrates near its ceiling, whereas GenEval 2 retains a broad dynamic range. (b) The contrast persists within every model’s generations. (c) Preference agreement between PickScore and GenEval 2 is only slightly above the 50% random baseline (dashed line); error bars show prompt-level standard deviations.

Figure 3: Overview of CAST. A Causal Scene Graph (CSG) decomposes the prompt into $K$ verifiable-atoms. Three representative types of facts are shown: attributes (red), counts (blue), and spatial relations (green). A frozen VLM uses separate passes to score each atom and extract its teacher-forced attention map. The resulting atom advantages and maps are combined into a signed spatial map that weights the SDE policy objective.

Figure 4: Training dynamics on FLUX.2-dev. Left: SFT self-distillation loss under three learning rates; the reported setting is emphasized, and all runs plateau within the budget. Right: mean Soft-TIFA reward on the shared RL manifest; the CAST curve is an atom-averaged diagnostic rather than its optimized objective, and the aesthetic reward (right axis) is on a separate scale.

Variant	Advantage	Spatial map	G2 Overall $\uparrow$	QIB Overall $\uparrow$
Sum-reward baseline	per-image	none	$84.54\pm 0.11$	$52.73\pm 0.06$
Per-atom advantage	per-atom	none	$85.12\pm 0.14$	$52.98\pm 0.18$
Spatial weighting	per-image	aggregated attention	$84.91\pm 0.16$	$53.06\pm 0.20$
Shuffled correspondence	per-atom	shuffled atom maps	$85.00\pm 0.15$	$52.89\pm 0.22$
CAST	per-atom	matched atom maps	$\mathbf{85.95\pm 0.10}$	$\mathbf{53.44\pm 0.29}$

Table 2: Ablation of atom-level credit assignment and spatial weighting on FLUX.2-dev. All variants use the same training prompts, SDE window, generated-image budget, and optimization settings. Values are mean

\pm

sample standard deviation over four matched generation-seed replicates.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Additional qualitative comparisons on FLUX.2-dev (part I). GenEval 2 prompts with multi-object counts, attributes, and chained spatial relations. GenEval 2 score reported below each image.

Figure 6: Additional qualitative comparisons on FLUX.2-dev (part II). GenEval 2 score reported below each image.

Figure 7: Image structure forms in early denoising steps. Left: intermediate decoded images at five timesteps for three diffusion models. Dashed borders mark the knee point where CLIP similarity plateaus. Right: CLIP-ViT-L/14 similarity between Tweedie estimates and the text prompt. The knee point (triangle) occurs before $\tau{=}0.25$ across all models, indicating that the main layout and composition are established early.

Setting	Value
Image resolution	$1024\times 1024$
Sampling steps	50
Guidance	Backbone default
LoRA rank	64
LoRA scaling factor	$\alpha=64$
Target modules	Matched within each backbone

Appendix

Table 3: Settings shared by all compared methods.

Figure 8: Attention-head sensitivity across denoising steps. Each heatmap cell averages a $4\!\times\!4$ group of transformer layers and attention heads. The side panels show the layer-wise mean sensitivity at each step, and the dashed line separates the dual-stream layers (0–7) from the single-stream layers (8–55).

Explore similar work

Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models
May 11, 2026Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken +2Generative Flow Networks Diffusion Models
1TU Munich, MCML · 2MIT CSAIL · 3King’s College London +2
Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
Sep 29, 2025Shuchen Xue, Chongjian Ge, Shilong Zhang +2Diffusion Models Diffusion Policies
UCAS · Adobe Research · HKU +1
DRM: Diffusion-based Reward Model With Step-wise Guidance
May 25, 2026Jaxon Zhang, Binxin Yang, Hubery Yin +2Diffusion Alignment Diffusion Models
Peking University · Work done during an internship at WeChat Vision, Tencent Inc. · WeChat Vision, Tencent Inc.

May 11, 2026, cs.LGJ/K move · Enter open · S save

Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models

Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich+1

1TU Munich, MCML · 2MIT CSAIL · 3King’s College London+2

Generative Flow Networks Diffusion Models Reinforcement Learning Post-Training Reward Gradients Adjoint Matching Post-Training+1

Open paper AskPDF arXiv

Diffusion and flow-matching models scale because pretraining is supervised regression: a clean sample is noised analytically, and a model regresses against a closed-form target. RL post-training aligns the model with a reward. In image generation, this makes samples compose objects correctly, render text legibly, and match human preferences. Existing methods rely on costly SDE rollouts, reward gradients, or surrogate losses, sacrificing pretraining's regression structure. We show that the structure extends to RL post-training. Under KL-regularized reward maximization, the optimal generative process tilts the clean-endpoint distribution towards samples with higher reward and leaves the noising law unchanged. Combining this with the adjoint-matching optimality condition and a REINFORCE identity, we derive Reinforce Adjoint Matching (RAM): a consistency loss that corrects the pretraining target with the reward. At each step, we draw a clean endpoint from the current model, evaluate its reward, noise it as in pretraining, and regress. No SDE rollouts, backward adjoint sweeps, or reward gradients are required. Like the pretraining objective, RAM is simple and scales. On Stable Diffusion 3.5M, RAM achieves the highest reward on composability, text rendering, and human preference, reaching Flow-GRPO's peak reward in up to

50\times

fewer training steps.

Abstract

Explore similar work

Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

DRM: Diffusion-based Reward Model With Step-wise Guidance