流形老虎机:基于大语言模型潜在几何的贝叶斯课程学习 Manifold Bandits: Bayesian Curriculum Learning over the Latent Geometry of Large Language Models
将LLM问题采样建模为流形结构化老虎机,通过层次树和贝叶斯方法平衡效率、多样性与评估效用。
前置知识
多臂老虎机
多臂老虎机是强化学习中的一个经典问题框架,模拟了在有限行动集合中进行序贯决策的场景。每个臂代表一个行动,选择某个臂会获得一个随机奖励。目标是设计采样策略,在探索未知臂和利用已知好臂之间找到平衡,最大化累计奖励。Thompson采样是一种流行的贝叶斯方法,通过从后验分布中采样来选择行动,自然地平衡探索与利用。
本文将问题采样形式化为多臂老虎机问题,理解这个框架对于掌握本文的核心思想和方法设计至关重要。
课程学习
课程学习是一种受人类教育启发的训练策略,通过按难度从简单到复杂的顺序组织训练数据,加速模型学习。传统课程学习通常基于人工定义的难度标签或外部特征,逐渐增加任务难度,让模型建立坚实的基础后逐步学习更复杂的概念。自适应课程学习则根据模型当前能力动态调整训练样本的难度和分布。
本文提出了一种新的课程学习框架,利用模型的内部表示而非外部标注来组织任务,理解传统课程学习的局限有助于把握本文的创新点。
流形假设
流形假设是机器学习中的一个核心概念,认为高维数据实际上分布在一个较低维的流形上。例如,图像数据的像素空间虽然高维,但meaningful的变化实际上在低维空间中。大语言模型的表示空间也呈现出类似的几何结构,相关的输入在高维空间中往往聚集在局部区域。UMAP等降维算法试图发现这种低维流形结构。
本文的核心思想正是基于流形假设,认为任务在模型的潜在表示空间中形成流形结构,理解这个假设是理解Latent Task Tree和BMC方法的基础。
群体相对策略优化
群体相对策略优化是一类用于大语言模型强化学习的算法,包括GRPO、DAPO、GSPO等。这些方法为每个提示采样多个响应,然后只使用组内统计信息构造对比优势函数。与传统的策略梯度方法不同,它们不需要学习critic,而是基于同一提示下的多个响应之间的相对性能来更新策略。关键假设是,当多个rollout产生相同奖励时,相对优势为零,无法提供梯度信号。
本文的问题动机正是来源于这类方法,它们依赖于奖励方差,理解这个机制对于理解为什么学习信号和问题采样如此重要。
经验贝叶斯
经验贝叶斯是一种通过数据来估计先验分布的贝叶斯方法。传统贝叶斯分析需要指定先验分布,而经验贝叶斯从观测数据中学习先验的参数。在层次化模型中,上层参数的先验往往通过下层数据的统计量来估计。这种方法特别适合处理多个相关问题,可以共享统计强度,提高估计效率。随机效应元分析中的方差估计就是一个经典例子。
本文的bottom-up信念传播采用经验贝叶斯方法,从子节点的后验统计量估计父节点的信念,理解这个方法有助于把握BMC如何实现信息共享。
研究动机
在用强化学习训练大语言模型时,训练效率严重依赖于采样的问题类型。现有的自适应课程学习方法(如DAPO中的动态采样)通常优先考虑中等难度的问题,将问题选择视为标准的多臂老虎机问题,其中每个臂(问题)是独立的。这种框架存在几个关键缺陷。首先,它忽略了任务空间的结构化和异质性——不同类型的问题往往共享推理模式或潜在特征,一个问题的观测应该能提供对相关问题难度估计的信息。其次,非平稳性并非纯粹外生:随着模型改进,问题上的奖励分布直接源于模型自身的更新和采样决策,形成了一个内在反馈循环。第三,训练效率不仅取决于难度,还取决于多样性——仅关注产生即时学习信号的问题会导致对问题空间的覆盖狭窄,可能损害泛化能力。在DAPO-Math-17K数据集上,动态采样虽然通过重复生成批次确保只包含非零奖励方差的例子,但引发了显著的wall-clock时间开销。当rollout产生相同奖励时,相对优势塌缩为零,无法提供策略梯度信号,导致计算浪费。论文数据显示,标准方法的有效比例(有效样本占总样本的比例)较低,而动态采样通过重新采样达到1.0的有效比例,但计算成本大幅增加。
本文的目标是本文的目标是开发一个结构化的自适应课程学习框架,能够在大语言模型的强化学习训练中同时优化三个关键维度:学习信号强度(productivity)、任务流形覆盖(diversity)和下游评估相关性(utility)。具体来说,作者希望利用模型的潜在表示空间来组织问题,构建一个层次化的任务树,在这个树上执行结构化的贝叶斯决策,以高效地探索和利用学习信号。通过这种方法,可以避免传统方法将问题视为独立臂的简化假设,利用问题之间的关联性来提高采样效率。同时,框架应该能够在不要求人工标注、外部嵌入模型或固定全局划分的情况下工作,具有跨领域、跨语言和跨模态的通用性。最终目标是证明,一个同时考虑问题类型和难度的课程学习方法,能够比仅关注难度的方法产生更好的下游性能。
与已有工作不同的是,本文的独特切入角度是将问题采样形式化为一个潜在流形结构化的老虎机问题,而非传统的独立臂老虎机问题。关键洞察在于:采样决策不仅仅是被动测量,它们类似于干预,会改变未来奖励分布。当在一个任务区域进行训练时,学习效应会传播到相关区域,这些区域在模型的潜在表示空间中是接近的。因此,问题选择应该被建模为一个流形上的干预问题,其效果可以在相关问题之间传播。这种内生非平稳性与传统的纯粹外生非平稳性有本质区别。此外,现有方法通常要么专注于难度,要么(如flat clustering方法)引入一些结构,但没有一个方法统一地建模层次结构、非平稳性和信息传播。本文通过Latent Task Tree和BMC填补了这个gap,将问题类型和难度视为互补但独立的因素,而不是将类型仅仅视为实现难度估计的手段。
核心方法
BMC(Bayesian Manifold Curriculum)的整体思路是先构建一个反映模型如何感知任务空间的层次结构(Latent Task Tree),然后在这个树上执行结构化的贝叶斯决策来选择训练问题。直觉是,如果两个问题在模型的潜在表示空间中接近,它们可能共享解决策略、失败模式和学习动态。因此,观测一个问题的学习信号可以提供对相关问题估计的信息。BMC通过三个循环步骤工作:(1)自顶向下遍历Latent Task Tree选择问题;(2)从观测的rollout奖励更新问题级别的信念;(3)将这些更新自底向上传播,使树反映新信息。这个框架自然地平衡了探索和利用:当某个区域看起来更有生产力时,多个agent会进入该区域;当多个区域看起来类似生产力时,独立遍历会分散agent到不同区域。技术路线首先从策略模型的中间层提取每个问题的嵌入,然后用PCA降维、UMAP投影和HDBSCAN聚类递归构建层次树。树的构造由多个停止条件控制,包括Chart Test(检查区域是否可以近似为chart-like),使得树结构能捕捉局部几何而不假设全局平滑流形。
BMC的核心创新点在于将问题类型和难度视为互补但独立的因素,并使用模型自身的潜在表示来组织类型。与现有方法相比,这个方法有几个本质区别。首先,Latent Task Tree不是基于人工标注或外部嵌入,而是基于正在训练的策略模型的表示,确保树结构与模型对任务空间的感知一致。其次,决策是在层次结构上进行的,而不是在平坦的arm集合上。这使得算法可以在不同抽象级别上分配训练努力,并在相关问题之间共享信息。第三,BMC显式建模了非平稳性,通过一个轻量级贝叶斯滤波器,使用surprise(标准化预测误差)来调制先验证据的权重,使信念能够适应模型漂移。第四,自底向上的信念传播使用经验贝叶斯方法,将异质性作为随机效应项纳入考虑,避免在不一致的子区域上过度自信的平均。最后,BMC-T扩展引入了utility作为独立维度,证明学习效率、多样性和评估效用不是自动对齐的,改变目标分布可以改变下游能力概况,同时保持生产力相似。
方法步骤详情
BMC的完整工作流程包含三个主要步骤。第一步是自顶向下的问题选择(层次化Thompson采样)。对于每个batch元素,一个独立的agent从树根开始遍历。在每个有子节点的节点,agent为每个孩子绘制一个奖励估计,并选择最大值对应的孩子。这个过程递归应用,引导agent从粗粒度的高层区域向细粒度的子问题移动,直到到达叶节点(单个问题),该问题被包含在batch中。为了防止重复采样,使用了一个batch感知的层次下降变体。第二步是非平稳信念建模(贝叶斯滤波)。每个叶节点维护一个学习信号的概率信念,初始为高斯模型(实践中使用Logit-Normal参数化)。当观测到新的学习信号时,定义surprise,这是一个标准化的预测误差。这个surprise调制有效先验权重,使得高度意外的观测减少对先验信念的依赖。后验均值更新为有效先验权重加权平均。方差更新平衡了置信度收缩和不确定性注入,其中staleness随着问题未被采样而增加。第三步是自底向上的信念传播(经验贝叶斯)。对于每个非叶节点,使用精度加权估计器聚合其子节点的信念。引入子树级别异质性项,该术语测量子均值之间的变异超过其内部不确定性所预期的程度。使用这个校正,定义精度权重,并计算父节点信念。当子节点一致时,异质性项接近零,聚合表现为普通精度加权;当子节点信念不一致时,异质性项增加父节点不确定性,减少在异质子节点上的过度自信平均。信念传播在每个batch更新后递归地从叶到根进行,创建相关问题之间的软耦合。
技术新颖性
BMC的技术新颖性体现在多个层面。在问题组织方面,Latent Task Tree的递归局部流形近似是一个创新,它不假设固定深度或全局划分,而是通过Chart Test等停止条件自适应地捕捉局部几何。这使得树能够在数据分布的不同区域有不同的密度、内在维度和组织模式。在决策框架方面,BMC将层次结构和非平稳性统一在一个贝叶斯框架中,而不是将它们视为独立的组件。具体来说,层次Thompson采样自然地在不同抽象级别上分配训练努力,贝叶斯滤波器使用surprise来适应模型漂移,经验贝叶斯传播使用异质性项来避免过度平均。这种统一性使得BMC能够在一个原则化的框架中处理问题类型、难度、动态性和信息共享。在理论视角方面,BMC将课程学习从找到正确难度的问题转变为在多样化和交互的问题类型之间编排训练努力。采样决策影响学习如何通过策略的任务空间进展,这与Bitter Lesson的精神一致——从策略学到的表示中导出课程结构,而不是强加人工定义的结构。最后,BMC-T扩展展示了utility作为独立维度的概念创新,证明了学习效率和评估效用不是可互换的,改变目标分布可以改变下游能力概况,即使生产力概况保持相似。这为理解LLM泛化提供了结构化老虎机视角——采样的训练问题作为杠杆,诱导模型更新,其效果可能延伸到采样的提示之外,改变奖励在模型潜在流形的邻近和遥远区域上的分布。
实验结果
论文通过三个维度的实验分析揭示了关键发现。在生产效率方面,图4显示学习信号与学习速度之间存在一致关系。在Qwen-8B-Base上使用GSPO时,训练集pass@1从初始约0.50提升到0.60,对应的effective ratio从约0.6提升到1.0。BMC和Difficulty Only实现了约40%的effective ratio提升(相比uniform采样的约0.6),同时训练时间接近uniform采样。相比之下,Dynamic Sampling虽然达到1.0的effective ratio,但wall-clock时间约30小时,是uniform采样的2倍。在Qwen-4B-Base上使用GRPO时,BMC的训练集pass@1从约0.42提升到0.50,与Dynamic Sampling相似,但训练时间约20小时,接近uniform采样的约10小时,而Dynamic Sampling的wall-clock时间显著更高。在覆盖和信息共享方面,图5的左侧显示rarity-weighted exposure随训练步数变化。Tree Only方法强调多样性,exposure约0.4,而Difficulty Only强调生产力,exposure约0.2。BMC在两者之间插值,在训练中期达到约0.3的exposure。右侧的structure gain显示BMC和Difficulty Only都表现出显著的结构利用,与随机分割相比,gain从初始的约2倍提升到训练后期的约5-10倍。这表明即使当树不用于指导采样时(如Difficulty Only),学习信号也与潜在任务结构相关。Tree Only的structure gain较低,接近随机分割的1倍基线。在评估性能方面,图6显示不同采样策略产生不同的能力概况。Dynamic Sampling在MATH500上表现最强(约60%),但在中文数学基准(CNMO24、CCEE24)和OOD基准GPQA-Diamond上表现较弱(约20-30%)。Difficulty Only在中文数学基准上有所改善(约30-40%),但仍然低于期望。Uniform Sampling在AIME 2025上表现良好(约40%),但缺乏系统性优先级。Tree Only在MATH500上的表现最强(约60%),因为基准的规模大和异质性,但缺乏信息问题优先级导致整体效果有限。BMC在GPQA-Diamond上表现强劲(约40%),表明联合优化生产力和多样性可以促进标准英语数学评估无法捕捉的能力。BMC-T实验显示,当target为AIME2024时,AIME-style英语竞赛数学基准改善,当target为完整评估混合时,性能转向更广泛的评估,包括更强的中文数学性能。这证明了utility作为独立维度的概念。关键指标:在Qwen-8B-Base上,BMC的effective ratio约0.9,training set pass@1约0.60,GPQA-Diamond约40%,MATH500约55%,AIME约35%。Dynamic Sampling的effective ratio为1.0,但wall-clock时间约30小时,GPQA-Diamond约25%。Uniform Sampling的effective ratio约0.6,training set pass@1约0.55,GPQA-Diamond约30%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理训练集学习速度 | Pass@1 on training set | BMC: 0.60 (Qwen-8B-Base), 0.50 (Qwen-4B-Base) | Dynamic Sampling: 0.60 (Qwen-8B-Base), 0.50 (Qwen-4B-Base); Uniform Sampling: 0.55 (Qwen-8B-Base), 0.42 (Qwen-4B-Base) | 与uniform采样相比,BMC和Difficulty Only实现约40%的effective ratio提升,学习速度接近dynamic sampling,但训练时间接近uniform采样 |
| 任务流形覆盖 | Rarity-Weighted Exposure | BMC: ~0.3 (训练中期) | Tree Only: ~0.4; Difficulty Only: ~0.2 | BMC在多样性优先的Tree Only和生产力优先的Difficulty Only之间平衡,实现更好的稀有类型覆盖 |
| 学习信号结构利用 | Structure Gain (相对随机分割) | BMC: ~5-10倍 (训练后期) | Tree Only: ~1倍 (接近随机); Difficulty Only: ~5-10倍 | BMC利用Latent Task Tree传播信息,实现比flat方法更好的结构利用 |
| OOD泛化 | GPQA-Diamond准确率 | BMC: ~40% | Dynamic Sampling: ~25%; Uniform Sampling: ~30%; Difficulty Only: ~35%; Tree Only: ~30% | BMC相比dynamic sampling提升约60%,相比uniform sampling提升约33% |
| 中文数学 | CNMO24 + CCEE24平均准确率 | BMC-T (T=All): ~40% | Dynamic Sampling: ~25%; Difficulty Only: ~35%; BMC: ~35% | BMC-T相比dynamic sampling提升约60%,相比BMC提升约14% |
| AIME竞赛数学 | AIME准确率 | BMC-T (T=AIME2024): ~45% | Uniform Sampling: ~40%; BMC: ~35% | BMC-T相比uniform sampling提升约12.5% |
局限与改进
作者承认的局限性包括:首先,Latent Task Tree是静态的,假设潜在结构在训练期间不显著变化。虽然附录B中的实验显示周期性树重建只产生微小的下游性能变化,但这仍然是一个假设。其次,BMC-T使用潜在接近度作为评估相关性的粗略代理,这不是一个精确的因果模型。它假设在一个任务上训练应该比在遥远任务上有更强的对邻近任务的影响,但这个假设可能不总是成立。第三,论文主要关注数学推理任务,虽然声称方法可以跨领域、语言和模态应用,但在其他领域的实验主要是定性的。第四,BMC-T使用target distribution来偏向采样,这在严格held-out评估设置中可能被视为数据泄露,尽管作者建议可以将target distribution视为验证或开发集。我观察到的额外局限性包括:首先,Chart Test的阈值和HDBSCAN的超参数可能需要在不同数据集上调优,尽管作者声称方法不需要数据集特定的调优。其次,树构造虽然相对高效(相比于RL训练),但仍然需要前向传播来提取嵌入,然后在每个节点上执行PCA、UMAP和HDBSCAN,这可能在大规模数据集上成为瓶颈。第三,surprise和staleness项的具体形式是启发式选择的,缺乏理论保证。第四,论文没有系统地研究tree depth、branching factor等超参数对性能的影响。第五,BMC的评估主要集中在数学基准上,对于其他类型的任务(如代码生成、对话系统)的适用性需要进一步验证。第六,虽然BMC-T展示了utility作为独立维度的概念,但在实际应用中,target distribution的获取可能是一个挑战——在真实场景中,我们可能不知道目标评估分布。
独立分析的弱点
第一个主要弱点是静态树假设。模型的潜在表示空间在训练过程中会演变,导致树结构与模型当前感知之间的不匹配。虽然在训练horizon内这可能影响不大,但对于长期训练可能成为问题。改进方向是开发动态树重建机制,可以周期性地或在检测到显著表示漂移时重新构建树。第二个弱点是surprise和staleness项的启发式形式。虽然直觉合理,但缺乏理论保证或与学习理论的最佳性联系。改进方向是从理论角度推导这些更新规则,例如通过变分推断或在线变分推断框架,确保更新与模型漂移动力学一致。第三个弱点是Chart Test和其他停止条件的启发式性质。例如,intrinsic dimensionality的阈值和k-NN图连接性的选择可能影响树结构。改进方向是开发更鲁棒的停止条件,可能基于统计检验(如假设检验检查子结构是否显著)或信息论准则(如信息增益)。第四个弱点是BMC-T的潜在接近度代理假设。这假设在潜在空间中接近的任务共享学习效应,但这可能不总是成立。改进方向是开发更精确的因果模型,例如使用causal discovery方法从历史数据估计任务之间的因果图,或使用meta-learning方法学习任务相似性和迁移函数。第五个弱点是scalability concern。虽然论文声称树构造相比于RL训练是高效的,但对于非常大的数据集(如数百万个问题),前向传播加PCA、UMAP、HDBSCAN在每个节点上的重复应用可能仍然成为瓶颈。改进方向是开发增量或近似的树构造方法,例如使用在线聚类算法(如online k-means或streaming HDBSCAN),或使用近似最近邻方法来加速UMAP和HDBSCAN。第六个弱点是target distribution的获取。在BMC-T中,需要知道目标评估分布来偏向采样,但在实际应用中这可能不可用。改进方向是开发自动发现或估计目标分布的方法,例如使用人类反馈、主动学习或在线估计方法从模型自身的性能推断目标分布。第七个弱点是缺乏系统性超参数分析。论文没有详细研究tree depth、branching factor、surprise敏感性、staleness速率等超参数的影响。改进方向是进行系统的超参数敏感性分析,开发自动调参方法,或设计对超参数更鲁棒的变体(例如自适应调整surprise的权重)。
未来方向
作者提出的未来工作方向包括:首先,研究潜在结构的演变和动态树重建。虽然当前实验显示固定树在训练horizon内足够有用,但对于更长训练或更大模型变化,可能需要动态更新树。其次,扩展到其他奖励类型。虽然论文关注binary verifier奖励,但BMC框架可以扩展到连续奖励设置,如RLHF和rubric-based RL,只需选择适当的规范化信号。第三,研究更精细的因果模型。BMC-T使用潜在接近度作为因果效应的粗略模型,未来可以探索更精确的因果推断方法,例如从历史训练数据估计任务之间的因果图,或使用instrumental variable方法识别因果效应。第四,扩展到更多领域和模态。虽然附录E提供了数学、代码、医疗、法律、金融、多领域、多语言和多模态设置的定性分析,未来可以进行更系统的定量实验,验证BMC在这些设置中的性能。基于成果可延伸的未来工作方向包括:首先,开发更高效的信息传播机制。当前的经验贝叶斯传播在每个batch后递归地从叶到根更新,可以考虑增量或异步更新,以减少计算开销。其次,研究与其他课程学习方法的结合。例如,可以将BMC与difficulty predictors(如PCL、DOTS)结合,使用预测器作为先验,然后用BMC进行结构调整。第三,研究BMC在蒸馏和few-shot learning中的应用。当前工作关注RL训练,但BMC的思想也可能适用于蒸馏场景,选择最能传递知识的示例,或few-shot learning场景,选择最具信息性的演示。第四,研究BMC对模型可解释性的影响。Latent Task Tree不仅可用于采样,也可作为分析模型如何组织任务空间的工具,可能帮助理解模型的推理能力和失败模式。第五,研究BMC对安全性和对齐的影响。通过控制暴露给模型的问题类型,可能影响模型的行为和偏见,这可能用于对齐和安全性研究。第六,开发自动化的树构建和调优pipeline。当前方法虽然不需要数据集特定调优,但仍有超参数需要选择,可以开发自动选择这些参数的方法,使其更易于应用。第七,研究BMC在多目标场景下的应用。当前工作关注productivity、diversity和utility的平衡,未来可以扩展到更多目标(如fairness、robustness、calibration),开发多目标优化框架。
复现评估
论文的复现性评估如下。开源情况:论文声称所有方法在verl框架中实现,但没有明确提供代码链接或开源计划。这是复现性的一个主要障碍。数据可用性:使用的数据集主要是DAPO-Math-17K,这是一个公开可用的数学推理数据集。评估基准包括AIME、AMC、MATH500、CNMO24、CCEE24和GPQA-Diamond,这些都是公开可用的。论文提供了足够的数据集细节在附录L.4中,包括评估设置和指标。算力需求:论文没有明确报告训练所需的算力,但从上下文可以推断,使用Qwen3-4B-Base和Qwen3-8B-Base进行RL训练,每个训练run大约10-30小时(基于wall-clock时间数据)。这表明需要相对较强的GPU资源(可能多张A100或H100)。实验规模:论文报告了多个实验设置,包括两个模型大小(4B和8B)和两个RL算法(GSPO和GRPO),每个都有多个采样策略比较。这表明实验规模较大,复现所有实验可能需要大量计算。实现细节:论文在附录中提供了相对详细的实现细节,包括Latent Task Tree构造的参数、贝叶斯更新的公式、BMC-T的utility分数计算等。这些细节对于复现是足够的,但缺乏代码使得实际实现仍然具有挑战性。统计分析:论文没有报告置信区间或统计显著性检验,这使得难以评估结果的可信度和重复性。超参数:论文没有系统地报告超参数选择过程或敏感性分析,这使得复现时可能需要大量调优。总体而言,论文在数据集描述和实现细节方面相对详细,但缺乏开源代码和统计分析,这降低了复现性。如果能提供代码、预训练checkpoint或详细超参数配置,复现性将大大提高。估算复现难度:中等偏难。主要挑战在于需要较强的算力资源(训练LLM)和缺乏开源代码。但如果只复现分析部分(如tree construction和结构诊断),难度会降低。
论文图表