Aligning Transformers with Continuous Feedback via Energy Rank Alignment
作者: Shriram Chennakesavalu, Frank Hu, Sebastian Ibarraran, Grant M. Rotskoff
分类: cs.LG, cs.AI, physics.chem-ph, q-bio.QM
发布日期: 2024-05-21 (更新: 2025-10-22)
💡 一句话要点
提出能量排序对齐(ERA)算法，用于优化自回归策略以生成具有指定属性的分子和蛋白质序列。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 分子生成 蛋白质序列生成 自回归模型 能量排序对齐 奖励函数
📋 核心要点
- 化学空间搜索面临组合爆炸问题，现有方法难以有效生成具有特定属性的分子。
- 提出能量排序对齐（ERA）算法，利用奖励函数优化自回归策略，无需强化学习。
- 实验表明，ERA算法能有效对齐分子Transformer和蛋白质语言模型，生成具有指定属性的序列。
📝 摘要（中文）
在化学空间中搜索极具挑战性，因为可能分子的数量随原子数量呈组合增长。基于化学化合物数据库训练的大型自回归模型已经产生了强大的生成器，但我们仍然缺乏生成具有所需属性分子的稳健策略。这种分子搜索问题与大型语言模型的“对齐”问题非常相似，但对于许多化学任务，我们有一个明确且易于评估的奖励函数。本文介绍了一种名为能量排序对齐（ERA）的算法，该算法利用显式奖励函数来产生基于梯度的目标，我们使用该目标来优化自回归策略。理论上，我们证明该算法与近端策略优化（PPO）和直接偏好优化（DPO）密切相关，但其极小值收敛到理想的吉布斯-玻尔兹曼分布，其中奖励充当能量函数。此外，该算法具有高度可扩展性，不需要强化学习，并且在每次配对的偏好观察数量较少时，相对于DPO表现良好。我们部署这种方法来对齐分子Transformer和蛋白质语言模型，以分别生成具有外部指定属性的分子和蛋白质序列，并发现它能够稳健地做到这一点，从而搜索化学空间的不同部分。
🔬 方法详解
问题定义：论文旨在解决分子和蛋白质序列生成中，如何高效地生成具有特定属性的序列的问题。现有方法，特别是基于大型自回归模型的方法，虽然能够生成序列，但缺乏有效的策略来引导生成过程，使其满足外部指定的属性要求。传统的强化学习方法虽然可以用于优化生成过程，但训练不稳定且计算成本高昂。
核心思路：论文的核心思路是将奖励函数视为能量函数，并利用能量排序对齐（ERA）算法来优化自回归策略。ERA算法的目标是使生成序列的概率分布接近于一个理想的吉布斯-玻尔兹曼分布，其中序列的能量由其奖励值决定。通过最小化生成序列的能量与奖励之间的差异，可以有效地引导生成过程，使其生成具有更高奖励值的序列。
技术框架：ERA算法的整体框架包括以下几个主要步骤：1) 使用自回归模型生成序列；2) 使用外部奖励函数评估生成序列的质量；3) 计算生成序列的能量，能量与奖励相关；4) 使用ERA算法更新自回归模型的参数，使得生成序列的概率分布更接近于理想的吉布斯-玻尔兹曼分布。该框架不需要强化学习，可以直接使用梯度下降等优化方法进行训练。
关键创新：ERA算法的关键创新在于将奖励函数与能量函数联系起来，并提出了一种新的优化目标，使得自回归模型能够直接学习如何生成具有高奖励值的序列。与传统的强化学习方法相比，ERA算法更加稳定和高效，因为它避免了策略梯度估计的方差问题。此外，ERA算法还具有良好的可扩展性，可以应用于各种不同的分子和蛋白质序列生成任务。
关键设计：ERA算法的关键设计包括：1) 能量函数的选择，通常选择与奖励函数成比例的函数；2) 优化目标的构建，目标是最小化生成序列的能量与奖励之间的差异；3) 学习率的设置，需要根据具体任务进行调整，以保证训练的稳定性和收敛速度。论文中还提到，ERA算法与PPO和DPO算法密切相关，但在每次配对的偏好观察数量较少时，ERA算法表现更好。
🖼️ 关键图片
📊 实验亮点
实验结果表明，ERA算法能够有效地对齐分子Transformer和蛋白质语言模型，生成具有外部指定属性的分子和蛋白质序列。在分子生成任务中，ERA算法能够生成具有更高药物活性的分子。在蛋白质序列生成任务中，ERA算法能够生成具有更高稳定性和功能的蛋白质序列。ERA算法在少量偏好观察下优于DPO。
🎯 应用场景
该研究成果可广泛应用于药物发现、材料设计和蛋白质工程等领域。通过优化分子和蛋白质序列的生成过程，可以加速新药物和新材料的开发，并提高蛋白质的功能和稳定性。该方法还可以应用于其他序列生成任务，例如自然语言生成和音乐生成。
📄 摘要（原文）
Searching through chemical space is an exceptionally challenging problem because the number of possible molecules grows combinatorially with the number of atoms. Large, autoregressive models trained on databases of chemical compounds have yielded powerful generators, but we still lack robust strategies for generating molecules with desired properties. This molecular search problem closely resembles the "alignment" problem for large language models, though for many chemical tasks we have a specific and easily evaluable reward function. Here, we introduce an algorithm called energy rank alignment (ERA) that leverages an explicit reward function to produce a gradient-based objective that we use to optimize autoregressive policies. We show theoretically that this algorithm is closely related to proximal policy optimization (PPO) and direct preference optimization (DPO), but has a minimizer that converges to an ideal Gibbs-Boltzmann distribution with the reward playing the role of an energy function. Furthermore, this algorithm is highly scalable, does not require reinforcement learning, and performs well relative to DPO when the number of preference observations per pairing is small. We deploy this approach to align molecular transformers and protein language models to generate molecules and protein sequences, respectively, with externally specified properties and find that it does so robustly, searching through diverse parts of chemical space.