知识蒸馏中推理习惯比分数更重要：中间层特征与注意力对齐

我在基于一个已有的大模型做知识蒸馏时，踩过一个相当隐蔽的坑：学生模型在验证集上的准确率很漂亮，但一旦进入带推理链的任务，比如数学解题、逻辑推导，它就表现得像一个只会背答案的机器。后来我才意识到，问题不是蒸馏流程不对，而是我从一开始就盯错了东西——我盯着教师模型的最终分数，却忽略了教师模型的推理习惯。

“分数”在这里指的是经典知识蒸馏里的软标签，也就是教师模型输出的概率分布。很多人会把知识蒸馏理解成“让教师模型给答案，学生模型照着答案学”，这在简单任务上确实有效，但在推理类任务里，这条路径的尽头很窄。因为教师模型的最终答案只是一个高度压缩的结果，真正的推理过程隐藏在所有中间层里。如果你只拿结果去训练学生模型，学生模型学到的是“猜答案”，而不是“会思考”。

这篇文章想讨论的，正是“推理习惯比分数更重要”这句话到底意味着什么，以及在实际训练里，你该怎么把教师的推理习惯蒸馏给学生模型。

1. 先看清一个误区：知识蒸馏并不只是“学习教师模型的答案”

1.1 软标签蒸馏的基本逻辑和它为什么有效

经典蒸馏方法通常分两步：先用一个大模型作为教师，在样本上推理，记录它的输出概率分布；然后把这个概率分布作为监督信号，去训练一个小模型作为学生。教师模型的输出不是非0即1的硬标签，而是带有“知识”的软标签。比如在图像分类任务里，一张猫的图片，教师模型可能会给出0.7猫、0.2狗、0.1狐狸的概率。这种分布包含了类别之间的关系，比硬标签信息量大得多。

在分类、检索、语义匹配这类任务上，软标签确实能显著加速学生模型收敛，也能提升表现。原因是教师模型把类别之间的相似性都编码进了概率分布，学生模型可以从中学到“猫和狗比较像，猫和汽车不太像”。这种相对关系是硬标签给不了的，也是知识蒸馏最有价值的地方之一。

但问题在于，这个价值在很大程度上依赖于一个前提：教师模型最终输出的分布仍然保留着足够的“结构”。一旦这个前提不成立，软标签蒸馏就会退化成一种接近硬标签学习的操作。

1.2 在推理类任务中，答案本身并不等于推理能力

推理任务和识别任务有一个本质差异：识别任务的输出空间通常很大，而且输出和输入之间存在一种“整体映射”；推理任务则往往最终只有一个很小的答案空间，比如一个数字、一个布尔值、一个代码片段。拿数学题举例：“小明有3个苹果，小华有5个苹果，总共有几个？”教师模型可能直接输出“8”，而且概率分布中“8”占绝对优势，其他数字几乎忽略不计。这种情况下，软标签携带的“知识”量非常少，学生模型学到的是“看到这一题就输出8”，却没有学到为什么是8，也没有学到遇到类似问题时应该如何拆解。

我自己做过一个很直接的对比：分别用硬标签、软标签、中间层特征对齐来训练同一个学生模型，训练数据是同一批数学题。结果在常见题目上，三种方式差异不大；但在新题型上，中间层特征对齐的学生模型表现明显更稳，失败率低很多。这个实验结果让我意识到，对于推理类任务，最终分数并不是教师模型中最值得蒸馏的资源。

1.3 一个关键转折：学生模型学到的是“模仿结果”还是“模仿思考”

那为什么大部分人的第一反应仍然是用最终分数做蒸馏？因为分数最容易获取。你只需要让教师模型做一次前向推理，就能拿到概率分布，不需要访问中间层，也不需要担心结构差异。而中间层特征、注意力分布这些“推理习惯”信息，获取成本更高，并且要设计对齐方式，甚至要处理教师和学生模型结构不一致的问题。

但推理任务恰恰不能只看最终分布。推理的多步累积特性决定了，任何一个中间步骤的错误都可能被放大到最终结果。如果学生模型没有学会教师的中间思考方式，它只能靠记忆大量“问题-答案”对来拟合训练集，一旦遇到稍微改变条件的题目，它就会暴露出泛化能力不足。所以这里的关键转折是：你到底是在让学生模型模仿教师模型的“答案”，还是在让学生模型模仿教师模型的“思考路径”。答案只能靠中间层信息传递，而不是靠最终概率分布。

