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

arXiv:2609.39441 · cs.CV, cs.AI, cs.LG · Submitted 2026-09-30 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models".

Jane: Online reinforcement learning has been successfully extended to flow matching for diffusion model (DM) image generation, but existing methods suffer from limitations regarding window selection, reward saturation, and sample inefficiency.

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, we're diving into this paper today titled "CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models." Basically, the authors are tackling some real headaches in using reinforcement learning for diffusion models.

Jane: That sounds complex, Tom. Can you break down what they're saying about the main problems they're trying to solve? I want to make sure our listeners get the big picture right away.

Lu: Well, the core thesis seems to be that existing methods struggle with window selection, where people have to guess when exploration noise should be added into the denoising steps.

Lu: They also found that reward saturation is a big issue because preference models like PickScore give very high scores on strong diffusion models, and a simple scalar reward just mixes up different kinds of failures.

Meng: From an engineering standpoint, I’m curious about how they propose solving that sample inefficiency where a single reward doesn't tell you which part of the image generation failed.

Lalam: I think this is really significant because if we can separate the rewards for different parts of a prompt, it could lead to much more nuanced and controllable image generation capabilities in the future.

Tom: Exactly, that’s what they are aiming for with CAST: identifying the optimal denoising step based on layout formation and rewarding distinct semantic units independently through pixel-space grounding. It seems like a really smart way to guide the training process.

Jane: So, instead of one big score for the whole image, they are breaking down prompts into smaller, verifiable atoms—like an object or a spatial relation—and scoring those separately before applying the reinforcement learning signal. That makes a lot more sense for debugging generation issues.

Lu: They use something called a Causal Scene Graph, which decomposes the prompt so that every node and edge has a verifiable-atom probe attached to it, essentially turning the prompt into these independently checkable semantic units.

Lu: This structure lets them build data where they can check for causal links between these atoms.

Meng: That sounds like a way to create highly structured training data, which is crucial for keeping the AI consistent as models get bigger and more complex. How does this structural approach translate into a practical optimization step?

Lalam: It moves the focus from just getting a generally good image to ensuring every specific component of that image meets its required criteria, which really elevates the quality control aspect of diffusion model training.

Tom: And they tie all this together in three stages: first, constructing the CSG; second, per-atom scoring and grounding using models like Qwen3-VL; and finally, spatially weighted optimization where those atom advantages are projected onto a per-pixel advantage map to weight the SDE policy objective.

Paper summary: Jane: That mapping step sounds particularly innovative. So they aren't just rewarding the atom score; they are turning that into a map showing exactly which pixels need attention during optimization.

Lu: Precisely, that per-pixel advantage map is key because it projects those atom-level advantages into pixel space via VLM attention heatmaps, which allows gradients to reach the specific image regions responsible for success or failure.

Lu: This is how they ensure the model learns where to focus its effort.

Meng: If we can pinpoint exactly where the optimization needs to happen based on structural correctness rather than just a general preference score, that drastically simplifies our iterative development process. It moves us closer to predictable control over complex outputs.

Lalam: And for me, this means the future of AI image generation isn't just about producing pretty pictures; it’s about building systems where we can guarantee specific compositional elements are rendered correctly every single time.

Tom: The paper also points out that existing methods fail because a scalar reward collapses different failure modes into identical scores, meaning they miss distinct issues entirely. They explicitly show that this is true when comparing CAST to the Flow-GRPO with a scalar reward on the Hard Case subset, where CAST sees improvements of up to three point zero seven times over those baseline results.

Jane: That comparison number is quite compelling; showing a three point zero seven times improvement in those challenging cases suggests that this compositional approach yields much stronger guidance than the standard methods we see today.

Lu: They also addressed window selection by using a Tweedie estimate to track when objects and their spatial layout become established in each model, which they then use to set the SDE sampling window where exploration noise is injected.

Lu: This tackles that issue of manually setting those steps effectively.

Tom: So, to wrap up this summary of "CAST: Causal Advantage-Structured Training with Spatially Grounded Compositional Rewards for Diffusion Models," we see a method that systematically decomposes generation into verifiable parts, scores those parts separately, and then uses that structural information to guide the optimization process directly into the image pixels.

Jane: It seems like they've moved RL fine-tuning away from relying on general preference models toward a more grounded, atom-level understanding of what makes an image successful.

Lu: The implication for research is that we can now build better ways to structure our prompts for training, leveraging causal graphs instead of just raw text.

Meng: For practical AI development, this suggests we can create more robust training pipelines where the model is explicitly told *how* to handle composition rather than implicitly hoping it figures it out from a single reward signal.

Lalam: If this technique becomes standard, we could see an explosion in the complexity and reliability of generative AI systems across all modalities, making them far more trustworthy for complex tasks.

Conclusion: Tom: So, to recap this session, we’ve been dissecting CAST—Causal Advantage-Structured Training—and how it uses atom-level rewards to fine-tune diffusion models by grounding those rewards in pixel space.

