STAR: 时空自适应奖励分配用于文本到图像强化学习后训练

# STAR: 文本到图像强化学习后训练中的时空自适应奖励分配
来源: https://arxiv.org/html/2606.17979
11institutetext:###### 摘要

现有的文本到图像生成的RL后训练方法通常将最终图像的奖励转换为单个标量优势，并以相同的强度应用于整个生成轨迹。然而，文本到图像生成天然具有时间和空间结构：不同的去噪步骤负责不同的生成阶段，而真正决定文本对齐的内容往往只出现在图像的一部分。这种粒度不匹配使得策略更新难以聚焦于实际影响奖励的生成组件。为了解决这个问题，我们提出**时空自适应奖励（STAR）分配**，用于文本到图像扩散和流模型的RL后训练。STAR利用生成模型内部的文本-图像注意力，并从用户提示中真正关心的核心内容出发，构建随去噪步骤和采样轨迹动态变化的空间分配图，将相同的组相对优势分配给更相关的潜在区域，且几乎不增加额外计算开销。然后，STAR通过空间解析的策略目标对这些区域施加更强的策略更新。我们使用Stable Diffusion 3.5 Medium作为基础模型，并在三个任务上评估：GenEval、OCR文本渲染和PickScore。实验结果表明，STAR在不改变外部奖励源的情况下，改进了组合语义对齐、文本渲染和偏好优化，在GenEval、OCR和PickScore上分别达到**0.9759**、**0.9757**和**23.60**。

参见说明图 1: STAR（时空自适应奖励分配）概览。给定相同的图像级组相对优势，STAR从生成骨干网络中提取文本-图像注意力图，并构建特定时间步的空间分配图，突出显示与关键提示组件对应的潜在区域。这些图将标量优势路由到空间解析的训练信号中，使得策略更新选择性地聚焦于决定文本对齐的生成区域。## 1 引言

扩散模型已成为高保真文本到图像生成的主流框架[ho2020denoising, dhariwal2021diffusion, song2020score, ddim, rombach2022high]。最近的文本到图像系统通过潜在扩散、修正流变换器和大规模提示-图像训练进一步提高了分辨率和语言理解能力[podell2023sdxl, drawbench, sd3]。尽管这些模型在视觉质量上取得了显著进步，但在复杂提示下的组合泛化仍然是一个主要弱点：模型经常在物体计数、属性绑定和空间关系上出错。由于这类错误通常只能由自动评估器、偏好模型[xu2024imagereward, kirstain2023pick]或人工反馈在完成的图像上判断，RL风格的post-training为直接优化图像级目标提供了有效途径。

现有文本到图像生成的RL后训练的一个关键限制在于信用分配的粒度。最终奖励通常是图像级标量，而扩散/流生成轨迹具有明确的时间和空间结构：采样跨越多个去噪/流步骤，每个步骤包含许多潜在空间标记。然而，图像级奖励并不意味着整个图像对奖励承担同等责任。即使是简单的提示如“一辆自行车”，文本对齐主要取决于包含自行车本身的区域，而不是天空或道路等背景区域。当提示进一步包含颜色、数量和空间关系时，奖励更强烈地取决于少数语义相关区域是否被正确生成。这个问题在在线扩散/流RL后训练方法中尤为明显，包括Flow-GRPO[liu2025flowgrpo]，其中组相对奖励仍被转换为标量优势用于策略优化。现有的策略梯度训练通常将该优势以相同强度应用于所有时间步和空间位置，使得难以区分主体区域和背景区域的奖励责任。

这种不匹配源于扩散和流生成轨迹的时空结构[lipman2022flow, rectified_flow]。不同的去噪阶段分别影响全局布局、物体形成和细节细化；在每个阶段内，只有一部分潜在区域与当前提示中的主体、属性和关系直接相关。我们研究这一结构化过程中的时空信用分配：给定相同的图像级结果反馈，我们将其分配到去噪时间×潜在空间上，以便策略更新聚焦于真正决定文本对齐的生成组件。

我们提出**时空自适应奖励（STAR）分配**，一种基于注意力的T2I扩散和流后训练框架。STAR利用生成模型内部的文本-图像交互来构建特定时间步的路由图，并将其转换为潜在空间分配图。给定相同的标量优势，STAR将其训练强度分布到去噪步骤和潜在区域上。结合空间解析的似然比目标，对应于关键提示组件的区域获得更强的更新，而背景和弱相关区域保持较弱的影响。

