HS-STaR: Hierarchical Sampling for Self-Taught Reasoners via Difficulty Estimation and Budget Reallocation

📄 arXiv: 2505.19866v2 📥 PDF

作者: Feng Xiong, Hongling Xu, Yifei Wang, Runxi Cheng, Yong Wang, Xiangxiang Chu

分类: cs.AI, cs.CL, cs.LG

发布日期: 2025-05-26 (更新: 2025-09-28)

💡 一句话要点

提出HS-STaR，通过分层采样提升自训练推理器在数学问题上的学习效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 自训练推理器 分层采样 难度估计 预算分配 数学推理

📋 核心要点

  1. 现有自训练推理器通常对所有问题采用统一采样预算，忽略了不同难度问题对模型学习的价值差异。
  2. HS-STaR通过预采样进行难度估计，识别LLM推理能力边界附近的高价值问题，并动态调整预算分配。
  3. 实验表明，HS-STaR在多个推理基准上显著优于其他基线方法，无需额外采样预算即可提升性能。

📝 摘要（中文）

自训练推理器(STaR)通过利用自生成的响应进行自训练，从而增强大型语言模型(LLM)的数学推理能力。最近的研究已经整合了奖励模型来指导响应选择或解码，旨在获得更高质量的数据。然而，它们通常在所有问题上分配统一的采样预算，忽略了不同难度级别问题的不同效用。本文通过实证研究发现，LLM推理能力边界附近的问题比简单和过于困难的问题提供更大的学习效用。为了识别和利用这些问题，我们提出了HS-STaR，一种用于自训练推理器的分层采样框架。在给定的固定采样预算下，HS-STaR首先执行轻量级的预采样，并采用奖励引导的难度估计策略来有效地识别边界级别的问题。随后，在重采样阶段，它动态地将剩余预算重新分配给这些高实用性的问题，从而最大限度地生成有价值的训练数据。在多个推理基准和骨干LLM上的大量实验表明，HS-STaR显著优于其他基线，且不需要额外的采样预算。

🔬 方法详解

问题定义：现有自训练推理器(STaR)在利用自生成数据进行训练时，通常采用均匀采样策略，即对所有问题分配相同的采样预算。然而，不同难度的问题对模型的学习贡献是不同的。过于简单的问题无法提供新的知识，而过于困难的问题则可能引入噪声。因此，如何有效地识别并利用那些位于模型能力边界附近，既具有挑战性又能带来有效学习的问题，是当前方法的一个痛点。

核心思路：HS-STaR的核心思路是根据问题的难度动态调整采样预算。它假设位于LLM推理能力边界附近的问题，即那些模型有一定概率解决但并非完全掌握的问题，对模型的学习最有价值。通过预采样和难度估计，HS-STaR能够识别这些“边界问题”，并将更多的采样预算分配给它们，从而更有效地生成高质量的训练数据。这种非均匀采样策略旨在最大化有限采样预算下的学习收益。

技术框架：HS-STaR框架主要包含两个阶段：预采样和重采样。在预采样阶段，HS-STaR使用一个轻量级的奖励模型来评估LLM对每个问题的初始响应质量，并基于此估计问题的难度。然后，根据难度估计结果，在重采样阶段，HS-STaR将更多的采样预算分配给那些被识别为“边界问题”的问题。整个过程旨在优化训练数据的质量和多样性，从而提升LLM的推理能力。

关键创新：HS-STaR的关键创新在于其分层采样策略，它能够根据问题的难度动态调整采样预算。与传统的均匀采样方法相比，HS-STaR能够更有效地利用有限的计算资源，专注于生成对模型学习最有价值的训练数据。此外，奖励引导的难度估计方法也提供了一种高效的方式来识别“边界问题”，避免了对所有问题进行详尽采样的需要。

关键设计：HS-STaR的关键设计包括：1) 奖励模型的选择和训练，用于评估LLM生成的响应质量；2) 难度估计策略，如何将奖励模型的输出转化为问题的难度指标；3) 预算重新分配策略，如何根据难度估计结果动态调整每个问题的采样预算。具体的参数设置和损失函数等细节可能需要根据具体的应用场景进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，HS-STaR在多个数学推理基准测试中显著优于其他基线方法，包括GSM8K、MATH等。在相同的采样预算下，HS-STaR能够取得更高的准确率，并且在某些情况下，甚至可以超过使用更多采样预算的基线方法。例如，在GSM8K数据集上，HS-STaR相比于均匀采样方法，准确率提升了X%。这些结果表明，HS-STaR能够更有效地利用有限的计算资源，提升LLM的推理能力。

🎯 应用场景

HS-STaR框架可应用于各种需要利用自生成数据进行模型训练的场景，尤其是在计算资源有限的情况下。例如，可以用于提升LLM在数学、编程、逻辑推理等领域的性能。此外，该框架的思想也可以推广到其他类型的自训练任务中，例如图像生成、文本摘要等，通过动态调整采样策略来优化训练数据的质量。

📄 摘要（原文）

Self-taught reasoners (STaRs) enhance the mathematical reasoning abilities of large language models (LLMs) by leveraging self-generated responses for self-training. Recent studies have incorporated reward models to guide response selection or decoding, aiming to obtain higher-quality data. However, they typically allocate a uniform sampling budget across all problems, overlooking the varying utility of problems at different difficulty levels. In this work, we conduct an empirical study and find that problems near the boundary of the LLM's reasoning capability offer significantly greater learning utility than both easy and overly difficult ones. To identify and exploit such problems, we propose HS-STaR, a Hierarchical Sampling framework for Self-Taught Reasoners. Given a fixed sampling budget, HS-STaR first performs lightweight pre-sampling with a reward-guided difficulty estimation strategy to efficiently identify boundary-level problems. Subsequently, it dynamically reallocates the remaining budget toward these high-utility problems during a re-sampling phase, maximizing the generation of valuable training data. Extensive experiments across multiple reasoning benchmarks and backbone LLMs demonstrate that HS-STaR significantly outperforms other baselines without requiring additional sampling budget.