约束流优化：基于序列微调的分子设计方法

# 约束流优化：通过序列微调进行分子设计  
来源：https://arxiv.org/html/2605.30610  

###### 摘要  

调整生成式基础模型（特别是扩散模型和流模型）以优化给定奖励函数（如结合亲和力），同时满足约束条件（如分子可合成性），对于其在分子设计或蛋白质工程等现实科学发现应用中的采用至关重要。尽管近期工作已通过强化学习和控制方案引入了可扩展的方法，用于对此类模型进行奖励引导的微调，但如何以可靠且可预测的方式算法性地权衡奖励最大化与约束满足仍是一个开放问题。受此挑战启发，我们首先提出一个严格的框架——*约束生成式优化*，该框架从优化视角审视所引入的适应问题，并将约束生成的相关任务作为子案例纳入其中。接着，我们介绍**约束流优化（CFO）**，一种通过将原始问题简化为基于现有可扩展方法的序列微调，自动且可证明地平衡奖励最大化与约束满足的算法。我们提供了CFO在约束生成式优化和约束生成方面的收敛性保证。最后，我们在合成但具有说明性的设置以及分子设计任务上对CFO进行了实验评估。在这些评估中，CFO在确保高约束满足率的同时实现了奖励的一致提升，展示了其在约束生成式优化中的实用价值。  

机器学习，ICML  

## 1 引言  

生成式建模的最新进展，特别是扩散模型（Ho et al., 2020；Song et al., 2020, 2022）和流模型（Lipman et al., 2022）的出现，已在图像合成（Rombach et al., 2022）、生物学（Corso et al., 2022；Wohlwend et al., 2025）和化学（Hoogeboom et al., 2022）等多个领域取得了最先进的性能。特别是，它们已被应用于蛋白质结构（Wu et al., 2024）、类药分子（Dunn and Koes, 2024）和DNA序列（Stark et al., 2024）等设计任务。这些生成式模型擅长捕捉复杂的数据分布并生成逼真的样本。然而，对于大多数现实世界的发现应用而言，仅近似地从数据分布中采样是不够的。在这些应用中，通常希望生成最大化特定任务*奖励*的候选样本，这一问题近期被称为*生成式优化*（De Santi et al., 2025a；Li et al., 2024）。感兴趣的奖励示例包括药物发现中的结合亲和力（Pantsar and Poso, 2018）或药物相似性（Bickerton et al., 2012）。  

为解决生成式优化问题，近期工作引入了可扩展的微调方法，这些方法通过强化学习或控制理论方法（例如，Domingo-Enrich et al., 2024；Uehara et al., 2024b；Tang and Zhou, 2024），在KL正则化下调整预训练的流模型或扩散模型以最大化给定奖励函数。  

**已知约束在生成式优化中的重要性。** 许多生成式设计和科学发现问题要求生成的样本满足明确的、领域特定的约束，例如毒性上限（Amorim et al., 2024）、合成可及性（Ertl and Schuffenhauer, 2009；Neeser et al., 2024）或对接姿态的生物物理合理性（Buttenschoen et al., 2024）。尽管当前的微调方案通过正则化向预训练模型（Domingo-Enrich et al., 2024；Uehara et al., 2024b；Tang and Zhou, 2024）限制分布漂移，但它们无法保证硬约束的满足（Uehara et al., 2024a）。这一局限性源于任务特定的约束可能未在原始数据集中编码，或仅能从有限的训练数据中不完美地学习到。一种处理此类显式约束的朴素方法是将它们作为奖励纳入，即作为手动加权目标函数中的另一项。然而，这种方法在实践中不可靠，因为奖励与约束之间的适当权重因任务和训练阶段而异，需要通过低效的试错来确定。此外，随着优化探索高奖励区域，所选权重可能意外地倾向于奖励而牺牲约束满足，从而产生具有诱人奖励但违反领域特定约束的样本。  

