Distilling System 2 into System 1

📄 arXiv: 2407.06023v3 📥 PDF

作者: Ping Yu, Jing Xu, Jason Weston, Ilia Kulikov

分类: cs.CL, cs.AI

发布日期: 2024-07-08 (更新: 2024-07-24)

💡 一句话要点

通过自监督蒸馏将System 2推理能力迁移至System 1，提升LLM效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 大型语言模型 推理加速 自监督学习 System 1 System 2 计算效率 模型压缩

📋 核心要点

  1. 现有System 2方法通过增加计算量提升LLM性能，但推理效率较低，限制了实际应用。
  2. 论文提出自监督蒸馏方法，将System 2的推理能力迁移到System 1，无需中间推理步骤。
  3. 实验表明，该方法在提升性能的同时，降低了推理成本，优于原始System 1模型。

📝 摘要（中文）

大型语言模型(LLMs)可以在推理过程中花费额外的计算资源来生成中间思考步骤，从而产生更好的最终结果。自从Chain-of-Thought (Wei et al., 2022) 提出以来，许多System 2技术，如Rephrase and Respond (Deng et al., 2023a)，System 2 Attention (Weston and Sukhbaatar, 2023)和Branch-Solve-Merge (Saha et al., 2023)相继被提出。本文研究了自监督方法，将System 2技术产生的高质量输出“编译”（蒸馏）回LLM的生成过程中，无需中间推理token序列，因为这种推理已经被蒸馏到System 1中。我们证明了几种这样的技术可以成功地被蒸馏，与原始System 1的性能相比，结果得到了改善，并且推理成本低于System 2。我们认为，这种System 2蒸馏将是未来持续学习AI系统的重要特征，使其能够将System 2的能力集中在它们还不能很好地完成的推理任务上。

🔬 方法详解

问题定义：论文旨在解决大型语言模型推理过程中，System 2方法（如CoT）虽然能提升性能，但计算成本高、推理速度慢的问题。现有方法难以在效率和性能之间取得平衡，限制了LLM在资源受限场景下的应用。

核心思路：论文的核心思路是通过知识蒸馏，将System 2模型的推理能力“压缩”到System 1模型中。具体来说，利用System 2模型生成的高质量输出作为训练数据，训练System 1模型，使其能够直接生成高质量的答案，而无需显式的中间推理步骤。这样，System 1模型就能在保持较高性能的同时，显著降低推理成本。

技术框架：整体框架包含两个阶段：1) System 2模型生成阶段：利用各种System 2技术（如CoT、Rephrase and Respond等）对输入问题进行推理，生成高质量的答案。2) System 1模型蒸馏阶段：使用System 2模型生成的答案作为监督信号，训练System 1模型。训练目标是使System 1模型能够直接生成与System 2模型相似的答案，从而学习到System 2模型的推理能力。

关键创新：论文的关键创新在于提出了一种自监督的蒸馏方法，能够有效地将System 2模型的推理能力迁移到System 1模型中。与传统的知识蒸馏方法不同，该方法不需要人工标注的中间推理步骤，而是直接利用System 2模型生成的答案作为监督信号，降低了标注成本。此外，该方法还能够灵活地集成各种System 2技术，进一步提升System 1模型的性能。

关键设计：论文中关键的设计包括：1) 使用System 2模型生成的高质量答案作为训练数据；2) 采用合适的损失函数，例如交叉熵损失或序列到序列损失，来训练System 1模型；3) 可以探索不同的网络结构和训练策略，以进一步提升System 1模型的性能。具体的参数设置和网络结构选择可能需要根据具体的任务和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文实验表明，通过System 2蒸馏，System 1模型在多个任务上取得了显著的性能提升，同时推理成本远低于System 2模型。具体性能数据未知，但强调了在性能提升的同时降低了推理成本，实现了效率和准确率的平衡。

🎯 应用场景

该研究成果可广泛应用于各种需要快速推理和低计算成本的场景，例如移动设备上的智能助手、边缘计算环境下的自然语言处理应用等。通过将复杂的推理过程蒸馏到更小的模型中，可以实现更高效、更节能的AI系统，加速人工智能在各个领域的普及。

📄 摘要（原文）

Large language models (LLMs) can spend extra compute during inference to generate intermediate thoughts, which helps to produce better final responses. Since Chain-of-Thought (Wei et al., 2022), many such System 2 techniques have been proposed such as Rephrase and Respond (Deng et al., 2023a), System 2 Attention (Weston and Sukhbaatar, 2023) and Branch-Solve-Merge (Saha et al., 2023). In this work we investigate self-supervised methods to ``compile'' (distill) higher quality outputs from System 2 techniques back into LLM generations without intermediate reasoning token sequences, as this reasoning has been distilled into System 1. We show that several such techniques can be successfully distilled, resulting in improved results compared to the original System 1 performance, and with less inference cost than System 2. We posit that such System 2 distillation will be an important feature of future continually learning AI systems, enabling them to focus System 2 capabilities on the reasoning tasks that they cannot yet do well.