用蒸馏让轻量模型在有限时间内跑出比大模型更准的推理结果。
Thinking Slow, Fast: Scaling Inference Compute with Distilled Reasoners
- 从Transformer蒸馏纯Mamba和混合Mamba模型，提升生成效率。
- 在固定时间预算下，轻量模型推理覆盖率达92.3%，准确率超教师模型。
- 适合追求高吞吐、低延迟推理的部署场景，如在线问答系统。
近期研究表明，通过在测试时扩展计算资源可显著提升大语言模型性能。常见策略是生成多个思维链（CoT）轨迹，并通过不同选择机制聚合输出。这引出一个核心问题：复杂度更低的模型能否凭借更高的生成吞吐量，在固定计算预算下超越同等规模的Transformer？为解决此问题并克服缺乏强亚二次复杂度推理模型的瓶颈，我们从预训练Transformer中蒸馏出纯Mamba和混合Mamba模型。仅用80亿词训练后，这些蒸馏模型在数学推理数据集上表现优异且具备良好可扩展性，同时在大批量和长序列推理中速度更快。尽管蒸馏导致零样本性能下降，但在固定时间预算下，纯与混合Mamba模型仍能实现超过其Transformer教师模型的覆盖率与准确率，开辟了推理计算扩展的新路径。
原文摘要 · Abstract (English)
Recent advancements have demonstrated that the performance of large language models (LLMs) can be significantly enhanced by scaling computational resources at test time. A common strategy involves generating multiple Chain-of-Thought (CoT) trajectories and aggregating their outputs through various selection mechanisms. This raises a fundamental question: can models with lower complexity leverage their superior generation throughput to outperform similarly sized Transformers for a fixed computational budget? To address this question and overcome the lack of strong subquadratic reasoners, we distill pure and hybrid Mamba models from pretrained Transformers. Trained on only 8 billion tokens, our distilled models show strong performance and scaling on mathematical reasoning datasets while being much faster at inference for large batches and long sequences. Despite the zero-shot performance hit due to distillation, both pure and hybrid Mamba models can scale their coverage and accuracy performance past their Transformer teacher models under fixed time budgets, opening a new direction for scaling inference compute.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。