奖励传输用噪声空间对齐实现流匹配属性控制

背景

流匹配（Flow Matching）是一类生成模型，通过学习从噪声分布到数据分布的连续确定性变换来生成样本。其中的关键设计是“耦合”（coupling），即如何将噪声向量与训练数据点配对。以往的研究通常将耦合视为纯粹的计算选择，主要影响训练效率和样本质量，而非生成过程的可控性。与此同时，分子生成领域迫切需要能够根据特定性质（如脂溶性 logP 或药物相似性 QED）定向调控生成结果的机制。现有的方法往往依赖额外的奖励模型、梯度指导或无分类器引导（Classifier-Free Guidance），这些方法要么引入额外计算开销，要么需要预先训练的评分函数。本文提出一种全新的视角：将耦合本身用作对齐接口（alignment interface），通过最优传输（Optimal Transport）在训练时将噪声空间的一个标量坐标与目标分子奖励对齐，从而实现推理时无需任何附加模型即可连续调控生成分布的性质。

核心内容

论文的标题为 Reward Transport: Property Control in Flow Matching via Noise-Space Alignment，发表于 arXiv cs.AI（2026年6月13日提交）。核心思想是：流匹配中的耦合规则（将噪声向量与数据点配对）不应仅看作计算选择，而可以作为“对齐接口”。具体来说，通过在训练时根据目标分子属性匹配噪声和数据，可以直接将可控结构嵌入学习到的流场中。基于此，作者提出 Reward Transport 方法：

• 训练阶段：使用最优传输耦合来建立噪声空间中的一个标量坐标与分子奖励（例如 logP 值）之间的对齐。该标量坐标取自噪声向量（通常是高斯噪声）的某一维度或线性组合，经过最优传输后，该坐标的取值与训练数据对应的奖励值形成单调对应关系。
• 推理阶段：在生成过程中，通过改变这个对齐后的标量坐标（即“噪声空间坐标”），就能连续地引导生成分子的性质分布，而无需调用任何 oracle、奖励模型、梯度指导或额外计算。该坐标本质上成为一个分布级的“控制旋钮”。
• 理论联系：在保持该耦合结构的极限情况下（即耦合完全保留），对该坐标进行阈值化操作可以恢复交叉熵方法（Cross-Entropy Method, CEM）中的截断奖励分布。这为方法提供了理论依据：它原理性地提供了一个连续可调的分布级控制旋钮，而不仅仅是一个样本级修改。
• 实验验证：在标准分子生成基准 ZINC-250K 和 GuacaMol 上，通过扫描该标量坐标，实现了对 logP 的单调控制（随坐标增加 logP 单调增大）以及对 QED 在其操作范围内的一致控制。特别值得注意的是，相同的控制旋钮对不同目标产生相反的结构响应：对于 logP，分子尺寸随奖励增加而增大；对于 QED，分子尺寸却随奖励增加而减小。这排除了该控制单纯来源于通用尺寸偏差的可能性，说明该方法真正捕捉到了不同性质对应的结构特征。
• 兼容性与对比：Reward Transport 与无分类器引导和条件流匹配（Conditional Flow Matching）在原理上是互补的，可以结合使用。同时，作者在 epsilon-prediction 扩散模型（一种常用的扩散模型训练范式）上进行了实验，得到负面结果——在该设置下，耦合级别的对齐在结构上缺失，无法实现类似的控制，这从反面印证了流匹配耦合本身作为对齐接口的独特性。
• 代码可用：论文提供了开源代码（URL 在摘要中）。

关键要点

• 耦合的新角色：将流匹配中的耦合从计算选择提升为一种可学习的对齐接口，直接嵌入属性信息到流场中。
• Reward Transport 方法：利用最优传输在训练时将噪声空间标量坐标与分子奖励对齐，推理时通过调整该坐标即可连续控制生成分布的性质，无需外部模型或额外计算。
• 与 CEM 的理论联系：在保持耦合的极限下，阈值化该坐标等价于交叉熵方法的截断奖励分布，提供了理论根基和可解释性。
• 实验控制能力：在 ZINC-250K 和 GuacaMol 上对 logP 实现单调控制，对 QED 实现范围内一致控制；同一旋钮对不同性质产生相反的结构响应（如分子尺寸），证明控制并非源自尺寸偏置。
• 互补性与负结果：该接口与无分类器引导、条件流匹配互补；在 epsilon-prediction 扩散下该对齐失败，凸显了流匹配耦合结构的必要性和独特性。
• 零额外推理代价：推理时仅需调整一个标量，无需梯度计算、奖励模型评估或重复采样。

意义与影响

Reward Transport 为生成模型中的属性控制提供了一种新颖而简洁的范式。它打破了“耦合仅用于训练效率”的传统认知，揭示了噪声-数据配对方式本身可以编码任务相关的结构信息。与需要额外训练奖励模型或使用梯度指导的方法相比，该方法实现了零额外推理开销的连续控制，且控制粒度是分布级的而非样本级的。这一特性在药物分子设计等数据稀缺或奖励函数难以定义的任务中尤为重要——因为它只需要在训练时用目标属性值指导最优传输配对，推理时便不再需要任何外部监督。

此外，论文通过对比实验（扩散模型上的负结果）清晰界定了该方法适用的模型架构范围：它依赖于流匹配中显式的耦合结构，而非扩散模型中的隐式噪声预测。这为未来设计可控制生成模型提供了明确的设计原则——如果想要实现类似的对齐控制，应优先考虑流匹配框架并精心设计耦合策略。

该工作也开放了若干后续方向：例如将标量坐标扩展为多维控制向量（同时调控多个性质），或者将最优传输耦合与其他生成范式（如 score-based models）结合。总体而言，Reward Transport 提出了一种原理清晰、实现简单且效果显著的属性控制方法，可能对计算化学、材料设计以及更一般的可控生成领域产生重要影响。