MiniRBT中文小型预训练模型:结合了全词掩码(Whole Word Masking)技术和两段式知识蒸馏

举报
汀丶 发表于 2023/08/06 20:53:33 2023/08/06
【摘要】 MiniRBT中文小型预训练模型:结合了全词掩码(Whole Word Masking)技术和两段式知识蒸馏(Knowledge Distillation)技术,加快推理速度

MiniRBT中文小型预训练模型:结合了全词掩码(Whole Word Masking)技术和两段式知识蒸馏(Knowledge Distillation)技术,加快推理速度

在自然语言处理领域中,预训练语言模型(Pre-trained Language Models)已成为非常重要的基础技术。为了进一步促进中文信息处理的研究发展,哈工大讯飞联合实验室(HFL)基于自主研发的知识蒸馏工具TextBrewer,结合了全词掩码(Whole Word Masking)技术和知识蒸馏(Knowledge Distillation)技术推出中文小型预训练模型MiniRBT


中文LERT | 中英文PERT | 中文MacBERT | 中文ELECTRA | 中文XLNet | 中文BERT | 知识蒸馏工具TextBrewer | 模型裁剪工具TextPruner

1.简介

目前预训练模型存在参数量大,推理时间长,部署难度大的问题,为了减少模型参数及存储空间,加快推理速度,我们推出了实用性强、适用面广的中文小型预训练模型MiniRBT,我们采用了如下技术:

  • 全词掩码技术:全词掩码技术(Whole Word Masking)是预训练阶段的训练样本生成策略。简单来说,原有基于WordPiece的分词方式会把一个完整的词切分成若干个子词,在生成训练样本时,这些被分开的子词会随机被mask(替换成[MASK];保持原词汇;随机替换成另外一个词)。而在WWM中,如果一个完整的词的部分WordPiece子词被mask,则同属该词的其他部分也会被mask。更详细的说明及样例请参考:Chinese-BERT-wwm,本工作中我们使用哈工大LTP作为分词工具。

  • 两段式蒸馏:相较于教师模型直接蒸馏到学生模型的传统方法,我们采用中间模型辅助教师模型到学生模型蒸馏的两段式蒸馏方法,即教师模型先蒸馏到助教模型(Teacher Assistant),学生模型通过对助教模型蒸馏得到,以此提升学生模型在下游任务的表现。并在下文中贴出了下游任务上两段式蒸馏与一段式蒸馏的实验对比,结果表明两段式蒸馏能取得相比一段式蒸馏更优的效果。

  • 构建窄而深的学生模型。相较于宽而浅的网络结构,如TinyBERT结构(4层,隐层维数312),我们构建了窄而深的网络结构作为学生模型MiniRBT(6层,隐层维数256和288),实验表明窄而深的结构下游任务表现更优异。

MiniRBT目前有两个分支模型,分别为MiniRBT-H256MiniRBT-H288,表示隐层维数256和288,均为6层Transformer结构,由两段式蒸馏得到。同时为了方便实验效果对比,我们也提供了TinyBERT结构的RBT4-H312模型下载。

我们会在近期提供完整的技术报告,敬请期待。

2.模型下载

模型简称 层数 隐层大小 注意力头 参数量 Google下载 百度盘下载
MiniRBT-h288 6 288 8 12.3M [PyTorch] [PyTorch]
(密码:7313)
MiniRBT-h256 6 256 8 10.4M [PyTorch] [PyTorch]
(密码:iy53)
RBT4-h312 (TinyBERT同大小) 4 312 12 11.4M [PyTorch] [PyTorch]
(密码:ssdw)

也可以直接通过huggingface官网下载模型(PyTorch & TF2):https://huggingface.co/hfl

下载方法:点击任意需要下载的模型 → 选择"Files and versions"选项卡 → 下载对应的模型文件。

3.快速加载

3.1使用Huggingface-Transformers

依托于🤗transformers库,可轻松调用以上模型。

from transformers import BertTokenizer, BertModel

tokenizer = BertTokenizer.from_pretrained("MODEL_NAME")
model = BertModel.from_pretrained("MODEL_NAME")

注意:本目录中的所有模型均使用BertTokenizer以及BertModel加载,请勿使用RobertaTokenizer/RobertaModel!

对应的MODEL_NAME 如下所示:

原模型 模型调用名
MiniRBT-H256 “hfl/minirbt-h256”
MiniRBT-H288 “hfl/minirbt-h288”
RBT4-H312 “hfl/rbt4-h312”

3.2模型对比

模型结构细节与参数量汇总如下。

