模型蒸馏任务

任务简介

• 基于ERNIE预训练模型效果上达到业界领先，但是由于模型比较大，预测性能可能无法满足上线需求。
• 直接使用ERNIE-Tiny系列轻量模型finetune，效果可能不够理想。如果采用数据蒸馏策略，又需要提供海量未标注数据，可能并不具备客观条件。
• 因此，本专题采用主流的知识蒸馏的方案来压缩模型，在满足用户预测性能、预测效果的需求同时，不依赖海量未标注数据，提升开发效率。
• 文心提供三种不同大小的基于字粒度的ERNIE-Tiny学生模型，满足不同用户的需求。

注：知识蒸馏（KD）是将复杂模型（teacher）中的dark knowledge迁移到简单模型（student）中去，teacher具有强大的能力和表现，而student则更为紧凑。通过知识蒸馏，希望student能尽可能逼近亦或是超过teacher，从而用更少的复杂度来获得类似的预测效果。

代码结构

.
├── data
│   ├── dev_data_distill
│   ├── dict
│   ├── download_data.sh
│   ├── predict_data
│   ├── test_data_distill
│   └── train_data_distill
├── examples
│   ├── cls_ernie_2.0_L3H128A2_ft_for_distill.json
│   ├── cls_ernie_2.0_L3H128A2_task_distill_step1.json
│   └── cls_ernie_2.0_L3H128A2_task_distill_step2.json
├── __init__.py
├── model
│   ├── base_cls.py
│   ├── bow_classification.py
│   ├── ernie_classification_distill.py
│   ├── ernie_classification.py
│   ├── __init__.py
├── output
└── trainer
    ├── custom_dynamic_trainer.py
    ├── custom_dynamic_trainer_step1.py
    ├── custom_dynamic_trainer_step2.py
    ├── __init__.py

模型蒸馏原理

• 知识蒸馏是一种模型压缩常见方法，指的是在teacher-student框架中，将复杂、学习能力强的网络(teacher)学到的特征表示“知识”蒸馏出来，传递给参数量小、学习能力弱的网络(student)。
• 在训练过程中，往往以最优化训练集的准确率作为训练目标，但真实目标其实应该是最优化模型的泛化能力。显然如果能直接以提升模型的泛化能力为目标进行训练是最好的，但这需要正确的关于泛化能力的信息，而这些信息通常不可用。如果我们使用由大型模型产生的所有类概率作为训练小模型的目标，就可以让小模型得到不输大模型的性能。这种把大模型的知识迁移到小模型的方式就是蒸馏。
• 基本原理可参考Hinton经典论文：https://arxiv.org/abs/1503.02531

ERNIE模型蒸馏

• 基于transformer结构的知识蒸馏，目的是将教师模型学习到的知识迁移到学生模型中，训练过程中，选用教师模型的每一层transformer layer的输出作为监督信息；
• 教师模型：基于ERNIE Large模型Fine-tune得到的模型，24层transformer。特点：效果好，预测性能差；
• 学生模型：只有几层transformer。特点：预测性能好，效果略低于教师模型；

• 两阶段蒸馏

  • GD：general distillation通用蒸馏，在预训练阶段训练。使用大规模无监督的数据， 帮助学生网络学习到尚未微调的教师网络中的知识，有利于提高泛化能力。
  • TD：task specific distillation，使用具体任务的数据，学习到更多任务相关的具体知识。
  • 可参考：https://arxiv.org/abs/1909.10351

蒸馏步骤

• FT阶段：基于ERNIE Large模型Fine-tune得到教师模型；

• TD1阶段：蒸馏transformer layer输出和attention矩阵

  1. 热启动Fine-tune的教师模型和通用学生模型。其中，通用的学生模型由文心平台提供，从bos上下载。
  2. 使用教师模型的transformer层的输出与学生模型的transformer层的输出的mse loss作为最终的loss。
  3. 如果教师模型与学生模型的transformer头数相等，可以增加attention蒸馏，即将教师模型与学生模型的attention矩阵的mse loss作为最终的loss。
  4. 由于教师模型层数较多，教师模型每隔几层与学生模型做MSE loss。假设 教师有M层，学生模型有N层(N<M), n = g(m) 是学生模型到教师模型的映射函数，表示学生模型的第m层是从教师模型中的第n层学习知识的，映射函数可表示为n = g(m)。比如教师模型有24层，学生模型4层，这里的映射函数为：g(m) = 6 * m。
  5. 反向传播阶段只更新学生模型参数，教师模型参数不更新；

