先预测下一个词再用强化学习,能让模型更高效学会复杂逻辑推理。
How Reinforcement Learning After Next-Token Prediction Facilitates Learning
- 在短长序列混合数据中,强化学习提升模型泛化能力
- 仅靠预测下一个词需极端资源才能解决比特奇偶性问题
- 适合研究大模型推理能力的提升机制与训练方法
近期神经网络在推理任务中的进展主要依赖于一种训练方法:先用预训练语言模型预测序列中的下一个词,再通过强化学习进行优化。本文提出一个框架,理论分析了该范式为何有效。研究聚焦于包含短和长“思维链”序列的混合分布,以单个任务为例。当任务为预测d个比特的奇偶性且长序列稀少时,我们证明:经过强化学习微调后的自回归Transformer可实现泛化,而仅依赖下一个词预测则需极高的统计或计算资源。此外,强化学习利用测试时增加的计算量(表现为更长的回答)促进学习。在简化设定下,我们理论上证明:只要长示范在数据混合中的比例不随输入维度d呈指数级减小,自回归线性模型就能高效学习比特奇偶性。最后,我们在其他场景中验证了这些现象,包括在常见数学推理基准数据集的混合变体上对Llama系列模型进行后训练。
原文摘要 · Abstract (English)
Recent advances in reasoning domains with neural networks have primarily been enabled by a training recipe that optimizes Large Language Models, previously trained to predict the next-token in a sequence, with reinforcement learning algorithms. We introduce a framework to study the success of this paradigm, and we theoretically expose the optimization mechanisms by which reinforcement learning improves over next-token prediction in this setting. We study learning from mixture distributions of short and long ``chain-of-thought'' sequences encoding a single task. In particular, when the task consists of predicting the parity of $d$ bits and long sequences are rare, we show how reinforcement learning after next-token prediction enables autoregressive transformers to generalize, whereas mere next-token prediction requires extreme statistical or computational resources to do so. We further explain how reinforcement learning leverages increased test-time computation, manifested in longer responses, to facilitate this learning process. In a simplified setting, we theoretically prove that autoregressive linear models following this training recipe can efficiently learn to predict the parity of $d$ bits as long as the proportion of long demonstrations in the data mix is not exponentially small in the input dimension $d$. Finally, we demonstrate these same phenomena in other settings, including the post-training of Llama-series models on mixture variations of common mathematical reasoning benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。