Efficient Adjoint Matching for Fine-tuning Diffusion Models¶

会议: NeurIPS 2026
arXiv: 2605.11480
代码: https://github.com/swswss/EAM
领域: 图像生成
关键词: 扩散模型微调, 随机最优控制, 伴随匹配, 奖励微调, 线性基础漂移

一句话总结¶

通过将随机最优控制框架的基础漂移重新设计为满足无记忆性的线性形式并推导修正终端代价，消除了传统伴随匹配中高开销的全轨迹前向 SDE 模拟与反向伴随 ODE 模拟，实现收敛速度提升高达 4 倍且理论最优路径完全等价的高效扩散模型奖励微调。

研究背景与动机¶

扩散模型与流匹配模型已成为现代大规模文生图生成任务的核心骨架，但在标准预训练目标下生成的图像往往难以完美契合人类偏好、美学质量以及复杂的 Prompt 描述指令。为了解决这一错位，基于奖励函数的后训练微调逐渐成为主流范式。在各类奖励微调方法中，基于奖励梯度的微调技术能够直接利用可微偏好模型（如 PickScore、HPSv2、ImageReward）对样本输入的局部梯度指引更新方向，相较于需要对每个 Prompt 密集采样 12 到 24 张图像的无梯度强化学习策略，展现出极高的样本利用效率。

在此技术脉络下，伴随匹配（Adjoint Matching, AM）通过将奖励微调建立为严格的随机最优控制（Stochastic Optimal Control, SOC）问题，奠定了具有分布收敛保证的严谨理论基准。然而，AM 在训练迭代效率上面临严重的计算瓶颈。其根源在于：AM 继承了预训练模型的非平凡漂移项作为基础动力学，为了满足无记忆性条件以获得无偏控制估计，必须使用前向随机 SDE 求解器生成完整的长轨迹（通常需要 40 步函数评估 NFE），并在显存中缓存整条轨迹；随后还必须沿着该轨迹执行反向伴随 ODE 积分，每一步都需要反向传播计算穿透预训练网络权重的雅可比-向量积（JVP），导致极高的时间与显存开销。

这种效率瓶颈的核心矛盾在于：现有 SOC 形式强行将预训练生成漂移绑定为参考基准过程，导致前向与反向动力学均陷入对网络昂贵微分与多步随机仿真的依赖。本文的切入角度是打破参考基准漂移与预训练模型强绑定的思维定势——只要终端受控边缘分布与最优路径测度保持等价，基础动力学可以完全重构。核心 idea：将 SOC 的基础漂移重新设计为满足无记忆性与无偏采样的线性形式，并相应推导补偿终端代价，在前向阶段仅需少步数确定性 ODE 采样终点加噪即可构造轨迹，在反向阶段导出解析闭式解彻底消除伴随状态模拟，在保证最优路径测度完全等价的同时实现至多 4 倍的微调加速。

方法详解¶

整体框架¶

Efficient Adjoint Matching（EAM）的整体微调流程由终端样本生成、解析单步中间态加噪、解析终端代价梯度推导与控制匹配损失优化四个阶段组成。算法首先使用高效的确定性 ODE 求解器（仅需 10 NFE）直接生成终端样本 $X_1$，利用基础线性漂移的无记忆性与高斯摄动核在单步内直接合成中间时间步状态 $X_t$；接着通过重构的终端代价函数解析计算目标梯度，最终通过闭式缩放的伴随匹配目标直接优化速度模型的 LoRA 权重，全程完全规避了轨迹缓存与反向神经网络积分。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始噪声 X0 ~ N(0, I)"] --> B["确定性少步 ODE 采样<br/>生成终端样本 X1 (10 NFE)"]
    B --> C["解析单步中间态加噪<br/>从原摄动核采样 Xt = t*X1 + (1-t)*ε"]
    C --> D["线性基础漂移闭式伴随映射<br/>解析缩放因子替代反向 ODE 模拟"]
    B --> E["修正终端代价与得分估计<br/>结合 Tweedie 公式与奖励梯度 ∇r(X1)"]
    D --> F["伴随匹配损失优化<br/>更新速度网络 LoRA 参数"]
    E --> F
    F --> G["优化后的对齐生成模型 v_ft"]

