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

arXiv: https://arxiv.org/abs/2609.39441 · HTML: https://mlsanity.com/paper/2609.39441/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 30 of its 33 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 30)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Qwen3-VL Technical Report](https://mlsanity.com/paper/2511.21631/) (2025)
- [Prompt-to-Prompt Image Editing with Cross Attention Control](https://mlsanity.com/paper/2208.01626/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Qwen-Image Technical Report](https://mlsanity.com/paper/2508.02324/) (2025)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [Evaluating Text-to-Visual Generation with Image-to-Text Generation](https://mlsanity.com/paper/2404.01291/) (2024)
- [TIFA: Accurate and Interpretable Text-to-Image Faithfulness Evaluation with Question Answering](https://mlsanity.com/paper/2303.11897/) (2023)
- [Generic Attention-model Explainability for Interpreting Bi-Modal and Encoder-Decoder Transformers](https://mlsanity.com/paper/2103.15679/) (2021)
- [Perception Prioritized Training of Diffusion Models](https://mlsanity.com/paper/2204.00227/) (2022)
- [Z-Image: An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer](https://mlsanity.com/paper/2511.22699/) (2025)
- [MixGRPO: Unlocking Flow-based GRPO Efficiency with Mixed ODE-SDE](https://mlsanity.com/paper/2507.21802/) (2025)
- [VisionReward: Fine-Grained Multi-Dimensional Human Preference Learning for Image and Video Generation](https://mlsanity.com/paper/2412.21059/) (2024)
- [Learning Multi-dimensional Human Preference for Text-to-Image Generation](https://mlsanity.com/paper/2405.14705/) (2024)
- [GenEval 2: Addressing Benchmark Drift in Text-to-Image Evaluation](https://mlsanity.com/paper/2512.16853/) (2025)
- [HiDream-O1-Image: A Natively Unified Image Generative Foundation Model with Pixel-level Unified Transformer](https://mlsanity.com/paper/2605.11061/) (2026)
- [Qwen-Image-Bench: From Generation to Creation in Text-to-Image Evaluation](https://mlsanity.com/paper/2605.28091/) (2026)
- [DenseGRPO: From Sparse to Dense Reward for Flow Matching Model Alignment](https://mlsanity.com/paper/2601.20218/) (2026)
- [Neighbor GRPO: Contrastive ODE Policy Optimization Aligns Flow Models](https://mlsanity.com/paper/2511.16955/) (2025)
- [Stepwise Credit Assignment for GRPO on Flow-Matching Models](https://mlsanity.com/paper/2603.28718/) (2026)
- [Focus-N-Fix: Region-Aware Fine-Tuning for Text-to-Image Generation](https://mlsanity.com/paper/2501.06481/) (2025)
- [Alleviating Sparse Rewards by Modeling Step-Wise and Long-Term Sampling Effects in Flow-Based GRPO](https://mlsanity.com/paper/2602.06422/) (2026)
- [AEGPO: Adaptive Entropy-Guided Policy Optimization for Diffusion Models](https://mlsanity.com/paper/2602.06825/) (2026)
- [LINA: Learning INterventions Adaptively for Physical Alignment and Generalization in Diffusion Models](https://mlsanity.com/paper/2512.13290/) (2025)

## 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.