PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Shu Yu, Chaochao Lu

arXiv:2609.39441v1cs.CVcs.AIcs.LG

TL;DR

Strong diffusion models expose saturated preference scores, while scalar rewards obscure which compositional units fail. CAST measures a model-specific layout window and spatially grounds independently scored atom advantages; at equal budget, it outperforms Flow-GRPO on the hardest GenEval 2 prompts while preserving overall quality.

  • Problem

    Diffusion RL lacks a model-specific basis for choosing the SDE window, preference scores saturate on strong open-source models, and scalar rewards obscure which prompt atoms fail.

    Source: §1 · pp. 1–2

  • Method

    CAST measures each model’s layout-formation transition, scores CSG-derived verifiable atoms independently, and uses teacher-forced VLM attention to map their advantages onto pixels for spatially weighted SDE optimization.

    Source: §1 · p. 2

  • Results

    On the GenEval 2 Hard Case subset, CAST improves over the base model by 1.93×–3.07× as much as Flow-GRPO at the same budget and reward signal, while achieving the best overall quality on Qwen-Image-Bench.

    Source: §1 · p. 3

  • Takeaways & Limitations

    On the evaluated strong diffusion models, CAST provides finer-grained credit assignment than scalar rewards while preserving overall generation quality.

    Source: §7 · p. 10

    The motivating diagnostics are supported by empirical evidence on state-of-the-art open-source diffusion models.

    Source: §3 · p. 3

Abstract

from arXiv · show

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.

1 INTRODUCTION

CAST addresses model-specific SDE window selection and weak scalar rewards with independently scored, spatially grounded atom advantages. On challenging GenEval 2 prompts, it outperforms Flow-GRPO at similar budget while improving overall quality.

1 INTRODUCTION

  • Because layout is established at different denoising steps across models, CAST measures each model’s transition from layout formation to detail refinement to set its SDE window.Objects and spatial layout form early, while later steps mainly refine low-level visual details.Source: §1 · pp. 1–2
  • Preference scores nearly indistinguishable on strong open-source diffusion models and scalar rewards that conflate distinct failures provide little signal for targeted improvement.Prompts contain 3–10 independently verifiable atoms, with a median of 6.5; even summed atom rewards barely reveal which atom failed.Source: §1 · p. 2
  • CAST builds training prompts from Causal Scene Graphs, scores each verifiable atom independently, and uses its advantage and attention heatmap to spatially weight the SDE policy objective.The atom heatmaps come from teacher-forced attention in a frozen vision-language model.Source: §1 · p. 2
  • On the GenEval 2 Hard Case subset, CAST improves over the base model by 1.93×–3.07× as much as Flow-GRPO at the same budget and reward signal, while achieving the best overall quality on Qwen-Image-Bench.The base models already score ≥0.95 on 11.6%–18.8% of prompts, so the study focuses on cases where they still fail.Source: §1 · p. 3
  • The per-pixel advantage map directs gradients toward image regions responsible for each atom’s success or failure instead of weighting the whole image uniformly.CAST projects atom-level advantages into pixel space through vision-language-model attention heatmaps.Source: §1 · p. 3

2 PRELIMINARIES

The preliminaries describe rectified-flow sampling and the marginal-preserving SDE used for stochastic exploration, then introduce GRPO’s group-normalized advantages. CAST uses these foundations with per-atom, spatially weighted advantages.

2.1 FLOW MATCHING AND MARGINAL-PRESERVING SDE SAMPLING

  • Rectified-flow sampling interpolates data and Gaussian noise along a linear path, with the velocity network defining the reference flow.The setup samples data and conditioning, Gaussian noise, and a time uniformly from [0, 1].Source: §2.1 · p. 3
  • Flow-GRPO replaces the deterministic ODE rollout with a reverse-time SDE that preserves continuous-time marginals; its Gaussian transition enables transition ratios and KL regularization.For the rectified-flow path, the score is computed without a separate score model, and CAST uses the SDE inside its calibrated window and the ODE outside.Source: §2.1 · p. 3

2.2 GROUP RELATIVE POLICY OPTIMIZATION (GRPO)

  • GRPO normalizes rewards within a group to form advantages for a clipped policy-gradient objective; CAST replaces the image-level advantage with per-atom, spatially weighted advantages.The group advantage is each output’s reward minus the group mean, divided by the group standard deviation plus ϵ.Source: §2.2 · p. 3