关键设计¶

1. 线性基础漂移与解析闭式伴随状态解：彻底消除反向伴随模拟

在原生伴随匹配中，由于基础漂移包含复杂的预训练神经网络 $v^{\text{pt}}(x, t)$，伴随状态的逆向 ODE $\frac{\mathrm{d}}{\mathrm{d}t}a(t; X_t) = -\nabla_x b(X_t, t)^\top a(t; X_t)$ 在每一步积分时都必须对网络执行 JVP 传播，导致显存消耗暴增且耗时巨大。EAM 证明，若将基础漂移约束为关于状态 $x$ 的纯标量线性形式 $b(x, t) = D(t)x$，则其空间雅可比矩阵直接简化为纯对角矩阵 $\nabla_x b(X_t, t) = D(t)I$。伴随 ODE 由此退化为标量微分方程，其沿轨迹的状态转移矩阵可直接积分写为解析闭式解： $$a(t; X_t) = \exp\left(\int_t^1 D(\tau)\,\mathrm{d}\tau\right) a(1; X_1)$$ 伴随状态在任意时刻 $t$ 仅仅是终端代价梯度 $a(1; X_1) = \nabla g(X_1)$ 乘以一个仅与时间 $t$ 相关的标量系数。这彻底消除了反向伴随 ODE 模拟以及伴随路径中的所有神经网络梯度回传。

2. 保持原摄动核的单参数容许线性漂移族：使能少步 ODE 端点生成与瞬时加噪

为了使 SOC 问题的解具有无偏可学习性，基础动力学必须具备时间点 $t=0$ 与 $t=1$ 的无记忆性。进一步地，为了不改变扩散模型本身建立的加噪几何结构，EAM 强制要求基础动力学的后验加噪核 $p^{\text{base}}_{t|1}(\cdot|X_1)$ 严格等于原始扩散过程的高斯摄动核 $q_t(X_t|X_1) = \mathcal{N}(tX_1, (1-t)^2 I)$。论文通过变分约束证明，在固定的前向扩散扩散系数 $\sigma(t) = \sqrt{\frac{2(1-t)}{t}}$ 条件下，满足上述条件的线性漂移 $D(t)$ 由一个单自由度常数 $C > \frac{1}{2}$ 唯一定义： $$D(t) = \frac{2Ct^2 - 1}{t(2Ct^2 - 2t + 1)}$$ 在此线性漂移下，在最优受控动力学处联合概率分布满足 $p^u(X_t, X_1) = q_t(X_t|X_1)p^u_1(X_1)$。因此，在训练生成前向样本时，模型无需再运行 40 步昂贵的高方差随机 SDE 生成整条轨迹，而是直接使用现成的高效确定性 ODE 求解器（如 10 步 DPM-Solver++）采样终点 $X_1 \sim p_1^{\bar{u}}$，并仅凭一步公式 $X_t = tX_1 + (1-t)\epsilon$ 瞬间构造中间状态，轨迹采样时间减少 73% 且无需暂存整条轨迹。

3. 修正终端代价与重用中间态的高效得分估计：保障目标分布严格等价

