WTF?! Simulation-Free Reinforcement Learning with Wasserstein-Tilted Flow Maps
Abbas Mammadov, Jerry Y. Huang, Justin Lin, Partha Kaushik, Sheel Shah, Kartik Nair, Yee Whye Teh, Nicholas M. Boffi
cs.LG, cs.CV, stat.ML
2026-09-23
KL 正则换成基于预训练漂移的最优传输后，流模型微调把样本直接搬向高奖励区，文生图 4 步推理拿下最高 HPSv2，算力比 Flow-GRPO 省 280 倍。
奖励微调是生成模型后训练的常规动作：拿一个打分模型当奖励，把预训练生成器往高分方向拉。主流配方把目标写成 KL 正则的奖励最大化，最优解是对基分布按 e^(λr) 加权的「倾斜分布」。这套框架有两个长期的别扭。第一，KL 倾斜只能在基分布已有的支撑上重新分配概率质量：奖励相同的两个点，无论 λ 多大，相对密度都不变；把质量搬进基分布本就稀薄的高奖励区，它在机制上做不到。第二，能落地的训练目标依赖 Girsanov 定理，那是随机扩散过程的工具。用确定性流（flow matching）的团队只好绕路：把流按扩散处理、微调、再转回去，训练中还要滚动采样估计价值函数，方差高、代价大。至于已经蒸馏成一步生成的 flow map，现有流程根本用不上。
WTF（Wasserstein-Tilted Flow Maps）把 KL 正则整个换掉，改成建立在预训练漂移上的最优传输代价：候选速度的成本按它偏离基流速度的能量（偏差范数平方对时间的积分）计算，只对偏离基流的部分计费，搬得越远付得越多。
理论上有三步，每步都把目标往能跑的算法推一格：
flow map 要单独说一句。流匹配模型学的是 ODE 的速度场，采样要逐步积分，几十上百次网络调用；flow map 学的是这个 ODE 的解算子，一次调用跨任意时间区间，多步采样靠半群性质拼接。WTF 直接微调 flow map，产出的模型天生少步可用，不需要「先微调速度场、再蒸馏」的二次工序。大规模实验里奖励梯度用欧氏形式，直接取终点处的奖励梯度，实测比精确的雅可比回拉更稳。
ImageNet-256，基座 DMF XL/2，HPSv2 作训练奖励，三个种子，各方法同一预算：一块 8×H100 节点上约 5 小时。
| 方法 | NFE | HPSv2 | DreamSim 多样性 |
| 基座 | 1 | 0.214 | 0.790 |
| Adjoint Matching | 250 | 0.243 | 0.453 |
| VFM | 1 | 0.307 | 0.482 |
| WTF | 1 | 0.319 | 0.482 |
文生图，基座 TiM-T2I，微调约 48 GPU 时：
| 方法 | NFE | HPSv2 | PickScore | ImageReward |
| 基座 | 1 | 0.228 | 20.36 | 0.182 |
| Flow-GRPO | 50 | 0.303 | 22.06 | 1.210 |
| Adjoint Matching | 50 | 0.357 | 22.75 | 1.380 |
| WTF | 4 | 0.401 | 22.97 | 1.383 |
算力对照按「追平对手峰值奖励」计：比 Adjoint Matching 省 47 倍，比 Flow-GRPO 省 280 倍训练算力。一步生成的 WTF 拿 0.319，已经高于基座 250 步的 0.218 和 Adjoint Matching 250 步的 0.243。
RL 后训练正在变成生成模型的标准开支，rollout 在 4 步上做还是在 50 步上做，成本差出数量级，280 倍的差距就是这么来的。方法论上的提示更根本：KL 倾斜不是奖励微调唯一合法的正则，对确定性流来说，从传输代价出发的正则才是原生的。少步加速采样器不再只是推理端的技巧，而是后训练可以直接瞄准的基础设施。
作者自述四条：固定点结论只在总体层面成立，没有有限样本保证；大规模实验用的欧氏奖励梯度与理论分析的精确形式不一致；奖励升、多样性降，靠 λ 调节取舍；方法假设手头已有预训练 flow map，从速度模型出发要先蒸馏，这笔成本没有算进 47 与 280 倍的对照。读下来另有两点。HPSv2 同时是训练目标，PickScore 与 ImageReward 的涨幅里有多少来自奖励模型之间的相关性，论文没有拆开。每个基准只测一个基座，换基座能否复现，没有证据。多样性也要诚实读：WTF 的 0.482 对基座 250 步的 0.841 是明显下降，只是比同类微调方法好（MFM 是 0.379）。