2. 为什么推理任务中“分数”会失效

2.1 推理任务的特殊性：过程可解释、步骤依赖强、错误可传播

推理任务里的错误往往是累积的。比如在代码生成任务中，判断一段代码是否正确，不仅要看最终输出，还要看变量之间的依赖关系、控制流逻辑、边界条件处理。如果学生模型只是记住了“输入A返回B”的映射，当输入中的全局变量名发生变化，或者某个函数调用顺序调整，它就很难正确产出代码。

这种“步骤依赖”特性，在识别任务中并不常见。识别一张图片是不是猫，图片里的每个像素都对最终类别有贡献，但不存在严格的先后依赖。推理任务则不同，中间步骤是有顺序的，需要不断在已有中间结果的基础上继续推导。教师模型的中间层，正是在记录这些推导步骤产生的表征。只蒸馏最终分数，相当于把一条完整的推理链压缩成一个点，学生模型拿到的只是一张“结果摘要”，而非推演过程。

2.2 教师模型的最终分数（概率分布）丢失了哪些信息

最终概率分布到底丢失了什么？这里可以列一个清单：

• 中间步骤的表示状态，比如模型在第5层已经知道“A是数量，B是操作符”；
• 每一步决策时的注意力焦点，比如教师模型在读到哪几个token时开始跨句建立关系；
• 遇到不确定性时识别风险的能力，比如某个条件不充分时，教师模型内部会有多个候选路径，而最终分布可能只留下其中一个；
• 对不同推理路径的偏好，比如教师模型倾向于先做条件判断、再执行运算，而不是反过来。

这些信息都存在于隐藏层里，而不是最后的softmax层。软标签蒸馏只能传递“最终答案的倾向”，而这些倾向一旦被压缩，很多推理路径上的细节就被抹掉了。

2.3 反直觉现象：教师分数越高，蒸馏效果不一定越好

还有一个反直觉的现象：教师模型越强，最终分数越集中，软标签越容易退化成硬标签。如果教师模型在训练集上准确率接近100%，且置信度接近1，那么它输出的概率分布几乎就是“正确答案概率1，其他答案概率0”。这时候软标签携带的类间关系信息非常少，学生模型得到的监督信号，和直接用硬标签训练没有本质区别。

这也是为什么很多蒸馏实践里需要把温度参数调高。温度的作用是让分布更平缓，把中间类别的关系重新展露出来。但温度调高只是在最终结果上做软化，并不能还原推理过程中被压缩的中间状态。换句话说，你给一个已经压缩成二维码的图片做模糊化处理，也无法变回原始高清图。

3. 推理习惯到底是什么：从中间层到注意力分布

3.1 推理习惯不是你想象的“思维链文本”

很多人一听到“推理习惯”，第一反应是思维链（CoT）文本，也就是让模型输出一步步的解释。这其实是一个误解。思维链文本是模型生成的自然语言解释，它确实可以反映一部分推理过程，但它本身仍然是生成的产物，并不等价于模型内部的表征状态。一个模型可能输出了很合理的思维链，但真正决定答案正确性的，是它在每一步编码中隐藏的信息。

换句话说，思维链文本可能是“推理习惯”的表层投影，但它不完整，也不一定可靠。比如教师模型可能用错误的内部逻辑生成了一段看起来很合规的思维链，最后答案碰巧是对的。学生模型如果只模仿这段文本，就很可能学到错误的推理路径。

3.2 中间层特征和特征关系

推理习惯更重要的部分，是教师模型在每一层变换中如何处理和组合输入信息。在数学推理任务中，较早层可能负责抽取“数量”和“操作符”信息，中间层负责建立它们之间的关系，高层负责选择运算顺序。如果学生模型能在中间层上模拟这种动态变化，它就等于继承了教师的推理路径。