由于改变了参考基础漂移，基础过程的终端边缘分布转变为 $p_1^{\text{base}} = \mathcal{N}(0, (2C-1)I)$，不再自然等价于真实数据分布 $p_1^{\text{pt}} = p_{\text{data}}$。为了确保最优受控动力学的终端分布仍然严格收敛到奖励倾斜的目标分布 $p_1^\star(x) \propto e^{\beta r(x)} p_{\text{data}}(x)$，EAM 重新推导了终端代价函数及其梯度： $$\nabla g(X_1) = -\frac{X_1}{2C-1} - \nabla \log p_1^{\text{pt}}(X_1) - \beta \nabla r(X_1)$$ 为了避免在终端 $t=1$ 附近计算预训练数据得分 $\nabla \log p_1^{\text{pt}}(X_1)$ 时引入数值奇异与高方差，EAM 巧妙复用已采样的匹配中间态 $(X_t, t)$ 与 Tweedie 公式进行得分估计： $$\nabla g(X_1) \approx -\frac{1}{2C-1}X_1 + \frac{1}{1-t}X_t - \frac{t}{1-t}v^{\text{pt}}(X_t, t) - \beta \nabla r(X_1)$$ 理论分析（Theorem 3.3）严格证明，尽管 EAM 与 AM 定义在不同的参考过程上且最优控制函数不同，但二者对应的最优随机过程的漂移、扩散系数与路径测度完全重合，EAM 在数学本质上收敛到与 AM 严格同一的最优解。

损失函数 / 训练策略¶

为了使微调从预训练权重平滑启动，EAM 将可学习控制项 $u(x, t)$ 参数化为包含微调速度场 $v^{\text{ft}}(x, t) = v^{\text{pt}}(x, t) + \text{LoRA}_\theta(v^{\text{pt}})$ 的形式： $$u(x, t) = \frac{1}{\sigma(t)}\left(-D(t)x - \frac{1}{t}x + 2v^{\text{ft}}(x, t)\right)$$ 当 LoRA 初始化为 0 时，受控动力学精确退回原预训练生成动力学。EAM 最终优化的回归匹配损失为： $$\mathcal{L}_{\text{EAM}}(\theta) = \mathbb{E}_{X_1 \sim p_1^{\bar{v}^{\text{ft}}}, t \sim \mathcal{U}[0, 1], \epsilon \sim \mathcal{N}(0, I)}\left[\frac{1}{2}\left\|u(X_t, t) + \sigma(t)\frac{(2C-1)t}{2Ct^2-2t+1}\nabla g(X_1)\right\|^2\right]$$ 在训练实践中，为防止极端时间步的数值不稳定，对奖励梯度项引入时间步加权系数 $w(t) = (1-t)^{0.9}/t^{1.5}$，并使用其倒数对最终损失进行加权。超参数常数默认设定为 $C = 0.51$。

实验关键数据¶

主实验¶

实验基于 Stable Diffusion 3.5-Medium（SD3.5-M）在 4 块 NVIDIA A100 GPU 上训练，采用 LoRA（rank=32）微调，训练提示词来自 Pick-a-Pic，并在 DrawBench 基准上使用 10 步 DPM-Solver++ 评估。下表给出了在单奖励（PickScore）及多奖励（PickScore + HPSv2.1 + Aesthetics）设定下的生成质量主对比：

方法	优化奖励设定	PickScore ↑	CLIPScore ↑	ImageReward ↑	HPSv2.1 ↑	Aesthetics ↑
SD3.5-M (基线)	无微调	20.34	0.226	-0.58	0.214	5.29
SD3.5-M + CFG	无微调 (推理引导)	21.82	0.275	0.65	0.265	5.41
AM (Adjoint Matching)	PickScore	23.00	0.259	0.97	0.292	6.25
EAM (本文方法)	PickScore	23.16	0.252	0.95	0.297	6.27
AM (Adjoint Matching)	Pick + HPS + Aesthetic	23.06	0.253	0.83	0.303	6.21
EAM (本文方法)	Pick + HPS + Aesthetic	22.93	0.261	0.98	0.306	6.56

消融实验¶

为了量化 EAM 训练效率与显存优势的来源，论文对前向轨迹生成与反向伴随计算进行了匹配 NFE 与单步开销的系统分解：

