64 lines
2.3 KiB
Markdown
64 lines
2.3 KiB
Markdown
---
|
||
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|论文页面]]
|