# MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers

以下是本例的简要目录结构及说明：
```
.
├── general_distill.py       # 通用蒸馏脚本
├── run_clue.py              # 在下游任务上的微调脚本
└── README.md                # 文档，本文件
```
## 简介
本目录下的实验主要参考论文[《MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers》](https://arxiv.org/abs/2012.15828)实现。

MiniLMv2也是从层数深的 Transformer 类模型到层数较浅的 Transformer 类模型的蒸馏策略。它的优势是只需要取教师模型和学生模型中的各一层进行蒸馏训练，而不像其他方法需要蒸馏更多的层，避免面对更加复杂的 layer mapping 问题，并且效果优于 TinyBert 的蒸馏策略。

MiniLMv2蒸馏的目标是教师模型某层的 q 与 q, k 与 k, v 与 v 的矩阵乘结果和学生模型最后一层的 q 与 q, k 与 k, v 与 v 的矩阵乘之间的 kl 散度 loss。其中教师模型是 large size 时，选择实验并选取倒数某一层，当教师模型是 base size 时，选择最后一层进行蒸馏即可。

为了防止教师模型是 large size 时，head size 与学生模型不同，蒸馏目标的 shape 无法匹配，MiniLMv2还需要对 head 进行重组，先合并再按 relation_head_num 重新分割 head_num 和 head_size。

## 数据、预训练模型介绍及获取

### 数据获取
由于本实验是通用场景下的蒸馏，因此数据和预训练类似。可以参考[NLP Chinese Corpus](https://github.com/brightmart/nlp_chinese_corpus)中提供的数据。
数据下载完成后，需要将所有数据集整理成每行一条文本数据，再将数据切分成多个小文件，并放在一个目录下，以便使用多卡并行训练。

### 训练启动方式

假设我们把切分好的预训练数据文件都放在`${dataset}`下，那么我们可以运行如下命令用单机八卡进行预训练蒸馏：
```shell

dataset=/PaddleNLP/dataset
output_dir=./pretrain

python -m paddle.distributed.launch --gpus "0,1,2,3,4,5,6,7" general_distill.py \
    --student_model_type tinybert \
    --num_relation_heads 48 \
    --student_model_name_or_path tinybert-6l-768d-zh \
    --init_from_student False \
    --teacher_model_type bert \
    --teacher_model_name_or_path bert-base-chinese \
    --max_seq_length 128 \
    --batch_size 256 \
    --learning_rate 6e-4 \
    --logging_steps 20 \
    --max_steps 100000 \
    --warmup_steps 4000 \
    --save_steps 5000 \
    --teacher_layer_index 11 \
    --student_layer_index 5 \
    --weight_decay 1e-2 \
    --output_dir ${output_dir} \
    --device gpu \
    --input_dir ${dataset} \

```

其中参数释义如下：

- `student_model_type` 学生模型的类型
- `num_relation_heads` head 重新组合之后的 head 数
- `student_model_name_or_path` 学生模型的名字（需要与学生模型类型对应），或者是学生模型的路径
- `init_from_student` 本次蒸馏的学生模型是否用`student_model_name_or_path`中的参数进行初始化，是个 bool 类型的参数。默认是 False
- `teacher_model_type bert` 教师模型的类型
- `teacher_model_name_or_path`  教师模型的名字
- `max_seq_length 128` 表示最大句子长度，超过该长度将被截断。
- `warmup_steps` 学习率 warmup up 的步数
- `save_steps` 保存模型的频率
- `teacher_layer_index`表示学生模型从教师模型学习的教师层
- `student_layer_index` 表示学生模型从教师模型学习的学生层
- `output_dir` 模型输出的目录
- `device gpu` 表示运行该程序的设备，默认是 gpu
- `input_dir` 预训练数据的存放地址



### 评价方法

假设预训练完成后的模型存储在`${pretrained_models}`下，这里也提供了我们已经预训练完成的一版[模型](https://bj.bcebos.com/paddlenlp/models/general_distill/minilmv2_6l_768d_ch.tar.gz)可供参考，模型与`tinybert-6l-768d-zh`结构相同，因此可以使用`TinyBertForSequenceClassification.from_pretrained()`对模型直接进行加载。
本示例训练出的通用模型需要在下游任务上 Fine-tuning，利用下游任务上的指标进行评价。
我们可以运行如下脚本在单卡上进行 Fine-tuning：

```shell

export CUDA_VISIBLE_DEVICES="0"

python -u ./run_clue.py \
    --model_type tinybert  \
    --model_name_or_path ${pretrained_models} \
    --task_name ${TASK_NAME} \
    --max_seq_length ${max_seq_len} \
    --batch_size 16   \
    --learning_rate ${learning_rate} \
    --num_train_epochs ${num_train_epochs} \
    --logging_steps 100 \
    --seed 42  \
    --save_steps  100 \
    --warmup_proportion 0.1 \
    --weight_decay 0.01 \
    --adam_epsilon 1e-8 \
    --device gpu  \

```


其中不同的任务下，`${learning_rate}`、`${num_train_epochs}`、`${max_seq_len}`，我们推荐不同的 Fine-tuning 的超参数，可以参考以下配置：

| TASK_NAME        | AFQMC | TNEWS | IFLYTEK | OCNLI | CMNLI | CLUEWSC2020 | CSL  |
| ---------------- | ----- | ----- | ------- | ----- | ----- | ----------- | ---- |
| learning_rate    | 2e-5  | 2e-5  | 2e-5    | 3e-5  | 3e-5  | 1e-5        | 1e-5 |
| num_train_epochs | 3     | 3     | 6       | 6     | 3     | 50          | 8    |
| max_seq_len      | 128   | 128   | 128     | 128   | 128   | 128         | 256  |


### 蒸馏实验结果

本示例选择的是 CLUE 中的分类任务，以`bert-base-chinese`作教师模型，利用 MiniLMv2策略对6层模型进行蒸馏，可以得到的通用模型在 CLUE 上的指标为：

| CLUE    | AFQMC | TNEWS | IFLYTEK | CMNLI | OCNLI | CLUEWSC2020 | CSL   |
| ------- | ----- | ----- | ------- | ----- | ----- | ----------- | ----- |
| Acc (%) | 71.38 | 56.46 | 58.87   | 79.01 | 73.02 | 68.42       | 77.73 |


## 参考文献

Wang W, Bao H, Huang S, Dong L, Wei F. [MiniLMv2: Multi-Head Self-Attention Relation Distillation for Compressing Pretrained Transformers](https://arxiv.org/abs/2012.15828)[J]. arXiv preprint arXiv:2012.15828v2, 2021.
