Improving Reasoning Capabilities in Small Models through Mixture-of-Layers Distillation with Stepwise Attention on Key Information

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

一句话定位

提出 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
• 源材料判断：论文结构完整，方法描述清晰，实验充分，包含消融、超参分析和案例分析，数据可信。