Bengio团队提出TBA，解耦采样与训练，让LLM后训练最高提速50倍。
原文标题：无惧Off-Policy偏移！Bengio团队解绑后训练，大模型RL提速50倍
原文作者：数据派THU
冷月清谈：
怜星夜思：
2、文章里说旧轨迹也能继续学习，这会不会让模型学到过时甚至错误的东西？
3、TBA最高提速50倍，这个数字在实际训练大模型时有多大参考价值？
4、Trajectory Balance相比PPO、GRPO这类方法，真正优势是在算法本身还是系统架构？
原文内容
本文约2500字，建议阅读5分钟
采样和学习终于分家。
大模型 RL 后训练最慢的环节，往往不是更新参数，而是苦等一轮轮 rollout。Bengio 团队用异步框架把采样和训练拆开，让旧轨迹也能继续变成学习信号。
大模型 RL 后训练最慢的环节，往往不是更新参数，而是苦等一轮轮 rollout。Bengio 团队用异步框架把采样和训练拆开，让旧轨迹也能继续变成学习信号。
LLM 后训练的慢，很多时候不在学习本身，而在等模型一条条生成样本。
PPO、RLOO、GRPO 这类主流方法通常依赖 on-policy 数据，模型先生成回答，再计算奖励，最后更新策略。问题在于，rollout 是逐 token 解码，慢；训练是对完整序列做并行计算，快。
为了让样本尽量来自当前策略，训练环节只能等生成环节先跑完，集群算力很难真正跑满。策略一更新，前面生成的样本又会变成 off-policy。采样规模越大，数据反而越难复用。
为了解决这一瓶颈，Yoshua Bengio 团队在 NeurIPS 2025 提出了 TBA（Trajectory Balance with Asynchrony）。
这套异步框架把后训练拆成两条流水线，Searcher 负责生成轨迹，Trainer 持续从 replay buffer 中抽样更新模型。
TBA 引入 Trajectory Balance 目标来处理这些 off-policy 轨迹，让原本容易被丢掉的旧样本继续参与学习。
实验显示，在数学推理、偏好微调、自动红队等任务上，TBA 在保持甚至提升性能的同时，将大模型 RL 后训练最高提速 50 倍。
论文标题：Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training
论文链接：https://arxiv.org/abs/2503.18929
代码链接：https://github.com/bbartoldson/TBA
01 架构解耦：Searcher与Trainer独立运转
TBA 把计算集群里的节点分成两类：一类是负责探索采样的 Searcher，另一类是负责模型更新的 Trainer。
实际运行中，每个 Searcher 都维护一份相对滞后的模型权重，从数据集中取出 prompt 并生成回答。
生成的回答和奖励会先存入 Searcher 本地的 replay buffer，再在周期性同步时汇总到 Trainer 端的全局 buffer。Trainer 不必等待每一轮 rollout 完成，而是主要从全局 buffer 中抽取 batch 数据更新策略。
为了避免 Searcher 与 Trainer 的策略版本相差过远，系统设置了同步周期 k。每隔 k 个优化步，系统会把 Trainer 的最新权重同步给 Searcher，同时把各个 Searcher 的本地经验汇总到全局 buffer。
这种周期性同步的异步设计，减少了训练等待 rollout 的时间，也让集群资源利用率更高。
02 用轨迹平衡处理off-policy轨迹
打破同步约束后，Trainer 从全局 buffer 中拿到的数据，往往不再严格来自当前策略，而是带有不同程度的 off-policy 性质。
对于传统 on-policy 方法而言，直接使用这些旧轨迹，可能带来策略偏移和训练不稳定。这也是 TBA 选择轨迹平衡（Trajectory Balance）目标的原因。
TB 来源于 GFlowNet，一个关键性质是 off-policy：只要采样分布具有 full support，轨迹不必来自当前策略，也可以用于训练。
在 on-policy 数据上，VarGrad TB 等价于带均值基线的 REINFORCE，并使用 KL 正则化后的奖励。而在 off-policy 的异步环境中，它展现出远超传统方法的鲁棒性。
传统方法处理这类数据，往往只能硬套复杂且极易失效的重要性采样机制。TB 为 TBA 提供了一种更直接利用 off-policy 轨迹的训练目标，不必完全依赖重要性采样修正。
03 经验池的动态采样调控
Buffer 变大之后，问题不只是有没有数据，而是应该从里面挑哪些数据。纯随机采样效率低下，而单纯依据奖励分数优先采样则容易导致模型输出同质化。
TBA 设计了一套混合调节方案，引入超参数 （Most-On-Policy Probability）。
Trainer 在构建训练 Batch 时，会以概率 挑选那些在最近一次同步中新加入 Buffer 的数据。这部分数据离当前策略最近，可以给训练提供更稳定的参照。
对于剩下的 1-m 概率，系统会在整个历史库中游走，混合使用基于奖励分数的 Softmax 采样和均匀采样。这种设计既不浪费高质量的历史样本，又守住了探索的多样性。
这套采样机制试图在数据新鲜度和探索广度之间做平衡。消融结果显示，数学推理任务对 m 较敏感，较高的最近样本比例通常更稳。
04 核心任务的效能提升
论文在三个后训练任务上验证了 TBA 的速度和性能优势：
1. 数学推理：提速与精度并存
在 4x A100 的计算环境下，TBA 在 GSM8K 任务中表现出显著的效率优势。
相比 VinePPO，TBA 将训练时间缩短了近 50 倍，同时 Pass@1 准确率高出了约 1.2% 至 1.8%。相比 Online DPO，TBA 实现了 1.6 倍加速并提升了 2.0% 的准确率。
2. 偏好微调：重塑帕累托前沿
在 TL;DR 摘要任务中，TBA 在 KL/perplexity 与 win-rate 的权衡上形成了更好的 Pareto 前沿。
在 410M 到 2.8B 等不同模型规模下，TBA 均能在更短时间内达到更高胜率，比经过优化的异步 DPO 基线快了约 3.8 到 5.3 倍。
05 自动红队：稀疏奖励下的扩展能力
自动红队高度依赖采样广度和 prompt 多样性。在非等算力设置下，TBA 相比非分布式同步 GFlowNet 基线显著缩短 wall-clock time，最快约 7×。
结果显示，随着 Searcher 数量增加，攻击成功率和 prompt 多样性整体上升。
TBA 在更大模型和高度 off-policy 设置下是否仍然稳定？
论文进一步在 Qwen 2.5 7B 上测试了变体 TBA′，并与 Dr. GRPO 对比。
这里使用的是 TBA′，也就是基于 PRIME-RL 的简化版本。它基于 PRIME-RL 构建，超参数更少，也支持多 GPU 训练，因此更适合测试较大模型和更长上下文。
在 10 steps off-policy 设置下，Dr. GRPO 的曲线出现明显波动，而 TBA′ 整体更稳定。这进一步说明，TB 目标在高度 off-policy 的异步训练中更有稳定性优势。
07 结语
TBA 的价值，不只是让后训练更快，而是重新组织了采样和学习的关系。
通过引入 Trajectory Balance，TBA 让异步产生的 off-policy 轨迹也能参与训练，使大规模并行探索不再只是“采样更多”，而是能真正转化为学习信号。
不过，trajectory-level 目标也会带来更高的梯度方差。论文的做法是为每个 query 使用更多 responses 来降低方差，例如在 GSM8K 消融中比较了 K=20 和 K=40。
这也意味着，TBA 对 batch 构造和采样策略提出了更高要求。
如果未来的 LLM 后训练越来越依赖大规模 rollout 和稀疏奖励，TBA 这类异步搜索-学习框架，可能会在自我改进训练、搜索式推理等方向变得更重要。
编辑：文婧
校对：林亦霖