Jane: That’s right, Tom; essentially, they figured out a way to give the AI feedback not just on the final image quality, but on *why* certain parts of that image succeeded or failed during the generation process.

Lu: I think what’s really striking is how they build this causal scene graph to systematically probe every semantic unit in the prompt, which opens up so many avenues for structured data creation.

Meng: From my side, it’s fascinating how they manage to translate those high-level semantic scores down into a per-pixel map that directly influences the optimization steps of the diffusion model.

Lalam: For me, this means we’re moving toward a future where AI can be trained with much more fine-grained control over its compositional logic, which has huge implications for how we think about digital culture.

Tom: Speaking of control, the title itself is pretty descriptive; CAST really lays out exactly what they did—causal advantage structure and spatial grounding—to solve those old problems with reward saturation and sample inefficiency.

Jane: And the authors’ work on breaking down these complex generation processes into these verifiable atoms seems like a really practical step toward making AI outputs more predictable.

Lu: The implication here is that we can start treating prompts not just as long strings of text, but as structured blueprints where every component has its own measurable contribution to the final result.

Meng: That structured blueprint idea is what interests me; it suggests a path for building training pipelines that are less reliant on broad preference models and more focused on explicit compositional correctness.

Lalam: If we can get this level of atom-level control, I see AI systems becoming far more reliable for tasks requiring strict adherence to complex visual rules, which could really reshape creative and technical fields.

Tom: It’s clear that CAST tackles the core issues of reward ambiguity and sample inefficiency by making the reinforcement learning signal specific to the spatial layout decisions happening inside the model.

Jane: So, we’re talking about moving from broad success metrics to a detailed map showing exactly which pixels need attention during training, which is a big conceptual step.

Lu: The future work they hinted at involves expanding this causal framework beyond simple scene graphs to handle more intricate, long-range dependencies in the generated images.

Meng: I’m looking forward to seeing how this translates into real-world deployment costs; if the training time remains comparable to existing methods, that makes it immediately interesting for practical application.

Lalam: It really feels like a step toward building AI that understands not just what an object is, but precisely where and how that object relates spatially to everything else in the frame.

Tom: It’s wild stuff, folks; CAST gives us a concrete mechanism for training models with far more structured and traceable compositional awareness.

Jane: It really puts the focus squarely on the mechanics of image formation rather than just chasing a high aggregate score.

Lu: The way they grounded those atom advantages into pixel space via VLM attention heatmaps is particularly elegant, showing how language understanding can directly guide physical optimization in the latent space.

Meng: I’m genuinely excited to see if this approach scales effectively across different types of generative models and data sets.

Lalam: This kind of structured guidance means we can start designing AI for specific aesthetic or functional requirements with much greater precision than we have now.

Tom: Next up, we’ll be looking at the detailed experimental results showing those performance improvements on challenging datasets.

Shu Yu, Chaochao Lu

Shanghai Artificial Intelligence Laboratory · Shanghai Innovation Institute · Fudan University

cs.CV, cs.AI, cs.LG

Submitted: 2026-09-30

Updated: 2026-09-30

Comments: Project page: https://opencausalab.github.io/CAST

Project page: https://opencausalab.github.io/CAST

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 92/100

The gist: Online reinforcement learning has been successfully extended to flow matching for diffusion model (DM) image generation, but existing methods suffer from limitations regarding window selection,

Key concepts

Causal Scene Graph (CSG)
This is a structure built from prompts where every element—like an object or relation—is paired with a verifiable probe question. This decomposes the complex prompt into 'verifiable-atoms,' which are minimal, independently checkable semantic units like an attribute or a spatial relation.
Per-pixel Advantage Map
This map projects the calculated advantages of individual semantic atoms onto the image pixels. It uses attention heatmaps from a Vision-Language Model to show exactly which pixel regions are most important for success or failure, allowing optimization to target specific areas.
Verifiable-Atoms
These are the minimal semantic units derived from the CSG, such as an object's count or a spatial relationship. Each atom can be checked independently using a probe question. They allow the model to be rewarded for satisfying specific, distinct parts of the prompt.
Spatially Weighted Optimization
This is the final stage where advantages from different atoms are combined and projected onto a per-pixel map. This map is then used to weight the diffusion model's objective function, ensuring training focuses on denoising steps that correctly form the desired spatial layout.

Terminology

Summary

Online reinforcement learning has been successfully extended to flow matching for diffusion model (DM) image generation, but existing methods suffer from limitations regarding window selection, reward saturation, and sample inefficiency. This paper introduces CAST (Causal Advantage-Structured Training), a new RL fine-tuning method that addresses these issues by identifying the optimal denoising step based on layout formation and rewarding distinct semantic units independently through pixel-space grounding.

How it works

CAST is an RL fine-tuning method for pretrained DMs consisting of three stages: (1) CSG probe construction, (2) per-atom scoring and grounding, and (3) spatially weighted optimization. The first stage involves constructing a Causal Scene Graph (CSG), where prompts are synthesized from these graphs such that every node and edge is paired with a verifiable-atom probe. This decomposes the prompt into verifiable-atoms, i.e., minimal semantic units such as an object, count, attribute, or spatial relation that can each be checked independently.

