Reward Transport: 通过噪声空间对齐实现流匹配中的属性控制

# 通过噪声空间对齐实现流匹配中的属性控制 来源：https://arxiv.org/html/2607.08781

Kehan Guo¹, Yili Shen¹, Yujun Zhou¹, Yue Huang¹, Chujie Gao¹, Shiyi Du², Xiangliang Zhang¹  
¹圣母大学，²卡内基梅隆大学  
[email protected], [email protected]  

###### 摘要  
流匹配中的耦合——即噪声向量与数据点的配对规则——通常被视为一个计算选择。我们证明，这种耦合可以充当一个对齐接口：通过根据目标分子属性匹配噪声和数据，它直接将可控结构嵌入到学习到的流场中。基于这一观点，我们引入了奖励传输（Reward Transport），它在训练时使用最优传输耦合来将标量噪声空间坐标与分子奖励对齐；在推理时，改变这个坐标即可引导生成分布，而无需预言机、奖励模型、梯度引导或额外计算。在耦合保持极限下，对该坐标进行阈值化操作可恢复交叉熵方法（Cross-Entropy Method）的截断奖励分布，从而提供一个原理性、可连续调节的分布级控制旋钮。实验上，在ZINC-250K和GuacaMol上，扫描该标量可实现对logP的单调控制和QED在其工作范围内的稳定控制；最值得注意的是，相同的旋钮对不同的目标产生相反的结构响应——对logP分子增大，对QED分子缩小——这排除了通用尺寸偏差的可能性。该接口与无分类器引导和条件流匹配互为补充，而ε-预测扩散下的负面结果则阐明了耦合级对齐在结构上缺失的情况。我们的代码见：https://github.com/KehanGuo2/reward-transport。

## 1 引言

