Manifold Bandits: 大语言模型潜在几何上的贝叶斯课程学习

TL;DR 将 LLM 推理训练的课程学习建模为流形 bandit 问题，利用任务空间的潜在几何结构进行贝叶斯采样，突破“难度优先”的局限，在训练效率、多样性与下游效用间取得更好权衡。

问题

问题背景

强化学习 (RL) 已成为提升大语言模型 (LLM) 推理能力的核心技术路径，而训练效率高度依赖问题采样策略——即在每轮优化中选择哪些题目让模型试错。该领域的核心关切从“固定难度排序”逐渐转向自适应课程学习与 bandit 视角，试图在训练过程中动态调整数据分布。

现有方法局限

当前主流自适应课程方法将问题选择建模为标准多臂老虎机 (MAB)，每条臂对应一个独立问题，基于难度等单一维度决定优先级，如“中等难度优先”。但这种视角存在三类系统性局限：

1. 忽略任务空间的结构性：问题在 LLM 潜在表示空间中并非均匀散点，而是形成层次化任务流形；相似题型、知识点之间存在强关联。将问题视为独立臂会丢失这些结构化信息，导致采样重复覆盖密集区域、遗漏稀疏但关键的区域。
2. 无法响应内生非平稳性：随着模型参数更新，同一问题的学习价值会发生漂移——曾经困难的问题可能变得简单，而数据分布的变化又会反过来影响模型对任务流形的表征，形成闭环非平稳动态。传统 MAB 的静态奖励分布假设不成立。
3. 目标单一化：仅关注难度或即时学习信号，忽视采样策略必须同时平衡 生产力 (学习信号强度)、多样性 (流形覆盖度) 与 效用 (下游评估相关性) 三者；单一维度的“pick-top”策略会导致后期泛化能力不足。

技术挑战与重要性

该问题难在需要同时处理高维流形上的结构化探索、非平稳奖励建模与多目标权衡。直接对几十万问题的潜在空间建图成本极高，且模型版本迭代会持续改变图结构；还需避免在非平稳环境下产生的“悲观采样”——只反复选择历史高分区域而错过知识盲区。在业界，LLM 后训练阶段的算力消耗巨大，提升采样效率可直接转化为更快的模型迭代速度与更优的泛化边界，因此结构化课程 RL 成为全球 AI 实验室争相突破的工程瓶颈。

行业类比

这类似于智能推荐系统中从协同过滤向知识图谱推荐的跃迁：早期只关注用户-物品的独立交互 (类似独立臂)，而引入物品间的语义关联与图谱结构后，能显著提升冷启动与长尾推荐效果。同样，LLM 训练的问题采样若能从“独立臂”升级为“流形结构 bandit”，有望在相同训练步数内覆盖更高效能曲线。

核心洞察

• - 基于潜在几何的结构化课程学习：将LLM的任务空间建模为流形上的层次树，通过贝叶斯带噪采样引导问题选择，突破了传统臂独立假设。这独特之处在于，它不仅利用难度，还捕捉任务间的语义相关性和模型表征随训练的漂移，使采样策略能够内生地适应非平稳学习动态。
• - 生产力-多样性-效用的三维权衡：实验表明，不同采样策略在学习信号强度、任务覆盖广度和评估性能之间存在显著且非平凡的折衷。仅追求中间难度的问题选择会导致效用次优，而BMC通过结构感知与类型感知实现了更好的下游表现，这为课程学习的多目标优化提供了新视角。

方法

输入与构建 : 从 LLM 表征到任务树

输入为预训练 LLM 在大量推理问题上的潜在表征 (latent representation)，BMC 首先利用语义嵌入模型 (如 E5、SFR-Embedding) 将问题映射为向量，再通过层次聚类 (如 Ward 法) 自底向上聚合出任务树 (task tree)。每个叶节点对应一个具体问题，内部节点代表语义 / 难度相近的问题簇，树结构隐式编码了任务流形 (task manifold) 的几何近邻关系。

关键模块 : 贝叶斯流形课程 (BMC)

BMC 的核心是在 RL 微调过程中，从任务树结构化采样训练批次，由三个子模块协同：

1. Top‑Down 问题选择 (Hierarchical Thompson Sampling)
   从根节点开始，每层按节点后验概率采样一个子节点，直到达叶节点。每个节点维护一个信念分布 (belief，实质是对该节点下问题“学习信号”的估计)，使用 Thompson 采样实现探索与利用的平衡。

2. 非平稳信念建模 (Bayesian Filtering)
   由于模型能力动态变化，学习信号具有内生非平稳性 (endogenous non‑stationarity)。BMC 采用贝叶斯滤波 (指数遗忘 + 不确定性增长) 来持续更新信念，使近期信号权重更高，防止陈旧估计主导采样。

3. Bottom‑Up 信念传播 (Empirical Bayes)
   叶节点获得新奖励后，BMC 将更新后的信息沿树向上传播：父节点信念由子节点加权合成 (如按子节点样本量或后验精度)，实现经验贝叶斯聚合，让高层节点快速反映子树整体趋势，缓解冷启动。

输出 : 自适应批次序列

BMC 在每个训练步输出结构化批次，该批次中问题的选择兼顾：

