Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information
作者: Yao Chen, Jiawei Sheng, Wenyuan Zhang, Tingwen Liu
分类: cs.CL
发布日期: 2026-04-17
备注: Accepted at EMNLP 2025
DOI: 10.18653/v1/2025.emnlp-main.250
💡 一句话要点
提出基于分层混合与逐步注意力蒸馏的小模型推理能力提升方法
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 思维链蒸馏 注意力机制 知识迁移 小模型推理 分层混合 语言模型 逐步注意力
📋 核心要点
- 现有思维链蒸馏方法未能充分利用教师模型在推理过程中对关键信息的动态关注。
- 提出一种新颖的CoT蒸馏框架，通过分层混合模块将教师模型对关键信息的逐步注意力转移到学生模型。
- 实验结果表明，该方法在多个数学和常识推理数据集上实现了持续的性能提升。
📝 摘要（中文）
大型语言模型的巨大计算需求激发了人们对将推理能力提炼到小型模型中的兴趣，特别是通过思维链（CoT）蒸馏。现有的CoT蒸馏方法主要集中于转移教师模型生成的复杂推理过程，但未能充分利用教师模型在推理过程中对关键信息的动态关注。我们发现，语言模型在推理过程中会逐步将注意力转移到关键信息上，这为得出结论提供了重要线索。基于此，我们提出了一种新颖的CoT蒸馏框架，将教师模型对关键信息的逐步注意力转移到学生模型，从而为学生模型在推理过程中逐步关注关键信息建立结构化指导。更重要的是，我们开发了一个分层混合模块，实现动态对齐，以适应教师和学生模型之间的不同层。我们的方法在多个数学和常识推理数据集上实现了持续的性能提升。据我们所知，这是第一个在CoT蒸馏中利用逐步注意力来提高小型模型推理能力的方法。
🔬 方法详解
问题定义：论文旨在解决如何将大型语言模型的推理能力有效地迁移到小型模型的问题。现有基于思维链（CoT）蒸馏的方法主要关注于迁移教师模型生成的推理过程，但忽略了教师模型在推理过程中对关键信息的动态关注，导致学生模型难以捕捉到推理的关键线索。
核心思路：论文的核心思路是利用教师模型在推理过程中逐步将注意力转移到关键信息上的特性，通过蒸馏的方式将这种逐步注意力的模式迁移到学生模型。通过让学生模型学习教师模型在不同推理步骤中对关键信息的关注程度，从而引导学生模型更好地进行推理。
技术框架：整体框架包含教师模型、学生模型以及一个分层混合（Mixture of Layers）模块。教师模型负责生成推理过程和注意力权重，学生模型负责学习推理能力。分层混合模块用于动态对齐教师模型和学生模型之间的不同层，从而更好地进行注意力蒸馏。具体流程为：首先，教师模型对输入问题进行推理，生成思维链和注意力权重；然后，分层混合模块将教师模型的注意力权重映射到学生模型的对应层；最后，学生模型学习模仿教师模型的推理过程和注意力权重。
关键创新：论文的关键创新在于提出了基于逐步注意力的思维链蒸馏方法，并设计了分层混合模块来实现教师模型和学生模型之间的动态对齐。与现有方法相比，该方法能够更有效地将教师模型的推理能力迁移到学生模型，从而提高学生模型的推理性能。
关键设计：分层混合模块的设计是关键。该模块允许学生模型的每一层学习教师模型不同层的注意力模式，从而实现更灵活的知识迁移。损失函数包括两部分：一部分是模仿教师模型推理过程的损失，另一部分是模仿教师模型注意力权重的损失。具体来说，可以使用交叉熵损失来衡量学生模型推理过程与教师模型推理过程之间的差异，使用KL散度来衡量学生模型注意力权重与教师模型注意力权重之间的差异。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在多个数学和常识推理数据集上取得了显著的性能提升。例如，在GSM8K数据集上，该方法将学生模型的准确率提高了超过5个百分点，超过了现有的CoT蒸馏方法。此外，消融实验表明，逐步注意力和分层混合模块都对性能提升起到了关键作用。
🎯 应用场景
该研究成果可应用于各种需要小型化语言模型的场景，例如移动设备、嵌入式系统和边缘计算设备。通过将大型语言模型的推理能力提炼到小型模型中，可以在资源受限的环境下实现高效的推理，从而扩展语言模型的应用范围。此外，该方法还可以用于提高现有小型语言模型的推理能力，使其能够更好地解决复杂的推理问题。
📄 摘要（原文）
The significant computational demands of large language models have increased interest in distilling reasoning abilities into smaller models via Chain-of-Thought (CoT) distillation. Current CoT distillation methods mainly focus on transferring teacher-generated rationales for complex reasoning to student models. However, they do not adequately explore teachers' dynamic attention toward critical information during reasoning. We find that language models exhibit progressive attention shifts towards key information during reasoning, which implies essential clues for drawing conclusions. Building on this observation and analysis, we introduce a novel CoT distillation framework that transfers the teacher's stepwise attention on key information to the student model. This establishes structured guidance for the student's progressive concentration on key information during reasoning. More importantly, we develop a Mixture of Layers module enabling dynamic alignment that adapts to different layers between the teacher and student. Our method achieves consistent performance improvements across multiple mathematical and commonsense reasoning datasets. To our knowledge, it is the first method to leverage stepwise attention within CoT distillation to improve small model reasoning.