WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps¶

会议: NeurIPS 2026
arXiv: 2609.27033
领域: 强化学习
关键词: 强化学习微调, 流图模型, 最优传输正则化, 确定性最优控制, 无模拟策略迭代

一句话总结¶

针对流生成模型奖励微调依赖扩散代理与多步慢速采样的难题，本文提出基于基础漂移场的最优传输正则化与确定性最优控制框架 WTF，利用流图单步跨时间步转移实现无需轨迹模拟且无 Critic 的端到端微调，在大幅降低训练算力的同时保持少步数生成质量。

研究背景与动机¶

流匹配与扩散模型在图像、视频、生物大分子与文本等连续和离散生成任务中取得了最先进的生成保真度。然而在实际部署中，无条件或仅依赖数据分布的采样往往不足以满足下游需求，需要结合下游奖励函数（例如人类美学偏好评分、安全性约束或分子结合亲和力）进行奖励微调（Reward Fine-Tuning）。当前的主流微调理论通常基于 KL 正则化的奖励最大化范式，将微调目标形式化为奖励倾斜分布（Reward-Tilted Distribution）。为了求解该目标，现有算法普遍借助 Doob $h$-变换与 Girsanov 定理，把确定性流转换为具有独立前向/反向过程的随机扩散模型，在扩散路径空间上进行随机最优控制求解，最后再转换回等效流。

这种“流转扩散、微调扩散、再转回流”的管线在计算效率与模型形态上存在严重的内在矛盾。一方面，扩散微调的核心依赖于对辅助随机过程的终端奖励进行蒙特卡洛值函数评估，不仅需要数十至数百步昂贵的数值积分，而且其指数奖励期望在高维下方差极大；另一方面，生成模型领域正加速向少步甚至单步推断的流图（Flow Maps）蒸馏演进，而现有的扩散微调方法根本无法直接利用流图的长距离跳跃特性来加速训练中的样本生成与梯度回传。此外，KL 正则化本质上只在既有数据分布的支撑集内对样本权重进行重标度，无法将远离高概率密度的低频样本直接搬运到高奖励区域。

为了打破这种算法与推断结构的脱节，本文重新审视了微调的正则化机制：不再度量分布间的相对熵，而是直接衡量确定性动力系统偏离预训练漂移场的控制能耗。核心 idea：引入由预训练漂移场诱导的最优传输正则化，将奖励微调严格等价为终端奖励的确定性最优控制问题，并利用流图单次前向直接抵达任意时间步的能力，构建无模拟（Simulation-Free）、无 Critic 且无偏的蒙特卡洛值梯度回归微调算法。

方法详解¶

整体框架¶

WTF 的微调流程直接在确定性流图参数化模型 $\hat{w}_{s,t}(x)$ 上进行，其对应的状态转移为 $X_{s,t}^{\hat{w}}(x) = x + (t-s)\hat{w}_{s,t}(x)$。在任意训练步中，模型使用冻结参数的参考流图 $\bar{w}$ 生成在策略状态轨迹点与终端预测样本，通过对标量时间积分的单点时间采样计算值函数及解析值梯度，在对角线 $s=t$ 上指导瞬时速度场的策略改进更新；同时在非对角线 $s < t$ 上联合执行流图自蒸馏损失，确保流图在微调过程中持续保持半群性质与少步数采样能力。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["先验采样<br/>x₀ ~ 𝒩(0, I) 及时间 t, τ ~ [0, 1]"] --> B["流图跳跃采样<br/>冻结参考 w̄ 评估 x̄_t, x̄_τ, x̄_1"]
    B --> C["单样本无模拟值函数估计<br/>结合终端奖励 r(x̄_1) 与控制耗费"]
    C --> D["对角线策略改进更新<br/>ŵ_t,t 拟合预训练漂移与值梯度之和"]
    D --> E["非对角线自蒸馏约束<br/>保持流图半群性质与跨时间步跳跃"]
    E --> F["输出微调后流图<br/>任意 NFE 预算下均可直接部署采样"]

关键设计¶

