通过逐步关注关键信息,提升小模型的推理能力。
Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information

- 让小模型学习大模型在推理中逐步聚焦关键信息的注意力模式
- 在多个数学和常识推理数据集上实现稳定性能提升
- 适合希望提升小模型推理能力的研究者和开发者
大语言模型的高昂计算成本促使研究者关注将复杂推理能力蒸馏到小型模型中。现有链式思维(CoT)蒸馏方法主要传递教师生成的推理步骤,但未充分挖掘教师在推理过程中对关键信息的动态注意力变化。我们发现语言模型在推理时会逐步聚焦关键信息,这一过程蕴含重要推理线索。基于此,我们提出一种新型CoT蒸馏框架,将教师在推理过程中对关键信息的逐步注意力转移给学生模型,为学生提供结构化引导,使其在推理中逐步集中于关键信息。更重要的是,我们设计了分层混合模块(Mixture of Layers),实现教师与学生不同层间的动态对齐。该方法在多个数学和常识推理数据集上均取得一致性能提升。据我们所知,这是首个在CoT蒸馏中利用逐步注意力来增强小模型推理能力的方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。