但这里有一个容易被忽略的点：中间层特征不仅仅是单个token的表征，还包括特征之间的关系。比如在“小明有3个苹果，小华有5个苹果”这句话里，“3”和“5”之间是什么关系？“小明”和“小华”是否在同一个语义槽位？这些关系正是通过注意力机制和前后层变换逐步建立的。如果只对齐单点特征，而不对齐特征间的相对关系，学生模型仍然可能学到“各自背下来的向量”，而不是“相互关联的推理”。

3.3 注意力模式与token贡献

在Transformer架构中，注意力分布是推理习惯最直观的体现。教师模型在处理一个长文本时，每一步会把注意力放在哪几个token上，体现了它对信息的依赖关系。比如在代码理解任务中，教师模型几乎一定会重点关注变量定义处、函数调用处、边界条件判断处。如果学生模型的注意力分布和教师模型不一致，即使它最终预测的token是正确的，也很可能在泛化时出现偏差。

注意力对齐的一层特殊价值在于，它不需要对特征维度做精确匹配，只需要对齐每个位置上“注意力权重的相对大小”。这在一定程度上降低了对模型结构的要求。但要注意，不同模型层的注意力含义可能差异很大，简单地对全层注意力做对齐，并不一定都有正收益。

3.4 错误模式和边界感知

推理习惯还包括“知道自己不确定什么”。教师模型在面对模糊问题时，内部表征会在多个可能路径之间犹豫，这其实是一种有价值的边界信息。比如一道题缺少了一个关键条件时，教师模型内部的注意力可能会在两个候选条件之间反复横跳，最终答案置信度较低。如果学生模型只看到最终答案和置信度，就很容易学习到过于自信的映射，它不知道哪些情况是风险点。

通过蒸馏教师模型的中间层不确定性，或者对齐教师模型在低置信度样本上的特征分布，可以让学生模型获得更好的边界感知。这一点在实际部署中尤其重要，因为推理模型的错误往往不像识别模型那样“温和”，它可能在一个很自然的推导链里突然出现逻辑断点，很难被事后校验察觉。

4. 如何把推理习惯“蒸馏”给学生模型：三种可落地的路径

4.1 路径一：中间层特征对齐（Feature Alignment）

中间层特征对齐是最直接的做法。先选择一个教师模型和共有的中间层，提取特征，然后计算两者之间的L2距离或余弦距离，把这个距离作为辅助损失函数，加在总损失里。

下面是一个简化示例：

# 示例结构：特征对齐损失 teacher_feature = teacher_model(x, output_hidden_states=True).hidden_states[6] student_feature = student_model(x, output_hidden_states=True).hidden_states[6] # 对特征做L2归一化，避免量纲影响 teacher_feature = F.normalize(teacher_feature, dim=-1) student_feature = F.normalize(student_feature, dim=-1) loss_feature = F.mse_loss(student_feature, teacher_feature)

这里的关键是选择一个有代表性的层。更建议先选接近输出但又不是最后几层的隐藏层，因为太早的层太通用，太晚的层已经接近答案预测，推理信息被压缩得比较严重。实际落地时，可以先试教师模型的倒数第2层或倒数第3层，映射到学生模型的最后1层或最后2层。如果层数和维度差异很大，就要用投影层把学生特征升维或降维，但这种做法会增加训练参数，要谨慎。

4.2 路径二：注意力分布对齐（Attention Distillation）

注意力分布对齐的目标是让教师模型和学生模型在相同的输入上产生相似的注意力权重。它的好处是，即使两个模型的隐藏层维度不同，只要多头注意力的头数一致，或者能通过平均操作把头数对齐，就可以计算损失。

一个常见写法：

# 注意力对齐损失（简化版） teacher_attn = teacher_model(x, output_attentions=True).attentions[-1] student_attn = student_model(x, output_attentions=True).attentions[-1] # 平均头数，将维度统一为 (batch, seq, seq) teacher_attn = teacher_attn.mean(dim=1) student_attn = student_attn.mean(dim=1) loss_attn = F.kl_div(student_attn.log(), teacher_attn, reduction='batchmean')

在真实项目里，注意力对齐需要特别留意数值稳定性。注意力权重是归一化后的概率值，直接算KL散度时，学生模型概率为零会导致log算出负无穷。所以通常会在学生注意力输出上加一个很小的平滑项，比如student_attn + 1e-8。