受当前流适应方法在约束满足方面的这些局限性的驱动，我们提出以下问题：*如何微调预训练的流模型或扩散模型，以可靠且可预测地权衡奖励优化与约束满足？*  

**我们的方法。** 越来越多的研究表明，经典的优化思想可以有意义地应用于流模型和扩散模型的微调，包括受镜像下降（Nemirovskij and Yudin, 1983；De Santi et al., 2025a）、机会约束（Ben-Tal and Nemirovski, 2000；Zhang et al., 2025a）和双层优化（Bracken and McGill, 1973；Xiao et al., 2025）启发的公式。类似地，在本工作中，我们旨在通过引入一个正式的框架——*约束生成式优化*（第3节），通过流模型微调来解决这一问题，该框架涉及调整预训练流模型以生成最大化奖励函数同时满足任意约束的样本。此外，所提出的公式将约束生成式建模的相关任务作为奖励函数为常数时的子案例纳入。接下来，我们介绍**约束流优化（CFO）**，这是一种基于增广拉格朗日方案（Birgin and Martínez, 2014）的对偶方法，它将约束目标转化为一系列普通的生成式优化子问题。在高层面上，CFO交替执行两个步骤：求解KL正则化的微调问题（Domingo-Enrich et al., 2024；Uehara et al., 2024b）以最大化增广奖励函数，以及利用生成样本上估计的约束违反来更新增广奖励的参数（见第4节）。这一过程顺序调整了对约束违反的惩罚，从而避免了手动权衡权重选择的需要。CFO使得调整预训练流模型以最大化期望奖励，同时强制满足任意约束并保持与预训练模型的接近性成为可能。我们在近似求解器的现实假设下提供了保证约束满足的保障，并在更理想化的设置下实现了奖励最大化（第5节）。最后，我们在约束生成式优化和建模问题上评估了CFO，展示了其在视觉上可解释的设置以及分子设计任务中的性能，包括量子力学性质的约束优化（第6节）。  

**我们的贡献。** 我们提出以下贡献：  

- • 我们通过流微调形式化了*约束生成式优化*，捕捉了在给定约束下进行奖励引导适应的实际相关任务（第3节）。  
- • 我们介绍了**约束流优化（CFO）**，一种基于增广拉格朗日的方法，通过序列微调可证明地解决上述问题（第4节）。  
- • 我们在两种预言机假设下，利用增广拉格朗日理论，提供了CFO在约束生成和优化方面的保证（第5节）。  
- • 我们在视觉上可解释的设置以及高维分子设计任务中展示了CFO权衡奖励最大化与约束满足的能力（第6节）。  

参见图注(a) 约束生成式优化微调问题。参见图注(b) 适应黑线内的低成本区域。  

图1：(1(a)) 预训练和微调策略诱导密度 \(p_1^{\text{pre}}\) 和相对于奖励 \(r\) 的最优密度 \(p_1^{*}\)，奖励向下增加，红色区域为高成本区域。(1(b)) 预训练模型 \(p_1^{\text{pre}}\) 适应为 \(p_1^{*}\) 以最大化 \(r\) 并保持在黑线内的约束区域内。  

## 2 背景与符号  

**流模型。** 基于流的生成式模型构成了将简单基分布 \(p^{\text{base}}\)（如高斯分布）变换为复杂数据分布 \(p_{\rm{data}}\) 的一类重要方法（Song et al., 2022, 2020；Lipman et al., 2022）。形式上，流是一个时间依赖的映射 \(\psi: [0,1] \times \mathbb{R}^d \to \mathbb{R}^d\)，其中 \(\psi_t(x_0)\) 表示在时间 \(t\) 起始于 \(x_0\) 的样本的位置。\(x_t\) 的轨迹由一个时间依赖的速度场 \(u: [0,1] \times \mathbb{R}^d \to \mathbb{R}^d\) 通过常微分方程（ODE）控制：  

