可微奖励就该直接微调！用HJB方程颠覆流匹配对齐

用强化学习微调扩散模型，还有更好的办法吗？

来自港中深、微软研究院等机构的 VGG-Flow 团队给出了一个新思路：既然奖励函数本身是可微的，为什么非要绕弯路用 PPO、GRPO。

在大规模生成模型的对齐任务中，通常依赖强化学习，在某个奖励函数上微调模型以贴近人类偏好。而事实上，大部分奖励模型本身是在偏好数据集上训练过的神经网络。既然奖励是可微的，能否直接利用 " 可微性 " 本身，高效而稳定地微调流匹配模型？

主流做法主要分为两类路径：一条路是把模型当作黑盒，通过像 Flow-GRPO 那样，把原本确定性的 ODE 采样过程强行转为随机 SDE，适配经典的强化学习框架来采用高方差的策略梯度方法（如 PPO、GRPO）。

另一条路则更加直接，如 ReFL 等方法，通过近似手段优化某些取样步对应的奖励值，但这种做法在目标层面上缺乏严格的理论支撑，也往往容易导致过拟合与模式坍塌。那么是否可以走一条新路线？

VGG-Flow团队回归第一性原理，将奖励微调重新表述为一个连续时间最优控制问题。通过 Hamilton – Jacobi – Bellman（HJB）方程，直接将 " 可微奖励 " 转化为价值梯度，为流匹配对齐提供了一条更稳定、更鲁棒的路径。目前该项目已被 NeurIPS 2025 接收。

VGG-Flow 的核心思路：最大化 " 终态奖励−累计代价 "

流匹配模型通过在随机取样的 x ₀上模拟时间 t=0 到 t=1 的轨迹ẋ =v ( x,t ) 来生成样本，其中 v ( x,t ) 是流匹配模型的速度场。

微调后的速度场可以被写成预训练模型与残差的和：v θ ( x,t ) =vbase ( x,t ) + ṽ θ ( x,t ) ，其中预训练模型是 vbase ( x,t ) ，残差是ṽ θ ( x,t ) 。

直观来看，为了避免模型在微调过程中过度偏离原有分布，微调在最大化样本奖励的同时，需要约束预训练模型与微调模型在取样路径上的差：

从最优控制的角度看，这就是一个终态目标加上一段路径累计代价（cost-to-go）。

HJB 方程：从价值到梯度的解析直觉

在最优控制理论中，价值函数 V ( x,t ) 描述了从状态 ( x,t ) 出发的最优预期成本。根据定义，上述目标对应如下的价值函数：

其演化满足以下 Hamilton – Jacobi – Bellman（HJB）方程（强化学习中贝尔曼方程的连续时间形式）：

由此可以得到最优修正项的解析形式：

这得到一个非常直接的结论：最优微调方向 = 价值函数的梯度。

不需要采样优势函数，不需要计算对数概率比，也不需要进行策略比值裁剪。只需估计价值梯度，即可直接、可微地更新流匹配模型。

这个价值函数如何得到？将最优速度场代回 HJB 方程，可以得到如下的价值一致性关系：

通过求解满足该一致性关系的价值函数，即可得到用于训练速度场的目标梯度。

参数化：将先验写入梯度

为了使价值梯度∇ V ( x,t ) 在训练初期具备合理的引导方向，VGG-Flow 引入了 Forward-looking 参数化方法：

1. 预估终点：在 xt 处进行一步 Euler 前推，得到预估终点

2. 参数化引导：利用一步前推的奖励梯度对价值梯度∇ V ( x,t ) 进行参数化：

其中 η t 为随时间减小而减小的标量权重，为可学习的残差项。

在这种设计下，模型仅需学习残差项即可实现对齐。团队在实验中发现，即使不学习残差，仅依赖奖励梯度的参数化引导也能实现明显的对齐效果。这一技巧可以显著降低流匹配模型微调的计算成本。

优化目标：在该框架下，VGG-Flow 的损失函数可以写为：

1. 梯度匹配损失：使速度场修正项拟合价值梯度

2. 价值一致性损失：最小化 HJB 方程的残差

3. 终端边界损失：根据价值函数定义，确保 t=1 时的边界条件成立

注：若采用上一节所述的设置（即不学习价值梯度中的残差项），仅需优化梯度匹配损失。

实验效果

在 Stable Diffusion 3 上的实验中，仅需 400 次更新，VGG-Flow 即可实现奖励信号的稳定提升。在 Aesthetic 和 PickScore 指标上，方法展现出较高的收敛效率与良好的多样性保持能力。相比 ReFL、DRaFT 等方法，其表现更为稳健，更不易遗忘预训练模型中的先验，生成结果更加自然。其收敛也快，并且直接作用于流匹配模型本身，无需额外将 ODE 转换为 SDE。