1. 基于预训练漂移的最优传输正则化：将分布重标度转变为连续空间粒子搬运 传统微调采用 $\mathrm{KL}(\nu \parallel \rho_1)$ 正则化，其最优解 $\tilde{\rho}_1(x) \propto e^{\lambda r(x)}\rho_1(x)$ 严格受限于基分布的支撑集，对具有相同奖励的样本只能按原密度同比例调整。本文构建了以先验分布 $\rho_0$ 为源分布、以候选分布 $\nu$ 为目标的最优传输代价 $\mathcal{T}_b(\rho_0, \nu) = \inf_{\pi \in \Pi(\rho_0, \nu)} \int c_b(x, y) \pi(dx, dy)$。其中粒子点对代价 $c_b(x, y) = \inf_{\omega_0=x, \omega_1=y} \int_0^1 \frac{1}{2} \|\dot{\omega}_t - b_t(\omega_t)\|^2 dt$ 衡量将粒子从 $x$ 驱动至 $y$ 时偏离预训练漂移场 $b_t$ 所需的最小残差控制能量。理论证明，该正则化目标在 Kantorovich 对偶下可解耦为单样本层面的近端传输极值问题： $$ \sup_\nu \left{ \lambda \mathbb{E}{x \sim \nu}[r(x)] - \mathcal{T}b(\rho_0, \nu) \right} = \mathbb{E} \left[ \sup { \lambda r(y) - c_b(x, y) } \right] $$ 这意味着每个先验样本能够自主寻找满足奖励增益与动力学偏离惩罚最佳权衡的终端位置，使模型能够跨越低密度空隙、主动将质量输运至高奖励模态。}^d

2. 确定性最优控制等价性与值梯度回归：消除 Critic 网络的策略提升 通过 Benamou-Brenier 动力学形式转化，本文证明上述最优传输正则化问题与状态方程 $\dot{x}_t^u = b_t(x_t^u) + u_t(x_t^u)$ 下的终端奖励确定性最优控制问题完全等价。针对直接端到端回传终端奖励易引发“对角动力学与非对角流图解耦脱节”的奖励黑客（Reward Hacking）现象，WTF 建立在策略迭代框架上。定义参考控制 $\bar{u}$ 下的值函数 $V_t^{\bar{u}}(x) = \lambda r(X_{t,1}^{\bar{u}}(x)) - \frac{1}{2}\int_t^1 \|u_\tau(X_{t,\tau}^{\bar{u}}(x))\|^2 d\tau$，本文证明了策略差分恒等式（Performance Difference Identity）：将当前可学习控制拟合至参考值函数梯度 $\bar{g}_t(x) = \nabla_x V_t^{\bar{u}}(x)$，保证在总体分布水平上实现单调策略提升，且不动点唯一对应全局最优控制 $u_t^*(x) = \nabla_x V_t^*(x)$。

