Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners

作者: Daniele Paliotta, Junxiong Wang, Matteo Pagliardini, Kevin Y. Li, Aviv Bick, J. Zico Kolter, Albert Gu, François Fleuret, Tri Dao

发布时间: 2025-02-28

来源: arxiv

研究方向: 机器学习，特别是大语言模型（LLMs）的推理计算优化和推理能力提升

主要内容

本文研究了如何通过提升推理计算能力来优化LLMs的推理性能，特别关注于利用低复杂度模型通过其更快的生成吞吐量来超越相同规模的Transformers，从而在固定计算预算下提升性能。

1. 提出了从预训练Transformers中提取知识以训练混合和纯Mamba模型的方法，以探索低复杂度模型的推理能力。

2. 开发了一种将推理时间扩展到推理任务的策略，通过生成多个思维链（CoT）来提高推理性能。

3. 实现了高效的推理计算优化，使低复杂度模型在固定时间内实现更高的覆盖率和准确率。

4. 展示了Mamba模型在数学推理任务中的优势，特别是在处理大型批次和长序列时。

5. 提出了监督微调（SFT）技术，以进一步提高通过蒸馏得到的模型性能。

1. 知识蒸馏

2. Mamba架构设计

3. CoT生成和选择机制

4. 评估推理时间的覆盖率和准确性

5. 监督微调（SFT）

实验结果表明，蒸馏后的Mamba模型在大多数时间预算下在数学推理任务中（如MATH和GSM8K）比Transformers教师模型具有更高的覆盖率和准确率。此外，这些模型在生成大量完成项方面具有更高的吞吐量，从而在相同时间内提供更多样本。

未来的工作将集中在进一步探索Mamba和其他低复杂度模型在推理计算扩展方面的潜力，以及开发更有效的知识蒸馏和微调技术。此外，将研究这些模型在处理对话和主观任务时的表现，以实现更高的性能和速度。