2.2 KiB
2.2 KiB
title, created, updated, type, tags, sources
| title | created | updated | type | tags | sources | |||||
|---|---|---|---|---|---|---|---|---|---|---|
| Trajectory Balance with Asynchrony (TBA) | 2026-05-12 | 2026-05-12 | concept |
|
|
Trajectory Balance with Asynchrony (TBA)
TBA 是一个分布式异步 RL 框架,将 trajectory-balance-objective 目标应用于 LLM 后训练,实现 4×–50× 加速。
核心思想
将 LLM 后训练解耦为两个独立并行的角色:
- S EARCHER:生成响应、评估奖励、存入 replay buffer
- T RAINER:从 buffer 采样、计算 TB loss、更新策略
关键突破:TB 目标天然 off-policy 兼容,使得 Trainer 可以在 Searcher 持续产生数据的同时学习——无需等待 on-policy rollouts。
两种实现
| 变体 | 代码 | 特点 |
|---|---|---|
| TBA | 从 scratch 实现 | buffer 采样、β 退火、多 searcher |
| TBA′ | 基于 PRIME-RL | 简单、多 GPU 训练、ref-policy reset |
核心流程
- 每 k 步同步:Searcher ↔ Trainer 交换权重和 buffer 数据
- Buffer 采样:概率 m 采样最近数据(recency),1−m 采样高奖励数据
- TB 梯度更新:使用 VarGrad 估计 Z(x),避免学习 value network
关键结果
- GSM8K (RhoMath-1B): 55% 准确率 + 50× 加速 vs VinePPO
- PFT (TL;DR): 在 16 步 off-policy 下超越 on-policy Online DPO
- RT (GPT-2): 达到 diversity-toxicity Pareto 前沿
- MATH (Qwen 2.5 7B): TBA′ 在高度 off-policy 下超越 Dr. GRPO
为什么 TB 适合异步?
TB 的 VarGrad 梯度等效于 mean-baseline REINFORCE + KL 正则化奖励,但关键在于 不要求 on-policy 数据。只要数据有 full support,TB 保证收敛——使异步产生的 stale 数据仍然可被高效利用。
相关概念
- trajectory-balance-objective — TB 目标详解
- asynchronous-rl-llm — 异步 RL 范式
- searcher-trainer-decoupling — 架构模式
- replay-buffer-rl-llm — Buffer 设计
- reward-recency-sampling — 采样策略
- gflownet-fine-tuning — GFlowNet 基础
- bartoldson-tba-2025