The second stage focuses on scoring and grounding these atoms. Each verifiable-atom is scored using the GenEval 2 Soft-TIFA protocol with a frozen Vision-Language Model (VLM), such as Qwen3-VL. This yields a score, but to provide spatial information, we run a separate teacher-forced decoding pass. This process extracts evidence by mapping each atom to the words in the prompt that states the corresponding fact and normalizing it to unit spatial mean at pixel space, yielding a heatmap for every atom.

The third stage involves spatially weighted optimization. CAST normalizes each atom’s score within a group into an advantage, projects these advantages into a per-pixel advantage map, and uses this map to weight the SDE policy objective. The final loss function is defined as:

LCAST = − Ei, t∈W, ph

min r(i)t,pM(i) p

This approach ensures that optimization targets the steps where layout is decided by using a model-specific early candidate range derived from the measured structure-formation point.

Diagnostic of Limitations

The paper identifies three key limitations in current DM RL paradigms: (1) Window selection, which is addressed by using a Tweedie estimate to track when objects and their spatial layout become established in each model; (2) Reward saturation, where preference models like PickScore are found to be extremely high and nearly indistinguishable on the latest SOTA open-source DMs; and (3) Sample inefficiency, where a scalar reward collapses different failure modes into identical scores. The paper demonstrates that a scalar reward is blind to structural correctness because it fails to separately check each verifiable-atom, leading to credit assignment ambiguity.

Compositional Generation and Evaluation

The method relies on the Causal Scene Graph (CSG) for data construction, where every node and edge is paired with a VQA probe question. This allows for scaling datasets while maintaining a clean causal structure among its verifiable-atoms. The evaluation utilizes GenEval 2 to check adherence to all verifiable-atoms, and QwenImage-Bench to measure overall generation quality. The results show that CAST produces images that satisfy all verifiable-atoms where existing methods fail, achieving improvements of up to 3.07× over Flow-GRPO on the most challenging GenEval 2 prompts while maintaining overall quality.

Key Contributions

The main contributions are: (1) showing that preference rewards saturate on strong DMs and that a scalar reward conflates distinct failures; (2) proposing CAST, which applies independently normalized verifiable-atom advantages within a per-model SDE window; and (3) introducing the per-pixel advantage map, which projects atom-level advantages into pixel space via VLM attention heatmaps. This enables gradients to reach the specific image regions responsible for success or failure.

Experimental Results

When applied to FLUX.2-dev and Qwen-Image-2512, CAST improves over the base model by 1.93×–3.07× as much as scalar-reward Flow-GRPO on the Hard Case subset, while also surpassing high-selectivity distillation SFT (3/64) in terms of hard case gain and achieving the highest Qwen-Image-Bench Overall quality on both backbones. The training cost for CAST is comparable to Flow-GRPO variants within the same budget. The ablation studies confirm that Per-atom advantage per-atom none provides a stronger signal than sum-reward baseline, and CAST per-atom matched atom maps yields the highest overall performance metrics.

The gist: CAST trains diffusion models with spatially grounded atom-level rewards by identifying the optimal denoising step based on layout formation and rewarding distinct semantic units independently through pixel-space grounding.

Improvements for AI systems

Here are the specific improvements that can be made to AI systems using the CAST (Causal Advantage-Structured Training) method, based on this research:

  1. The model will gain a mechanism for structural credit assignment by learning which specific components of a complex prompt (objects, counts, attributes, spatial relations) caused success or failure during image generation.

  2. The system will be able to dynamically adjust the exploration noise injection window during the diffusion process; it will stop exploring when it determines that the core objects and their spatial arrangement have been fixed in the early denoising steps (the structural formation point).

  3. The system will use a highly discriminative, atom-level reward signal instead of relying on saturated preference scores, allowing for targeted training on specific compositional errors (e.g., fixing the relative position of a cow and a clock) rather than receiving an ineffective global score.

  4. The optimization process will be spatially aware; it will project the learned atom-level advantages onto a per-pixel advantage map using teacher-forced attention heatmaps, ensuring that gradient updates are applied precisely to the image regions responsible for fixing or breaking each specific object relationship.

  5. The system's training efficiency will improve by focusing optimization efforts on the most critical early steps (the structural formation window), preventing late-stage refinement from disrupting established layouts.

By implementing these changes, the improved AI system can:

  1. Generate images that satisfy complex, multi-object prompts with high fidelity across all required attributes and spatial arrangements (e.g., correctly placing five striped monkeys in front of three pink cows to the left of four clocks).

  2. Achieve a significant improvement in adherence to compositional constraints on challenging benchmarks like GenEval 2, where current models often fail despite high overall quality scores.

  3. Be fine-tuned with a much more precise and informative reward signal, leading to faster convergence on specific compositional tasks compared to methods relying on scalar or preference rewards alone.

Sources

Related papers