河南梦之网网络科技有限公司
梦之网科技出品
扫描关注梦之网科技微信公众账号

扫小程序码联系客服

TinyBERT:模型小7倍,速度快8倍,华中科大、华为出品-郑州小程序开发

梦之网科技2019-09-30文章动态

参与:张倩、杜伟

BERT 等大模型性能强大,但很难部署到算力、内存有限的设备中。为此,来自华中科技大学、华为诺亚方舟实验室的研究者提出了 TinyBERT,这是一种为基于 transformer 的模型专门设计的知识蒸馏方法,模型大小还不到 BERT 的 1/7,但速度是 BERT 的 9 倍还要多,而且性能没有出现明显下降。目前,该论文已经提交机器学习顶会 ICLR 2020。

论文:https://arxiv.org/abs/1909.10351

在自然语言处理(NLP)领域,BERT 等预训练语言模型极大地提升了诸多 NLP 任务的性能。但是,这类预训练语言模型通常计算开销大,内存占用也大,因此很难在一些资源紧张的设备上有效执行。

为了在加快推理速度和降低模型大小的同时保持准确率,来自华中科技大学和华为诺亚方舟实验室的研究者提出了一种新颖的 transformer 蒸馏法,这是为基于 transformer 的模型专门设计的知识蒸馏(knowledge distillation,KD)方法。通过这种新的 KD 方法,大型 teacherBERT 模型中编码的大量知识可以很好地迁移到小型 student TinyBERT 模型中。

此外,研究者还提出了一种专门用于 TinyBERT 的两段式学习框架,从而分别在预训练和针对特定任务的学习阶段执行 transformer 蒸馏。这一框架确保 TinyBERT 可以获取 teacherBERT 的通用知识和针对特定任务的知识。

除了提出新的 transformer 蒸馏法之外,研究者还提出了一种专门用于 TinyBERT 的两段式学习框架,从而分别在预训练和针对特定任务的具体学习阶段执行 transformer 蒸馏。这一框架确保 TinyBERT 可以获取 teacherBERT 的通用和针对特定任务的知识。

实证研究结果表明,TinyBERT 是有效的,在 GLUE 基准上实现了与 BERT 相当(下降 3 个百分点)的效果,并且模型大小仅为 BERT 的 13.3%(BERT 是 TinyBERT 的 7.5 倍),推理速度是 BERT 的 9.4 倍。此外,TinyBERT 还显著优于当前的 SOTA 基准方法(BERT-PKD),但参数仅为为后者的 28%,推理时间仅为后者的 31%左右。

Transformer 蒸馏

研究者提出的 Transformer 蒸馏是专门为 Transformer 网络设计的知识蒸馏方法,下图 1 为本文提出的 Transformer 蒸馏方法概览图:

TinyBERT:模型小7倍,速度快8倍,华中科大、华为出品-郑州小程序开发

图 1:Transformer 蒸馏概览图。

在这篇论文中,student 和 teacher 网络都是通过 Transformer 层构建的。为了表述清楚,研究者在详解 TinyBERT 之前阐述了以下问题。

假定 student 模型有 M 个 Transformer 层,teacher 模型有 N 个 Transformer 层,从 teacher 模型中选择 M 个 Transformer 层用于 Transformer 层蒸馏。n=g(m) 是 student 层到 teacher 层的映射函数,这意味着 student 模型的第 m 层从 teacher 模型的第 n 层开始学习信息。嵌入层蒸馏和预测层蒸馏也考虑进来,将嵌入层的指数设为 0,预测层的指数设为 M+1,并且对应的层映射分别定义为 0 = g(0) 和 N + 1 = g(M + 1)。下文实验部分将探讨不同的映射函数对性能的影响。在形式上,通过最小化以下目标函数,student 模型可以获取 teacher 模型的知识:

TinyBERT:模型小7倍,速度快8倍,华中科大、华为出品-郑州小程序开发

其中 L_layer 是给定模型层(如 Transformer 层或嵌入层)的损失函数,λ_m 是表征第 m 层蒸馏重要度的超参数。

研究者提出的 Transformer 层蒸馏包含基于注意力的蒸馏和基于隐状态的蒸馏,具体可参考上图 1(b)。基于注意力的蒸馏是为了鼓励语言知识从 teacherBERT 迁移到 student TinyBERT 模型中。具体而言,student 网络学习如何拟合 teacher 网络中多头注意力的矩阵,目标函数定义如下:

TinyBERT:模型小7倍,速度快8倍,华中科大、华为出品-郑州小程序开发

其中,h 是注意力头数。A_i ∈ R^( l×l) 是与 teacher 或 student 的第 i 个注意力头对应的注意力矩阵。

此外,(非归一化)注意力矩阵 A_i 用作拟合目标,而不是其 softmax 输出的 softmax(A_i),因为实验表明前者的设置呈现更快的收敛速度和更佳的性能。除了基于注意力的蒸馏之外,研究者还对 Transformer 层输出的知识进行蒸馏处理(具体可参考上图 1(b)),目标函数定义如下:

其中 H^S∈R^l×d'和 H^T∈R^l×d 分别表示 student 和 teacher 网络的隐状态,由方程式 4 计算得到。标量值 d 和 d'分别表示 teacher 和 student 模型的隐状态,并且 d'通常小于 d,以获得更小的 student 网络。

研究者还执行了嵌入层的蒸馏,与基于隐状态的蒸馏类似,定义如下:

其中矩阵 E^S 和 H^T 分别表示 student 和 teacher 网络的嵌入。在论文中,这两种矩阵的形状与隐状态矩阵相同。矩阵 W_e 表示线性变化,它起到与 W_h 类似的作用。

文章关键词
华为
华中科技大学
transformer
bert