Dual-Space Knowledge Distillation for Large Language Models

目前llm白盒蒸馏往往是对齐教师学生最后输出的概率分布，本文希望在隐藏状态层面进行对齐。同时目前的llm kd局限于同词汇表的架构之间的蒸馏，对上一个点的改进也有助于解决这个问题。

考虑隐藏状态\(h_{1:n}^t\),其过一个linear层\(P^{t \rightarrow s}\)，映射到学生的隐藏状态空间中，表示为 \[ h_{1:n}^{t \rightarrow s} = P^{t \rightarrow s}(h_{1:n}^t;\theta_{P}^{t \rightarrow s}) \in \mathbb{R}^{n\times d} \] 其中\(d\)是学生emd size 从而教师映射之后得到的概率分布可以表示为 \[ p_{1:n}^{t \rightarrow s} = \text{softmax}(h_{1:n}^{t \rightarrow s}W^s) \] 其中\(W^s\)就是学生的映射头，此时双方可以共享一个映射头。这里唯一的参数是\(P^{t \rightarrow s}\),其训练时通过将\(p_{1:n}^{t \rightarrow s}\)与ground truth进行ce loss对齐实现的。当然后面做教师学生KLD对齐的时候，这里的\(W^s\)并不会做反向传播。 类似的，学生的隐藏状态也可以映射到教师空间中，然后通过教师映射头得到概率分布，并与教师进行对齐。 最终训练的时候，两个对齐的方向都会使用。

llm kd中另一个问题是，由于模型输出的概率分布是定义在词汇表中的token上的，如果某个token在词汇表上不存在，那么其出现的概率也是没有定义的，从而在使用KLD对齐教师学生概率分布的时候，不同的词汇表就会导致不对应的token对齐错误。因此目前llm 白盒蒸馏仅仅局限于同词汇表的模型之间。 尝试通过attention来解决这个问题。每次对齐的时候，考虑学生向教师对齐，则将Q设计为学生相关emd，K设计为教师相关emd，然后设计一个V，即可得到学生映射后的emd，即隐藏状态，然后仿照上文进行概率分布对齐即可。教师向学生对齐的部分也是同理。

具体来说，对于当前句子在学生中的embeding表示，记为\(e_{1:n}^s\)，其目标tokens的embeding表示记为\(e_{2:n+1}^s\)，则 \[ Q = P^q([e_{1:n}^s; e_{2:n+1}^s];\theta_{P}^q)\in R^{n\times 2D} \] 其中\(D\)是教师emd size \(K\)设计为 \[ K = [e_{1:m}^t;e_{2:m+1}^t]\in R^{m\times 2D} \] \(V\)设计为 \[ V = P^v(e_{2:m+1}^t+h_{1:m}^t;\theta_{P}^v)\in R^{m\times d} \] 关键主要在于\(V\)的设计，这里设计为目标tokens在教师中的embedding表示，再加上教师最后一个隐藏层emd，不懂为什么是这样。 最终映射结果表示为 \[ softmax(\frac{QK^T}{\sqrt{ 2D }}V) \]