Improved techniques for fine-tuning flow models via adjoint matching: a deterministic control pipeline¶
会议: NeurIPS 2026
arXiv: 2605.06583
OpenReview: https://openreview.net/forum?id=9ye0wbdavj
论文: https://arxiv.org/abs/2605.06583
领域: 其他
关键词: 流匹配 (Flow Matching), 伴随匹配 (Adjoint Matching), 确定性控制 (Deterministic Optimal Control), 人类偏好对齐, 截断反向传播 (Truncated Adjoint)
一句话总结¶
针对流生成模型后训练偏好对齐中全轨迹伴随方程计算沉重且容易诱发奖励黑客与模式坍塌的问题，提出基于确定性速度场控制的伴随匹配管线，通过理论证明控制增益集中于去噪终盘并引入时间截断与广义高阶多项式正则化，在 FLUX.2-Klein-4B 与 SiT-XL/2 上取得超过 10 倍训练加速并显著保持样本多样性。
研究背景与动机¶
流匹配模型（Flow Matching）与连续时间修正流（Rectified Flow）因其推断轨迹平直、采样高效且扩展性优异，已成为 Stable Diffusion 3、FLUX.2 以及 WAN 等前沿大规模图像和视频生成系统的核心骨干。然而，仅在大规模无监督数据集上预训练的基座模型往往无法稳定契合人类主观审美偏好或精确遵循复杂的长文本指令。为了使模型输出契合人类期望，借鉴大语言模型 RLHF 的后训练对齐技术被广泛引入扩散与流模型中。其中，基于最优控制理论的后训练对齐因其原则性强、梯度估计方差低而备受关注；近期提出的伴随匹配（Adjoint Matching）方法通过逆时间求解伴随微分方程，将控制场对齐至奖励诱导的目标值梯度，提供了一种无需重要性加权即可直接回归的对齐范式。
然而，现存的伴随匹配方法绝大多数建立在无记忆随机最优控制（Stochastic Optimal Control）和标准 KL 散度正则化的框架下。这种设计在实际部署中暴露出了严重的理论与工程瓶颈：首先，当前顶尖的流模型本质上是以确定性常微分方程（ODE）为首选采样流形，强制引入随机朗之万扩散不仅破坏了高效直线的流动力学，而且在去噪初始的高噪声阶段注入了极大扰动，使前期梯度方差剧增、优化极其缓慢；其次，沿着整个扩散轨迹全步数逆向积分伴随 ODE 计算成本极度昂贵，在多达数十步的基座大模型（如 4B 参数量的 FLUX.2-Klein）上单步迭代耗时长达数百秒，严重限制了工业级规模的扩展；最后，传统的二次 L2 控制代价对应于 KL 正则化，其理论目标是将输出分布偏转至指数倾斜分布（exponentially tilted distribution），但在实际有限步对齐中极易过拟合于极端奖励峰值，导致生成样本色彩过饱和、多样性急剧萎缩甚至出现严重模式坍塌。
面对上述困境，本文从确定性动力学视角重新审视流匹配对齐任务。理论分析表明，在确定性去噪轨迹中，奖励信号对控制梯度的决定性影响高度集中在靠近干净数据的尾部时间步，且这一峰值特性由高信噪比阶段的动力学几何决定。核心 idea：将流模型偏好对齐重构为速度场上的确定性最优控制问题，推导无随机扰动的精简伴随匹配目标，并结合尾部时间截断与超越 KL 的高阶多项式控制正则化，在计算量缩减一个数量级的同时兼顾对齐质量与生成分布的多样性保持。
方法详解¶
整体框架¶
所提框架将预训练流模型的参数微调建模为一个在速度场上施加控制扰动的确定性最优控制系统。给定预训练速度场 \(v^{\mathrm{base}}(x, t)\)，引入可学习参数化控制量 \(u_\theta(x, t) = v_\theta(x, t) - v^{\mathrm{base}}(x, t)\)，驱动状态变量 \(X_t\) 遵循确定性控制受控常微分方程演化。为了将末端状态 \(X_1\) 对齐至最大化奖励 \(r(X_1)\)，系统在目标泛函中引入对控制场幅值的惩罚项 \(f(\|u\|)\) 与末端奖励损失 \(g(X_1) = -r(X_1)\)。
在求解控制梯度时，框架利用庞特里亚金极大值原理（PMP）推导出伴随状态。由于一阶最优性条件使控制量与伴随向量代数共线，原本包含复杂控制项雅可比矩阵的伴随微分方程可大幅化简为仅依赖于冻结基座模型速度梯度的“精简伴随方程”（Lean Adjoint ODE）。在此基础上，算法引入截断参数 \(\tau\)，仅在后半程关键时步内逆向积分伴随状态并计算匹配损失，前序时段则完全遵循确定性前向传播。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["初始高斯噪声 X_0 ~ N(0, I)"] --> B["前向确定性 ODE 采样<br/>当前模型 v_θ(X_t, t) 积分生成轨迹"]
B --> C["末端生成样本 X_1 与奖励评估<br/>计算奖励梯度 a_1 = -∇_X1 r(X_1)"]
C --> D["精简伴随逆向积分 (Lean Adjoint)<br/>仅在截断区间 [τ, 1] 逆向回传 ã_t"]
D --> E["高阶多项式最优控制目标构造<br/>计算代数闭式 u*(ã_t) 作为回归目标"]
E --> F["确定性伴随匹配损失优化<br/>更新流模型速度场参数 θ"]关键设计¶
1. 确定性最优控制与精简伴随方程：彻底剥离随机扰动与雅可比计算开销
现有伴随匹配严重依赖反向 SDE 采样与无记忆布朗运动噪声计划，在去噪初始阶段注入巨大随机噪声导致梯度震荡。本文直接基于确定性常微分方程构建控制问题：
$\(\dot{X}_t = v^{\mathrm{base}}(X_t, t) + u(X_t, t), \quad t \in [0, 1]\)$
优化目标为最小化总代价 \(\mathcal{L}(u) = \int_0^1 f(\|u(X_t, t)\|) dt + g(X_1)\)。定义伴随状态 \(a(t) = \nabla_{X_t} J(u; X_t, t)\)。传统连续伴随方程需要对总受控速度场 \(v^{\mathrm{base}} + u\) 计算雅可比矩阵向量积，计算开销极其昂贵。本文利用哈密顿函数在一阶最优点处的一阶条件 \(\nabla_u f(\|u\|) + a = 0\)，将控制项的偏导与正则梯度的乘积相互抵消，推导出仅依赖基座速度场雅可比的精简伴随方程（Lean Adjoint）：
$\(\frac{d}{dt}\tilde{a}_t = - \left[\tilde{a}_t^\top \nabla_{X_t} v^{\mathrm{base}}(X_t, t)\right], \quad \tilde{a}_1 = \nabla_{X_1} g(X_1) = -\nabla_{X_1} r(X_1)\)$
在训练过程中，前向采样得到的轨迹 \(X_t\) 和逆向积分得到的伴随向量 \(\tilde{a}_t\) 均通过 stopgrad 冻结，将复杂的控制优化问题直接转化为了一个针对目标控制值的确定性逐点回归问题，彻底摆脱了在线强化学习的高方差弊病。
2. 终局时间截断机制：聚焦高信噪比末端时步实现数量级提速
全轨迹伴随积分在多步大模型中存在两大严重缺陷：一是计算时间随着轨迹步数严格线性倍增，二是逆向微分方程在长时程积分中容易积累离散误差。论文通过对高斯初值与目标测度的整流场进行解析推导，提出了关键理论命题（Proposition 4.1）：当初始先验方差 \(\sigma > 1\) 时，最优控制强度 \(\|u_t^\star\|\) 在去噪时间轴上具有唯一的极值点 \(t^\star = \frac{\sigma^2}{1 + \sigma^2}\)。随着去噪进程推进至高信噪比阶段，控制强度单调剧烈升高，而在前期的低信噪比纯噪声阶段控制增益几乎趋近于零。
基于这一洞察，作者提出了时间截断伴随匹配策略（Truncated AM）：将有效控制区间限制在靠近数据端的终盘窗口 \([\tau, 1]\)（其中 \(\tau \in [0, 1)\)）。在前向过程 \(t \in [0, \tau)\) 阶段，模型按照基座或当前速度正常推进；而在逆向伴随积分中，伴随 ODE 仅从 \(t=1\) 逆向回传至 \(\tau\) 处即终止。实验证明，即使将伴随计算压缩至最后 1 到 3 步（如 FLUX.2 上 \(n_{\mathrm{truncate}}=1\) 或 \(3\)），模型依然能够捕获 95% 以上的核心奖励梯度方向，同时将每次迭代更新的耗时从 345 秒骤降至 32 秒。
3. 超越 KL 的高阶多项式正则化：展宽有效控制窗口并防御模式坍塌
传统强化学习和随机控制对齐普遍默认采用二次惩罚项 \(f(r) = \frac{1}{2} r^2\)，其对偶于逆向 KL 散度约束。然而，二次正则化在数学上会导致控制响应对伴随向量的极值过度敏感，梯度更新会被极少数过饱和维度主导，从而引发严重的模式坍塌（Mode Collapse）和多样性损失。本文将控制正则化推广至任意连续可微、严格凸且单调递增的函数 \(f: \mathbb{R}_+ \to \mathbb{R}\)。
根据庞特里亚金极大值原理，最优控制 \(u_t^\star\) 与伴随向量具有严格的对偶反演形式： $\(u_t^\star = - \frac{(f')^{-1}(\|a_t\|)}{\|a_t\|} a_t\)$ 作者选取引入尺度加权的高阶多项式正则化族 \(f(x) = \frac{1}{p\lambda} x^p\)（其中 \(p > 1\) 为阶数，\(\lambda > 0\) 为代价权重），求得解析闭式解： $\(u_t^\star = - \lambda^{\frac{1}{p-1}} \|a_t\|^{\frac{2-p}{p-1}} a_t\)$ 当 \(p=2\) 时退化为经典的二次形式；而当 \(p=4\) 或 \(p=6\) 时，相对控制强度剖面 \(R_p(t) = \left(\frac{D(t^\star)}{D(t)}\right)^{\frac{1}{2(p-1)}}\) 在时间维度上显著平展（Proposition 4.2）。高阶正则化使得更广泛的末端时步能够维持接近峰值的有效控制强度，避免了单一时步的过冲激化，从而在强力提升奖励得分的同时大幅提升了生成分布的样本多样性（LPIPS / MS-SSIM）与数据覆盖率（Coverage / Recall）。
损失函数 / 训练策略¶
在每个训练迭代循环中，首先在无梯度模式下利用当前模型从标准高斯噪声前向积分生成 \(m\) 条轨迹： $\(X_{t+h} = X_t + h v_\theta(X_t, t), \quad X_0 \sim \mathcal{N}(0, I)\)$ 随后在末端图像计算奖励梯度 \(\tilde{a}_1 = -\nabla_{X_1} r(X_1)\)，并沿着轨迹逆向执行 \(T = n_{\mathrm{truncate}}\) 步欧拉离散积分： $\(\tilde{a}_{t-h} = \tilde{a}_t + h \tilde{a}_t^\top \nabla_{X_t} v^{\mathrm{base}}(X_t, t)\)$ 固定轨迹状态与逆向伴随状态，参数优化目标通过截断区间的均方误差匹配损失计算： $\(\mathcal{L}(\theta) = \frac{1}{m T} \sum_{i=1}^m \sum_{t \in \{1, 1-h, \ldots, 1-(T-1)h\}} \left\| \big(v_\theta(X_t, t) - v^{\mathrm{base}}(X_t, t)\big) - \left(-\lambda^{\frac{1}{p-1}} \|\tilde{a}_t\|^{\frac{2-p}{p-1}} \tilde{a}_t\right) \right\|^2\)$ 训练过程采用 Adam 优化器，FLUX.2-Klein-4B 采用学习率 \(1 \times 10^{-5}\) 及 Cosine 衰减，批大小 16，奖励缩放因子 \(\lambda = 2 \times 10^5\)；SiT-XL/2 采用学习率 \(6 \times 10^{-5}\)，批大小 8，常数学习率。
实验关键数据¶
主实验¶
实验分别在类别条件生成模型 SiT-XL/2（675M 参数，ImageNet 256×256）与文本到图像流模型 FLUX.2-Klein-4B（4B 参数）上展开，统一使用 HPSv2 奖励模型进行后训练，并在 Aesthetic Score、ImageReward、PickScore、CLIPScore 以及壁钟训练耗时等维度全面评估。
| 模型骨干 | 方法 | 截断步数 \(n_{\mathrm{truncate}}\) | 正则阶数 \(p\) | HPSv2 ↑ | ImageReward ↑ | CLIPScore ↑ | PickScore ↑ | 迭代耗时 (s) ↓ |
|---|---|---|---|---|---|---|---|---|
| SiT-XL/2 | Base Model | – | – | \(0.183 \pm 0.038\) | \(-0.658 \pm 0.915\) | \(0.239 \pm 0.049\) | \(0.190 \pm 0.011\) | – |
| DRaFT-1 | – | – | \(0.296 \pm 0.041\) | \(0.266 \pm 0.909\) | \(0.260 \pm 0.041\) | \(0.200 \pm 0.011\) | 26.7 | |
| ReFL-10 | – | – | \(0.297 \pm 0.034\) | \(0.282 \pm 0.885\) | \(0.265 \pm 0.038\) | \(0.200 \pm 0.011\) | 26.4 | |
| 2nd SDE-AM-Full | 50 (Full) | 2nd | \(0.279 \pm 0.038\) | \(-0.084 \pm 0.989\) | \(0.244 \pm 0.052\) | \(0.197 \pm 0.012\) | 83.8 | |
| 2nd ODE-AM-Full | 50 (Full) | 2nd | \(0.290 \pm 0.037\) | \(0.185 \pm 0.967\) | \(0.243 \pm 0.049\) | \(0.199 \pm 0.012\) | 83.8 | |
| 2nd ODE-AM | 10 | 2nd | \(0.316 \pm 0.040\) | \(0.450 \pm 0.903\) | \(0.266 \pm 0.040\) | \(0.204 \pm 0.011\) | 47.3 | |
| 4th ODE-AM | 10 | 4th | \(0.328 \pm 0.038\) | 0.574 ± 0.834 | \(0.269 \pm 0.040\) | \(0.206 \pm 0.011\) | 47.3 | |
| 6th ODE-AM | 10 | 6th | 0.329 ± 0.039 | \(0.535 \pm 0.862\) | \(0.269 \pm 0.039\) | 0.207 ± 0.011 | 47.7 | |
| FLUX.2-4B | Base Model | – | – | \(0.290 \pm 0.039\) | \(1.074 \pm 0.645\) | – | \(0.221 \pm 0.014\) | – |
| DRaFT-1 | – | – | \(0.458 \pm 0.045\) | \(1.556 \pm 0.455\) | – | \(0.235 \pm 0.017\) | 24.7 | |
| ReFL-5 | – | – | \(0.466 \pm 0.046\) | \(1.473 \pm 0.521\) | – | \(0.234 \pm 0.018\) | 24.1 | |
| 2nd ODE-AM-Full | 20 (Full) | 2nd | \(0.382 \pm 0.044\) | \(1.160 \pm 0.679\) | – | \(0.227 \pm 0.017\) | 345.5 | |
| 2nd ODE-AM-1 | 1 | 2nd | \(0.413 \pm 0.047\) | \(1.470 \pm 0.503\) | – | \(0.235 \pm 0.017\) | 32.2 | |
| 4th ODE-AM-1 | 1 | 4th | \(0.449 \pm 0.043\) | \(1.522 \pm 0.503\) | – | \(0.235 \pm 0.017\) | 32.2 | |
| 6th ODE-AM-1 | 1 | 6th | 0.470 ± 0.046 | \(1.531 \pm 0.514\) | – | \(0.235 \pm 0.018\) | 32.2 | |
| 4th ODE-AM-3 | 3 | 4th | \(0.449 \pm 0.043\) | 1.557 ± 0.475 | – | \(0.235 \pm 0.017\) | 65.3 |
消融实验¶
针对 FLUX.2-Klein-4B 在微调后容易发生的生成模式坍塌现象，深入测试了生成样本的组内多样性（LPIPS、1-MS-SSIM）以及与基座模型真实数据流形契合度的分布保持指标（Coverage、Recall）。
| 模型配置 | 截断步数 \(n_{\mathrm{truncate}}\) | 正则阶数 \(p\) | LPIPS ↑ | 1-MS-SSIM ↑ | Coverage ↑ | Recall ↑ |
|---|---|---|---|---|---|---|
| Base Model | – | – | \(0.505 \pm 0.073\) | \(0.857 \pm 0.113\) | – | – |
| DRaFT-1 | – | – | \(0.401 \pm 0.069\) | \(0.757 \pm 0.102\) | 0.080 | 0.010 |
| ReFL-5 | – | – | \(0.431 \pm 0.070\) | \(0.779 \pm 0.103\) | 0.047 | 0.002 |
| 2nd ODE-AM-1 | 1 | 2nd | \(0.448 \pm 0.078\) | \(0.832 \pm 0.111\) | 0.376 | 0.129 |
| 4th ODE-AM-1 | 1 | 4th | \(0.420 \pm 0.074\) | \(0.801 \pm 0.109\) | 0.204 | 0.037 |
| 6th ODE-AM-1 | 1 | 6th | \(0.424 \pm 0.073\) | \(0.804 \pm 0.102\) | 0.139 | 0.012 |
| 2nd ODE-AM-3 | 3 | 2nd | 0.454 ± 0.074 | 0.843 ± 0.104 | 0.306 | 0.057 |
| 6th ODE-AM-3 | 3 | 6th | \(0.450 \pm 0.073\) | \(0.813 \pm 0.115\) | 0.396 | 0.083 |
| 2nd ODE-AM-Full | 20 (Full) | 2nd | \(0.449 \pm 0.065\) | \(0.817 \pm 0.098\) | 0.262 | 0.106 |
关键发现¶
- 截断步数饱和规律：在 SiT 实验中将截断步数从 1 步逐渐增加至 15 步时，对齐奖励稳步上升，但在 \(n_{\mathrm{truncate}} \ge 10 \sim 15\) 步后增益完全饱和；全轨迹 50 步由于累积截断误差不仅耗时增加一倍以上，最终 HPSv2 仅有 0.290，反而显著落后于 10 步截断（0.316）。在 FLUX.2 上仅截断最后 1 步（32s/iter）即可匹敌甚至超越 20 步全轨迹（345s/iter）。
- 多样性与模式坍塌防御：DRaFT 和 ReFL 等直接一阶反馈方法虽然训练收敛快，但 Coverage 和 Recall 均跌破 0.08 和 0.01，视觉上呈现出刺眼高亮、纹理千篇一律的奖励黑客现象；而 2nd ODE-AM-1 与 6th ODE-AM-3 的 Coverage 达到 0.376 与 0.396，Recall 维持在 0.08~0.13，几乎完整保全了原始底模的语义分布丰富度。
- 高阶多项式正则化的优势：提高多项式阶数至 \(p=4\) 和 \(p=6\) 能将控制强度曲线沿时间轴均匀拉平，使得截断时步内的梯度更平缓稳定，在 SiT 上 ImageReward 显著提升至 0.574（相比 2 阶的 0.450），证明高阶控制正则化能克服传统 KL 散度引起的单一极值敏感问题。
亮点与洞察¶
- 精简伴随消去雅可比交叉项：利用一阶极大值最优性条件，将受控速度场与正则项的导数直接对消，使反向伴随积分仅需评估冷冻底模的向量-雅可比积（VJP），无需计算复杂的可学习网络控制项导数，大幅降低内存开销。
- 高信噪比驱动的计算资源再分配：从概率流积分的解析高斯模型中严格证明了控制强度在去噪终盘存在唯一极值点，以此为基础大刀阔斧截断低信噪比冗余步数，打破了扩散与流模型后训练“必须全轨迹反向传播”的思维定势。
- 超越 KL 的通用控制代价设计：突破了传统生成模型 RLHF 机械套用逆向 KL 散度的理论桎梏，通过任意单调凸函数建立可解析反演的控制惩罚，为生成模型的后训练约束设计开辟了更广阔的数学工具箱。
局限与展望¶
- 截断时步超参仍需手工经验调优：目前截断时钟 \(\tau\) 针对不同模型需要通过网格搜索或经验设定（如 SiT 选 10 步、FLUX 选 1~3 步），尚未实现基于在线瞬时控制模长估计的自适应动态时钟截断。
- 仅在确定性动力学下完成高阶推广：高阶多项式正则化推导目前仅适用于确定性常微分控制系统，在随机最大值原理（SMP）下的 BSDE 求解中第二阶噪声扩散校正项极其复杂，尚未建立完备的随机高阶理论。
- 应用模态集中于二维图像生成：当前实证仅在类别条件与文生图模型中验证，尚未扩展至对时空因果连续性要求更高的高维长视频生成（如 Sora、WAN 架构）或离散流匹配任务。
相关工作与启发¶
- vs. 随机伴随匹配 (Stochastic AM)：Domingo-Enrich 等人提出的原始伴随匹配强依赖无记忆布朗运动噪声，导致前期高噪声下梯度方差失控；本文采用确定性常微分动力学，不仅完美匹配现代流模型的高效确定性采样器，更大幅稳定了反向匹配训练。
- vs. 直接奖励求导方法 (DRaFT / ReFL)：DRaFT 直接对末端奖励通过采样图进行链式法则反向传播，极易受局部极端峰值诱导而出现饱和过曝和模式崩溃；本文基于控制论回归伴随目标，优化稳定且生成多样性与真实分布覆盖度完胜基线。
- vs. 价值梯度指引 (Value Gradient Guidance, VGG)：VGG 通过求解偏微分方程（PDE）来近似价值函数梯度，工程落地门槛高且高维难以精确近似；本文通过庞特里亚金极大值原理转化为 ODE 伴随匹配回归，结构更简洁直观。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立面向流匹配模型的确定性伴随控制体系并推广至超越 KL 的高阶多项式正则。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 675M 参数 SiT 与 4B 参数 FLUX.2，包含详实的消融分析、多样性量化与生成画质评估。
- 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，从玩具高斯动力学到大规模基础模型层层递进，逻辑严丝合缝。
- 价值: ⭐⭐⭐⭐⭐ 实现了伴随匹配算法逾 10 倍的训练加速，极大降低了大规模流模型后训练对齐的显存与计算门槛。