模型 层数 隐层大小 FFN大小 注意力头数 模型参数量 参数量(除去嵌入层) 加速比
RoBERTa 12 768 3072 12 102.3M (100%) 85.7M (100%) 1x
RBT6 (KD) 6 768 3072 12 59.76M (58.4%) 43.14M (50.3%) 1.7x
RBT3 3 768 3072 12 38.5M (37.6%) 21.9M (25.6%) 2.8x
RBT4-H312 4 312 1200 12 11.4M (11.1%) 4.7M (5.5%) 6.8x
MiniRBT-H256 6 256 1024 8 10.4M (10.2%) 4.8M (5.6%) 6.8x
MiniRBT-H288 6 288 1152 8 12.3M (12.0%) 6.1M (7.1%) 5.7x

括号内参数量百分比以原始base模型(即RoBERTa-wwm-ext)为基准

  • RBT的名字是RoBERTa三个音节首字母组成
  • RBT3:由RoBERTa-wwm-ext 3层进行初始化继续预训练得到,更详细的说明请参考:Chinese-BERT-wwm 小参数量模型
  • RBT6 (KD):助教模型,由RoBERTa-wwm-ext 6层进行初始化,通过对RoBERTa-wwm-ext蒸馏得到
  • MiniRBT-*:通过对助教模型RBT6(KD)蒸馏得到
  • RBT4-H312: 通过对RoBERTa直接蒸馏得到

3.3蒸馏设置

模型 Batch Size Training Steps Learning Rate Temperature Teacher
RBT6 (KD) 4096 100KMAX512 4e-4 8 RoBERTa-wwm-ext
RBT4-H312 4096 100KMAX512 4e-4 8 RoBERTa-wwm-ext
MiniRBT-H256 4096 100KMAX512 4e-4 8 RBT6 (KD)
MiniRBT-H288 4096 100KMAX512 4e-4 8 RBT6 (KD)

3.4中文基线系统效果

为了对比基线效果,我们在以下几个中文数据集上进行了测试。

经过学习率搜索,我们验证了小参数量模型需要更高的学习率和更多的迭代次数,以下是各数据集的学习率。

最佳学习率:

模型 CMRC 2018 DRCD OCNLI LCQMC BQ Corpus TNEWS ChnSentiCorp
RoBERTa 3e-5 3e-5 2e-5 2e-5 3e-5 2e-5 2e-5
* 1e-4 1e-4 5e-5 1e-4 1e-4 1e-4 1e-4

*代表所有小型预训练模型 (RBT3, RBT4-H312, MiniRBT-H256, MiniRBT-H288)

实验结果(开发集):

Task CMRC 2018 DRCD OCNLI LCQMC BQ Corpus TNEWS ChnSentiCorp
RoBERTa 87.3/68 94.4/89.4 76.58 89.07 85.76 57.66 94.89
RBT6 (KD) 84.4/64.3 91.27/84.93 72.83 88.52 84.54 55.52 93.42
RBT3 80.3/57.73 85.87/77.63 69.80 87.3 84.47 55.39 93.86
RBT4-H312 77.9/54.93 84.13/75.07 68.50 85.49 83.42 54.15 93.31
MiniRBT-H256 78.47/56.27 86.83/78.57 68.73 86.81 83.68 54.45 92.97
MiniRBT-H288 80.53/58.83 87.1/78.73 68.32 86.38 83.77 54.62 92.83

相对效果:

Task CMRC 2018 DRCD OCNLI LCQMC BQ Corpus TNEWS ChnSentiCorp
RoBERTa 100%/100% 100%/100% 100% 100% 100% 100% 100%
RBT6 (KD) 96.7%/94.6% 96.7%/95% 95.1% 99.4% 98.6% 96.3% 98.5%
RBT3 92%/84.9% 91%/86.8% 91.1% 98% 98.5% 96.1% 98.9%
RBT4-H312 89.2%/80.8% 89.1%/84% 89.4% 96% 97.3% 93.9% 98.3%
MiniRBT-H256 89.9%/82.8% 92%/87.9% 89.7% 97.5% 97.6% 94.4% 98%
MiniRBT-H288 92.2%/86.5% 92.3%/88.1% 89.2% 97% 97.7% 94.7% 97.8%
  • 两段式蒸馏对比

我们对两段式蒸馏(RoBERTa→RBT6(KD)→MiniRBT-H256)与一段式蒸馏(RoBERTa→MiniRBT-H256)做了比较。实验结果证明两段式蒸馏效果较优。

模型 CMRC 2018 OCNLI LCQMC BQ Corpus TNEWS
MiniRBT-H256(两段式) 77.97/54.6 69.11 86.58 83.74 54.12
MiniRBT-H256(一段式) 77.57/54.27 68.32 86.39 83.55 53.94

:该表中预训练模型经过3万步蒸馏,不同于中文基线效果中呈现的模型。

4.预训练

