用记忆增强让小模型也能高效推理
Reasoning Under 1 Billion: Memory-Augmented Reinforcement Learning for Large Language Models
- 用记忆存储成功推理路径,引导小模型探索
- 在GSM8K和AI-MO上显著提升小模型样本效率
- 适合资源受限场景下提升小模型推理能力
近期基于强化学习微调大语言模型在复杂推理任务中表现优异,尤其结合思维链提示时效果显著。然而,这些成果主要集中在参数量达数十亿的大模型上,其强大的预训练基础保障了有效探索。相比之下,参数量在10亿以下的小型模型因缺乏足够预训练支持,难以有效探索,常导致次优推理模式。本文提出一种新型内在动机方法,利用情景记忆解决该问题,提升小型模型的思维链推理能力。受人类记忆驱动学习启发,该方法将成功的推理模式存入记忆,并允许可控探索生成新解法。通过kNN实现的情景记忆可高效计算内在奖励,使模型既能发现新策略,又能快速复用已有有效方案。在GSM8K和AI-MO数据集上的实验表明,该方法显著提升了小型模型的样本效率与泛化能力,使基于强化学习的推理优化更适用于低资源环境。
原文摘要 · Abstract (English)
Recent advances in fine-tuning large language models (LLMs) with reinforcement learning (RL) have shown promising improvements in complex reasoning tasks, particularly when paired with chain-of-thought (CoT) prompting. However, these successes have been largely demonstrated on large-scale models with billions of parameters, where a strong pretraining foundation ensures effective initial exploration. In contrast, RL remains challenging for tiny LLMs with 1 billion parameters or fewer because they lack the necessary pretraining strength to explore effectively, often leading to suboptimal reasoning patterns. This work introduces a novel intrinsic motivation approach that leverages episodic memory to address this challenge, improving tiny LLMs in CoT reasoning tasks. Inspired by human memory-driven learning, our method leverages successful reasoning patterns stored in memory while allowing for controlled exploration to generate novel responses. Intrinsic rewards are computed efficiently using a kNN-based episodic memory, allowing the model to discover new reasoning strategies while quickly adapting to effective past solutions. Experiments on fine-tuning GSM8K and AI-MO datasets demonstrate that our approach significantly enhances smaller LLMs' sample efficiency and generalization capability, making RL-based reasoning improvements more accessible in low-resource settings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。