我们将STAR实例化到组合T2I RL后训练中，并在相同外部奖励源下与标量分配基线进行比较。实验结果表明，STAR在GenEval、PickScore和OCR上分别达到0.9759、23.60和0.9757。

我们的贡献如下：

- • 我们提出**时空自适应奖励（STAR）分配**，一种用于T2I扩散/流RL后训练的方法。STAR将图像级奖励转换为去噪时间和潜在空间上的训练信号，使更新能够聚焦于提示中对应关键对象、属性和关系的生成区域。
- • 我们利用模型内部的文本-图像交互来构建特定时间步的空间分配图，将现有的RL后训练目标扩展为空间解析的策略目标，并为基于注意力的T2I扩散/流模型提供统一的奖励路由接口。
- • 我们在GenEval、PickScore和OCR上验证STAR，分别达到0.9759、23.60和0.9757，表明时空自适应奖励（STAR）分配能够同时改进组合语义对齐、偏好优化和文本渲染。

## 2 相关工作

**T2I的RL后训练**。近年来，直接优化下游奖励已成为改进扩散和流基T2I模型的重要方向，遵循强化学习的更广泛策略梯度观点[sutton1998reinforcement, williams1992simple]。DDPO[ddpo]将去噪过程表述为多步决策问题，并使用策略梯度方法优化扩散模型。DPOK[fan2024reinforcement]进一步将在线RL和KL正则化引入T2I扩散微调。Diffusion-DPO[wallace2024diffusion]建立在直接偏好优化[dpo]之上，而AlignProp[prabhudesai2023aligning]通过可微奖励反向传播优化人类偏好。最近，Flow-GRPO[liu2025flowgrpo]将在线策略梯度训练扩展到流匹配模型，并通过随机采样和去噪缩减实现高效后训练。这些方法证明了RL风格后训练对T2I模型的有效性，但它们通常将最终奖励视为应用于整个图像或整个轨迹的标量信号。STAR研究相同奖励信号应如何在结构化的去噪轨迹内分配。

**细粒度奖励与信用分配**。现有工作已开始认识到图像级标量奖励为扩散对齐提供了过于粗糙的反馈。密集奖励视角[yang2024dense]从序列决策角度重新审视偏好对齐，并推导出强调早期去噪步骤的密集目标。B2-DiffuRL通过向后渐进训练和基于分支的采样估计中间去噪动作的贡献，解决了稀疏终端奖励问题。另一条工作路线改进奖励源本身。ImageReward[xu2024imagereward]学习通用的T2I人类偏好奖励模型；ImageDoctor进一步提供多方面评分和接地热图，用于诊断生成图像中的局部错误。SpatialReward为空间一致性构建可验证的奖励建模。这些方法突显了T2I后训练中细粒度反馈的价值，同时留下一个互补方向：给定现有结果奖励，我们如何利用生成模型自身的文本-图像交互，在时空结构内分配该信号？

## 3 方法

STAR将现有的图像级奖励重新分配到扩散/流模型的结构化生成过程中。给定提示$c$，模型生成一个潜在轨迹

$$\tau = (z_T, z_{T-1}, \ldots, z_0), \qquad z_t \in \mathbb{R}^{C \times h_t \times w_t},$$
(1)

其中$z_0$对应最终图像。令$\Omega_t = \{1, \ldots, h_t\} \times \{1, \ldots, w_t\}$表示在步骤$t$时潜在特征图中的空间位置集合。外部奖励源仅对最终样本提供标量反馈，记为

$$R_i = R(z_0^i, c), \qquad A_i = \mathcal{B}(R_i, c),$$
(2)

其中$i$表示同一提示下的第$i$次采样轨迹，$A_i$是通过组相对归一化或基线减法获得的样本级优势。标准RL风格目标使用$A_i$作为轨迹级权重，如组相对策略优化[grpo]；STAR将$A_i$路由到去噪步骤和潜在区域。

### 3.1 注意力引导的分配

STAR将每个提示表示为一组奖励相关的文本单元：

$$\mathcal{K}(c) = \{q_1, \ldots, q_K\},$$
(3)