\[
\frac{\text{d}}{\text{d}t} \psi_t(x_0) = u_t(\psi_t(x_0)), \quad \psi_0(x_0) = x_0.
\]  

一个*生成式*流模型通过采样初始值 \(X_0 \sim p^{\text{base}}\) 并根据流映射演化 \(X_t = \psi_t(X_0)\)，定义了一个连续时间马尔可夫过程 \(\{X_t\}_{t \in [0,1]}\)。终端状态 \(X_1 = \psi_1(X_0)\) 需遵循目标分布，即 \(X_1 \sim p_{\rm{data}}\)。等价地，该流诱导了一族中间边际密度 \(p_t\)，描述在每个时间 \(t \in [0,1]\) 上 \(X_t\) 的分布。我们说一个速度场 \(u\) 生成了概率路径 \(\{p_t\}_{t \in [0,1]}\)，如果对于所有 \(t < 1\)，随机变量 \(X_t = \psi_t(X_0) \sim p_t\)。在实践中，选择 \(p^{\text{base}} = \mathcal{N}(0,I)\) 使得抽样易于处理，而 \(u_t\) 提供了达到 \(p_{\rm{data}}\) 所需的复杂性。  

**流匹配。** 流匹配（Lipman et al., 2022）是一种免模拟算法，用于学习一个向量场 \(u_\theta\)，使得其诱导的边际密度 \(p_t^{u_\theta}\) 与规定的概率路径 \(\{p_t\}_{t \in [0,1]}\) 一致，并满足 \(p_0^{u_\theta} = p^{\text{base}}\) 和 \(p_1^{u_\theta} = p_{\rm{data}}\)。Lipman等人（2022）证明，流匹配和条件流匹配目标具有相同的梯度，确保它们收敛到相同的最优向量场。在实践中，这通过引入一个参考流并将学习到的场 \(u_\theta(x_t,t)\) 对参考速度进行回归来实现：  

\[
\min_{\theta} \mathbb{E}_{t, p(x_0,x_1)} \left[ \left\lVert u_\theta(x_t,t) - \frac{\text{d}}{\text{d}t} \psi_t^{\text{ref}}(x) \right\rVert^2 \right].
\]  

通过适当选择参考流（特别是遵循扩散轨迹的参考流），流匹配框架将扩散模型恢复为一个特例，表明扩散训练目标可以被视为基于流学习的特殊实例（Lipman et al., 2022；Domingo-Enrich et al., 2024）。在实践中，\(u_\theta\) 由神经网络参数化，通过模拟公式(1)中的ODE来从 \(p_1^{u_\theta}\)（≈ \(p_{\rm{data}}\)）中采样。  

**连续时间强化学习。** 有限时域连续时间强化学习（RL）（Wang et al., 2020；Treven et al., 2023；Zhao et al., 2025）为动态系统中的决策制定提供了一个框架，并可视为最优控制的一个实例。状态空间为 \(\mathcal{X} \coloneqq \mathbb{R}^d \times [0,1]\)，动作取自动作空间 \(\mathcal{A}\)。策略 \(\pi: \mathcal{X} \to \mathcal{A}\) 为每个状态 \((x,t) \in \mathcal{X}\) 分配一个动作，产生动力学：  

\[
\frac{\text{d}}{\text{d}t} \psi_t(x) = a_t(\psi_t(x)), \quad a_t = \pi(X_t,t), \quad X_0 \sim p^{\text{base}}.
\]  

由此产生的过程 \(\{X_t\}_{t \in [0,1]}\) 诱导了一族边际分布 \(\{p_t^\pi\}_{t \in [0,1]}\)。目标是优化期望性能，通常通过沿轨迹积累的积分奖励和 \(t=1\) 时的终端奖励来表达（Wang et al., 2020）。在我们的设定中，我们仅关注终端奖励。我们使用RL符号以强调其通用性和与标准实践的联系，同时注意到该设定与确定性最优控制一致，因为动力系统和目标函数都是已知的。  

**预训练流模型。**