3 DIAGNOSTIC OF THREE LIMITATIONS IN CURRENT DM RL

The diagnostics identify three limitations in diffusion-model RL: layout timing varies by model, preference scores saturate on strong models, and scalar rewards obscure which compositional facts fail.

3.1 IMAGE STRUCTURE IS ESTABLISHED EARLY DURING DENOISING

  • Layout emerges early but at model-specific denoising points: maximum-curvature estimates occur at τ ≈0.11 for SD3.5-Large, τ ≈0.21 for FLUX.2-dev, and τ ≈0.14 for Qwen-Image-2512.The decoded estimates at those points already show the main layout and spatial composition; later SDE perturbations risk disrupting established content.Source: §3.1 · p. 4

3.2 PREFERENCE MODELS SATURATE ON STRONG DMS

  • PickScore has per-model means of 0.887–0.907 and CoVs of 0.045–0.048, while GenEval 2 means are 0.734–0.829 and CoVs are 0.177–0.225.The preference scores are highly saturated and weakly discriminative on strong DMs, whereas GenEval 2 preserves more score variation for advantage estimation.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3 SCALAR REWARD IS BLIND TO STRUCTURAL CORRECTNESS

  • PickScore and GenEval 2 agree on pairwise image rankings only 52.5%–56.7% of the time, near the 50% random baseline.The figure compares score distributions pooled across four DMs and within each model, then shows preference agreement with prompt-level standard-deviation error bars.Source: §3.3 · p. 4, Figure 2 · p. 4
  • A scalar reward can assign nearly identical values to images with different failures, such as incorrect spatial relations versus missing objects.Collapsing GenEval 2’s separate atom scores into an aggregate would recreate this credit-assignment ambiguity, so the structured outputs must be preserved through optimization.Source: §3.3 · pp. 4–5

4 METHOD

CAST decomposes prompts into verifiable facts, scores and spatially grounds each fact, then uses the resulting atom-level advantages to weight policy optimization over a model-specific SDE window.

4 METHOD

  • CAST uses independently verifiable atoms to direct training gradients toward image regions containing their evidence through probe construction, scoring and grounding, and spatially weighted optimization.Its pipeline combines independently normalized atom advantages within a per-model SDE window.Source: §4 · p. 5, §4.3 · p. 7

4.1 VQA PROBE CONSTRUCTION VIA CAUSAL SCENE GRAPHS

  • A Causal Scene Graph represents entities as nodes and relations as edges, mapping node existence, attributes, and edges to object, property, and relational probes.CAST samples the graph first and renders it into a prompt, pairing each atom with a probe and the prompt words that express it before image generation.Source: §4.1 · p. 5

4.2 PER-ATOM SCORING AND SPATIAL GROUNDING

  • A frozen Qwen3-VL-8B scorer scores each atom independently by summing softmax probabilities for accepted answer variants.A separate teacher-forced pass uses atom-linked prompt tokens to produce a unit-spatial-mean attention heatmap for each atom.Source: §4.2 · p. 6

4.3 SPATIALLY WEIGHTED SDE POLICY OPTIMIZATION

  • Each trajectory samples a contiguous two-transition SDE window from a model-specific early candidate range defined by the measured structure-formation point.Sampling outside the window follows the original deterministic ODE; transition probabilities inside the window are retained for policy optimization.Source: §4.3 · p. 6
  • CAST combines each atom’s group-normalized advantage with its heatmap to create a signed patch-level advantage map for the clipped spatial policy objective.If score variation is below the scorer’s numerical resolution, the corresponding advantages are set to zero; signed contributions can offset at a location before clipping.Source: §4.3 · pp. 6–7

5 EXPERIMENTS

Experiments compare CAST with base, distillation, and Flow-GRPO baselines on compositional adherence and overall quality. On hard prompts, CAST outperforms matched-budget scalar-reward Flow-GRPO on both backbones.