注意：采样前先用一个batch验证注意力对齐损失不会出现NaN。如果出现NaN，先检查是否需要对注意力权重做平滑，再检查教师模型是否已经开启eval模式，训练模式会引入随机dropout，导致注意力分布不稳定。

4.3 路径三：关系蒸馏（Relational Distillation）

有时候教师模型和学生模型的架构差异很大，比如教师是Decoder-only大模型，学生是MLP或轻量Transformer。直接对齐特征维度既不现实，也不稳定。这时可以改用关系蒸馏：把同一个batch内的所有样本分别输入两个模型，得到对应的特征表示，然后计算样本间的相似度矩阵，再对齐两个相似度矩阵。

示例结构：

# 关系蒸馏损失 teacher_sim = torch.cosine_similarity( teacher_feature.unsqueeze(1), teacher_feature.unsqueeze(0), dim=-1 ) student_sim = torch.cosine_similarity( student_feature.unsqueeze(1), student_feature.unsqueeze(0), dim=-1 ) loss_rel = F.mse_loss(student_sim, teacher_sim)

关系蒸馏的核心思路是“不要求点对点一致，只要求样例间的相对关系一致”。这对学生模型来说比较友好，因为它给了学生模型更多自由度去用自己的方式表达推理状态。代价是对batch size有一定要求，batch太小的话，关系矩阵本身的信息量不足。一般建议batch size至少设为32，才比较稳定。

4.4 一个附加建议：借助少量推理步骤文本作为辅助监督

尽管前面提到思维链文本不等于推理习惯，但少量高质量的推理步骤文本仍然可以作为辅助监督信号。原因在于，思维链文本可以给输出层提供一个显式的“推理路径约束”，帮助学生模型在生成时更稳定地按步骤推进。但它不能替代中间层特征对齐，更多是作为加速收敛的补充。

具体做法是：先用教师模型为训练集生成少量带推理步骤的样本，然后让学生模型也生成推理步骤，与教师生成的推理步骤计算语言模型损失。注意，这里不需要依赖大量CoT数据，甚至100~200条就够，因为核心监督信号仍然来自中间层特征或注意力对齐。如果数据量太大，反而会让学生模型过度关注“模仿文本格式”，忽略真正的表征学习。

5. 实际落地中的流程、参数和排查链路

5.1 一个建议的最小实验流程

不要一开始就把所有技巧都加进来。我更建议从一套最小可运行流程开始，跑通后再逐步扩展。

1. 准备教师模型和学生模型，确认两者都可以输出中间层特征或注意力。
2. 准备一个包含推理任务的数据集，拆成训练集和验证集，并单独保留一组“新题型”测试集。
3. 先只跑一个batch，确认教师模型前向、学生模型前向、特征提取、损失计算都没有问题。
4. 从单层特征对齐开始，温度设为1，特征对齐权重设为0.1，跑10步观察loss变化。
5. 效果稳定后，再逐步加入注意力对齐、关系蒸馏和辅助CoT损失。
6. 对照组始终保留：使用纯软标签蒸馏作为baseline。最后对比在新题型测试集上的表现，而不仅仅是验证集准确率。

5.2 参数设置：温度、层数、归一化方式、权重比例

参数设置往往直接决定蒸馏效果，但这里没有绝对正确的值，只能给出一个常见起点。

权重比例是最难调的部分。一个稳妥做法是先用很小的权重跑一个版本，比如0.1，然后观察学生模型在验证集上的表现。如果学生模型已经能达到教师模型90%以上的效果，再逐步增大特征对齐权重，看是否还能继续提升。如果一上来就1:1，很容易让学生模型被教师特征“绑架”，失去自己学习数据规律的能力。

5.3 排查链路：为什么蒸馏后学生模型效果不佳

蒸馏后效果不佳，原因可能来自多个层次。建议按下述顺序排查：

