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

2026-06-16

AI Summary

Existing reinforcement learning (RL) post-training methods for text-to-image generation collapse rewards into a single scalar, failing to align with the spatiotemporal structure of the generative process and resulting in inefficient policy updates. This work proposes STAR, the first approach to introduce a spatiotemporally adaptive reward allocation mechanism in RL post-training for both diffusion and flow models. STAR dynamically constructs spatial reward maps based on text-image attention and applies stronger policy updates to critical latent regions during key denoising steps. By integrating spatial reward allocation, group-relative advantage estimation, and spatially resolved policy objectives, STAR achieves substantial performance gains with negligible computational overhead. Evaluated on Stable Diffusion 3.5 Medium, STAR significantly improves semantic alignment (GenEval: 0.9759), text rendering quality (OCR score: 0.9757), and human preference (PickScore: 23.60).

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 propose \textbf{SpatioTemporal Adaptive Reward (STAR) Allocation} for 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, achieving $\mathbf{0.9759}$, $\mathbf{0.9757}$, and $\mathbf{23.60}$ on GenEval, OCR, and PickScore, respectively.

Problem

• text-to-image generation
• reinforcement learning post-training
• spatiotemporal structure
• reward allocation
• semantic alignment

Innovation

• SpatioTemporal Reward Allocation
• Text-to-Image Generation
• Reinforcement Learning Post-Training
• Diffusion Models
• Attention-based Reward Shaping