Knowledge Distilling
知识蒸馏实践
1.什么是蒸馏
Hinton大佬的一篇论文《Distilling the Knowledge in a Neural Network》中提出一种模型压缩的新思路：知识蒸馏。通过训练好的Teacher model（通常是大参数模型或者ensemble model）,将知识蒸馏到Student model。这样的做法比起直接使用Student model在数据集上训练，能够取得更接近Teacher model的效果。简而言之，就是从Teacher model中学习到的信息提取精华后，灌输到Student model中去，Student model将同时具备低复杂度和高精度的优势，这个过程就叫做知识蒸馏。
2. 为什么要蒸馏
知识蒸馏的做法在15年就被提出，当时主要应用于一些集成学习模型系统。随着18年bert的出现，大规模预训练语言模型广泛应用于业务中，但是这类模型往往参数量巨大，所消耗的计算资源巨大，直接将bert模型部署在线上环境时，响应时间不能够满足线上要求，因此对bert模型的压缩、加速推理提出了要求。
3. 如何蒸馏
蒸馏的方式比较灵活多变，需要根据具体的任务，以及Teacher model 和 Student model的模型结构作调整。在这里我以中文命名实体任务为例，将训练好的roberta+crf作为Teacher model，bilstm+crf作为Student mdoel。（roberta预训练模型采用在米读小说数据集上继续预训练后的模型，相关细节参考：Fiction-Bert ）
蒸馏的通用流程如下所示：
训练Teacher模型：先用hard target，也就是正常的label训练Teacher模型。
计算soft target：利用训练好的大模型来计算soft target。也就是Teacher模型通过温度T参数进行“软化”再经过softmax的output。
训练Student模型，在小模型的基础上再加一个额外的soft target的loss function，通过lambda参数来调节两个loss functions的比重。
预测时，将训练好的Student模型按常规方式使用。
对于ner任务，在实践中有两个难点：1.最后一层不是softmax，而是crf层，需要自行实现crf层，蒸馏过程视作softmax用法一致。2.对于序列标注问题，学生模型和教师模型的序列长度需要对齐，而bert与lstm在处理序列问题上通常的编码方式是不一致的，bert依赖于词表，以字为单位；lstm会先将句子用jieba进行分词，再通过词向量矩阵，以词为单位。这就导致了在处理soft target时出现维度不一致，不容易对应上的问题。为了解决这个问题，我将lstm的词编码格式与bert保持一致，词向量矩阵使用bert的token-embedding参数进行初始化。
下图是teacher模型（部分）和student模型信息，soft target加入了两部分的中间层信息，对应关系如下所示：
roberta -> bilstm
Embedding-Token (Embedding) -> embedding_1 (Embedding) （固定参数）
Transformer-11-FeedForward-Norm -> bidirectional_1 (Bidirection)
dense_72 (Dense) -> logit (Dense)
4. 实践细节
温度参数T：也许是配合蒸馏的叫法，因此将该变量命名为温度。从直观意义上来理解，通过升温，将教师模型的知识进行蒸馏，灌输到学生模型中。从数学角度来分析，T越大，输出的概率值越平缓。软化后，输出的概率值所包含的信息熵更大，对于学生模型来说更具学习价值。在我们的蒸馏过程中，T为10（通常为1-100）。
Last updated