• TD2阶段：蒸馏预测层，产出最终的学生模型。

  1. 热启动Fine-tune的教师模型和TD1阶段训练的学生模型；
  2. loss包括两部分：

  a) pred_loss：学生模型的预测层输出与教师模型的预测层输出（软label）的交叉熵；

  b) student_ce_loss：学生模型的预测层输出与真实label的交叉熵；

  3. pred_loss与student_ce_loss之和作为最终loss，进行反向传播；loss = pred_loss + student_ce_loss；
  4. 反向传播阶段只更新学生模型参数，教师模型参数不更新；

蒸馏策略

文心提供两种不同结构的预训练学生模型，学生模型采用的蒸馏策略见下表：

• hidden loss：将学生模型每一层transformer的输出与教师模型对应的层做MSE loss。
• attention loss：将学生模型transformer的attention矩阵与教师模型的transformer的attention矩阵做MSE loss，注意：只有教师模型与学生模型的头数一致的时候，才可以用attention蒸馏。
• pred_loss：学生模型的预测层输出与教师模型的预测层输出（软label）的交叉熵；
• student_ce_loss：学生模型的预测层输出与真实label的交叉熵；

开始使用

数据准备

使用ChnSentiCorp数据集。ChnSentiCorp 是一个中文情感分析数据集，包含酒店、笔记本电脑和书籍的网购评论。

• 训练集：wenxin_appzoo/tasks/model_distillation/data/train_data_distill/part.0，9600条数据；
• 验证集：wenxin_appzoo/tasks/model_distillation/data/dev_data_distill/part.0, 1200条数据；
• 测试集：wenxin_appzoo/tasks/model_distillation/data/test_data_distill/part.0, 1200条数据。

分阶段运行

这里以分类任务、学生模型ERNIE-Tiny-L3H128A2为例进行说明：

1. 准备预训练教师模型和学生模型

cd wenxin_appzoo/models_hub
bash download_ernie_2.0_base_ch.sh
bash download_pretrained_student_model.sh L3H128A2
cd wenxin_appzoo/tasks/model_distillation

2. fine-tune 教师模型

python run_trainer.py --param_path ./examples/cls_ernie_2.0_L3H128A2_ft_for_distill.json

3. 任务蒸馏一阶段：TD1

python run_trainer.py --param_path ./examples/cls_ernie_2.0_L3H128A2_task_distill_step1.json

4. 任务蒸馏二阶段：TD2

python run_trainer.py --param_path ./examples/cls_ernie_2.0_L3H128A2_task_distill_step2.json

效果验证

• 教师模型与学生模型的效果对比

• 与其他模型的效果对比

常见问题

• 为确保效果，模型蒸馏任务必须使用单卡训练，如果使用GPU多卡，模型效果有所降低；

• 问题1：怎么修改学生模型的层数？上面提供了三个不同的学生模型，但每个学生模型的层数、hidden size等都是固定的，如果想更改，需要在哪些地方更改？

  1. 文心提供了三种不同结构的预训练学生模型，如果修改层数、hidden size等，会导致预训练学生模型参数无法加载，在此情况下，蒸馏出来的学生模型效果无法保证，建议用户还是使用文心提供的预训练模型，不更改模型结构；

  2. 如果用户认为更改学生模型的结构非常有必要，需要对文心做以下改动：

     1. 修改TD1阶段的pre_train_model配置项，去掉预训练学生模型的加载，只保留微调后的教师模型：
     2. 预训练学生模型的目录里params子目录删除，只保留ernie_distill_student.json配置文件和vocab.txt词典文件；

     3. 修改 cls_ernie_2.0_L3H128A2_task_distill_step1.json 和 cls_ernie_2.0_L3H128A2_task_distill_step1.json 配置文件，更改模型结构：

            "pre_train_model": [
            # 热启动fine-tune的teacher模型
              {
                "name": "finetuned_teacher_model",
                "params_path": "./output/cls_ernie_2.0_base_fc_chn/save_checkpoints/checkpoints_step_6000"
              }
            ]

     4. 修改 student_embedding 配置项，改成对应学生模型的配置：

      "student_embedding": {
        "config_path": "../../models_hub/pretrained_student_model/your_student_model/ernie_distill_student.json"
      },

  3. 再次强调，上述修改后，由于无法加载预训练学生模型，蒸馏出来的学生模型效果无法保证。（用户训练数据量到百万样本以上可以考虑尝试一下）