Technical Summary: STAR (SpatioTemporal Adaptive Reward Allocation)

Problem Statement

Existing Reinforcement Learning (RL) post-training methods for text-to-image (T2I) generation, such as Flow-GRPO, typically convert the final image-level reward into a single scalar advantage. This scalar is then applied with uniform strength across the entire generative trajectory (all denoising steps and all spatial latent tokens).

The paper identifies a fundamental granularity mismatch in this approach:

1. Temporal Structure: Different denoising steps are responsible for distinct generation stages (e.g., global layout vs. detail refinement).
2. Spatial Structure: The content that determines text alignment (e.g., specific objects, attributes, or spatial relations) often appears only in specific regions of the image, while other regions (e.g., background) contribute less to the prompt's semantic requirements.

When a scalar reward is applied uniformly, the policy update struggles to distinguish between regions that are critical for satisfying the prompt and those that are not. This makes it difficult for the model to focus learning on the specific generative components that actually drive the reward, leading to suboptimal performance in compositional generalization, object counting, and text rendering.

Methodology: STAR

The authors propose SpatioTemporal Adaptive Reward (STAR) Allocation, a framework designed to reallocate image-level rewards across the structured denoising trajectory of diffusion and flow models. STAR operates without modifying the external reward source, instead leveraging the internal mechanics of the generative model.

1. Attention-Guided Allocation

STAR constructs dynamic spatial allocation maps that vary across denoising steps ($t$) and rollouts ($i$). The process involves:

• Text-Unit Extraction: The prompt $c$ is decomposed into a set of core text units $K(c)$ (e.g., subjects, styles, attributes).
• Attention Mapping: The model utilizes text-image attention layers (specifically from the joint-attention backbone of Stable Diffusion 3.5) to extract routing signals. For each latent spatial position $u$ and text token $j$, an attention weight $\alpha(u, j)$ is computed.
• Heatmap Construction: For each text unit, a spatial map is generated by aggregating attention weights over the corresponding text tokens. These maps are aggregated into a trajectory-specific heatmap $H_t^i$ that highlights latent regions most active with respect to the prompt's key components.
• Normalization: The heatmap is normalized and scaled to produce a coefficient map $C_t^i(u)$. This map assigns higher weights to prompt-relevant regions and lower (but non-zero) weights to background regions.

2. Spatially Resolved Policy Objective

STAR converts the scalar group-relative advantage $A_i$ into a spatially resolved training signal $A_t^i(u) = A_i \cdot C_t^i(u)$.

• The policy objective is modified to operate on spatial log-probability contributions rather than a scalar trajectory probability.
• A normalized spatial ratio $\rho$ is computed for each latent position.
• The final objective function (based on PPO-style clipping) sums the clipped ratios weighted by the spatial advantage map $A_t^i(u)$ across all time steps and spatial locations. This ensures that policy updates are stronger in regions corresponding to key prompt components.

Key Contributions

1. STAR Framework: A method for T2I diffusion/flow RL post-training that converts image-level rewards into training signals distributed over denoising time and latent space, allowing updates to focus on generative regions corresponding to key objects, attributes, and relations.
2. Unified Interface: The use of internal text-image interactions to construct timestep-specific spatial allocation maps, extending existing RL objectives into a spatially resolved policy objective without requiring new reward models.
3. Empirical Validation: Demonstration that STAR improves compositional semantic alignment, text rendering, and preference optimization across three distinct tasks (GenEval, OCR, PickScore) using Stable Diffusion 3.5 Medium as the base.

Experimental Results

The authors evaluated STAR on Stable Diffusion 3.5 Medium (SD3.5-M) compared against the Flow-GRPO baseline and other state-of-the-art models.

• GenEval (Compositional Generation): STAR achieved an overall score of 0.9759, outperforming Flow-GRPO (0.95) and the base SD3.5-M (0.63). Notable improvements were seen in counting (0.98), colors (0.96), and attribute binding (0.94).
• OCR Text Rendering: STAR achieved an accuracy of 0.9757, surpassing the Flow-GRPO baseline (0.92).
• PickScore (Preference Optimization): STAR achieved a score of 23.60, compared to 23.31 for Flow-GRPO.
• Efficiency: The method introduces minimal computational overhead, adding less than 2% to training time and increasing peak memory usage by only 1.56%, as it reuses attention maps from the forward pass.

Significance and Claims

The paper claims that STAR addresses the "credit assignment" problem in T2I RL by acknowledging the spatiotemporal nature of diffusion generation. By routing the same scalar reward signal to the specific latent regions that determine text alignment, STAR enables more selective and effective policy updates.

The authors emphasize that STAR achieves these gains without changing the external reward source or requiring additional reward models. It provides a "simple, lightweight, and effective interface" for allocating image-level feedback to structured trajectories. The results suggest that spatiotemporal adaptive allocation is a critical factor in optimizing fine-grained semantic understanding, local content control, and preference alignment in generative models.