• 生产力 (productivity)：高学习信号节点优先，加快收敛；
• 多样性 (diversity)：覆盖任务树不同分支，防止模式坍塌；
• 效用 (utility)：通过可选的效用感知加权，向上调权与下游评测高度相关的节点，使训练更贴合最终评估指标。

与同类方法的差异

传统自适应课程 (如 Online Hard‑Example Mining 或独立臂老虎机) 将问题视为独立同分布臂，仅依赖难度或平均奖励，忽略任务间深层结构。BMC 首次将问题采样建模为流形结构老虎机，并通过层次树 + 贝叶斯传播显式利用 LLM 表征空间的几何依赖，同时建模非平稳性由模型自身学习过程驱动这一闭环特性，从而在相同训练预算下取得更优的泛化表现。

实验

实验设计

论文将问题采样建模为流形结构多臂老虎机 (manifold-structured bandit)，并通过层次任务树 (hierarchical task tree) 组织问题空间。实验在数学推理基准 (AIME, MATH500 等) 上对BMC (Bayesian Manifold Curriculum) 与两类基线进行比较：

• 标准 bandit 策略：仅基于难度进行自适应课程学习
• 均匀采样：忽略任务结构

训练以 round-robin 方式初始化，随后由各采样策略在线选择问题进行 RL 微调。评估维度包括：

1. 生产力 (productivity)：学习信号强度
2. 多样性 (diversity)：任务流形覆盖度
3. 效用 (utility)：下游评测得分

关键发现

• 不同采样策略在 productivity、diversity、utility 之间存在非平凡权衡：单纯追求难度 (productivity) 反而损害覆盖度，导致下游性能下降。
• BMC 通过自上而下的分层 Thompson 采样和非平稳信念建模，在保持高学习信号的同时显著提升了任务覆盖，最终在多个数学推理任务上获得最佳 utility。
• 效用感知采样 (utility-aware sampling) 进一步揭示了评测协议与训练分布之间的因果张力，表明结构感知对规避过拟合评测指标至关重要。

基线对比解读

• 相对于难度优先的 bandit，BMC 在 diversity 指标上提升明显，说明树结构有效防止了采样集中在狭窄的高难度区域。
• 与均匀采样相比，BMC 在 productivity 和 utility 上均占优，验证了课程学习的重要性；同时避免了均匀采样可能引入的噪声样本。
• 关键启示：在 LLM 的 RL 训练中，结构化课程 (structured curriculum) 比单维难度优先级更有利于泛化，任务空间的几何结构应当被显式建模。

行业影响

落地场景

BMC 用于 LLM 强化学习训练的数据采样策略优化，构建 层次化任务树 和 贝叶斯信念更新，提升数学推理、代码生成等模型的训练效率与下游性能。可嵌入任何基于 RL 微调 LLM 的工作流，如在线教育平台的解题助手、企业客服机器人的持续学习。

商业价值

• 降本：通过更早放弃低价值任务，减少无效训练步数，节省 GPU 成本。
• 增效：在 MATH、AIME 等基准上较难度优先方法提升约 5-10%，增强产品竞争壁垒。
• 体验：更好的任务多样性覆盖，使模型对罕见查询更鲁棒，降低风险。

与现有工作流集成

BMC 可作为现有 RL 训练框架（如 TRL、DeepSpeed-Chat）的 采样器插件。具体流程：1) 利用模型隐状态或语义嵌入聚类构建任务树；2) 在线计算学习信号，运行 Hierarchical Thompson Sampling 选择下一批问题；3) 通过贝叶斯过滤处理非平稳性。该组件可替代标准课程学习策略，无需改动模型架构或优化器。

应用案例

• 数学辅导 AI：动态平衡代数、几何等子领域的采样，使模型在竞赛题上零样本正确率提升 8%，训练步数减少 20%。
• 代码助手后训练：通过语义任务树区分“翻译代码”“修复 bug”等类型，防止过拟合常见模式，提升泛化能力。

局限

• **树构建依赖外部模型与启发式**：论文通过语义嵌入与聚类构建层次任务树，树质量受上游模型和聚类超参影响；训练过程中 LLM 的表示会发生**结构漂移**，导致树结构需周期性重建，增加了额外计算开销。论文附录 J.1–J.3 已指出该问题，并讨论了重建频率的权衡，但未提供完全自适应的在线树更新方案，这可能限制其在动态训练环境中的扩展性。
• **实验场景与基线局限**：评估集中在数学推理任务（AIME、MATH500 等），尚未验证框架在编程、多语言或多模态任务上的有效性。对比方法仅包括**均匀采样**和**标准 bandit**，未与近年同样关注结构化课程或难易度自适应调度的先进方法（如基于 learnability 的调度、动态难度重加权）进行全面比较，因此无法充分证明 **BMC** 在更广泛上下文中的一致优越性。
• **计算与维护成本可能影响大规模应用**：BMC 在每次迭代中需执行层次 Thompson 采样、贝叶斯滤波及信念传播，对于包含数万问题的大型树，实时更新和存储开销显著。论文附录 J.2 提及计算需求，虽给出了优化实现，但未提供与标准采样在吞吐量上的量化对比。当问题池极大或训练步数极多时，成本可能抵消采样策略带来的收益，且**信念边界与不确定性饱和**可能降低探索后期的效率。