3. 流图赋能的无模拟值函数估计：$\mathcal{O}(1)$ 评估替代 $\mathcal{O}(T)$ 轨迹展开 在常规方法中，计算值梯度需要沿整个未来轨迹积分并进行跨时间步反向传播，面临高昂的显存与计算开销。而在流图参数化下，从当前状态 $x_t$ 到最终时刻 $t=1$ 的状态转移可由流图在 $\mathcal{O}(1)$ 次网络调用内直接闭式计算。值函数中的连续时间控制能耗积分可转化为对采样时间标量 $\tau \sim \mathrm{Unif}[t, 1]$ 的一维期望，从而构建无偏单样本蒙特卡洛值函数估计： $$ \widehat{V}(t, x) = \lambda r(\bar{x}1) - \frac{1-t}{2} |\bar{w}}(\bar{x\tau) - b\tau(\bar{x}_\tau)|^2 $$ 其中 $\bar{x}_\tau = X_{t, \tau}^{\bar{w}}(x)$ 且 $\bar{x}_1 = X_{\tau, 1}^{\bar{w}}(\bar{x}_\tau)$。该估计器方差仅源于标量时间 $\tau$ 的采样，彻底规避了高维布朗运动积分的高方差问题。通过自动微分直接求导即可获得无偏回归靶标 $\hat{g} = \nabla_{\bar{x}_t} \widehat{V}$，无需任何预先训练或同步更新的价值网络（Critic-Free）。

4. 欧氏奖励梯度与规范化尺度缩放：稳定大尺度图像合成 在处理高分辨率文生图等视觉任务时，精确奖励梯度项包含流图状态转移的雅可比逆转置 $\nabla X_{t,1}^\top \nabla r(x_1)$。针对流图雅可比可能引入的病态曲率，WTF 在大规模图像实验中引入欧氏奖励梯度更新，将生成终点改写为 $x_1 = x_t + \mathrm{sg}(X_{t,1}(x_t) - x_t)$，使梯度直接以 $\nabla r(x_1)$ 反传。同时，针对预训练漂移模长与外部奖励模型输出尺度不一致的问题，动态引入范数比率缩放因子 $\kappa = \mathbb{E}[\|b_t(\bar{x}_t)\|_2 / \|\nabla_x r(\bar{x}_1)\|_2]$，使用自适应有效参数 $\lambda_{\mathrm{eff}} = \lambda \cdot \kappa$ 消除不同奖励模型的手动调参成本。

损失函数与训练策略¶

WTF 的联合优化目标由对角线值梯度回归损失与非对角线欧拉自蒸馏（Eulerian Self-Distillation, ESD）损失构成： $$ \mathcal{L}{\mathrm{WTF}}(\hat{w}) = \mathbb{E}} \left[ |\hat{w}_{t,t}(\bar{x}_t) - b_t(\bar{x}_t) - \hat{g}(t, \bar{xt)|^2 \right] + \beta \mathcal{L}) $$ 其中 }}(\hat{w$\bar{w} = \mathrm{sg}(\hat{w})$ 为冻结参考（在 ImageNet 实验中使用 EMA 权重，文生图实验中直接使用当前步的 stop-gradient 副本）。自蒸馏损失形式为 $\mathcal{L}_{\mathrm{dist}}(\hat{w}) = \mathbb{E}[\|\partial_t X_{t,\tau}^{\hat{w}}(x_t) + \mathrm{sg}(\nabla X_{t,\tau}^{\hat{w}}(x_t) \hat{w}_{t,t}(x_t))\|^2]$。理论证明，在充分支撑与适定正则性条件下，该损失的总体驻点 $\hat{w}$ 严格收敛至最优漂移场 $b_t + u_t^*$，且其非对角映射完整保持为最优控制动力学的真流图。

实验关键数据¶

主实验¶

在文生图任务上，研究者以 TiM-T2I 为基础流图，以 HPSv2 人类偏好评分作为训练奖励，在相同的单台 8×H100 节点（约 48 GPU-hours）计算预算下与一阶代表方法 Adjoint Matching 及零阶强化学习方法 Flow-GRPO 进行对比。WTF 单次训练所得流图可原生评估于 1、4、8 及 50 NFE（网络评估步数）。

模型与方法	NFE	HPSv2 (微调奖励) ↑	PickScore ↑	ImageReward ↑	DreamSim (多样性) ↑	CLIP 多样性 ↑
TiM-T2I Base	1	0.228	20.36	0.182	0.548	0.413
TiM-T2I Base	4	0.255	21.06	0.665	0.426	0.346
TiM-T2I Base	8	0.264	21.29	0.801	0.361	0.296
TiM-T2I Base	50	0.274	21.58	0.997	0.281	0.239
Flow-GRPO	50	0.303	22.06	1.210	0.250	0.215
Adjoint Matching	50	0.357	22.75	1.380	0.210	0.188
WTF (Ours)	1	0.347	22.31	1.069	0.325	0.289
WTF (Ours)	4	0.401	22.97	1.383	0.250	0.220
WTF (Ours)	8	0.400	22.93	1.394	0.233	0.212
WTF (Ours)	50	0.398	22.87	1.416	0.217	0.205

在 ImageNet-256 类条件图像生成任务上，以 DMF XL/2 为底模，微调 HPSv2 奖励 12,500 步：

模型与方法	产物类型	评估 NFE	HPSv2 ↑	PickScore ↑	ImageReward ↑	DreamSim ↑	CLIP 多样性 ↑
DMF XL/2 Base	基础流图	1	0.214	19.21	-0.271	0.790	0.350
DMF XL/2 Base	基础流图	250	0.218	19.45	-0.252	0.841	0.340
Adjoint Matching	速度场流	250	0.243	20.26	0.005	0.453	0.237
MFM	速度场流	250	0.291	19.83	0.544	0.379	0.189
VFM	单步采样器	1	0.307	20.91	0.860	0.482	0.194
WTF (Ours)	任意步流图	1	0.319	20.24	0.612	0.482	0.201
WTF (Ours)	任意步流图	250	0.330	20.45	0.718	0.460	0.177

梯度设计与算力消融分析¶

文生图消融对比了流图雅可比回传（精确梯度）与欧氏奖励梯度的性能差异，以及达成相同奖励水平所需的算力对比：

评估配置 / 消融项	HPSv2 ↑	PickScore ↑	ImageReward ↑	DreamSim 多样性 ↑	达到 Baseline 峰值奖励算力节约倍数
精确梯度 ($J^\top \nabla r$) @ 800步	0.377	22.59	1.340	0.163	—
欧氏梯度 ($\nabla r$) @ 800步	0.395	22.83	1.400	0.195	—
对比 Adjoint Matching 峰值	0.357	22.75	1.380	0.210	47× 训练算力缩减
对比 Flow-GRPO 峰值	0.303	22.06	1.210	0.250	280× 训练算力缩减

关键发现¶

训练收敛速度呈现数量级飞跃：在文生图模型后训练中，WTF 达到 Adjoint Matching 峰值奖励所需算力减少 47 倍，达到 Flow-GRPO 峰值奖励所需算力减少 280 倍，彻底打破了多步轨迹反传与强化学习在线采样的算力壁垒。
天然继承并保持少步数生成质量：WTF 在仅需 4 步推断预算下即可取得高达 0.401 的 HPSv2 分数，甚至微弱优于 50 步的 0.398，单步推断（1 NFE）亦可达到 0.347，无需后续进行繁琐复杂的两阶段事后蒸馏。
欧氏奖励梯度兼顾高奖励与高多样性：省去流图雅可比矩阵回传的欧氏奖励梯度在保持终点不变的前提下，大幅抑制了高维雅可比诱发的过拟合模式坍缩，在所有对比维度上均取得了更高的人类偏好与多样性指标。

亮点与洞察¶

范式创新：从“概率分布重标度”迈向“微分动力学最优输运”：传统微调将生成模型死板绑定在 KL 散度与薛定谔桥框架下，而 WTF 从控制论本质出发，将偏离惩罚建立在预训练漂移场的二次控制泛函上，使粒子能够主动跨过低密度区域迁移至高奖励区域。
架构同源：流图既是推断加速器，更是后训练的基础设施：以往工作仅把流图作为事后提速的蒸馏手段，而 WTF 首次揭示出流图对角/非对角解耦性质是实现“无模拟 Monte Carlo 梯度回归”的数学基石，将生成效率直接反哺至后训练优化。
全栈无 Critic 设计：通过严格的动力学值函数时间积分变换，WTF 用单次随机标量时间抽样取代了昂贵复杂的判别器网络训练，规避了强化学习中常见的值函数高估与拟合不稳定性。

局限与展望¶

依赖预训练流图初始权重：WTF 默认系统已拥有可用的预训练流图检查点；若仅有传统速度场模型，则必须先经历一步流图蒸馏阶段，该阶段的预处理成本未计入微调算力中。
固定点定理的理论与实操间隙：命题 4.2 的无偏收敛性是在总体连续极限下成立的，在大尺度视觉实践中采用的欧氏梯度与 LoRA 低秩近似引入了微小的动力学修正，未来可进一步完善经验流形上的有限样本泛化界。
奖励与多样性的内在权衡：随着温度系数 $\lambda$ 增大，生成样本的偏好评分上升伴随一定程度的多样性下降，自适应根据局部能量景观动态调节输运强度将是值得探索的演进方向。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次跳出 KL 框架，建立面向流图的原生最优输运与确定性最优控制微调理论体系。
实验充分度: ⭐⭐⭐⭐⭐ 兼顾一维理论基准、ImageNet 类条件生成与大规模文生图，指标覆盖三项偏好模型与两项嵌入多样性。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，从变分法、HJB 方程到策略提升恒等式层层递进，设计动机极其清晰。
价值: ⭐⭐⭐⭐⭐ 将扩散与流后训练开销降低数十至数百倍，且输出原生少步流图，极具工业级应用前景。

评估配置 / 消融项	HPSv2 ↑	PickScore ↑	ImageReward ↑	DreamSim 多样性 ↑	达到 Baseline 峰值奖励算力节约倍数
精确梯度 (\(J^\top \nabla r\)) @ 800步	0.377	22.59	1.340	0.163	—
欧氏梯度 (\(\nabla r\)) @ 800步	0.395	22.83	1.400	0.195	—
对比 Adjoint Matching 峰值	0.357	22.75	1.380	0.210	47× 训练算力缩减
对比 Flow-GRPO 峰值	0.303	22.06	1.210	0.250	280× 训练算力缩减