5.1 EXPERIMENTAL SETUP

  • The trainable methods share 1,024 synthesized prompts, while the Flow-GRPO variants and CAST also share a 128-prompt manifest, prompt order, seed, SDE window, and learning-rate schedule.The matched runs each use 2,048 generated images and 128 optimizer updates; baselines include two SFT variants and two Flow-GRPO reward variants.Source: §5.1 · p. 7

5.2 EVALUATION

  • Evaluation covers all 800 GenEval 2 prompts and 1,000 Qwen-Image-Bench prompts, with a 50-prompt Hard Case subset reported alongside full-set averages.Hard Case prompts are selected using frozen FLUX.2-dev base generations and verifier scores; evaluation uses matched seeds and four generation-seed replicates.Source: §5.2 · pp. 7–8

5.3 MAIN RESULTS

  • CAST’s Hard Case gain over the base model is 1.93× that of scalar-reward Flow-GRPO on FLUX.2-dev and 3.07× on Qwen-Image-2512, with the same 2,048-image and 128-update budget.CAST also exceeds SFT (3/64) on both backbones and attains the highest Qwen-Image-Bench Overall on each.Source: §5.3 · p. 8, Table 1 · p. 8
  • The ablations isolate atom-level credit and spatial routing under the same on-policy training setting, SDE window, budget, and optimization settings.Figure 4 also presents training dynamics for SFT and RL runs on FLUX.2-dev.Source: §5.3 · p. 8

6 RELATED WORK

Prior diffusion-RL work extends policy optimization to flow matching and improves temporal credit assignment, while related work develops compositional evaluation and different forms of reward guidance. CAST is distinguished by combining temporal and spatial credit assignment with atom-specific evidence.

RL for Diffusion and Flow Matching

  • CAST complements temporal credit assignment with spatial credit assignment by weighting each atom’s evidence region within a calibrated early SDE window.Earlier methods use terminal rewards, per-step reward gains, adaptive timestep-wise stochasticity, or initial-noise perturbations.Source: §6 · p. 9

Compositional Generation and Evaluation

  • GenEval 2 evaluates prompt atoms across object, count, attribute, spatial-relation, and verb skills, while CAST uses causal scene graphs to train and localize atom advantages.TIFA uses prompt-generated VQA questions; VQAScore uses a single global question; Soft-TIFA scores with VLM token probabilities.Source: §6 · p. 9

Dense Reward for Generation

  • Unlike global preference dimensions, scalar VQA-question scores, or a single global heatmap, CAST pairs each atom advantage with its own evidence map.MPS decomposes preference into dimensions without spatial localization; VisionReward scores each question globally; Focus-N-Fix uses one global heatmap.Source: §6 · p. 9

7 CONCLUSION

CAST combines a model-specific early SDE window with separate verifiable-atom advantages projected onto image regions through teacher-forced attention. On two strong diffusion models, it improves adherence on the hardest GenEval 2 prompts beyond Flow-GRPO at equal budget while preserving overall quality.

Conclusion

  • The diagnostics show that strong diffusion models fix their layout early and that preference rewards saturate on them, while scalar rewards conflate distinct failures.CAST responds with spatially grounded atom-level rewards.Source: §7 · pp. 9–10
  • CAST improves adherence on the hardest GenEval 2 prompts well beyond Flow-GRPO at equal budget while preserving overall generation quality.The paper presents this as finer-grained credit assignment than a scalar reward.Source: §7 · p. 10
  • CAST sets a model-specific early SDE window, computes a separate advantage for each verifiable atom, and projects it through teacher-forced attention onto its image region.The projected advantages weight the policy objective spatially.Source: §7 · p. 10

A ADDITIONAL QUALITATIVE COMPARISONS

The appendix provides additional FLUX.2-dev qualitative comparisons and details how the structure-formation point is detected and discretized into candidate SDE windows.

A ADDITIONAL QUALITATIVE COMPARISONS

  • Each qualitative-comparison row presents outputs from the base model, strongest SFT variant, both Flow-GRPO variants, and CAST for one GenEval 2 prompt, with a score for each image.Figure 6 is part II of the additional FLUX.2-dev comparisons.Source: A ADDITIONAL QUALITATIVE COMPARISONS · p. 15, Figure 6 · p. 16