一个流匹配模型既依赖于向量场训练目标，也依赖于*耦合*：即在训练过程中决定哪个噪声向量与哪个数据点配对的规则。尽管先前的工作仔细研究了训练目标（Karras et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib20); Esser et al., 2024 (https://arxiv.org/html/2607.08781#bib.bib12); Ma et al., 2024 (https://arxiv.org/html/2607.08781#bib.bib25)），但耦合通常被视为一个背景实现选择：为了简单而独立配对（Lipman et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib23)），或者使用小批量最优传输来平直轨迹并减少路径交叉（Tong et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib34); Pooladian et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib28)）。在每种情况下，选择耦合都是为了简化训练，而不是为了塑造模型最终学到的内容。

我们持有不同的观点。耦合不是一个训练细节，而是一个*对齐接口*：它决定了哪些噪声样本与哪些数据样本相关联，从而决定了有意义的数据属性是否会在噪声空间中组织起来。在大多数可控生成模型中（Ho and Salimans, 2022 (https://arxiv.org/html/2607.08781#bib.bib13); Dhariwal and Nichol, 2021 (https://arxiv.org/html/2607.08781#bib.bib10)），可控性是通过显式输入引入的，例如类别标签、属性嵌入或无分类器引导。在流匹配中，我们证明可控性也可以通过耦合本身引入。通过根据目标属性选择这种配对，学习到的流场在应用任何推理时引导之前就继承了一种属性对齐的组织（Peebles and Xie, 2023 (https://arxiv.org/html/2607.08781#bib.bib26); Zeng et al., 2025 (https://arxiv.org/html/2607.08781#bib.bib43)）。

分子是一个自然的试验场。它们具有重要的标量目标——logP（Wildman and Crippen, 1999 (https://arxiv.org/html/2607.08781#bib.bib38)）和QED（Bickerton et al., 2012 (https://arxiv.org/html/2607.08781#bib.bib3)）——离散的变长样本对朴素耦合构成挑战，并且属性无需预言机即可廉价验证。我们研究分子领域中的耦合作为对齐接口，并探究仅凭选择耦合是否足以实现可控生成，且仅在训练时使用属性标签来构建耦合。我们将该接口实例化为*奖励传输*（图1 (https://arxiv.org/html/2607.08781#S3.F1)）：一种属性对齐的单调耦合，它按标量坐标 \(s\) 对噪声向量排序，按目标属性 \(y\) 对分子排序，然后逐秩配对（Villani, 2003 (https://arxiv.org/html/2607.08781#bib.bib36)）。理论关键在于：在这种一维重排下，在推理时选择 \(s\)——在耦合保持极限下——会诱导出与交叉熵方法（Rubinstein, 1999 (https://arxiv.org/html/2607.08781#bib.bib29)）一步选择相同的截断分布，无需预言机、梯度和拒绝（命题1 (https://arxiv.org/html/2607.08781#Thmproposition1)）。单个标量旋钮，在生成前设置一次，就成为一个分布性陈述。

实验上，在ZINC-250K（Irwin et al., 2012 (https://arxiv.org/html/2607.08781#bib.bib17)）上，该单个旋钮在 \(3.14\) 单位范围内单调引导 logP（+137\% , \(\rho(s, \bar{y}) = 1.000\)），有效性为 100%，且推理开销为零；同一个在QED上训练的旋钮，以相同单调方式移动QED（表1 (https://arxiv.org/html/2607.08781#S5.T1)），并且该效果在GuacaMol（Brown et al., 2019 (https://arxiv.org/html/2607.08781#bib.bib5)）上得以复现。更重要的是，该旋钮并不编码通用的结构偏差。同一个坐标*增大* logP 分子（从 12 到 23 个重原子），而*缩小* QED 分子（从 23 到 16 个）。因此，\(s\) 不仅仅是一个尺寸坐标。它代表耦合与噪声空间对齐的任意属性：耦合写入了任务结构，流则对其进行传输。

我们在本工作中的贡献总结如下：
- **耦合作为对齐接口。** 我们将流匹配中的噪声-数据耦合重新定义为一种将任务结构注入学习到的流场的机制，而非训练时的启发式方法。在耦合保持极限下，属性对齐的单调耦合会诱导出与一步交叉熵方法选择相同的截断分布（命题1 (https://arxiv.org/html/2607.08781#Thmproposition1)），从而在推理时通过单个噪声标量实现分布级控制，且无附加推理成本。
- **单旋钮分子引导。** 我们在 ZINC-250K（Irwin et al., 2012 (https://arxiv.org/html/2607.08781#bib.bib17)）和 GuacaMol（Brown et al., 2019 (https://arxiv.org/html/2607.08781#bib.bib5)）上展示了 logP 和 QED 的分布级引导（表1 (https://arxiv.org/html/2607.08781#S5.T1)）。关键见解是，相同的旋钮对不同的目标产生相反的结构响应，排除了通用尺寸偏差的可能性，表明耦合将目标特定的结构写入了流中。
- **用于SELFIES生成的奖励传输。** 我们为变长SELFIES生成实现了奖励传输，其中连续流生成离散的分子令牌序列。我们识别并解决了两个失败模式：Pre-LayerNorm 下的幅度擦除（Xiong et al., 2020 (https://arxiv.org/html/2607.08781#bib.bib39)）以及填充MSE下的长度膨胀捷径。这可以推广到生成离散变长序列的其他流模型。

## 2 相关工作

#### 流匹配中的耦合。
流匹配（Lipman et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib23); Liu et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib24); Albergo and Vanden-Eijnden, 2023 (https://arxiv.org/html/2607.08781#bib.bib1)）学习一个从噪声到数据的ODE定义的传输，属于更广泛的扩散谱系（Sohl-Dickstein et al., 2015 (https://arxiv.org/html/2607.08781#bib.bib32); Ho et al., 2020 (https://arxiv.org/html/2607.08781#bib.bib14); Song and Ermon, 2019 (https://arxiv.org/html/2607.08781#bib.bib33); Karras et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib20); Chen et al., 2018 (https://arxiv.org/html/2607.08781#bib.bib6)）。该传输依赖于*耦合*：训练时使用的噪声和数据的联合分布。标准流匹配独立抽取噪声和数据。OT-CFM（Tong et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib34)）和多样本流匹配（Pooladian et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib28)）用一个小批量最优传输耦合取代独立耦合，使用噪声-数据对上的欧氏代价来平直轨迹并减少训练方差。Schrödinger桥视角（De Bortoli et al., 2021 (https://arxiv.org/html/2607.08781#bib.bib8)）将该族置于熵最优传输（Cuturi, 2013 (https://arxiv.org/html/2607.08781#bib.bib7); Peyré and Cuturi, 2019 (https://arxiv.org/html/2607.08781#bib.bib27); Villani, 2003 (https://arxiv.org/html/2607.08781#bib.bib36)）之中。在所有上述工作中，耦合是为优化而调整的；OT代价是一个距离。奖励传输使用相同的耦合机制，但采用属性值代价：噪声和数据按标量属性排序并逐秩配对。该耦合仍然解决一个一维最优传输问题，但现在的解编码的是可控属性结构，而非更短的轨迹。

#### 分子生成中的属性控制。
分子生成器使用多种表示——自回归SMILES（Weininger, 1988 (https://arxiv.org/html/2607.08781#bib.bib37); Segler et al., 2018 (https://arxiv.org/html/2607.08781#bib.bib30)）、接合树VAE（Jin et al., 2018 (https://arxiv.org/html/2607.08781#bib.bib18)）、图归一化流（Shi et al., 2020 (https://arxiv.org/html/2607.08781#bib.bib31); Zang and Wang, 2020 (https://arxiv.org/html/2607.08781#bib.bib42)）、图扩散（Vignac et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib35)）、3D等变扩散（Hoogeboom et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib15); Jing et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib19)）——在GuacaMol（Brown et al., 2019 (https://arxiv.org/html/2607.08781#bib.bib5)）上基准测试。我们使用SELFIES（Krenn et al., 2020 (https://arxiv.org/html/2607.08781#bib.bib21)）令牌来保证语法有效性。控制生成属性的方法在属性信息如何到达采样器方面有所不同。强化学习针对属性预言机对生成器进行微调（Blaschke et al., 2020 (https://arxiv.org/html/2607.08781#bib.bib4); You et al., 2018 (https://arxiv.org/html/2607.08781#bib.bib41); Zhou et al., 2019 (https://arxiv.org/html/2607.08781#bib.bib44); Bengio et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib2)）；预言机在训练期间以及通常在微调时采样期间被查询。引导方法在每个采样步骤评估一个属性预测器，并使用其梯度来引导轨迹（Ho and Salimans, 2022 (https://arxiv.org/html/2607.08781#bib.bib13); Lee et al., 2023 (https://arxiv.org/html/2607.08781#bib.bib22)）。条件流匹配在训练和推理时将属性值作为网络输入（Zeng et al., 2025 (https://arxiv.org/html/2607.08781#bib.bib43); Hou et al., 2024 (https://arxiv.org/html/2607.08781#bib.bib16); Deng et al., 2026 (https://arxiv.org/html/2607.08781#bib.bib9)）。这三种方法都需要在采样时提供属性信息——预言机、梯度或目标值。奖励传输则不需要：属性仅用于构建训练耦合，采样由一个标量噪声坐标驱动。关于分子图上贝叶斯流的并行工作（Xiong et al., 2025 (https://arxiv.org/html/2607.08781#bib.bib40)）也修改了耦合，使用一个拟Wasserstein几何代价与采样时条件相结合，这与我们的设计互补。

## 3 背景

参见图注  
图1：奖励传输概述。(A) 标准流匹配独立配对噪声和数据，在噪声空间中不产生属性结构。(B) 属性对齐的OT耦合按范数对噪声向量排序，按目标属性对分子排序，然后逐秩配对，在噪声坐标和分子属性之间创建单调映射。(C) 在推理时，从噪声分布的上尾采样（\(s \geq \tau\)）会得到一个截断的属性分布，等价于交叉熵方法的一步（命题1 (https://arxiv.org/html/2607.08781#Thmproposition1)），将生成的分子转向更高的属性值——无需任何预言机、奖励模型或额外计算。

#### 流匹配。
流匹配（Lipman et al., 2022 (https://arxiv.org/html/2607.08781#bib.bib23)）学习一个速度场 \(\mathbf{v}_\theta\)，它将先验分布 \(p_0 = \mathcal{N}(\mathbf{0}, \mathbf{I})\) 变换为数据分布 \(p_1\)。给定一个从耦合 \(\pi(\mathbf{z}, \mathbf{x})\) 中抽取的噪声-数据对 \((\mathbf{z}, \mathbf{x})\)，线性插值路径为
\[
\mathbf{x}_t = (1 - t) \, \mathbf{z} + t \, \mathbf{x}, \quad t \in [0, 1].
\tag{1}
\]
网络 \(\mathbf{v}_\theta\) 被训练来回归条件速度 \(\mathbf{u} = \mathbf{x} - \mathbf{z}\)：
\[
\mathcal{L}_{\mathrm{FM}} = \mathbb{E}_{t, \mathbf{z}, \mathbf{x}} \bigl\lVert \mathbf{v}_\theta(\mathbf{x}_t, t) - (\mathbf{x} - \mathbf{z}) \bigr\rVert^2.
\tag{2}
\]
样本通过积分 \(\mathrm{d}\mathbf{x}_t / \mathrm{d}t = \mathbf{v}_\theta(\mathbf{x}_t, t)\) 从 \(t = 0\) 到 \(t = 1\) 生成；我们将结果样本记为 \(\mathbf{x}_\theta(\mathbf{z})\)。耦合 \(\pi\) 决定了哪个噪声向量流向哪个数据点；标准FM独立抽取配对（\(\pi = p_0 \otimes p_1\)）。

## 4 奖励传输

### 4.1 属性对齐的最优传输耦合

我们将每个分子表示为一个由 \(L\) 个令牌嵌入组成的序列，位于 \(\mathbb{R}^d\) 中，因此噪声和数据都位于 \(\mathbb{R}^{L \times d}\) 中。耦合 \(\pi(\mathbf{z}, \mathbf{x})\) 将每个噪声序列与一个数据序列配对；我们用一种*单调*耦合替换标准的随机配对，该耦合将噪声的标量摘要与目标分子属性 \(y(\mathbf{x}) \in \mathbb{R}\) 对齐。设 \(\bar{\mathbf{z}} = \frac{1}{L} \sum_{i=1}^L \mathbf{z}_i \in \mathbb{R}^d\) 表示均值池化的噪声，并定义排序键\(^1\) [替代标量键（\(\|\cdot\|_\infty\)、随机线性投影或单个坐标）产生可比较的结果；见附录O (https://arxiv.org/html/2607.08781#A15)。]
\[
s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2.
\tag{3}
\]
给定一个训练集 \(\{(\mathbf{x}_n, y_n)\}_{n=1}^N\) 和一个新鲜采样的噪声 \(\{\mathbf{z}_n\}_{n=1}^N\)，我们分别按各自的键对两个序列进行排序，然后逐秩配对：
\[
\pi^* \colon \mathbf{z}_{\sigma(k)} \; \leftrightarrow \; \mathbf{x}_{\tau(k)}, \qquad \sigma = \operatorname{argsort}(s(\mathbf{z}_{1:N})), \quad \tau = \operatorname{argsort}(y(\mathbf{x}_{1:N})).
\tag{4}
\]
由于 \(s\) 和 \(y\) 是一维的，这是任何凸代价下的唯一最优传输方案——经典的单调重排（Villani, 2003 (https://arxiv.org/html/2607.08781#bib.bib36)）。我们现在形式化这个耦合在推理时带来的好处。

###### 命题1 (属性对齐的OT作为隐式CEM)。