我们使用了TextBrewer工具包实现知识蒸馏预训练过程。完整的训练代码位于pretraining目录下。

  • 代码结构
    • dataset:
      • train: 训练集
      • dev: 验证集
    • distill_configs: 学生模型结构配置文件
    • jsons: 数据集配置文件
    • pretrained_model_path:
      • ltp: ltp分词模型权重,包含pytorch_model.binvocab.txtconfig.json,共计3个文件
      • RoBERTa: 教师模型权重,包含pytorch_model.binvocab.txtconfig.json,共计3个文件
    • scripts: 模型初始化权重生成脚本
    • saves: 输出文件夹
    • config.py: 训练参数配置
    • matches.py: 教师模型和学生模型的匹配配置
    • my_datasets.py: 训练数据处理文件
    • run_chinese_ref.py: 生成含有分词信息的参考文件
    • train.py:预训练主函数
    • utils.py: 预训练蒸馏相关函数定义
    • distill.sh: 预训练蒸馏脚本

4.1环境准备

预训练代码所需依赖库仅在python3.8,PyTorch v1.8.1下测试过,一些特定依赖库可通过pip install -r requirements.txt命令安装。

  • 预训练模型准备

可从huggingface官网下载ltp分词模型权重与RoBERTa-wwm-ext预训练模型权重,并存放至${project-dir}/pretrained_model_path/目录下相应文件夹。

4.2数据准备

对于中文模型,我们需先生成含有分词信息的参考文件,可直接运行以下命令:

python run_chinese_ref.py

因为预训练数据集较大,推荐生成参考文件后进行预处理,仅需运行以下命令:

python my_datasets.py

4.3运行训练脚本

一旦你对数据做了预处理,进行预训练蒸馏就非常简单。我们在distill.sh中提供了预训练示例脚本。该脚本支持单机多卡训练,主要包含如下参数:

  • teacher_name or_path:教师模型权重文件
  • student_config: 学生模型结构配置文件
  • num_train_steps: 训练步数
  • ckpt_steps:每ckpt_steps保存一次模型
  • learning_rate: 预训练最大学习率
  • train_batch_size: 预训练批次大小
  • data_files_json: 数据集json文件
  • data_cache_dir:训练数据缓存文件夹
  • output_dir: 输出文件夹
  • output encoded layers:设置隐层输出为True
  • gradient_accumulation_steps:梯度累积
  • temperature:蒸馏温度
  • fp16:开启半精度浮点数训练

直接运行以下命令可实现MiniRBT-H256的预训练蒸馏:

sh distill.sh

提示:以良好的模型权重初始化有助于蒸馏预训练。在我们的实验中使用教师模型的前6层初始化助教模型RBT6(KD) ! 请参考scripts/init_checkpoint_TA.py来创建有效的初始化权重,并使用--student_pretrained_weights参数将此初始化用于蒸馏训练!

4.4 使用建议

  • 初始学习率是非常重要的一个参数,需要根据目标任务进行调整。
  • 小参数量模型的最佳学习率和RoBERT-wwm相差较大,所以使用小参数量模型时请务必调整学习率(基于以上实验结果,小参数量模型需要的初始学习率高,迭代次数更多)。
  • 在参数量(不包括嵌入层)基本相同的情况下,MiniRBT-H256的效果优于RBT4-H312,亦证明窄而深的模型结构优于宽而浅的模型结构
  • 在阅读理解相关任务上,MiniRBT-H288的效果较好。其他任务MiniRBT-H288MiniRBT-H256效果持平,可根据实际需求选择相应模型。

5.FAQ

Q: 这个模型怎么用?
A: 参考快速加载。使用方式和HFL推出的中文预训练模型系列如RoBERTa-wwm相同。

Q:为什么要单独生成含有分词信息的参考文件?
A: 假设我们有一个中文句子:天气很好,BERT将它标记为['天','气','很','好'](字符级别)。但在中文中天气是一个完整的单词。为了实现全词掩码,我们需要一个参考文件来告诉模型应该在哪个位置添加##,因此会生成类似于['天','##气','很','好']的结果。
注意:此为辅助参考文件,并不影响模型的原始输入(即与分词结果无关)。

Q: 为什么RBT6 (KD)在下游任务中的效果相较RoBERTa下降这么多? 为什么miniRBT-H256/miniRBT-H288/RBT4-H312效果这么低?如何提升效果?
A: 上文中所述RBT6 (KD)直接由RoBERTa-wwm-ext在预训练任务上蒸馏得到,然后在下游任务中fine-tuning,并不是通过对下游任务蒸馏得到。其他模型类似,我们仅做了预训练任务的蒸馏。如果希望进一步提升在下游任务上的效果,可在fine-tuning阶段再次使用知识蒸馏。

Q: 某某数据集在哪里下载?
A: 部分数据集提供了下载地址。未标注下载地址的数据集请自行搜索或与原作者联系获取数据。

更多优质内容请关注公号:汀丶人工智能;会提供一些相关的资源和优质文章,免费获取阅读。

【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息, 否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。