---
title: "Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information"
title_zh: "基于关键信息逐步注意力的混合层蒸馏提升小模型推理能力"
arxiv_id: "2604.15701"
paper_type: "method"
source_kind: "raw"
raw_path: "raw/papers/2026/04/2604.15701.md"
generated: "2026-05-11"
model: "LongCat-Flash-Chat"
quality: "ok"
---

# 基于关键信息逐步注意力的混合层蒸馏提升小模型推理能力

## 一句话定位

提出 MoLSAKI 框架，通过蒸馏教师模型在推理过程中对关键信息的逐步注意力，并结合混合层（MoL）动态对齐策略，提升小模型的复杂推理能力。

## 小学生也能听懂

这篇论文教小模型像大模型一样一步步思考，通过模仿大模型在解题时关注关键信息的方式，并用灵活的方法对齐不同层，让小模型变得更聪明，尤其在数学和常识题上表现更好。

## 为什么值得记录

- 首次将教师模型在推理过程中的逐步注意力模式引入 CoT 蒸馏，弥补了传统方法仅传递推理结果而忽略推理过程的缺陷
- 提出 Mixture-of-Layers (MoL) 模块，解决师生模型层数不匹配问题，实现自适应加权层对齐
- 在数学和常识推理任务上均取得一致性能提升，尤其在跨域泛化方面表现突出
- 方法兼容不同 tokenizer 和模型架构，无需投影层或共享词汇表，具备强实用性

## 核心方法

- 定义‘逐步注意力’为教师模型在每个推理步骤中对关键 token 的注意力分布，通过正则匹配自动提取关键 token（如数字或关键词）
- 在 CoT 蒸馏基础上引入新的损失项 L_att，强制学生模型学习教师对关键信息的动态关注模式
- 设计 MoL 模块：教师端基于各层注意力梯度分配权重，学生端通过可学习路由机制生成自适应层权重
- 使用温度控制的 softmax 对师生模型各层进行加权聚合，计算 KL 散度作为注意力蒸馏损失
- 总损失函数为预测损失、解释损失与注意力损失的加权和：L = αL_pre + (1−α)L_exp + βL_att

## 关键结果

- 在 Llama3-8B → GPT2-Large 设置下，MoLSAKI 在 SVAMP 上准确率达 49.5%，较 DSS 和 MMIloss 分别提升 1.5% 和 2.5%
- 在 Qwen2.5-32B → TinyLlama-1.1B 设置下，MoLSAKI 在 SVAMP 上达到 68.5% 准确率，相对提升 11.3%
- 在 out-of-domain 数据集（SingleEq, AsDiv, GSM8K）上均优于基线，显示强泛化能力
- MoL 对齐策略相比固定单层映射（SL）平均提升 2.0% 准确率，验证其有效性
- 即使使用不同模型生成 rationale 和注意力（如 PaLM-540B + Llama3-8B），MoLSAKI 仍保持优势

## 局限与风险

- 实验受限于计算资源，未探索更大范围模型规模与架构组合
- 当前方法聚焦短链推理任务，尚未验证在长链、复杂推理路径上的有效性
- 关键 token 提取依赖启发式规则（如正则匹配），可能遗漏语义关键但非显式的信息

## 适合沉淀的概念

`chain-of-thought-distillation`, `stepwise-attention`, `critical-tokens`, `mixture-of-layers`, `adaptive-layer-alignment`, `knowledge-distillation`, `reasoning-in-small-models`, `attention-pattern-transfer`

## 阅读注意

- 重点关注第 3.2 节中如何从自注意力矩阵中提取‘逐步注意力’的具体计算方式
- 注意 MoL 模块在师生模型中的差异化设计：教师端基于梯度统计，学生端基于可学习路由
- 附录 B 提供了关键 token 的可视化分析，有助于理解为何数字和关键词在推理中更受关注
- 实验部分应对比 Table 2 中 MoL 与多种 SL 映射的效果，理解动态对齐的价值
- 温度参数 τ₁ 和 τ₂ 的设置逻辑见附录 E，影响层权重的集中程度

## 质量说明

- 采用来源：`raw`，`raw/papers/2026/04/2604.15701.md`
- 生成模型：`LongCat-Flash-Chat`
- 源材料判断：论文结构完整，方法描述清晰，实验充分，包含消融、超参分析和案例分析，数据可信。