AIxiv专栏是AI在线发布学术、技术内容的栏目。过去数年，AI在线AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
本论文作者主要包括澳大利亚科学院院士、欧洲科学院外籍院士、IEEE Fellow陶大程，现任南洋理工大学杰出教授；新加坡工程院院士、IEEE Fellow文勇刚，现为南洋理工大学计算机科学与工程学院校长讲席教授，同时担任IEEE Transactions on Multimedia主编；张森，曾在悉尼大学从事博士后研究工作，现任TikTok机器学习工程师；詹忆冰，京东探索研究院算法科学家。本文的通讯作者是武汉大学计算机学院教授、博士生导师、国家特聘青年专家罗勇。第一作者为张子屹，目前在武汉大学计算机学院攻读博士二年级，研究方向为强化学习、扩散模型和大模型对齐。
OpenAI 最近发布的 o1 模型在数学、代码生成和长程规划等复杂任务上取得了突破性进展，据业内人士分析披露，其关键技术在于基于强化学习的搜索与学习机制。通过迭代式的自举过程，o1 基于现有大语言模型的强大推理能力，生成合理的推理过程，并将这些推理融入到其强化学习训练过程中。
强化学习技术能够显著提升模型解决复杂问题的能力，尤其是在 o1 所采用的细粒度奖励机制的加持下。这种奖励机制为模型的每一步推理提供细粒度的反馈，而不仅是依赖最终答案的正确性来评估模型的表现。通过精细化的控制，使模型能够不断优化其推理路径，有效应对复杂任务中的长程依赖问题，确保模型在推理过程中保持逻辑一致性。
前不久，来自武汉大学、悉尼大学、京东探索研究院和南洋理工大学的研究人员便在 ICML 2024 上发表了一项 “扩散模型对齐” 方向的研究，所提出的名为 TDPO-R 的强化学习算法与 o1 不谋而合地也采用了细粒度奖励机制。
论文链接：https://openreview.net/forum?id=v2o9rRJcEv
代码链接：https://github.com/ZiyiZhang27/tdpo
TDPO-R 在强化学习算法中引入了时间差分奖励机制，对文生图扩散模型的每一步生成过程提供细粒度的反馈，从而有效缓解了在扩散模型对齐时常见的奖励过优化问题。这项研究证实了细粒度奖励机制在扩散模型对齐中的关键性，而 o1 的最新技术同样揭示了这一机制在大模型领域中的广泛应用前景，有望推动生成模型在多样化、复杂任务中的持续发展与优化。
扩散模型（Diffusion Model）作为一种先进的生成式模型，通过学习并模拟自然界中的扩散过程来合成新数据，尤其在生成高质量逼真图像方面相较于其他模型架构具有显著优势。因此，尽管近年来也有基于 GPT 等其它架构的优秀图像、视频生成模型不断涌现，基于扩散模型的架构依然是当前大多数图像、视频生成任务的主流选择。
近期，为了确保生成的图像能够符合下游任务目标以及用户偏好，许多研究者开始探索如何对预训练扩散模型进行奖励或偏好驱动的微调训练 —— 这就是所谓的扩散模型对齐（Diffusion Model Alignment）。目前，最流行的对齐方法之一便是源于大语言模型领域的一项常用技术 —— 基于人类反馈的强化学习（Reinforcement Learning from Human Feedback, RLHF），即以强化学习方式依据人类反馈来不断地对已有的预训练模型参数进行迭代优化。
另外，也有其它一些不涉及强化学习的对齐方法，但它们大多也都需要以特定奖励函数作为优化目标，即将所对齐的下游任务目标（如图像美学质量、文图一致性等）构建为一个奖励函数 / 模型，并以最大化该奖励函数为目标对扩散模型进行迭代优化。还有一些方法，比如直接偏好优化（Direct Preference Optimization，DPO），虽然不显式地使用奖励函数来对齐模型，但其所采用的训练数据本身往往也隐含了某种奖励标准，间接引导模型朝着某种期望的目标优化。
扩散模型对齐中的奖励过优化
由于受奖励目标驱动，这些扩散模型对齐方法常常面临一个核心挑战 —— 奖励过优化（Reward Overoptimization），即经过微调后的模型可能会过度偏向于某一奖励目标，导致生成的图像丧失个性化和多样性、视觉保真度降低，最终偏离人类真实的审美偏好。如上图所示，当我们使用美学分数（Aesthetic Score）奖励模型，对 Stable Diffusion v1.4 模型进行对齐微调时，“奖励过优化” 这一现象在生成图像上肉眼可见为：图像风格趋于统一、与文本提示词的匹配度下降。
如何更加通俗地理解奖励过优化问题？想象你是一位学生，目标是通过考试拿到高分。你的老师设计了一套严格的评分标准（就像是奖励模型），目的是测试你的知识深度和理解能力。这套标准非常全面，涵盖了大量的知识点和难题，理论上应该能够准确衡量你的学习水平。
然而，你发现了一些 “捷径”：考试中虽然有很多问题，但也有着某些 “答题套路”，比如背诵标准答案，或者用某些关键词和答题模板。于是你并不真正花时间理解知识，而是把大量时间用在研究考试技巧上，专注于如何根据评分标准去 “迎合” 考试。结果，你每次考试都能拿到高分（就像模型获得了高奖励），但实际上你对知识的理解很肤浅。老师看到你的分数，误以为你学得非常好，但你自己知道，虽然分数高了，但你并没有真正掌握知识，甚至可能过于执着答题模板而错误地理解了一些问题，换另一位老师出题的话便可能原形毕露。
如何对奖励过优化问题进行定量分析？TDPO-R 算法的研究者发现，当扩散模型过度优化某一特定奖励函数时，其在域外奖励函数上的泛化能力随之下降。这种现象可以通过定量分析表现为：在以某一奖励函数为优化目标时，采用其他域外奖励函数（即不同于优化目标的奖励标准）对生成样本进行评分，会发现随着目标奖励逐渐上升，域外奖励评分的增长速度明显放缓，甚至出现域外评分逐渐下降的情况。这表明，模型在一个特定奖励函数上取得高分时，可能会以牺牲其在其他任务上的表现为代价，最终导致生成质量的不平衡。因此，研究者采用了这种跨奖励泛化度量，来定量地对比不同算法的奖励过优化程度。
研究者进一步指出，相较于基于监督学习的方法，基于强化学习的扩散模型对奖励过优化问题表现出更高的抗性。这是因为强化学习能够根据生成过程中的反馈信号不断调整策略，避免模型在单一目标上过度拟合。然而，这种抗性并非完全免疫。强化学习的这种优势往往是以牺牲样本效率为代价，即需要更多的样本和训练时间来达到较为平衡的效果。因此，在设计新的扩散模型对齐方法时，必须在提高样本效率与缓解奖励过优化之间做出权衡。
受到深度强化学习理论的启发，研究者首先从扩散模型对齐中的归纳偏置角度出发，分析了加重奖励过优化的潜在因素。归纳偏置（Inductive Bias）指的是机器学习算法在从有限的训练数据中学习时，基于特定假设提炼出的一般性模式。在深度强化学习的背景下，算法所采用的归纳偏置与其解决任务的契合程度，直接影响模型的泛化能力。简单来说，如果模型的偏置过于集中在某个特定目标上，就容易在泛化到新任务时出现问题。
然而，现有的扩散模型对齐方法往往过度依赖稀疏奖励 —— 也就是只根据最终生成图像的质量进行优化。这种做法隐含着一种假设，即最终图像的质量是最重要的评估标准，忽略了生成过程中的中间步骤（即每一步的去噪图像）。模型只在最后一步表现良好时被奖励，导致训练过程中没有充分利用多步去噪过程中的中间信息，而这部分信息恰恰包含了对生成过程的更细粒度反馈。因此，这种忽略中间步骤的做法，实际上引入了一种错位的归纳偏置 —— 即模型在训练时过于注重终极输出的奖励信号，而忽视了去噪过程中可能帮助模型提高质量的阶段性信号。这种偏置错位便是加重奖励过优化的潜在因素之一。
为了更好理解这种偏置错位，可以结合之前老师与学生的比喻：想象你是一名学生，你的老师只根据期末考试的成绩来评价你整个学期的表现（类似于稀疏奖励）。这种情况下，学生的学习过程（就像扩散模型中的中间去噪过程）没有被重视或衡量，学生很容易找到捷径（如通过背题库而非理解知识）来获得高分。这就类似于模型在训练中通过某种捷径，在最后生成的图像中获得高分，而忽视了生成过程中的细节和图像质量的演变。
然而，如果老师在整个学期中对你的每一次小测验、作业、甚至课堂表现都进行细粒度的评分和反馈（相当于细粒度奖励），那么你在每一步的学习过程中都会受到监控和指导。你无法依靠最后的考试来蒙混过关，必须在每个阶段都真正掌握知识。这与扩散模型对齐时类似，细粒度奖励机制能够让模型在每个生成步骤中都得到反馈，避免模型通过过度优化某一单一目标（如最终图像质量）来走捷径，确保生成过程中的每一步都保持质量和一致性。
TDPO-R 算法
面对以上问题，TDPO-R 通过引入时间差分奖励机制，为扩散模型的每一步去噪操作提供实时的奖励反馈，从而修正这种偏置错位，缓解奖励过优化问题。这种奖励机制将扩散过程中的每个时间步视为一个马尔科夫决策过程中的状态，模型在每一步的动作（即去噪操作）后，都会获得相应的奖励值。为了降低时间差分奖励函数的学习成本，TDPO-R 通过一个时间差分评判器（Temporal Critic）来近似估计每个时间步的奖励。这一评判器通过强化学习的方式与扩散模型一齐同步进行训练，逐步学习如何为每个时间步估算合适的奖励。具体来说，模型在每个时间步 t 做出去噪操作后，评判器会基于当前状态（中间去噪图像）与目标状态的差异，给出即时奖励反馈。
在策略更新时，TDPO-R 采用类似于策略梯度（Policy Gradient）的方法。特别是，得益于时间差分奖励机制提供了每个时间步的即时反馈，TDPO-R 在扩散模型去噪的每一步都可以立即进行策略更新，而无需等待所有步骤结束。这种即时更新让模型能够及时修正去噪过程中出现的偏差，避免问题在后续步骤中积累。同时，由于每一步都进行即时更新，模型可以在更短的时间内获得更好的优化效果，减少了传统全局更新方式下的延迟和不必要的计算开销，提高了训练的样本效率。
接着，从首要偏置（Primacy Bias）的角度出发，TDPO-R 的研究者进一步分析了奖励过优化的问题。首要偏置描述的是深度强化学习模型在训练中倾向于过度拟合早期的训练经验，从而影响后续训练阶段的优化效果的一种现象。他们发现，在扩散模型的对齐任务中，首要偏置与奖励过优化也有着密切的联系，特别是当模型在训练过程中过度倾向于某个早期的策略时，后续的生成质量可能会出现下降。
为了深入理解首要偏置与奖励过优化之间的关系，研究者引入了神经网络中神经元激活状态的概念。通常，深度神经网络中的神经元可以处于两种状态：
活跃神经元：在训练过程中频繁被激活的神经元，通常与当前任务的主要模式相关。
休眠神经元：在训练中很少被激活的神经元，通常被认为对当前任务贡献较少，甚至被认为会阻碍模型的学习能力。
在以往的研究中，活跃神经元往往被视为是模型的核心，因为它们反映了模型对当前数据的主要适应能力。而休眠神经元则通常被认为是冗余的或不必要的，因为它们没有参与主要的学习过程。
而在 TDPO-R 的研究中，研究者惊讶地发现，休眠神经元实际上在应对奖励过优化方面起到了重要作用。具体而言：
活跃神经元更易受到首要偏置的影响：由于活跃神经元频繁参与早期训练任务，它们倾向于过拟合早期的学习模式。因此，这些神经元在模型的优化过程中可能会过度强化某一特定奖励目标，导致奖励过优化问题的加剧。例如，模型可能因为过度优化早期阶段的奖励信号，导致后续的生成结果表现不佳，甚至在其他奖励函数上的泛化能力下降。
休眠神经元作为自适应正则化手段：与活跃神经元相反，休眠神经元由于较少参与训练任务，反而可以作为一种对抗奖励过优化的自适应正则化手段。这是因为，休眠神经元的低激活状态意味着它们没有过度依赖早期的奖励信号，从而在模型后期的训练中可以起到平衡的作用。另外，它们还能够提供一种类似于 “潜力储备” 的功能，在奖励函数发生变化或模型过拟合早期奖励信号时，休眠神经元可以重新被激活，以补充模型的适应能力，增强模型的泛化性能。
为了应对首要偏置，TDPO-R 引入了一种神经元重置机制。这种机制通过定期重置模型中的活跃神经元，将其恢复到类似 “休眠” 状态，从而打破首要偏置的影响，重新激发模型的学习能力。具体来说：
周期性神经元重置：在训练过程中，TDPO-R 会定期对评判器（critic）模型中过度活跃的神经元进行重置，降低它们的激活频率，从而打破它们在早期阶段对特定奖励信号的过拟合。通过这一操作，模型可以避免过度强化某一奖励目标，确保生成过程的多样性和泛化能力。
重新激活休眠神经元：随着活跃神经元被重置，模型的其他神经元，包括那些此前处于休眠状态的神经元，会被激活，以参与新的学习任务。这种神经元的 “轮替” 确保了模型的学习能力不会因为早期训练经验的固定化而受到限制，从而缓解了奖励过优化的问题。
实验评估 —— 跨奖励泛化度量
研究者使用提出的 TDPO-R 对 Stable Diffusion v1.4 模型进行微调训练，并采取不同文本提示词集和奖励函数进行评估实验。为了更直观地展现跨奖励泛化度量的效果，研究者通过图表展示了在优化某一特定目标奖励函数（如 Aesthetic Score）时，模型在其他多种奖励函数（如 ImageReward、HPSv2 和 PickScore）上的得分如何变化。如上图所示，其中每幅图表的横坐标代表训练保存的中间模型在目标奖励函数上的得分，而纵坐标代表对应的中间模型在域外奖励函数上的得分。通过这些图表，可以清楚地观察到 TDPO-R 和其他扩散模型对齐方法（如 DDPO 和 AlignProp）在跨奖励泛化能力上的差异。
在泛化较差的目标奖励函数（如 Aesthetic Score）上，TDPO-R 能够保持域外奖励函数的得分没有显著下滑，而其它方法则表现为明显的下降趋势；而在泛化较好的目标奖励函数（如 HPSv2 和 PickScore）上，TDPO-R 能够使域外奖励函数的得分更快地上涨。这表明 TDPO-R 在优化特定目标时，展现出更强的跨奖励泛化能力，有效缓解了奖励过优化的问题。
除了定量指标的对比，研究者还通过可视化图像样本的对比展示了 TDPO-R 相较于其他方法的优势。如下图所示对比了 TDPO-R 与其他方法，皆将目标奖励函数（Aesthetic Score）的得分提升至同一临界值时所生成的样本图像。可以看出，在其它方法中，不同文本提示词对应的图像样本风格趋同，尤其是在光线、色调和构图风格方面，缺乏多样性，甚至还存在图像失真的情况。相比之下，TDPO-R 不仅能够生成更具视觉自然感和细节保真的图像，还保留了风格的多样性，具有更平衡的美学表现，直观地体现了 TDPO-R 缓解奖励过优化的有效性。
如下图所示，在生成包含复杂场景描述的图像时，其它方法存在图像细节不准确的问题，忽略了关键的语义元素（如物体的数量、类型以及场景地点等），表现出典型的文图语义不一致问题。而 TDPO-R 生成的图像不仅视觉丰富，还能够准确反映文本描述的内容。这体现了 TDPO-R 的跨任务泛化能力，即在提升图像美学表现的同时，还能够有效保留文图的语义一致性。
除了上述实验结果之外，研究者还展示了其他一些主要实验内容，包括 “算法样本效率对比”、“对未见提示词的泛化能力”、“神经元状态的影响” 以及 “应对过优化的其它替代策略” 等方面。对于这些补充实验以及 TDPO-R 方法的更多技术细节，感兴趣的读者可以参阅论文原文进一步了解。