【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding

总述

TinyBert主要探究如何使用模型蒸馏来实现BERT模型的压缩。
主要包括两个创新点:

  1. 对Transformer的参数进行蒸馏,需要同时注意embedding,attention_weight, 过完全连接层之后的hidden,以及最后的logits。
  2. 对于预训练语言模型,要分为pretrain_model 蒸馏以及task-specific蒸馏。分别学习pretrain模型的初始参数以便给压缩模型的参数一个好的初始化,第二步学习pretrain model fine-tuning的logits让压缩模型再次学习。

模型

模型主要分成三个部分:

  1. Transformer Layer的蒸馏
    主要蒸馏两部分,第一是每一层的attention weight,第二是每一层输出的hidden。如下图所示。
    【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding_第1张图片

公式:
【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding_第2张图片
【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding_第3张图片
使用均方误差作为损失函数, 并且在hidden对比的时候引入了一个Wh,这是因为学生模型和老师模型的向量编码维度不一致(学生模型的向量维度要更小)
2. Embedding layer的蒸馏
在这里插入图片描述
E表示embeddign层的输出。
3. Predict logits的蒸馏
在这里插入图片描述
z表示老师模型与学生模型在task-specific任务上的预测概率。

此外还有一个细节便是数据增强,学生模型在task-specific任务上fine-tuning的时候,Tinybert对原数据集做了数据增强。(ps:这其实非常奇怪,因为后文实验中可以看到,去除数据增强之后,模型的效果比之之前的sota并无太大提升。而文章主要的卖点是模型蒸馏ummm)

实验和结论

  1. 蒸馏各个层次的重要性
    【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding_第4张图片
    可以看出,从重要性来说: Attn > Pred logits > Hidn > emb. 其中,Attn,Hidn以及emb在两个阶段的蒸馏中均有用到。

  2. 数据增强的重要性
    【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding_第5张图片
    GD (General Distillation)表示第一阶段蒸馏。
    TD (Task-specific Distillation)表示第二阶段蒸馏.
    and DA (Data Augmentation).表示数据增强。
    这张表得到的结论是,数据增强很重要 : (。

  3. 学生模型需要学习老师模型的哪些层
    【模型蒸馏】TinyBERT: Distilling BERT for Natural Language Understanding_第6张图片
    假设学生模型4层,老师模型12层
    top表示学生模型学习老师的后4层(10,11,12),bottom表示学习老师模型的前4层(1,2,3,4),uniform表示均匀学习(等间距,3,6,9,12)。
    可以看到,均匀学习各层的效果更好。

代码

# 此部分代码应该写在Trainer里面, loss.backward之前。
# 获取学生模型的logits, attention_weight以及hidden
 student_logits, student_atts, student_reps = student_model(input_ids, segment_ids, input_mask,
                                                            is_student=True)
# 在测试环境下获取老师模型的logits, attention_weight以及hidden
 with torch.no_grad():
     teacher_logits, teacher_atts, teacher_reps = teacher_model(input_ids, segment_ids, input_mask)

# 分为两步,一步是学习attentino_weight和hidden,还有一步是学习predict_logits。 总思想就是对学生模型的输出和老师模型的输出做loss,其中针对attention_weight和hidden是MSE loss, 针对logits是交叉熵。
 if not args.pred_distill:
     teacher_layer_num = len(teacher_atts)
     student_layer_num = len(student_atts)
     assert teacher_layer_num % student_layer_num == 0
     layers_per_block = int(teacher_layer_num / student_layer_num)
     new_teacher_atts = [teacher_atts[i * layers_per_block + layers_per_block - 1]
                         for i in range(student_layer_num)]

     for student_att, teacher_att in zip(student_atts, new_teacher_atts):
         student_att = torch.where(student_att <= -1e2, torch.zeros_like(student_att).to(device),
                                   student_att)
         teacher_att = torch.where(teacher_att <= -1e2, torch.zeros_like(teacher_att).to(device),
                                   teacher_att)

         tmp_loss = loss_mse(student_att, teacher_att)
         att_loss += tmp_loss

     new_teacher_reps = [teacher_reps[i * layers_per_block] for i in range(student_layer_num + 1)]
     new_student_reps = student_reps
     for student_rep, teacher_rep in zip(new_student_reps, new_teacher_reps):
         tmp_loss = loss_mse(student_rep, teacher_rep)
         rep_loss += tmp_loss

     loss = rep_loss + att_loss
     tr_att_loss += att_loss.item()
     tr_rep_loss += rep_loss.item()
 else:
     if output_mode == "classification":
         cls_loss = soft_cross_entropy(student_logits / args.temperature,
                                       teacher_logits / args.temperature)
     elif output_mode == "regression":
         loss_mse = MSELoss()
         cls_loss = loss_mse(student_logits.view(-1), label_ids.view(-1))

     loss = cls_loss
     tr_cls_loss += cls_loss.item()

你可能感兴趣的:(NLP,bert,transformer,人工智能)