1. 看现象：是训练loss不下降，还是验证集不掉点但新题型掉点？前者通常是优化问题，后者往往是监督信号不足。
2. 看输入：教师模型和学生模型的输入预处理是否一致？tokenizer是否一致？序列长度是否修剪一致？如果学生模型用的tokenizer和教师不一样，中间层特征对齐的意义就非常小。
3. 看环境：PyTorch版本、CUDA版本、模型库版本是否匹配？不同版本下hidden_states的返回格式可能有差异。
4. 看参数：温度是否过高导致软标签噪声过大？层数映射是否合理？特征对齐权重是否过大导致模型崩溃？
5. 看模型结构：教师和学生模型的层数、维度、头数是否差异过大？如果差异过大，直接特征对齐可能不适用，改成关系蒸馏更稳妥。
6. 看日志：中间层特征的梯度是否正常？是否有NaN？如果有NaN，检查归一化和损失计算，尤其是注意力对齐时的对数计算。

下面是一个常见的排查表：

6. 适用边界：不是所有任务都需要蒸馏推理习惯

6.1 什么时候分数蒸馏就够了

如果任务是简单分类、匹配、语义相似度，或者训练数据量非常大，任务模式相对固定，那么软标签蒸馏通常就够了。因为这类任务的“推理过程”被压缩在比较浅的模式中，学生模型可以通过大量数据自行学习条件映射。硬标签和软标签的差距，并没有多步推理任务里那么明显。

比如一个意图识别系统，输入一句话，输出“查天气”“订机票”“放音乐”等固定类别。教师模型在相似意图之间给出的软分布，确实能帮助学生模型更好地聚类意图。但这里不涉及严密的步骤依赖，所以不需要刻意去对齐中间层特征。只要把温度调合适，软标签蒸馏已经能带来不错的收益。

6.2 什么时候必须关注推理习惯

如果任务满足下面几个条件之一，建议认真考虑推理习惯蒸馏：

• 需要多步推理或逻辑推导，比如数学解题、规划、多轮对话中的条件判断；
• 输出空间很小，但输入空间很大且变化多，比如代码生成中的布尔条件判断；
• 需要较强的泛化能力，比如训练集之外的新题型、新组合、新坑位；
• 学生模型结构远小于教师模型，需要更压缩的表征知识。

一个很典型的场景是数学应用题。输入一句自然语言描述，输出一个数字。在这个任务里，软标签几乎无法传递任何有效信息，因为答案分布高度集中。如果不蒸馏中间层表征，学生模型就只能机械记忆“题目模式-答案”的对应关系，遇到稍微变换表述就会出错。

6.3 资源限制下的取舍

中间层特征对齐会明显增加显存和训练时间，因为你需要同时保存教师模型和学生模型的中间层特征，并且保留梯度。如果资源有限，可以采用“离线特征缓存”方案：先在训练集上跑一次教师模型，把选定层的特征、注意力权重全部存成文件，训练学生模型时不再做教师前向，只读取缓存。这个方案可以把额外的显存开销降到几乎为零，代价是磁盘存储和I/O。

另外，如果教师模型和学生模型的结构差异过大，比如一个是大模型，一个是非常小的MLP，直接特征对齐其实没有太大意义。这时关系蒸馏是更稳妥的选择，因为它不要求特征维度一致，只要教师模型能提取出样本级别的特征，就可以计算相对关系。总之，不要为了追新技巧而强行对齐特征，要衡量你手里的资源和最终任务是否真的需要这种复杂度。

7. 收尾：把蒸馏从“结果模仿”推进到“过程理解”

回头看踩过的那个坑，我最大的教训是：知识蒸馏里最值得注意的不是教师模型答对了什么，而是教师模型是如何想到答案的。最终分数只是推理过程的一个快照，如果你只把快照丢给学生，学生能学到的只有结果；如果你想让学生真正具有泛化能力，就必须把中间层的推理习惯也传递过去。

下一次做知识蒸馏实验时，建议你先问自己一句：我的教师模型在中间层到底在做什么？如果这个问题说不清楚，那单纯调温度、调权重比例，都只是在和最终分数的噪声搏斗。试着多看一眼隐藏状态，多对齐一层注意力，或者至少用关系蒸馏让学生模型模仿教师模型处理样本间关系的模式。你会发现，学生模型的泛化能力，往往就藏在你多看的这一层里。

这就如同带新人：如果你只给答案，新人只能照着做；如果你带他走一遍分析问题的思路，他才能真正独立处理那些没见过的题型。教师模型和学生模型之间，也是这个道理。