Video World Models · 2026-10-01
CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models
Shu Yu, Chaochao Lu
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
CAST is an RL fine-tuning method for pretrained diffusion image generation models (flow-matching). The authors (1) set the SDE stochastic window from each model’s denoising trajectory, (2) decompose prompts with Causal Scene Graphs into independently verifiable atoms and score them separately, and (3) project signed atom-level advantages into pixel space via teacher-forced attention to spatially weight the SDE policy objective.
Why it matters
The authors address three issues in diffusion RL: manually chosen SDE windows, preference-reward saturation on strong open-source DMs, and scalar rewards that conflate distinct failure modes. By using atom-level structured rewards and grounding them spatially, the authors aim to restore discriminative credit assignment and focus optimization on the image regions tied to compositional correctness.
Method
- Determine the SDE sampling window by measuring a model-specific structure-formation point from Tweedie-estimated clean images and CLIP-ViT-L/14 similarity curvature.
- Decompose prompts via Causal Scene Graphs into verifiable-atoms and score each atom independently using GenEval 2 Soft-TIFA (Qwen3-VL-8B) with teacher-forced attention heatmaps for pixel grounding.
- Normalize atom scores within each prompt group into z-score advantages, project them to a per-pixel advantage map, and use it to spatially weight a clipped spatial SDE policy objective; outside the window use the original deterministic ODE.
Limitation
The paper evaluates CAST on FLUX.2-dev and Qwen-Image-2512 and reports results on GenEval 2 and Qwen-Image-Bench; no other backbones or benchmarks are reported in the provided text.
Abstract (from arXiv)
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.
Related papers
- DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models
- Glance Forcing: Making Bidirectional Diffusion Autoregressive at a Glance
- Reinforcement Learning-driven Diffusion Models: Language, Vision, and Beyond
- Are We Overthinking One-Step Generation?
- LongForcing: Scaling Temporal Manifold Matching for Long-Horizon World Rollout
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-09-30 · Indexed by Pub-AI: 2026-10-01 UTC
CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models
Preprint posted: 2026-09-30 · Original record ↗
- arxiv
- 2609.39441
Explicit source versions (1)
- v1 · 2026-09-30
- Preprint available
Source date: 2026-09-30 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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.