通过关键信息的分步注意力混合层蒸馏提升小模型的推理能力
Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information
AI总结：
本文提出一种基于分步注意力的混合层蒸馏方法，通过引导学生模型逐步聚焦关键信息，提升小模型的推理能力，并在多个数学和常识推理数据集上取得一致性能提升。
摘要：
大规模语言模型显著的计算需求促使人们通过链式思维（CoT）蒸馏将推理能力转移到小模型中。当前CoT蒸馏方法主要关注将教师生成的复杂推理理由转移到学生模型，但未能充分探索教师在推理过程中对关键信息的动态注意力。我们发现语言模型在推理过程中表现出对关键信息的逐步注意力转移，这暗示了得出结论的重要线索。基于这一观察和分析，我们引入了一种新的CoT蒸馏框架，将教师对关键信息的分步注意力转移到学生模型中。这为学生模型在推理过程中逐步聚焦关键信息提供了结构化指导。更重要的是，我们开发了一个混合层模块，能够实现动态对齐，适应教师和学生之间不同的层。我们的方法在多个数学和常识推理数据集上实现了持续的性能提升。据我们所知，这是首次在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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