B STRUCTURE FORMATION POINT DETECTION

  • The structure-formation point is the maximum-curvature point of the rise-then-plateau CLIP-similarity trajectory, marking the knee between semantic alignment and refinement.The detected points are τ*≈0.11 for SD3.5-Large, τ*≈0.21 for FLUX.2-dev, and τ*≈0.14 for Qwen-Image-2512, all before τ=0.25.Source: B STRUCTURE FORMATION POINT DETECTION · pp. 16–17, Figure 7 · p. 17
  • For an N-step sampler, the measured point sets a candidate transition range, from which a length-2 window is sampled and only its two transitions are stochastic.For N=50, the ranges are [0,11) for FLUX.2-dev and [1,7) for Qwen-Image-2512; the other transitions follow the deterministic ODE.Source: B STRUCTURE FORMATION POINT DETECTION · p. 17

C ADDITIONAL ABLATION DETAILS

The ablation isolates atom-level credit assignment and spatial weighting by comparing reward aggregation and atom-to-map routing under a shared on-policy training setup.

C ADDITIONAL ABLATION DETAILS

  • The sum-reward baseline broadcasts one normalized per-image advantage uniformly across the image.The variants share prompts, SDE window, generated-image budget, and optimization settings.Source: C ADDITIONAL ABLATION DETAILS · p. 17
  • The per-atom variant normalizes each atom separately before broadcasting the sum of atom advantages.It uses the same on-policy setting as CAST.Source: C ADDITIONAL ABLATION DETAILS · p. 17
  • Spatial weighting combines the per-image advantage with an aggregate attention map, while shuffled correspondence permutes atom-to-map pairings.The ablation keeps the SDE window, budget, and optimization settings fixed across variants.Source: C ADDITIONAL ABLATION DETAILS · p. 17

D IMPLEMENTATION DETAILS

The appendix specifies shared training and cost-accounting conventions and reports a diagnostic showing skill-specific attention sensitivity timing alongside substantial overlap across skills.

D IMPLEMENTATION DETAILS

  • Training uses 1,024 graph-generated prompts with atom probes, disjoint from both evaluation sets, and 256 held-out development prompts.Each method runs 128 single-prompt collections with group size 16, totaling 2,048 generated images and 128 optimizer updates.Source: D IMPLEMENTATION DETAILS · p. 17
  • Methods share prompt order, seed, SDE window, schedule, and LoRA setup; Flow-GRPO G2 averages atom scores and uses neither atom advantages nor spatial maps.Generated-image and optimizer-update budgets are reported separately.Source: D IMPLEMENTATION DETAILS · pp. 17–18

E TRAINING COST ACCOUNTING

  • Reported H200-hours include generation, scoring, attention extraction when needed, and optimization; costs are compared only within the same backbone.Source: E TRAINING COST ACCOUNTING · p. 18

F ATTENTION-HEAD DIAGNOSTIC

  • The diagnostic compares correct and incorrect pre-projection activations across 56 blocks at steps 3, 7, 14, and 21 on 10 hard prompts, with permutation and FDR testing.It aggregates results by count, attribute, and spatial relation.Source: F ATTENTION-HEAD DIAGNOSTIC · p. 18
  • Spatial relations are most sensitive early, attributes from step 7 onward, and count at step 7; the most sensitive heads overlap substantially across skills.The attention-head diagnostic is not used by CAST during training.Source: F ATTENTION-HEAD DIAGNOSTIC · p. 18

G EVALUATION PROTOCOL

The evaluation protocol aggregates scores across fixed benchmark sets and defines a shared 50-prompt Hard Case subset using base-model results.

G EVALUATION PROTOCOL

  • GenEval 2 averages Overall and skill scores over 800 prompts, while Qwen-Image-Bench averages Overall over 1,000; reported results are means and sample standard deviations across four replicates.Source: G EVALUATION PROTOCOL · p. 18
  • Base-model scores were ≥0.95 on 18.8% of FLUX.2-dev prompts and 11.6% of Qwen-Image-2512 prompts.These figures average over four evaluation seeds.Source: G EVALUATION PROTOCOL · p. 18
  • The Hard Case subset retains the 100 prompts with highest within-prompt score variance, then selects the 50 with lowest mean score.The subset is fixed before training and shared across methods and both backbones.Source: G EVALUATION PROTOCOL · p. 19
Loading 2609.39441v1…