MiniRBT 中文小型预训练模型:结合了全词掩码技术和两段式知识蒸馏技术,加快推理速度
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-H256 和 MiniRBT-H288,表示隐层维数 256 和 288,均为 6 层 Transformer 结构,由两段式蒸馏得到。同时为了方便实验效果对比,我们也提供了 TinyBERT 结构的 RBT4-H312 模型下载。
我们会在近期提供完整的技术报告,敬请期待。
2.模型下载
也可以直接通过 huggingface 官网下载模型(PyTorch & TF2):https://huggingface.co/hfl
下载方法:点击任意需要下载的模型 → 选择"Files and versions"选项卡 → 下载对应的模型文件。
3.快速加载
3.1 使用 Huggingface-Transformers
依托于🤗transformers库,可轻松调用以上模型。
注意:本目录中的所有模型均使用 BertTokenizer 以及 BertModel 加载,请勿使用 RobertaTokenizer/RobertaModel!
对应的MODEL_NAME
如下所示:
3.2 模型对比
模型结构细节与参数量汇总如下。
括号内参数量百分比以原始 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 蒸馏设置
3.4 中文基线系统效果
为了对比基线效果,我们在以下几个中文数据集上进行了测试。
经过学习率搜索,我们验证了小参数量模型需要更高的学习率和更多的迭代次数,以下是各数据集的学习率。
最佳学习率:
*代表所有小型预训练模型 (RBT3, RBT4-H312, MiniRBT-H256, MiniRBT-H288)
实验结果(开发集):
相对效果:
两段式蒸馏对比
我们对两段式蒸馏(RoBERTa→RBT6(KD)→MiniRBT-H256)与一段式蒸馏(RoBERTa→MiniRBT-H256)做了比较。实验结果证明两段式蒸馏效果较优。
:该表中预训练模型经过 3 万步蒸馏,不同于中文基线效果中呈现的模型。
4.预训练
我们使用了TextBrewer工具包实现知识蒸馏预训练过程。完整的训练代码位于pretraining
目录下。
代码结构
dataset
:train
: 训练集dev
: 验证集distill_configs
: 学生模型结构配置文件jsons
: 数据集配置文件pretrained_model_path
:ltp
: ltp 分词模型权重,包含pytorch_model.bin
,vocab.txt
,config.json
,共计 3 个文件RoBERTa
: 教师模型权重,包含pytorch_model.bin
,vocab.txt
,config.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 数据准备
对于中文模型,我们需先生成含有分词信息的参考文件,可直接运行以下命令:
因为预训练数据集较大,推荐生成参考文件后进行预处理,仅需运行以下命令:
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
:设置隐层输出为 Truegradient_accumulation_steps
:梯度累积temperature
:蒸馏温度fp16
:开启半精度浮点数训练
直接运行以下命令可实现 MiniRBT-H256 的预训练蒸馏:
提示:以良好的模型权重初始化有助于蒸馏预训练。在我们的实验中使用教师模型的前 6 层初始化助教模型 RBT6(KD) ! 请参考scripts/init_checkpoint_TA.py
来创建有效的初始化权重,并使用--student_pretrained_weights
参数将此初始化用于蒸馏训练!
4.4 使用建议
初始学习率是非常重要的一个参数,需要根据目标任务进行调整。
小参数量模型的最佳学习率和
RoBERT-wwm
相差较大,所以使用小参数量模型时请务必调整学习率(基于以上实验结果,小参数量模型需要的初始学习率高,迭代次数更多)。在参数量(不包括嵌入层)基本相同的情况下,MiniRBT-H256 的效果优于 RBT4-H312,亦证明窄而深的模型结构优于宽而浅的模型结构
在阅读理解相关任务上,MiniRBT-H288 的效果较好。其他任务 MiniRBT-H288 和 MiniRBT-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: 部分数据集提供了下载地址。未标注下载地址的数据集请自行搜索或与原作者联系获取数据。
参考文献
[1] Pre-training with Whole Word Masking for Chinese BERT (Cui et al., IEEE/ACM TASLP 2021)
[2] TextBrewer: An Open-Source Knowledge Distillation Toolkit for Natural Language Processing (Yang et al., ACL 2020)
[3] CLUE: A Chinese Language Understanding Evaluation Benchmark (Xu et al., COLING 2020)
[4] TinyBERT: Distilling BERT for Natural Language Understanding (Jiao et al., Findings of EMNLP 2020)
更多优质内容请关注公号:汀丶人工智能;会提供一些相关的资源和优质文章,免费获取阅读。
版权声明: 本文为 InfoQ 作者【汀丶人工智能】的原创文章。
原文链接:【http://xie.infoq.cn/article/4b570791061c94aedbb07ac4f】。
本文遵守【CC-BY 4.0】协议,转载请保留原文出处及本版权声明。
评论