其中$q_k$表示提示中的一个核心组件。这些组件是最终图像需要满足的条件，包括主体、风格、属性、数量和空间关系。例如，对于提示“生成一座未来风格建筑”，$\mathcal{K}(c)$包含主体“建筑”和视觉风格要求“未来风格”。数据集特定的提取规则和用于文本单元提取的提示模板见附录A (https://arxiv.org/html/2606.17979#A1)。

假设完整提示$c$通过文本编码器分词器转换为$L_c$个词元。每个文本单元$q_k$对应若干文本词元，其词元索引记为

$$\mathcal{T}_k \subseteq \{1, \ldots, L_c\}.$$
(4)

STAR从选定的文本-图像注意力层集合$\mathcal{L}$中读取路由信号。在生成采样轨迹$i$时，基于注意力的T2I骨干网络在每个去噪步骤$t$建立图像潜在词元和文本词元之间的交互。对于任意$l \in \mathcal{L}$，该层提供一个注意力图$\alpha_t^{i,l} \in \mathbb{R}^{h_t \times w_t \times L_c}$；从潜在位置$u$到文本词元$j$的注意力权重记为

$$\alpha_t^{i,l}(u, j).$$
(5)

这里$u \in \Omega_t$是一个潜在空间位置，$j$是文本词元索引。在联合注意力骨干网络（例如Stable Diffusion 3.5, Qwen-Image）中，该值对应完整注意力矩阵的图像到文本块；在交叉注意力骨干网络（例如Stable Diffusion XL）中，它对应从图像查询到文本键的交叉注意力权重。STAR取选定层上最大的图像到文本注意力权重：

$$\alpha_t^{i}(u, j) = \max_{l \in \mathcal{L}} \alpha_t^{i,l}(u, j),$$
(6)

这给出了采样轨迹$i$在步骤$t$的特定时间步文本-图像路由图$\alpha_t^{i} \in \mathbb{R}^{h_t \times w_t \times L_c}$。

接下来，STAR提取每个文本单元在当前轨迹中出现的位置。给定文本单元$q_k$的词元集合$\mathcal{T}_k$，STAR构建其空间图$M_t^{i,k} \in \mathbb{R}^{h_t \times w_t}$，其中潜在位置$u$处的路由强度为

$$M_t^{i,k}(u) = \frac{1}{|\mathcal{T}_k|} \sum_{j \in \mathcal{T}_k} \alpha_t^{i}(u, j).$$
(7)

STAR将这些关键组件聚合成一个轨迹-时间步特定的热图$H_t^{i} \in \mathbb{R}^{h_t \times w_t}$：

$$H_t^{i}(u) = \frac{1}{K} \sum_{k=1}^{K} M_t^{i,k}(u).$$
(8)

该热图指定了与提示相关内容最活跃的位置。为了将其转换为归一化的奖励分配图，STAR将其映射到有界系数$C_t^{i} \in \mathbb{R}^{h_t \times w_t}$：

$$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) + \varepsilon}.$$
(9)

这里$\lambda \geq 1$控制分配强度，$\varepsilon$确保数值稳定性。较大的$C_t^{i}(u)$增强了对提示相关区域的更新；较小但非零的系数保留了背景、构图和整体质量对最终奖励的影响。

### 3.2 空间解析的策略目标

STAR使用分配图将样本级优势转换为空间化的训练信号。对于第$i$个采样轨迹和去噪步骤$t$，我们定义空间优势图为

$$A_t^{i}(u) = A_i \, C_t^{i}(u), \qquad A_t^{i} \in \mathbb{R}^{h_t \times w_t}.$$
(10)

该映射保留了$A_i$的符号，并通过$C_t^{i}(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)

其中$\ell_{\theta,t}^{i}(u)$表示潜在位置$u$处的对数概率贡献，$\bar{\ell}_{\theta,t}^{i}$对应标准策略梯度训练中使用的标量转移对数概率。STAR使用归一化的空间比率：

$$\rho_{\theta,t}^{i}(u) = \frac{\operatorname{sg}\!\left[\exp\!\left(\bar{\ell}_{\theta,t}^{i} - \bar{\ell}_{\text{old},t}^{i}\right)\right] \exp\!\left(\ell_{\theta,t}^{i}(u)\right)}{\operatorname{sg}\!\left[\exp\!\left(\ell_{\theta,t}^{i}(u)\right)\right]},$$