--- title: "Trajectory Balance (TB) 目标" created: 2026-05-12 updated: 2026-05-12 type: concept tags: ["gflownet", "off-policy-rl", "reinforcement-learning"] sources: ["arxiv:2201.13259", "arxiv:2503.18929"] --- # Trajectory Balance (TB) 目标 **Trajectory Balance** 是 [[gflownet-fine-tuning|GFlowNet]] 框架中的一种学习目标,具有 **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|Replay Buffer]] 中的 stale 数据仍然有效 - [[asynchronous-rl-llm|异步 RL]] 成为可能 - 支持多样化的探索策略 ### 梯度等价性 对于 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|TBA]] 利用 TB 的 off-policy 性质,使 Searcher 可以持续生成数据而 Trainer 持续训练,无需等待策略同步。这是 TBA 实现 4×–50× 加速的理论基础。 ## 相关概念 - [[tba|TBA]] — 框架应用 - [[gflownet-fine-tuning]] — GFlowNet 基础 - [[grpo]] — On-policy 对比 - [[off-policy-llm-post-training]] — Off-policy 范式 - [[bartoldson-tba-2025|论文页面]]