技术摘要：奖励传输 —— 通过噪声空间对齐实现流匹配中的属性控制

1. 问题陈述

在流匹配（Flow Matching, FM）及相关的基于扩散的生成模型中，耦合（coupling）——即决定在训练期间将哪个噪声向量与哪个数据点配对的规则——通常被视为一种计算启发式方法。标准方法为了简单起见通常采用独立配对，或使用小批量最优传输（minibatch OT）来使轨迹变直并减少方差。在这些范式中，耦合的优化仅是为了促进训练，而非为了将特定的结构属性编码进学习到的流场中。

因此，实现可控生成（例如，生成具有特定 logP 或 QED 分数的分子）通常需要在推理阶段引入显式机制，例如：

• 无分类器指导（Classifier-free guidance）： 需要将目标属性值作为输入。
• 梯度引导（Gradient guidance）： 在每个采样步骤中计算属性预测器的梯度。
• 强化学习（RL）： 针对预言机（oracle）对生成器进行微调，这通常需要在采样期间查询预言机。

这些方法引入了推理时的开销、对外部奖励模型或预言机的依赖，或者复杂的训练流程。本文探讨的是：耦合本身是否可以作为一个对齐接口（alignment interface），在训练期间将可控结构直接嵌入到流场中，从而使得在推理时无需额外的计算或预言机调用即可实现控制。

2. 方法论：奖励传输（Reward Transport）

作者提出了奖励传输（Reward Transport），该方法将噪声-数据耦合重新定义为一种向流场注入任务结构的机制。

2.1 属性对齐的最优传输耦合

奖励传输并非独立地配对噪声和数据，而是构建了一个基于目标分子属性 $y(x)$（如 logP, QED）的单调耦合（monotone coupling）。

1. 标量键（Scalar Keys）： 对于每个噪声向量 $z$（一系列 token 嵌入），计算一个标量排序键 $s(z)$。主要选择是均值池化后噪声的欧几里得范数：$s(z) = \|\bar{z}\|_2$。
2. 基于秩的配对（Rank-Based Pairing）： 给定一组噪声向量 $\{z_n\}$ 和具有属性 $\{y_n\}$ 的分子 $\{x_n\}$，两组数据分别根据其各自的键（$s$ 和 $y$）进行排序。然后它们按秩进行配对：具有第 $k$ 小 $s$ 的噪声向量与具有第 $k$ 小 $y$ 的分子配对。
3. 训练目标： 流匹配模型 $v_\theta$ 被训练以传输这些对齐的对。损失函数结合了在连续嵌入空间上的标准 MSE（端点预测）以及确保解码后的 token 形成有效分子的交叉熵（CE）损失。

2.2 方向条件化（Direction Conditioning）

为了确保模型在推理过程中能够访问噪声空间坐标 $s$，引入了一个方向 MLP（Direction MLP）。这个两层 MLP 将归一化的标量 $s$ 映射到一个 $d$ 维嵌入，并将其与时间及位置嵌入一起添加到 Transformer 的输入中。这使得模型能够根据用于单调配对的特定噪声坐标来对其速度场进行条件化。

2.3 推理机制

在推理时，用户指定一个标量方向信号 $s^*$（一个噪声空间坐标，而非目标属性值）。模型通过从满足 $s(z) \geq s^*$ 的噪声向量开始积分流场来生成分子。

• 理论保证（命题 1）： 在学习到的流完美保持 $s$ 与 $y$ 之间秩对应关系的极限情况下（耦合保持极限），对 $s \geq s^*$ 进行条件化等同于执行了**交叉熵方法（CEM）**的一次选择步骤。具体而言，它将生成的分布截断为属性分布的前 $\epsilon$ 部分，其中 $\epsilon$ 由阈值 $s^*$ 决定。
• 零开销： 这种控制不需要预言机、不需要奖励模型、不需要梯度计算，也不需要推理时的拒绝采样。

3. 核心贡献

