STAR: SpatioTemporal Adaptive Reward Allocation for Text-to-Image RL Post-Training

# STAR: SpatioTemporal Adaptive Reward Allocation for Text-to-Image RL Post-Training
Source: [https://arxiv.org/html/2606.17979](https://arxiv.org/html/2606.17979)
11institutetext:###### Abstract

Existing RL post\-training methods for text\-to\-image generation usually convert the final\-image reward into a single scalar advantage and apply it with the same strength to the entire generative trajectory\. However, text\-to\-image generation naturally has temporal and spatial structure: different denoising steps are responsible for different generation stages, and the content that truly determines text alignment often appears only in part of the image\. This granularity mismatch makes it difficult for policy updates to focus on the generative components that actually affect the reward\. To address this issue, we proposeSpatioTemporal Adaptive Reward \(STAR\) Allocationfor RL post\-training of text\-to\-image diffusion and flow models\. STAR uses text\-image attention inside the generative model and starts from the core content that the user truly cares about in the prompt\. It constructs spatial allocation maps that dynamically vary across denoising steps and rollouts, and allocates the same group\-relative advantage to more relevant latent regions with almost no additional computational overhead\. STAR then applies stronger policy updates to these regions through a spatially resolved policy objective\. We use Stable Diffusion 3\.5 Medium as the base model and evaluate on three tasks: GenEval, OCR text rendering, and PickScore\. Experimental results show that STAR improves compositional semantic alignment, text rendering, and preference optimization without changing the external reward source, achieving0\.9759\\mathbf\{0\.9759\},0\.9757\\mathbf\{0\.9757\}, and23\.60\\mathbf\{23\.60\}on GenEval, OCR, and PickScore, respectively\.

![Refer to caption](https://arxiv.org/html/2606.17979v1/x1.png)Figure 1:Overview of STAR \(SpatioTemporal Adaptive Reward Allocation\)\. Given the same image\-level group\-relative advantage, STAR extracts text\-image attention maps from the generative backbone and constructs timestep\-specific spatial allocation maps that highlight latent regions corresponding to key prompt components\. These maps route the scalar advantage into spatially resolved training signals, so policy updates selectively focus on the generative regions that determine text alignment\.## 1Introduction

Diffusion models have become mainstream frameworks for high\-fidelity text\-to\-image generation\[ho2020denoising,dhariwal2021diffusion,song2020score,ddim,rombach2022high\]\. Recent text\-to\-image systems further improve resolution and language understanding through latent diffusion, rectified\-flow transformers, and large\-scale prompt\-image training\[podell2023sdxl,drawbench,sd3\]\. Although these models have made substantial progress in visual quality, compositional generalization under complex prompts remains a major weakness: models often make errors in object counting, attribute binding, and spatial relations\. Since such errors can usually only be judged on the completed image by automatic evaluators, preference models\[xu2024imagereward,kirstain2023pick\], or human feedback, RL\-style post\-training provides an effective path for directly optimizing image\-level objectives\.

A key limitation of existing RL post\-training for text\-to\-image generation lies in the granularity of credit assignment\. The final reward is usually an image\-level scalar, whereas the diffusion/flow generative trajectory has explicit temporal and spatial structure: sampling spans multiple denoising/flow steps, and each step contains many latent spatial tokens\. However, an image\-level reward does not mean that the entire image bears equal responsibility for the reward\. Even for a simple prompt such as “a bicycle,” text alignment is mainly determined by the region containing the bicycle itself, rather than by background regions such as the sky or the road\. When the prompt further contains colors, counts, and spatial relations, the reward depends more strongly on whether a small number of semantically relevant regions are generated correctly\. This issue is particularly visible in online diffusion/flow RL post\-training methods, including Flow\-GRPO\[liu2025flowgrpo\], where a group\-relative reward is still converted into a scalar advantage for policy optimization\. Existing policy\-gradient training usually applies this advantage with the same strength to all timesteps and spatial locations, making it difficult to distinguish the reward responsibility of subject regions from that of background regions\.

This mismatch stems from the spatiotemporal structure of diffusion and flow generative trajectories\[lipman2022flow,rectified\_flow\]\. Different denoising stages affect global layout, object formation, and detail refinement, respectively; within each stage, only a subset of latent regions is directly related to the subjects, attributes, and relations in the current prompt\. We study spatiotemporal credit assignment in this structured process: given the same image\-level outcome feedback, we allocate it over denoising time×\\timeslatent space so that policy updates focus on the generative components that truly determine text alignment\.

We proposeSpatioTemporal Adaptive Reward \(STAR\) Allocation, a framework for attention\-based T2I diffusion and flow post\-training\. STAR uses text\-image interactions inside the generative model to construct timestep\-specific routing maps and converts them into latent\-space allocation maps\. Given the same scalar advantage, STAR distributes its training strength across denoising steps and latent regions\. Combined with a spatially resolved likelihood\-ratio objective, regions corresponding to key prompt components receive stronger updates, while background and weakly related regions maintain weaker influence\.

We instantiate STAR in compositional T2I RL post\-training and compare it with scalar\-allocation baselines under the same external reward source\. Experimental results show that STAR achieves0\.97590\.9759,23\.6023\.60, and0\.97570\.9757on GenEval, PickScore, and OCR, respectively\.

Our contributions are as follows:

- •We proposeSpatioTemporal Adaptive Reward \(STAR\) Allocation, a method for T2I diffusion/flow RL post\-training\. STAR converts image\-level rewards into training signals over denoising time and latent space, allowing updates to focus on the generative regions corresponding to key objects, attributes, and relations in the prompt\.
- •We use text\-image interactions inside the model to construct timestep\-specific spatial allocation maps, extend existing RL post\-training objectives into a spatially resolved policy objective, and provide a unified reward\-routing interface for attention\-based T2I diffusion/flow models\.
- •We validate STAR on GenEval, PickScore, and OCR, where it achieves0\.97590\.9759,23\.6023\.60, and0\.97570\.9757, respectively, showing that SpatioTemporal Adaptive Reward \(STAR\) Allocation can simultaneously improve compositional semantic alignment, preference optimization, and text rendering\.

## 2Related Work

RL post\-training for T2I\.In recent years, directly optimizing downstream rewards has become an important direction for improving diffusion and flow\-based T2I models, following the broader policy\-gradient view of reinforcement learning\[sutton1998reinforcement,williams1992simple\]\. DDPO\[ddpo\]formulates the denoising process as a multi\-step decision\-making problem and optimizes diffusion models with policy\-gradient methods\. DPOK\[fan2024reinforcement\]further introduces online RL and KL regularization into T2I diffusion fine\-tuning\. Diffusion\-DPO\[wallace2024diffusion\]builds on direct preference optimization\[dpo\], while AlignProp\[prabhudesai2023aligning\]optimizes human preferences through differentiable reward backpropagation\. More recently, Flow\-GRPO\[liu2025flowgrpo\]extends online policy\-gradient training to flow matching models and enables efficient post\-training through stochastic sampling and denoising reduction\. These methods demonstrate the effectiveness of RL\-style post\-training for T2I models, but they usually treat the final reward as a scalar signal applied to the whole image or the entire trajectory\. STAR studies how the same reward signal should be allocated within a structured denoising trajectory\.

Fine\-grained rewards and credit assignment\.Existing work has begun to recognize that image\-level scalar rewards provide overly coarse feedback for diffusion alignment\. A Dense Reward View\[yang2024dense\]revisits preference alignment from a sequential decision\-making perspective and derives a dense objective that emphasizes early denoising steps\. B2\-DiffuRL addresses sparse terminal rewards by estimating the contributions of intermediate denoising actions through backward progressive training and branch\-based sampling\. Another line of work improves the reward source itself\. ImageReward\[xu2024imagereward\]learns a general T2I human preference reward model; ImageDoctor further provides multi\-aspect scores and grounded heatmaps for diagnosing local errors in generated images\. SpatialReward builds verifiable reward modeling for spatial consistency\. These methods highlight the value of fine\-grained feedback in T2I post\-training, while leaving a complementary direction: given an existing outcome reward, how can we use the generative model’s own text\-image interactions to allocate this signal within the spatiotemporal structure?

## 3Method

STAR reallocates an existing image\-level reward to the structured generative process of diffusion/flow models\. Given a promptcc, the model generates a latent trajectory

τ=\(zT,zT−1,…,z0\),zt∈ℝC×ht×wt,\\tau=\(z\_\{T\},z\_\{T\-1\},\\ldots,z\_\{0\}\),\\qquad z\_\{t\}\\in\\mathbb\{R\}^\{C\\times h\_\{t\}\\times w\_\{t\}\},\(1\)wherez0z\_\{0\}corresponds to the final image\. LetΩt=\{1,…,ht\}×\{1,…,wt\}\\Omega\_\{t\}=\\\{1,\\ldots,h\_\{t\}\\\}\\times\\\{1,\\ldots,w\_\{t\}\\\}denote the set of spatial locations in the latent feature map at steptt\. The external reward source provides scalar feedback only on the final sample, denoted as

Ri=R​\(z0i,c\),Ai=ℬ​\(Ri,c\),R\_\{i\}=R\(z\_\{0\}^\{i\},c\),\\qquad A\_\{i\}=\\mathcal\{B\}\(R\_\{i\},c\),\(2\)whereiidenotes theii\-th rollout under the same prompt, andAiA\_\{i\}is the sample\-level advantage obtained by group\-relative normalization or baseline subtraction\. A standard RL\-style objective usesAiA\_\{i\}as a trajectory\-level weight, as in group\-relative policy optimization\[grpo\]; STAR routesAiA\_\{i\}to denoising steps and latent regions\.

### 3\.1Attention\-Guided Allocation

STAR represents each prompt as a set of reward\-relevant text units:

𝒦​\(c\)=\{q1,…,qK\},\\mathcal\{K\}\(c\)=\\\{q\_\{1\},\\ldots,q\_\{K\}\\\},\(3\)whereqkq\_\{k\}denotes a core component in the prompt\. These components are the conditions that the final image needs to satisfy, including subjects, styles, attributes, counts, and spatial relations\. For example, for the prompt “generate a futuristic\-style building,”𝒦​\(c\)\\mathcal\{K\}\(c\)contains the subject “building” and the visual style requirement “futuristic style\.” Dataset\-specific extraction rules and prompt templates for text\-unit extraction are provided in Appendix[A](https://arxiv.org/html/2606.17979#A1)\.

Suppose the full promptccis converted by the text encoder tokenizer intoLcL\_\{c\}tokens\. Each text unitqkq\_\{k\}corresponds to several text tokens, whose token indices are denoted as

𝒯k⊆\{1,…,Lc\}\.\\mathcal\{T\}\_\{k\}\\subseteq\\\{1,\\ldots,L\_\{c\}\\\}\.\(4\)
STAR reads routing signals from a selected set of text\-image attention layersℒ\\mathcal\{L\}\. When generating rolloutii, attention\-based T2I backbones establish interactions between image latent tokens and text tokens at each denoising steptt\. For anyℓ∈ℒ\\ell\\in\\mathcal\{L\}, this layer provides an attention mapαti,ℓ∈ℝht×wt×Lc\\alpha\_\{t\}^\{i,\\ell\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\\times L\_\{c\}\}; the attention weight from latent positionuuto text tokenjjis denoted as

αti,ℓ​\(u,j\)\.\\alpha\_\{t\}^\{i,\\ell\}\(u,j\)\.\(5\)Hereu∈Ωtu\\in\\Omega\_\{t\}is a latent spatial position andjjis a text token index\. In joint\-attention backbones \(e\.g\., Stable Diffusion 3\.5, Qwen\-Image\), this value corresponds to the image\-to\-text block of the full attention matrix; in cross\-attention backbones \(e\.g\., Stable Diffusion XL\), it corresponds to the cross\-attention weights from image queries to text keys\. STAR takes the maximum image\-to\-text attention weight over the selected layers:

αti​\(u,j\)=maxℓ∈ℒ⁡αti,ℓ​\(u,j\),\\alpha\_\{t\}^\{i\}\(u,j\)=\\max\_\{\\ell\\in\\mathcal\{L\}\}\\alpha\_\{t\}^\{i,\\ell\}\(u,j\),\(6\)which gives the timestep\-specific text\-image routing mapαti∈ℝht×wt×Lc\\alpha\_\{t\}^\{i\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\\times L\_\{c\}\}for rolloutiiat steptt\.

Next, STAR extracts where each text unit appears in the current trajectory\. Given the token set𝒯k\\mathcal\{T\}\_\{k\}of text unitqkq\_\{k\}, STAR constructs its spatial mapMti,k∈ℝht×wtM\_\{t\}^\{i,k\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\}, where the routing strength at latent positionuuis

Mti,k​\(u\)=1\|𝒯k\|​∑j∈𝒯kαti​\(u,j\)\.M\_\{t\}^\{i,k\}\(u\)=\\frac\{1\}\{\|\\mathcal\{T\}\_\{k\}\|\}\\sum\_\{j\\in\\mathcal\{T\}\_\{k\}\}\\alpha\_\{t\}^\{i\}\(u,j\)\.\(7\)STAR aggregates these key components into a trajectory\-timestep specific heatmapHti∈ℝht×wtH\_\{t\}^\{i\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\}:

Hti​\(u\)=1K​∑k=1KMti,k​\(u\)\.H\_\{t\}^\{i\}\(u\)=\\frac\{1\}\{K\}\\sum\_\{k=1\}^\{K\}M\_\{t\}^\{i,k\}\(u\)\.\(8\)This heatmap specifies where the prompt\-relevant content is most active\. To turn it into a normalized reward\-allocation map, STAR maps it to bounded coefficientsCti∈ℝht×wtC\_\{t\}^\{i\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\}:

Cti​\(u\)=1\+\(λ−1\)​Hti​\(u\)−minv∈Ωt⁡Hti​\(v\)maxv∈Ωt⁡Hti​\(v\)−minv∈Ωt⁡Hti​\(v\)\+ϵ\.C\_\{t\}^\{i\}\(u\)=1\+\(\\lambda\-1\)\\frac\{H\_\{t\}^\{i\}\(u\)\-\\min\_\{v\\in\\Omega\_\{t\}\}H\_\{t\}^\{i\}\(v\)\}\{\\max\_\{v\\in\\Omega\_\{t\}\}H\_\{t\}^\{i\}\(v\)\-\\min\_\{v\\in\\Omega\_\{t\}\}H\_\{t\}^\{i\}\(v\)\+\\epsilon\}\.\(9\)Hereλ≥1\\lambda\\geq 1controls the allocation strength, andϵ\\epsilonensures numerical stability\. LargerCti​\(u\)C\_\{t\}^\{i\}\(u\)strengthens updates on prompt\-relevant regions; smaller but nonzero coefficients preserve the influence of background, composition, and overall quality on the final reward\.

### 3\.2Spatially Resolved Policy Objective

STAR uses the allocation map to convert the sample\-level advantage into a spatialized training signal\. For theii\-th rollout and denoising steptt, we define the spatial advantage map as

Ati​\(u\)=Ai​Cti​\(u\),Ati∈ℝht×wt\.A\_\{t\}^\{i\}\(u\)=A\_\{i\}\\,C\_\{t\}^\{i\}\(u\),\\qquad A\_\{t\}^\{i\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\}\.\(10\)This mapping preserves the sign ofAiA\_\{i\}and increases the update magnitude for prompt\-relevant regions throughCti​\(u\)C\_\{t\}^\{i\}\(u\)\.

The policy update acts on spatial log\-probability contributions\. Let

ℓθ,ti∈ℝht×wt,ℓ¯θ,ti=1\|Ωt\|​∑u∈Ωtℓθ,ti​\(u\),\\ell\_\{\\theta,t\}^\{i\}\\in\\mathbb\{R\}^\{h\_\{t\}\\times w\_\{t\}\},\\qquad\\bar\{\\ell\}\_\{\\theta,t\}^\{i\}=\\frac\{1\}\{\|\\Omega\_\{t\}\|\}\\sum\_\{u\\in\\Omega\_\{t\}\}\\ell\_\{\\theta,t\}^\{i\}\(u\),\(11\)whereℓθ,ti​\(u\)\\ell\_\{\\theta,t\}^\{i\}\(u\)denotes the log\-probability contribution at latent positionuu, andℓ¯θ,ti\\bar\{\\ell\}\_\{\\theta,t\}^\{i\}corresponds to the scalar transition log\-probability used in standard policy\-gradient training\. STAR uses a normalized spatial ratio:

ρθ,ti​\(u\)=sg⁡\[exp⁡\(ℓ¯θ,ti−ℓ¯old,ti\)\]​exp⁡\(ℓθ,ti​\(u\)\)sg⁡\[exp⁡\(ℓθ,ti​\(u\)\)\],\\rho\_\{\\theta,t\}^\{i\}\(u\)=\\operatorname\{sg\}\\\!\\left\[\\exp\\\!\\left\(\\bar\{\\ell\}\_\{\\theta,t\}^\{i\}\-\\bar\{\\ell\}\_\{\\mathrm\{old\},t\}^\{i\}\\right\)\\right\]\\frac\{\\exp\(\\ell\_\{\\theta,t\}^\{i\}\(u\)\)\}\{\\operatorname\{sg\}\\\!\\left\[\\exp\(\\ell\_\{\\theta,t\}^\{i\}\(u\)\)\\right\]\},\(12\)whereℓ¯old,ti\\bar\{\\ell\}\_\{\\mathrm\{old\},t\}^\{i\}is the rollout log\-probability stored during sampling, andsg⁡\[⋅\]\\operatorname\{sg\}\[\\cdot\]denotes stop\-gradient\. This ratio preserves the overall scale of the scalar likelihood ratio while assigning gradients over spatial locations\.

Finally, STAR optimizes the clipped spatial objective, following the clipped\-ratio policy update used in PPO\-style training\[ppo\]:

𝒥STAR​\(θ\)=𝔼​\[1T​∑t=1T1\|Ωt\|​∑u∈Ωtmin⁡\(ρθ,ti​\(u\)​Ati​\(u\),clip⁡\(ρθ,ti​\(u\),1−ε,1\+ε\)​Ati​\(u\)\)\]−β​𝒟KL\.\\mathcal\{J\}\_\{\\mathrm\{STAR\}\}\(\\theta\)=\\mathbb\{E\}\\left\[\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\frac\{1\}\{\|\\Omega\_\{t\}\|\}\\sum\_\{u\\in\\Omega\_\{t\}\}\\min\\left\(\\rho\_\{\\theta,t\}^\{i\}\(u\)A\_\{t\}^\{i\}\(u\),\\operatorname\{clip\}\(\\rho\_\{\\theta,t\}^\{i\}\(u\),1\-\\varepsilon,1\+\\varepsilon\)A\_\{t\}^\{i\}\(u\)\\right\)\\right\]\-\\beta\\mathcal\{D\}\_\{\\mathrm\{KL\}\}\.\(13\)Here𝒟KL\\mathcal\{D\}\_\{\\mathrm\{KL\}\}is the KL regularization term between the current policy and the frozen reference policy over denoising transitions, which constrains the update magnitude and reduces the risk of reward hacking\.

## 4Experiments

### 4\.1Implementation Details

Evaluation tasks\.The experiments include three T2I post\-training tasks: GenEval, OCR text rendering, and PickScore\. GenEval evaluates compositional generation ability, including object counting, attribute binding, colors, and spatial relations\[geneval,T2i\-compbench\]; training prompts are generated from the official GenEval templates, and rewards are computed by the official rule\-based evaluation pipeline\. OCR text rendering requires the model to generate a specified text string in the image, and the reward is computed from the edit distance between the rendered text and the target text\. PickScore uses a learned preference reward model to score each prompt\-image pair\.

Baselines\.We compare against multiple representative T2I baselines, including diffusion models \(LDM, SD1\.5, SD2\.1, SD\-XL\[podell2023sdxl\], DALLE\-2\[ramesh2022hierarchical\], DALLE\-3\[betker2023improving\]\), autoregressive models \(Show\-o\[xie2024show\], Emu3\-Gen\[wang2024emu3\], JanusFlow\[ma2024janusflow\], Janus\-Pro\-7B\[chen2025janus\]\), and recent diffusion/flow models \(FLUX\.1 Dev\[flux2024\], SD3\.5\-L, SANA\-1\.5 4\.8B\[xie2025sana\], SD3\.5\-M\[sd3\], SD3\.5\-M\+Flow\-GRPO\[liu2025flowgrpo\]\)\.

Experimental setup\.We implement STAR based on the Flow\-GRPO\[liu2025flowgrpo\]training pipeline and use Stable Diffusion 3\.5 Medium \(SD3\.5\-M\) as the base model\. Training uses ODE\-to\-SDE rollouts for online sampling, with the number of sampling steps set toT=10T=10, the SDE noise level set toa=0\.7a=0\.7, and the image resolution set to512×512512\\times 512; evaluation usesT=40T=40steps\. For each prompt, we sampleG=24G=24images to form a group, and update the model with a group\-relative policy\-gradient objective, clipped ratio, and KL regularization\. The KL coefficient is set toβ=0\.04\\beta=0\.04for GenEval and OCR text rendering, andβ=0\.01\\beta=0\.01for PickScore\. We train the model with LoRA\[hu2021lora\], using rankr=32r=32and scaling factorα=64\\alpha=64\. STAR records text\-image attention during rollout and uses the spatial objective in Section 3 to allocate group\-relative advantages\. Attention routing uses the joint attention of SD3\.5\-M, reads the attention layers from transformer blocks 10–15, and sets the allocation strength toλ=1\.5\\lambda=1\.5\.

### 4\.2Main Results

Table 1:GenEval results\. We report the overall score and six compositional sub\-tasks\. Obj\. denotes object, and Attr\. denotes attribute\. Blue indicates the best result, and green indicates the second\-best result\.ModelOverallSingle Obj\.Two Obj\.CountingColorsPositionAttr\. BindingDiffusion ModelsLDM0\.370\.920\.290\.230\.700\.020\.05SD1\.50\.430\.970\.380\.350\.760\.040\.06SD2\.10\.500\.980\.510\.440\.850\.070\.17SD\-XL0\.550\.980\.740\.390\.850\.150\.23DALLE\-20\.520\.940\.660\.490\.770\.100\.19DALLE\-30\.670\.960\.870\.470\.830\.430\.45Autoregressive ModelsShow\-o0\.530\.950\.520\.490\.820\.110\.28Emu3\-Gen0\.540\.980\.710\.340\.810\.170\.21JanusFlow0\.630\.970\.590\.450\.830\.530\.42Janus\-Pro\-7B0\.800\.990\.890\.590\.900\.790\.66Diffusion / Flow ModelsFLUX\.1 Dev0\.660\.980\.810\.740\.790\.220\.45SD3\.5\-L0\.710\.980\.890\.730\.830\.340\.47SANA\-1\.5 4\.8B0\.810\.990\.930\.860\.840\.590\.65SD3\.5\-M0\.630\.980\.780\.500\.810\.240\.52SD3\.5\-M\+Flow\-GRPO0\.951\.000\.990\.950\.920\.990\.86SD3\.5\-M\+Flow\-GRPO\+STAR0\.97591\.001\.000\.980\.960\.980\.94

GenEval results\.Table[1](https://arxiv.org/html/2606.17979#S4.T1)presents the main results on GenEval\[geneval\]\. STAR achieves an overall score of0\.97590\.9759, outperforming SD3\.5\-M at0\.630\.63and Flow\-GRPO\[liu2025flowgrpo\]at0\.950\.95\. This improvement mainly comes from more fine\-grained compositional semantic constraints: on counting, colors, and attribute binding, STAR reaches0\.980\.98,0\.960\.96, and0\.940\.94, respectively, clearly improving over0\.950\.95,0\.920\.92, and0\.860\.86\. Two\-object generation further reaches1\.001\.00\. On single\-object generation and position, two sub\-tasks that are already close to saturation, STAR still maintains near\-optimal performance, reaching1\.001\.00and0\.980\.98, respectively\.

These results show that the gains of STAR come from more accurate spatiotemporal credit assignment\. By allocating the same image\-level advantage to prompt\-relevant latent regions, STAR can more selectively update the generative regions that affect object counts, colors, and attribute binding, thereby improving compositional generation performance\.

![Refer to caption](https://arxiv.org/html/2606.17979v1/x2.png)Figure 2:Main results on three T2I post\-training benchmarks\. STAR consistently improves over the Flow\-GRPO baseline on compositional generation, text rendering, and preference optimization\.Results across benchmarks\.Figure[2](https://arxiv.org/html/2606.17979#S4.F2)further compares Flow\-GRPO\[liu2025flowgrpo\]and STAR on the three post\-training tasks\. STAR achieves0\.97590\.9759on GenEval\[geneval\],0\.97570\.9757on OCR text rendering\[textdiffuser\], and23\.6023\.60on PickScore\[kirstain2023pick\], all outperforming the corresponding Flow\-GRPO baselines\. The improvements on GenEval and OCR indicate that spatiotemporal adaptive reward allocation can more effectively optimize verifiable text\-alignment objectives; the improvement on PickScore shows that this allocation scheme can also be used with continuous rewards provided by a preference model\. Overall, without changing the form of the task reward, STAR allocates the same sample\-level feedback to more relevant generative regions and thus obtains consistent gains in compositional generation, text rendering, and preference optimization\.

### 4\.3Ablation Study

![Refer to caption](https://arxiv.org/html/2606.17979v1/x3.png)\(a\)Prompt heatmap construction\.
![Refer to caption](https://arxiv.org/html/2606.17979v1/x4.png)\(b\)Attention layer selection\.
![Refer to caption](https://arxiv.org/html/2606.17979v1/x5.png)\(c\)Computational overhead\.

Figure 3:Ablation and efficiency analysis for STAR, covering reward routing design choices and additional training overhead\.#### 4\.3\.1Prompt Heatmap Construction

We compare two ways of constructing prompt heatmaps\. The single\-heatmap variant first aggregates the attention of multiple key components in the prompt into one heatmap and then obtains reward allocation coefficients; the component\-wise heatmaps variant first generates a heatmap for each key component separately and then stacks and normalizes them\. Both experiments use the same PickScore reward, training setting, and evaluation interval\.

Figure[3\(a\)](https://arxiv.org/html/2606.17979#S4.F3.sf1)shows that the training curves of the two variants are overall very close, with the single heatmap showing a slight advantage in the later stage of training\. This result indicates that directly aggregating multiple components that the user cares about in the prompt into a shared heatmap already provides a stable routing signal\. Generating heatmaps component by component and then stacking them does not bring additional gains and increases implementation complexity; therefore, we use single\-heatmap construction in the main experiments\.

#### 4\.3\.2Attention Layer Selection

We further analyze the choice of text\-image attention layers in STAR\. We compare four settings: front layers use transformer blocks0–99, STAR layers use blocks1010–1515, back layers use blocks1616–2323, and all layers use all attention layers\.

Figure[3\(b\)](https://arxiv.org/html/2606.17979#S4.F3.sf2)shows that layer selection affects the stability of attention routing\. Front layers and back layers both bring effective improvements, but their later\-stage curves are slightly lower; all layers reach the highest level in the final stage, but require reading more attention maps\. STAR layers remain in the best\-performing group for most of training and are overall close to all layers, while using only a small number of middle blocks\. This indicates that blocks1010–1515already contain sufficiently clear text\-image alignment signals\. Therefore, we use this layer set in the main experiments to achieve a better balance between effectiveness and implementation overhead\.

#### 4\.3\.3Computational Overhead

We further compare the training overhead of Flow\-GRPO\[liu2025flowgrpo\]and Flow\-GRPO\+STAR\. STAR reuses the text\-image attention from the forward pass of the generative model and only performs heatmap aggregation and advantage rescaling on the selected layers; it does not require training an additional reward model or adding new image evaluation steps\.

Figure[3\(c\)](https://arxiv.org/html/2606.17979#S4.F3.sf3)shows that STAR introduces less than2%2\\%additional training time overhead, and increases peak memory by only1\.56%1\.56\\%\. This result shows that the main computation of STAR still comes from the original diffusion/flow rollout and GRPO update; attention\-based allocation is only a lightweight reward\-routing step\. Therefore, we can obtain more fine\-grained spatiotemporal credit assignment with almost no change in training cost\.

### 4\.4Qualitative Comparison

To qualitatively examine the effect of post\-training with STAR, we compare the original SD3\.5\-M base model with SD3\.5\-M after Flow\-GRPO\+STAR post\-training on a diverse set of prompts\.

![Refer to caption](https://arxiv.org/html/2606.17979v1/x6.png)
![Refer to caption](https://arxiv.org/html/2606.17979v1/x7.png)

Figure 4:Qualitative comparison between the original SD3\.5\-M and SD3\.5\-M after Flow\-GRPO\+STAR post\-training\. Within every prompt pair, the left image is produced by SD3\.5\-M and the right image by SD3\.5\-M\+Flow\-GRPO\+STAR\.As shown in Figure[4](https://arxiv.org/html/2606.17979#S4.F4), both models generate images that capture the main concepts in the prompts\. Across the displayed examples, Flow\-GRPO\+STAR post\-training changes how the requested subjects, attributes, and surrounding scenes are realized\. For example, the post\-trained model depicts the black square building, reflective stainless\-steel Bitcoin logo, blue windmill, and compact starship with more complete prompt\-related elements and more coherent scene composition\.

These qualitative results complement the quantitative evaluations and illustrate the effect of applying Flow\-GRPO\+STAR post\-training to SD3\.5\-M across different types of prompts\. By allocating the reward signal to prompt\-relevant latent regions, STAR encourages the model to better incorporate the textual conditions during generation\.

## 5Conclusion

We proposeSpatioTemporal Adaptive Reward \(STAR\) Allocation, a spatiotemporal adaptive reward allocation method for RL post\-training of T2I diffusion and flow models\. STAR extracts prompt\-related spatiotemporal signals from text\-image interactions inside the generative model and uses them to allocate image\-level feedback more selectively to key regions of the generative trajectory, focusing policy updates on the parts that truly affect text alignment and generation quality\. Experimental results show that, without changing the external reward source and with almost no additional training cost, STAR stably improves existing RL post\-training methods, especially in scenarios requiring fine\-grained semantic understanding, local content control, and preference alignment\. Overall, STAR provides a simple, lightweight, and effective interface for allocating image\-level feedback to structured diffusion/flow trajectories\.

## Statements

Stable Diffusion 3\.5 Medium is used under the Stability AI Community License for academic research purposes\. The authors confirm that the use of this model is limited to academic research purposes and has not been used for any commercial activities\.

## References

## Appendix AText\-Unit Extraction Prompts

This appendix reports the task\-specific text\-unit extraction procedures used by STAR\. These text units are used only to locate the corresponding text\-token indices and aggregate text\-image attention maps; the external image\-level reward for each task remains unchanged\.

### A\.1GenEval Text Units

For GenEval, STAR does not require an LLM for text\-unit extraction\. Instead, it applies a deterministic focus\-span rule: remove a boilerplate image prefix such as “a photo of”, “an image of”, or “a painting of”, and use the remaining phrase as a single text unit\. This matches the GenEval prompt format, where the suffix already contains the objects, attributes, counts, and spatial relations to be grounded\.

``

`A\.2 PickScore Text Units For PickScore, STAR uses the same deterministic focus\-span rule by default\. We additionally consider component\-level decomposition, in which a vision\-language model extracts concise semantic units describing objects, attributes, quantities, and spatial relations\. The decomposition prompt is provided below\. A\.3 OCR Text Units OCR prompts contain exact quoted strings that must be rendered in the image\. We therefore extract visually localizable text units in advance and associate them with the corresponding training prompts\. These units include the exact rendered text, its carrier or surface, visible text attributes, and other concrete scene elements\. When pre\-extracted units are unavailable, STAR falls back to the deterministic focus\-span rule described above\. The extraction prompt is shown below\.`