C2KD: Cross-layer and Cross-head Knowledge Distillation for Small Language Model-based Recommendations

目前llm蒸馏主要有两个方向：prediction蒸馏与feature蒸馏，这其实与vanila KD是一样的 对于前者，由于教师与学生是在不同数据集上训练的，其输出头包含的信息实际上是不一样的，如果直接对齐两者的预测结果的话，并不能保证分布对齐。对于后者，教师的层实际上有某些是冗余的，在推荐场景下，有些层的贡献可能就会较小，那么层的选择也会是一个问题。

下图给出了到不同层为止的推荐能力，与SASRec进行对比。这里评估的方式是：到对应层停止推理，然后用可学习的预测头给出最终结果。背景是序列推荐。 ![[Pasted image 20250804131358.png|500]]

可以看到，不管是较大(8B)的教师还是较小的学生(1B)，其推荐能力普遍在十几层之后才超过SASRec，这可能是因为前面的层主要捕获低级文本线索，缺乏顺序推荐所需的上下文推理。另外注意到教师的表现最优的层并非最后一层，并且在不同数据集上，最优层对应的位置也在不断变化，这意味着，feature 蒸馏的时候，中间层的选择应当是要精心设计的。

具体选择的时候，选择教师具体层数时，首先得到教师参数\(H^l\)经过mlp变换后的结果\(g^l(H_{t}^l;W^l)\)，当然这个\(W^l\)是根据变换后结果在具体任务上学习的。之后通过如下式子得到最优层数 \[ l^* = argmax_{l\in \{1, 2, \dots L_{T}\}} \frac{<H_{s}^n, g^l(H_{t}^l;W^l)>}{||H_{s}^n||\ ||g^l(H_{t}^l;W^l)||} \] 即依照每一层变换后的教师特征与学生的相似度，取最相似的一层。之后直接MSE对齐即可，记为\(L_{feat}\)

对于另外一点，即教师和学生预测头分布不同的问题，作者选择对学生的最后一层的特征做一个变换，然后与教师使用同一个预测头。即对\(H_{s}^{L_{S}}\)乘上一个可学习的矩阵\(W_{proj}\)，进行维度对齐以及空间变换。为了避免特征投影中的信息丢失，对\(W_{proj}\)加上正交性限制 \[ L_{orth} = ||W_{proj}^TW_{proj} - I|| + ||W_{proj}W_{proj}^T - I|| \] 这一点其实与[1]的想法是很类似的，只是[1]不仅做了学生feature向教师的转换，还做了教师feature向学生的转换，当然省略了正交性限制。不管怎么样，之后对教师与学生在同一个预测头下的预测结果进行对齐即可，记为\(L_{logit}\)

那么最终Loss就表示为 \[ Loss = Loss_{feat} + L_{task} + L_{orth} + L_{logit} \]

Reference

[1] Zhang, S., Zhang, X., Sun, Z., Chen, Y., & Xu, J. Dual-Space Knowledge Distillation for Large Language Models. In EMNLP, 24