1. 耦合作为对齐接口： 本文将流匹配中的噪声-数据耦合从一种训练启发式方法重新定义为一种分布层级的对齐接口。它证明了通过基于目标属性选择耦合，可以组织学习到的流场，使得噪声空间坐标与属性值相对应。
2. 单旋钮分子转向（Single-Knob Molecular Steering）： 该方法利用单个标量噪声坐标实现了对分子属性（logP 和 QED）的单调控制。至关重要的是，同一个旋钮对不同的目标会产生相反的结构响应：增加 $s$ 会为 logP（亲脂性）生成较大的分子，但会为 QED（药物相似性）生成较小的分子。这排除了通用的尺寸偏差，并证实了耦合将目标特定的结构写入了模型中。
3. 针对离散序列的奖励传输： 作者将其应用于变长的 SELFIES token 序列，并识别并解决了流模型在生成离散数据时遇到的两个特定失效模式：
   • 隐式信号擦除（Implicit Signal Erasure）： 在 Pre-LayerNorm 架构中，幅度信息（由 $\|\bar{z}\|$ 编码）会被擦除。方向 MLP 通过注入学习到的方向模式解决了这一问题。
   • 长度混淆（Length Confound）： 如果不对填充（padding）token 使用未掩码（unmasked）的 MSE 损失，模型会学习到一种捷径，即通过增加序列长度来满足范数约束。未掩码 MSE 强制模型回归填充位置的原始噪声，从而锚定了幅度信号并防止了这种捷径。

4. 实验结果

该方法在 ZINC-250K 和 GuacaMol 数据集上进行了评估。

• 控制性能：
  • logP： 在 ZINC-250K 上，将 $s$ 从 -3 扫至 +7，使平均 logP 从 1.50 变为 5.44（+137%），且 $s$ 与平均属性之间的 Spearman 相关系数 $\rho=1.000$。效应量（Cohen's $d$）达到 1.63。
  • QED： 同一旋钮单调地改变了 QED（例如，在 ZINC 上当 $s=5$ 时为 +0.059），且 $\rho=1.000$。
  • 结构特异性： 对于 logP，增加 $s$ 会增加重原子数量（12 $\to$ 23）。对于 QED，增加 $s$ 会减少重原子数量（23 $\to$ 16）。
• 泛化能力： 结果在 Guacaicol (1.6M 分子) 上得到复现，显示出更强的 logP 效应量（$d=2.31$）以及在有效范围内的 QED 一致单调性。
• 基线对比：
  • 与条件流匹配（Conditional Flow Matching）（其中属性值作为显式输入）相比，奖励传输实现了相似的单调性（$\rho=1.000$），但它是通过分布截断而非显式目标命中来实现的。
  • 与 $\epsilon$-预测扩散（$\epsilon$-prediction diffusion） 相比，该方法表明在耦合层级的对齐在 $\epsilon$-预测中是结构性缺失的，因为贝叶斯最优目标与耦合是独立的，从而削弱了梯度信号。
• 训练动态： 秩保持（$s$ 与生成的 $y$ 之间的相关性）在微调期间迅速出现（在一个 epoch 内），并在 1.0 以下趋于平稳（例如 logP 为 0.57），这与命题 1 中描述的“分布近似”差距是一致的。

5. 重要性与主张

本文声称，奖励传输确立了流匹配中的耦合作为一个独特的、分布层级的控制通道，它可以独立携带对齐结构。

• 高效性： 它消除了在推理时需要调用预言机、奖励模型或进行迭代优化的需求，与基于引导或 RL 的方法相比，显著降低了计算成本。
• 机制： 它证明了生成模型的“可转向性”不一定依赖于显式的条件输入；它可以经由训练耦合被植入模型的几何结构中。
• 局限性与范围： 作者对其范围持谨慎态度。该方法提供的是一个 1D 分布旋钮（截断分布），而非精确的逐样本数值分配。它适用于与全局分子特征（如大小或极性）相关的标量属性，但对于像合成可及性（SA）这样复杂的拓扑属性则表现不佳，因为 1D 标量键无法区分具有相似大小的复杂结构。此外，研究表明该方法与 $\epsilon$-预测目标不兼容，强调了该接口取决于网络的特定预测目标。

总之，这项工作将可控生成范式从“在推理时引导轨迹”转向了“对齐训练分布”，为流匹配中的属性控制提供了一种原则性的、零开销的机制。