Newton Matching 在生成建模中的应用：微调与采样的统一框架

# 生成建模的牛顿匹配：微调与采样的统一框架  
来源：https://arxiv.org/html/2609.05727  
Zeyang Li [email protected] | 机构：Yunan Wang [email protected] | 机构：Paolo Giaretta [email protected] | 机构：Navid Azizan [email protected] | 机构：麻省理工学院  

###### 摘要  
扩散模型与流模型取得了显著成功，这很大程度上得益于条件匹配方法——当目标分布的样本可用时，它将分布学习转化为可扩展的监督回归问题。然而，许多重要任务（如微调和采样）属于相反的范畴：目标概率密度仅被确定到归一化常数，且无法直接获取样本。具体而言，我们考虑形式为 \(\pi(x) \propto \mu(x) e^{\tau r(x)}\) 的目标密度，其中 \(\mu\) 表示参考因子，\(r\) 表示奖励函数，\(\tau\) 控制对奖励的强调程度。在微调任务中，\(\mu = \rho^{\text{base}}\) 表示预训练模型的归一化终端密度；在采样任务中，\(\mu \equiv 1\)。早期方法将这些任务视为生成动力学上的强化学习，将终端密度问题提升到轨迹级优化，而很大程度上未利用条件匹配结构。近期的工作则试图保持监督回归的形式。基于重加权的方法保留了标准的匹配目标，但当提议分布无法充分覆盖目标密度时，可能会变得不稳定。另一类新兴方法从当前模型中采样，并使用基于奖励的反馈构建下一个回归更新。然而，这些方法源于不同的原理，且将模型表示、采样构建和反馈信号交织在一起。因此，它们的回归目标很少揭示底层的群体级更新，这掩盖了哪些差异是根本性的、哪些反映了不同的精确实现、哪些是由改变更新的近似方法产生的。更根本地，单个阶段如何转换终端密度、是否改进目标函数以及由此产生的迭代是否收敛，在很大程度上仍未被表征。在本文中，我们开发了**牛顿匹配（Newton Matching）**，这是一个用于微调和采样的统一框架。牛顿匹配基于两个视角的范式转变。首先，我们不将方法视为孤立的回归损失，而是将学习视为一个迭代优化过程：当前模型生成样本，每个阶段根据这些样本更新其终端密度，而由此产生的密度决定了下一个模型迭代。其次，我们利用标准匹配的条件期望结构来定义优化域。尽管一个终端密度可以通过无限多种生成动力学实现，但标准条件匹配选择了一个特殊的表示——如果能够从该密度中采样，则可以恢复的群体最小化器——我们称之为其规范模型。规范模型保留了使扩散和流训练可扩展的回归结构。在兼容的平滑实现假设下，它们的速度表示构成一个流形，与终端密度一一对应。在此共享域上，微调和采样成为迭代优化问题，我们的框架解耦了现有方法常常混淆的三个问题：哪个分布方向应主导迭代、沿着该方向的有限步长更新如何转换终端密度、以及这种转换如何通过匹配式的回归实现。为了指定一个有原则的方向，牛顿匹配选择反向Kullback-Leibler（KL）目标，为终端密度流形配备Fisher-Rao度量和混合连接，并将这两种几何结构传输到规范流形。值得注意的是，混合连接下的反向KL Hessian与Fisher-Rao度量一致，因此牛顿方向恰好是负的Fisher-Rao梯度。该方向由正则化奖励生成，该奖励结合了任务奖励和规定目标所需的密度比校正，并且允许等效的协方差和梯度表示，两者都可以通过样本级匹配学习，而无需重要性采样或通过完整生成轨迹的反向传播。牛顿匹配阶段首先执行阻尼或完整切向步长，然后将这个中间结果规范化，用与该终端密度相关的规范模型替换它。这种组合定义了一个规范回缩，并提供了所诱导的密度更新的精确有限步长刻画，从中我们建立了在全部允许步长范围内严格的反向KL下降、温和条件下的全局收敛，以及全步迭代的局部二次收敛。然后，我们将理想的切向更新转化为可扩展的算法。牛顿方向的协方差和梯度表示可以分别与正向或反向回归对构建相结合，产生具有相同群体最小化器的多种样本级回归目标。为了扩大计算设计空间，我们进一步开发了近似牛顿匹配，允许以各种有原则的方式近似切向更新。我们引入了临界点一致性，以区分那些当当前终端密度等于规定目标时切向更新恰好为零的近似方法。最后，这种模块化分解将现有的代表性方法恢复为精确实现、临界点一致近似或改变底层目标的变体，从而揭示了源于看似不同原理的方法的神秘性，并为新算法提供了有原则的设计空间。  
††首页：1††脚注文本：通讯作者为Zeyang Li和Yunan Wang。  

