56 lines
2.2 KiB
Markdown
56 lines
2.2 KiB
Markdown
---
|
||
title: "Trajectory Balance with Asynchrony (TBA)"
|
||
created: 2026-05-12
|
||
updated: 2026-05-12
|
||
type: concept
|
||
tags: ["reinforcement-learning", "llm-post-training", "gflownet", "asynchronous-rl"]
|
||
sources: ["arxiv:2503.18929"]
|
||
---
|
||
|
||
# Trajectory Balance with Asynchrony (TBA)
|
||
|
||
**TBA** 是一个分布式异步 RL 框架，将 [[trajectory-balance-objective|Trajectory Balance (TB)]] 目标应用于 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 |
|
||
|
||
## 核心流程
|
||
|
||
1. **每 k 步同步**：Searcher ↔ Trainer 交换权重和 buffer 数据
|
||
2. **Buffer 采样**：概率 m 采样最近数据（recency），1−m 采样高奖励数据
|
||
3. **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|论文页面]]
|