评估配置 / 阶段	采样方法与 NFE	训练耗时 (GPU hrs) ↓	单迭代耗时 (s) ↓	峰值显存 (GB) ↓	PickScore ↑
AM 基准	SDE, 40 NFE	64	-	-	23.00
AM 低步数尝试	SDE, 10 NFE	80	-	-	23.03
EAM (本文)	ODE, 10 NFE	16	-	-	23.16
前向轨迹阶段 (AM)	SDE 轨迹模拟	-	16.7	27.4	-
前向轨迹阶段 (EAM)	ODE 端点采样	-	4.5 (-73%)	27.2	-
伴随反向阶段 (AM)	ODE 伴随回传+JVP	-	34.7	73.6	-
伴随反向阶段 (EAM)	闭式解析计算	-	0.4 (-99%)	1.6 (-98%)	-

此外，针对平衡参数 $C$ 的取值消融表明：$C=0.51$ 取得最优表现（PickScore 23.16）；过小（$C=0.501$ 时为 22.53）减弱了对预训练速度的锚定作用；过大（$C=1.0$ 时为 21.97）导致噪声项主导并破坏训练稳定性。针对终端代价修正的消融表明，若去除修正项仅保留纯奖励梯度，PickScore 急剧恶化至 21.68，证明了修正理论的必要性。

关键发现¶

训练收敛加速 4 倍且性能无损：EAM 将 SD3.5-M 微调所需的 GPU 训练时间从 AM 的 64 小时骤降至 16 小时，并在 SANA 1.0（36h → 9h）和 SANA 1.5（120h → 30h）上展现出高度一致的 4 倍加速比。
显存与计算瓶颈完全解除：伴随计算耗时从单次迭代 34.7 秒降至 0.4 秒，伴随阶段峰值显存从 73.6 GB 压缩至 1.6 GB，消除了多步神经网络自动微分对于计算资源的严苛要求。
细节保真度显著提升：相较于 AM 倾向于过度平滑局部细节，EAM 在微调中保持了更好的精细纹理与复杂构图。这得益于训练时 10 步 ODE 采样与测试时完全一致，消除了训练/测试阶段动力学失配。

亮点与洞察¶

解耦 SOC 基准过程与生成主干：敏锐地洞察到伴随匹配的瓶颈本质上并非 SOC 范式本身，而是「机械采用预训练模型漂移作为基准」这一默认假设。通过换用线性漂移，巧妙将繁重微分转化为解析标量运算。
零代价的理论最优等价性：不仅提出了工程加速方案，更在严格测度论层面证明了新形式与原 AM 具有完全重合的最优受控路径测度与边缘漂移，做到了理论自洽。
训练与推理动力学一致性：利用前向 ODE 采样直接匹配了推理部署时的轻量求解器分布，有效避免了传统 SDE 微调模型在迁移至确定性快采样器时的分布偏移。

局限与展望¶

收敛性依赖不动点迭代机制：中间态采样公式在非最优解时属于近似，虽然最优解是算法的不动点，但由于漂移差异，有限网络容量下的局部优化景观与经验路径可能与理论最优点存在微小偏差。
超参数 C 的敏感性：常数 $C$ 需要精细设定在 $0.5$ 邻近区间（如 $0.51$），虽然在多个模型上具有良好通用性，但过大取值会导致梯度方差增大。
优化策略仍有扩展空间：当前算法单步直接更新，未来可结合经验回放池（Replay Buffer）、旧策略离线轨迹混合以及自适应退火学习率等强化学习常用技巧进一步挖掘极限潜力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙重构 SOC 基础漂移与终端代价，以纯解析形式彻底颠覆了伴随状态的反向仿真计算。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 SD3.5-M、SANA 1.0、SANA 1.5 跨架构实验，多奖励评测、时间显存细粒度分解与严谨消融非常详实。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密清晰，从痛点剖析、性质刻画到定理证明逻辑一气呵成。
价值: ⭐⭐⭐⭐⭐ 为扩散与流模型的高效对齐提供了坚实的理论工具与实用算法，极大降低了奖励梯度微调的算力门槛。