###### 关键词  
生成模型，扩散，流，强化学习，微调，采样  

###### 目录  
1. [引言](https://arxiv.org/html/2609.05727#S1)  
2. [预备知识](https://arxiv.org/html/2609.05727#S2)  
   1. [扩散与流](https://arxiv.org/html/2609.05727#S2.SS1)  
   2. [巴拿赫流形上的微分几何](https://arxiv.org/html/2609.05727#S2.SS2)  
3. [问题表述](https://arxiv.org/html/2609.05727#S3)  
   1. [采样与微调](https://arxiv.org/html/2609.05727#S3.SS1)  
   2. [速度场上的优化](https://arxiv.org/html/2609.05727#S3.SS2)  
4. [规范几何](https://arxiv.org/html/2609.05727#S4)  
   1. [整体几何](https://arxiv.org/html/2609.05727#S4.SS1)  
   2. [一阶几何](https://arxiv.org/html/2609.05727#S4.SS2)  
   3. [二阶几何](https://arxiv.org/html/2609.05727#S4.SS3)  
5. [规范回缩与值上升](https://arxiv.org/html/2609.05727#S5)  
   1. [速度与密度坐标中的规范回缩](https://arxiv.org/html/2609.05727#S5.SS1)  
   2. [值上升证书](https://arxiv.org/html/2609.05727#S5.SS2)  
6. [规范流形上的牛顿匹配](https://arxiv.org/html/2609.05727#S6)  
   1. [Fisher-Rao梯度](https://arxiv.org/html/2609.05727#S6.SS1)  
   2. [牛顿方向](https://arxiv.org/html/2609.05727#S6.SS2)  
   3. [带规范回缩的牛顿法](https://arxiv.org/html/2609.05727#S6.SS3)  
7. [牛顿匹配的收敛性保证](https://arxiv.org/html/2609.05727#S7)  
   1. [KL下降与全局收敛](https://arxiv.org/html/2609.05727#S7.SS1)  
   2. [局部二次收敛与延续法](https://arxiv.org/html/2609.05727#S7.SS2)  
8. [精确牛顿匹配](https://arxiv.org/html/2609.05727#S8)  
   1. [切向表示](https://arxiv.org/html/2609.05727#S8.SS1)  
   2. [正向与反向构造](https://arxiv.org/html/2609.05727#S8.SS2)  
   3. [协方差形式](https://arxiv.org/html/2609.05727#S8.SS3)  
   4. [梯度形式](https://arxiv.org/html/2609.05727#S8.SS4)  
   5. [通过控制变量统一](https://arxiv.org/html/2609.05727#S8.SS5)  
   6. [来自共享样本的多时间监督](https://arxiv.org/html/2609.05727#S8.SS6)  
   7. [规范化](https://arxiv.org/html/2609.05727#S8.SS7)  
9. [近似牛顿匹配](https://arxiv.org/html/2609.05727#S9)  
   1. [具有临界点一致性的近似协方差形式](https://arxiv.org/html/2609.05727#S9.SS1)  
   2. [具有临界点一致性的近似梯度形式](https://arxiv.org/html/2609.05727#S9.SS2)  
   3. [正则化权衡](https://arxiv.org/html/2609.05727#S9.SS3)  
   4. [通过零位移的固定点条件](https://arxiv.org/html/2609.05727#S9.SS4)  
10. [揭示现有工作](https://arxiv.org/html/2609.05727#S10)  
    1. [倾斜匹配](https://arxiv.org/html/2609.05727#S10.SS1)  
    2. [DiffusionNFT](https://arxiv.org/html/2609.05727#S10.SS2)  
    3. [Reinforce 伴随匹配](https://arxiv.org/html/2609.05727#S10.SS3)  
    4. [伴随匹配](https://arxiv.org/html/2609.05727#S10.SS4)  
    5. [流采样](https://arxiv.org/html/2609.05727#S10.SS5)  
    6. [伴随采样](https://arxiv.org/html/2609.05727#S10.SS6)  
11. [结论](https://arxiv.org/html/2609.05727#S11)  
12. 附录A [几何基础](https://arxiv.org/html/2609.05727#A1)  
    1. [环境类](https://arxiv.org/html/2609.05727#A1.SS1)  
    2. [集合论定义](https://arxiv.org/html/2609.05727#A1.SS2)  
    3. [平滑几何实现](https://arxiv.org/html/2609.05727#A1.SS3)  
    4. [规范流形 \(\mathscr{V}_{\mathrm{can}}\) 的非仿射性](https://arxiv.org/html/2609.05727#A1.SS4)  
13. 附录B [牛顿匹配的收敛性证明](https://arxiv.org/html/2609.05727#A2)  
    1. [全局收敛证明](https://arxiv.org/html/2609.05727#A2.SS1)  
    2. [局部收敛证明](https://arxiv.org/html/2609.05727#A2.SS2)  
14. 附录C [针对各向同性高斯族的牛顿匹配](https://arxiv.org/html/2609.05727#A3)  
    1. [高斯闭合](https://arxiv.org/html/2609.05727#A3.SS1)  
    2. [收敛性保证](https://arxiv.org/html/2609.05727#A3.SS2)  
15. 附录D [随机基础与计算工具](https://arxiv.org/html/2609.05727#A4)  
    1. [常见SDE设置与基本恒等式](https://arxiv.org/html/2609.05727#A4.SS1)  
    2. [后验值的后向方程](https://arxiv.org/html/2609.05727#A4.SS2)  
    3. [后验保持SDE](https://arxiv.org/html/2609.05727#A4.SS3)  
    4. [通用桥接](https://arxiv.org/html/2609.05727#A4.SS4)  
    5. [对数密度与对数密度比](https://arxiv.org/html/2609.05727#A4.SS5)  
    6. [初始状态敏感性与伴随演算](https://arxiv.org/html/2609.05727#A4.SS6)  
16. 附录E [后验Stein核](https://arxiv.org/html/2609.05727#A5)  
    1. [Langevin Stein恒等式](https://arxiv.org/html/2609.05727#A5.SS1)  
    2. [Stein控制变量](https://arxiv.org/html/2609.05727#A5.SS2)  
    3. [后验Stein核的协方差-梯度恒等式](https://arxiv.org/html/2609.05727#A5.SS3)  
    4. [后验敏感性核](https://arxiv.org/html/2609.05727#A5.SS4)  
    5. [高斯Stein核与后验矩](https://arxiv.org/html/2609.05727#A5.SS5)  
17. 附录F [近似牛顿匹配的证明](https://arxiv.org/html/2609.05727#A6)  
    1. [具有临界点一致性的近似协方差形式](https://arxiv.org/html/2609.05727#A6.SS1)  
    2. [具有临界点一致性的近似梯度形式](https://arxiv.org/html/2609.05727#A6.SS2)  
    3. [正则化权衡](https://arxiv.org/html/2609.05727#A6.SS3)  
18. 附录G [具有时间相关步长的规范更新](https://arxiv.org/html/2609.05727#A7)  
    1. [终端密度刻画](https://arxiv.org/html/2609.05727#A7.SS1)  
    2. [单调性保证](https://arxiv.org/html/2609.05727#A7.SS2)  
19. 附录H [扩展到替代坐标与插值器](https://arxiv.org/html/2609.05727#A8)  
    1. [得分与漂移坐标下的牛顿匹配](https://arxiv.org/html/2609.05727#A8.SS1)  
    2. [单边插值器下的牛顿匹配](https://arxiv.org/html/2609.05727#A8.SS2)  
20. 附录I [现有工作的详细推导](https://arxiv.org/html/2609.05727#A9)  
    1. [倾斜匹配](https://arxiv.org/html/2609.05727#A9.SS1)  
    2. [DiffusionNFT](https://arxiv.org/html/2609.05727#A9.SS2)  
    3. [Reinforce 伴随匹配](https://arxiv.org/html/2609.05727#A9.SS3)  
    4. [伴随匹配](https://arxiv.org/html/2609.05727#A9.SS4)  
    5. [流采样](https://arxiv.org/html/2609.05727#A9.SS5)  
    6. [伴随采样](https://arxiv.org/html/2609.05727#A9.SS6)  
21. [参考文献](https://arxiv.org/html/2609.05727#bib)  

## 1 引言  
扩散模型\[42 ([https://arxiv.org/html/2609.05727#bib.bib2](https://arxiv.org/html/2609.05727#bib.bib2)), 24 ([https://arxiv.org/html/2609.05727#bib.bib3](https://arxiv.org/html/2609.05727#bib.bib3))\]和流模型\[33 ([https://arxiv.org/html/2609.05727#bib.bib1](https://arxiv.org/html/2609.05727#bib.bib1))\]在图像和视频生成方面取得了显著成功\[41 ([https://arxiv.org/html/2609.05727#bib.bib4](https://arxiv.org/html/2609.05727#bib.bib4)), 25 ([https://arxiv.org/html/2609.05727#bib.bib5](https://arxiv.org/html/2609.05727#bib.bib5)), 16 ([https://arxiv.org/html/2609.05727#bib.bib6](https://arxiv.org/html/2609.05727#bib.bib6))\]，其应用正迅速扩展到更广泛的领域\[9 ([https://arxiv.org/html/2609.05727#bib.bib7](https://arxiv.org/html/2609.05727#bib.bib7)), 49 ([https://arxiv.org/html/2609.05727#bib.bib11](https://arxiv.org/html/2609.05727#bib.bib11)), 54 ([https://arxiv.org/html/2609.05727#bib.bib13](https://arxiv.org/html/2609.05727#bib.bib13))\]。这些模型的核心是学习一个时间依赖的生成过程，该过程将一个简单的先验分布（通常是高斯噪声）逐渐传输到一个复杂的目标分布中。然而，它们的成功不仅源于这些生成动力学的表达能力，还源于其学习方式异常简单且可扩展。当目标分布的样本可用时，两类模型都共享一种常见的条件匹配结构：中间状态是通过一个规定的加噪过程分析性地构建的，并且模型是通过针对易于处理的样本级目标进行监督回归来训练的。在群体层面，最优回归器正是驱动期望概率传输的条件期望。这种条件期望结构将学习复杂分布的难题转化为可扩展的监督回归问题。然而，许多重要任务恰恰违反了这种训练方法得以可能的假设：无法获取来自期望目标密度的样本。在基于奖励的微调中，从一个预训练的生成模型开始，寻求将概率质量向高奖励输出转移。在从未归一化密度中采样时，目标由一个能量函数指定。尽管通常分开研究，但这两个问题都可以通过共享目标 \(\pi(x) \propto \mu(x) e^{\tau r(x)}\) 来表达，其中 \(\mu\) 表示参考因子，\(r\) 是奖励函数，逆温度 \(\tau > 0\) 控制倾斜的强度。对于微调，\(\mu = \rho^{\text{base}}\) 是预训练模型诱导的终端密度；对于采样，\(\mu \equiv 1\)。这两种表述涵盖了广泛的应用范围。对于微调表述，即预训练模型的终端密度提供参考因子，代表性例子包括文本到图像模型的奖励对齐\[18 ([https://arxiv.org/html/2609.05727#bib.bib14](https://arxiv.org/html/2609.05727#bib.bib14)), 7 ([https://arxiv.org/html/2609.05727#bib.bib15](https://arxiv.org/html/2609.05727#bib.bib15))\]和离线强化学习\[48 ([https://arxiv.org/html/2609.05727#bib.bib9](https://arxiv.org/html/2609.05727#bib.bib9)), 47 ([https://arxiv.org/html/2609.05727#bib.bib10](https://arxiv.org/html/2609.05727#bib.bib10))\]。对于采样表述，其中目标是...