Trajectory Balance (TB) 目标

Trajectory Balance 是 gflownet-fine-tuning 框架中的一种学习目标，具有 off-policy 兼容 的关键性质，使其天然适用于异步分布式训练。

定义

对于 LLM 后训练，TB 目标定义为：

L_{TB}(y, x; \theta) = \left(\log \frac{Z(x)\pi_\theta(y|x)}{R(y; x)}\right)^2

其中：

R(y; x) = \pi_{ref}(y|x) \exp(\beta^{-1} r_\phi(y; x)) 是归一化的后验奖励
Z(x) 是 partition function（可用 VarGrad batch estimate 替代学习）
当 L_{TB} 最小化时，\pi_\theta \propto R

VarGrad 变体

用 K-sample batch estimate 替代学到的 Z(x)：

\log \hat{Z}(x^{(i)}) = \frac{1}{K}\sum_{j=1}^{K}\left[\log \pi_{ref}(y^{(i,j)}|x^{(i)}) - \log \pi_\theta(y^{(i,j)}|x^{(i)}) + \frac{1}{\beta}r_\phi\right]

使用 STOP-GRAD 确保估计值不回传梯度。

关键性质

Off-Policy 兼容

训练时 y 可以从任意分布采样（只要 full support），不需要来自当前策略 $\pi_\theta$。这使得：

replay-buffer-rl-llm 中的 stale 数据仍然有效
asynchronous-rl-llm 成为可能
支持多样化的探索策略

梯度等价性

对于 on-policy 数据，TB_VarGrad 的梯度等价于：

\nabla J_{TB}(\theta) = \hat{A}(y|x) \nabla \log \pi_\theta(y|x)

即 mean-baseline REINFORCE + KL 正则化奖励——与 RLOO 形式一致，但 off-policy 能力更强。

与 Kimi K1.5/K2 的关系

Kimi 的 RL 目标几乎匹配 TB，但排除了 log probability ratio 的平均值作为 control variate。TBA 的 ref-policy reset 策略也更保守（保持原始 base policy 作为参考更久）。

在 TBA 中的应用

tba 利用 TB 的 off-policy 性质，使 Searcher 可以持续生成数据而 Trainer 持续训练，无需等待策略同步。这是 TBA 实现 4×–50× 加速的理论基础。

2.3 KiB Raw Blame History Unescape Escape