慢速思考、快速思考:通过蒸馾推理者扩展推理计算
计算与语言
2025-02-28 v1 人工智能
摘要
最近的进展表明,扩大大型语言模型(LLM)在测试时间的计算资源可显著提升其性能。常见策略是生成多个链式思维(Chain-of-Thought, CoT)轨迹并通过各种选择机制对其输出进行聚合。这引出了一个根本性问题:对于固定的计算预算,规模较小的模型能否利用其更高的生成吞吐量超越同样大小的 Transformer?为回答此问题并克服缺乏强大亚二次方 reasoner 的局限,我们从预训练的 Transformer 中蒸馾出纯 Mamba 模型和混合 Mamba 模型。只需训练 80 亿个 token,我们的蒸馾模型在数学推理数据集上表现出强大的性能和扩展潜力,同时在大批量和长序列场景下的推理速度显著快于 Transformer。尽管蒸馾导致零样本性能下降,但无论是纯 Mamba 模型还是混合 Mamba 模型,都能在固定时间预算下将其覆盖范围和准确性性能超越其 Transformer 教师模型,为扩展推理计算开辟了新的方向。
引用
@article{arxiv.2502.20339,
title = {Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners},
author = {Daniele Paliotta and Junxiong Wang and Matteo Pagliardini and Kevin Y. Li and Aviv Bick and J. Zico Kolter and Albert Gu and François Fleuret and Tri Dao},
journal= {arXiv preprint arXiv:2502.20339},
year = {2025}
}