平衡长存：信息瓶颈驱动的树基策略优化

问题

问题背景

在线强化学习（online RL）已成为提升大语言模型复杂推理能力的重要范式，但探索与利用的平衡（exploration-exploitation trade-off）始终是核心痛点。

现有方法局限

当前主流方法（如 GRPO ）通常依赖经验性正则化（如 KL 散度约束）来避免策略偏离，但缺乏对探索-利用平衡的细粒度、可量化评估。

• 训练过程中，策略的推理多样性（探索）与正确答案互信息（利用）无法被实时监控，导致优化不稳定。
• 常见正则化器仅能抑制过度的策略更新，却无法主动引导模型在适当阶段增强探索或收敛。
• 采样效率低下：现有在线采样方式无法在固定token预算下生成足够多样且高回报的轨迹。

难点与重要性

关键在于如何定义并在线估计一个可优化的平衡度量指标。信息瓶颈理论提供了自然框架，但将其转化为可操作的树搜索目标面临两大挑战：

1. 量化难：需同时建模步骤级推理多样性（多样性熵）与正确性相关信息，计算复杂度高。
2. 采样效率瓶颈：在线训练中每个token预算都宝贵，如何低成本获取大量高质量样本并复用结构进行蒙特卡洛估计是工程核心。 该问题的解决直接关系到 LLM 在数学、编程等需要长链推理场景中的训练稳定性和最终表现，受到工业界和学术界高度关注。

行业类比

如同自动驾驶路径规划中，智能体需在尝试新路线（探索）与利用已知最短路径（利用）间动态权衡，LLM 的推理策略也需要可控的平衡调节器，这正是 IB-TPO 试图提供的。

核心洞察

• **IB-Score 作为探索-利用平衡的细粒度度量**：在线 RL 微调 LLM 时，现有方法（如 GRPO）常用 KL 散度或奖励裁剪等粗粒度正则，缺乏直接量化探索与利用平衡的工具。IB-Score 基于信息瓶颈理论，通过计算步骤级推理多样性与正确答案互信息之间的折衷，提供了可监控、可优化的数值指标，使训练中的失衡能被及时发现并矫正。
• **IB 引导的树采样提升采样效率与估计质量**：传统采样按 trajectory 独立生成，耗费大量 token 预算且无法复用中间步骤。IB-TPO 的树采样策略在相同 token 预算下生成 50% 更多的有效轨迹，并且树形结构天然支持 IB-Score 的蒙特卡洛估计，无需额外计算开销，既提升了数据多样性又保证了平衡优化的稳定性，在多个基准上显著超越 GRPO 与其他 SOTA 在线 RL 方法。

方法

IB-TPO 方法：信息瓶颈驱动的树策略优化

输入：语言模型策略 $\pi_\theta$，推理问题的提示 $x$ 与最终正确答案 $y$，以及用于采样的 token 预算。

核心流程按“采样 → 评估 → 优化”展开：

1. IB 引导的树采样 (Tree Sampling)：

   • 从提示 $x$ 出发，以当前策略 $\pi_\theta$ 生成中间推理步骤，构建一棵部分展开的蒙特卡洛树 (MCTS)，每个节点代表一个中间状态。
   • 利用信息瓶颈得分 (IB-Score) 指导树的扩展与剪枝：在每个决策点，优先选择既保持推理多样性又能有效传递与正确答案 $y$ 相关信息的路径。
   • 相比 GRPO 等方法的独立轨迹采样，该方式在相同 token 预算下可多生成 50% 的有效轨迹，并自然形成树状结构，便于后续重用。

2. IB-Score 计算：

   • 从信息瓶颈理论出发，将策略的探索-利用平衡量化为：
     • 步骤级推理多样性：不同分支覆盖的语义空间广度；
     • 与正确答案的互信息：生成的推理片段对预测正确答案 $y$ 的信息贡献。
   • 基于树结构，通过蒙特卡洛估计高效近似 IB-Score，避免对整个策略分布做精确计算。
   • 该指标能诊断训练各阶段是否出现探索过度（多样性高但信息量低）或利用过度（信息量高但多样性坍塌） 的问题。

3. 策略优化：

   • 将 IB-Score 转化为细粒度优化目标，直接作为强化学习过程中的额外信号，与原有的奖励函数结合，引导策略在提升任务性能的同时维持健康的探索-利用均衡。
   • 优化时重用树采样中的状态价值估计，减少重复前向传播，训练更稳定。

输出：经过多轮在线训练后的优化策略 $\pi_{\theta^*}$。

与同类方法的差异：IB-TPO 首次将信息瓶颈原理直接嵌入策略优化的采样与目标设计，而非像 GRPO 等仅依赖 KL 散度正则项来被动约束策略更新，从而更有效地动态维持探索-利用平衡，同时大幅提升采样效率。

实验

实验设计

本文在多个标准推理基准上对比了 IB-TPO 与 GRPO 等在线 RL 方法。实验统一控制 token 预算，将提出的 IB-Score 作为细粒度优化目标，并启用 IB 引导的树采样策略 生成轨迹。该策略在相同 token 预算下可多产生 50% 的有效轨迹，同时树结构本身被复用于 IB-Score 的蒙特卡洛估计，降低方差。评估指标聚焦于任务准确率与采样效率。

关键发现

• IB-TPO 能持续维持探索–利用平衡，训练过程更稳定，避免了 GRPO 等常见方法在中后期出现的性能退化。
• 在多个基准上，IB-TPO 的准确率较 GRPO 基线提升 2.9% 至 3.6%，并显著优于其他主流在线 RL 方法。
• 采样效率优势明显：相同 token 预算下轨迹数量增加 50%，表明信息瓶颈驱动的树采样既提升数据利用率，又保持推理多样性。

与基线的对比解读

传统在线 RL（如 GRPO）多依赖 KL 正则化 控制策略更新幅度，但难以自适应地平衡探索与利用，常陷入次优策略。IB-TPO 从信息瓶颈视角出发，将平衡性量化为 IB-Score，并直接作为优化目标，赋予了更精准的控制力。树采样策略不仅提高了在线采样的效率，其拓扑结构天然适合进行无偏的蒙特卡洛估计，避免了额外开销。这种系统性的联合设计使得 IB-TPO 在稳定性和最终性能上均明显超越现有方法，且无需复杂的超参数调节，工程部署友好。

局限

• - 计算与实现复杂度较高：IB-TPO 采用了树采样和 IB-Score 蒙特卡洛估计，相比 GRPO 等单轨迹采样方法，实现更为复杂，且在并行计算中可能增加显存开销。尽管在相同 token 预算下声称提升了 50% 的轨迹数量，但树结构的构建与维护在工程落地时仍需额外优化，对于资源受限的场景可能不够友好。
• - 泛化性有待验证：论文实验主要集中在标准推理基准（如数学或逻辑推理），缺少在开放式生成、多轮对话或代码生成等任务上的评估。信息瓶颈度量在步骤级推理多样性上的有效性，可能依赖于任务具有明确的多步推理结构，对于低结构化任务，该度量是否仍能指导最优探索-利用平衡尚不清楚。
• - 度量近似精度风险：IB-Score 的估计依赖于树采样策略获取的轨迹分布，蒙特卡洛近似可能存在偏差，尤其是在树结构未能充分覆盖推理空间时。若 IB-Score 本身误差较大，以此作为优化目标反而可能误导策略更新，导致训练不稳定或退化。论文未详细分析估计器的统计性质和偏差范围。