基于知识蒸馏的高效生物医学命名实体识别模型

删除或更新信息，请邮件至freekaoyan#163.com(#换成@)

本站小编 Free考研考试/2021-12-20

图 1 Transformer block的结构图

图选项

2.1.1 多头自注意力机制注意力机制可以有效地学习词与词之间的依赖关系，注意力函数的计算主要依赖于3个不同的向量，查询向量q、键向量k和值向量v。在Transformer block的多头自注意力机制中，q、k和v对应着同一个向量，即词在模型中间层的隐向量表示，注意力函数的形式表示如下：

$\boldsymbol{A}=\frac{Q \mathbf{K}^{\mathrm{T}}}{\sqrt{d_{K}}} , $

(1)

$\operatorname{Attention}(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})=\operatorname{softmax}(\boldsymbol{A}) \boldsymbol{V} . $

(2)

其中：Q、K、V分别是q、k和v在模型实际训练时对应的矩阵；A是由Q和K进行内积计算相似度后得到的注意力矩阵；d_K是K的维数，$ {\sqrt {{d_K}} }$作为惩罚因子，确保内积的结果不会过大。
函数的输出Q^′是对V的权重求和结果，权重是通过对A的每一列向量进行softmax()归一化后计算得到的。注意力机制的本质是一个查询到一系列键值对的映射，使得句子中的单词可以突破距离的限制，有效地捕捉其他单词的信息。多头自注意力机制即组合多个不同的注意力函数的结果，表示如下：

$\text { head }_{i}=\text { Attention }\left(\boldsymbol{Q} \boldsymbol{W}_{i}^{Q}, \boldsymbol{K} \boldsymbol{W}_{i}^{K}, \boldsymbol{V} \boldsymbol{W}_{i}^{V}\right) , $

(3)

$\begin{array}{*{20}{c}} \text { MultiHead }(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})=\\ \text { Concat(head } \left._{0}, \cdots, \text { head }_{h}\right) \boldsymbol{W}^{O} \text {. }\end{array} $

(4)

其中：h是多头注意力机制中头的数量，W_i^Q、W_i^K、W_i^V和W^O是可学习的线性变化参数。每一个注意力函数中的Q、K和V通过线性变化被投影到不同的子空间，这使模型具有能在不同的表示子空间里学习到相关信息的能力。
2.1.2 前馈神经网络层在Transformer block中，通过多头自注意力机制计算后的隐向量，将被输入前馈神经网络层。该前馈神经网络层的定义如下：

$\operatorname{FNN}\left(\boldsymbol{Q}^{\prime}\right)=\max \left(0, \boldsymbol{Q}^{\prime} \boldsymbol{W}_{1}+b_{1}\right) \boldsymbol{W}_{2}+b_{2}. $

(5)

前馈神经网络层在Transformer block中发挥了特征抽取器的作用，前馈神经网络层的输出H将作为输入进入下一层Transformer block进行计算。本文把多头自注意力机制层中输出的注意力矩阵A和前馈神经网络层的输出H引入到知识蒸馏的过程中来，方法的内容在节3中详细介绍。
2.2 知识蒸馏知识蒸馏是一种模型压缩技术，即通过优化知识蒸馏损失函数，让学生模型能够有效地模仿教师模型的行为。在神经网络模型的训练过程中，损失函数是用来估量模型的预测值与真实值的不一致程度，通过最小化损失函数使得模型的预测输出去拟合真实标签的分布。而知识蒸馏则需要使学生模型的预测输出去拟合教师模型的预测输出。对此，Hinton等^[18]提出了一种通用的方法：

$\operatorname{softmaxT}(\boldsymbol{Z})=\frac{\mathrm{e}^{\boldsymbol{Z} / \mathrm{T}}}{\sum\limits_{j} \mathrm{e}^{\mathbf{z}_{j} / T}} . $

(6)

通过在softmax()中引入温度参数T得到的软化标签soft label构造知识蒸馏损失函数。
2.3 基于动态权重的知识蒸馏方法尽管在知识蒸馏训练的过程中，蒸馏损失函数 $\mathcal{L} $_SL的目标是尽可能地使学生模型的预测输出去拟合教师模型的预测输出，但实际上，研究者更加关注的是这种模仿行为是否能够为学生模型对于真实标签预测的准确度带来提升，所以引入对于真实标签的损失函数$\mathcal{L} $_TL。
在以前的知识蒸馏工作^{[18, 21-22]}中，往往需要设置固定的权重参数，以此权衡$\mathcal{L} $_SL和$\mathcal{L} $_TL的重要程度。权重的设置是非常依赖于经验的，而在更加直观的理解下，知识蒸馏的过程更加接近于人类的学习行为，是一个从模仿到创新的过程。受这种“青出于蓝而胜于蓝”思想的启发，本文面向Transformer结构模型提出了一种基于动态权重的知识蒸馏方法:

$\mathcal{L}_{\mathrm{SL}}=\mathrm{CE}\left(\operatorname{softmax} \mathrm{T}\left(\boldsymbol{Z}^{\mathrm{t}}\right), \operatorname{softmax} \mathrm{T}\left(\boldsymbol{Z}^{\mathrm{S}}\right)\right), $

(7)

$\mathcal{L}_{\mathrm{TL}}=\mathrm{CE}\left(\operatorname{softmax}\left(\boldsymbol{Z}^{\mathrm{S}}\right), \boldsymbol{y}\right), $

(8)

$ \begin{gathered}\mathcal{L}_{\text {total }}=\varphi(s) \cdot \mathcal{L}_{\mathrm{SL}}+[1-\varphi(s)] \cdot \mathcal{L}_{\mathrm{TL}}+ \\\varphi(s) \cdot \sum\limits_{j=0}^{n} \mathcal{L}_{\mathrm{ML}}^{j} ,\end{gathered} $

(9)

$\varphi(s)=\max \left(1-\mathrm{e}^{\mathrm{s} / \mathrm{G}+b-1}, 0\right) . $

(10)

其中：CE()是交叉熵损失函数，Z^t和Z^S分别是教师模型和学生模型最后层输出的未归一化的概率logits，φ(s)是一个动态权重函数，s为训练的步长，G为训练的总步长，b为权重偏差参数。φ(s)的值分布于[0, 1]，从训练开始其值为1，然后逐渐加速下降到0。这使得总损失函数$\mathcal{L} $_total在训练开始阶段让学生模型更加倾向于模仿教师模型的行为，在训练后半段更加倾向于对真实标签分布的学习。b的引入使得$\mathcal{L} $_total能够有更多的训练步长去优化$\mathcal{L} $_HL，在实验中，b设置为0.3较为合理。
此外，$\mathcal{L}_{\mathrm{ML}}^{j} $是学生模型中的第j层对教师模型中间层的多头自注意力机制层中输出的注意力矩阵A和前馈神经网络层的输出H的学习，去拟合教师模型的输出。同时，是学生模型Transformer block层的层数，$\mathcal{L}_{\mathrm{ML}}^{j} $的定义如下：

$\begin{array}{*{20}{c}}\mathcal{L}_{\mathrm{ML}}=\operatorname{MSE}\left(\boldsymbol{H}^{\mathrm{S}} \boldsymbol{W}_{h}, \boldsymbol{H}^{\mathrm{t}}\right)+ \\\frac{1}{h} \sum\limits_{i=0}^{h} \operatorname{MSE}\left(\boldsymbol{A}_{i}^{\mathrm{S}}, \boldsymbol{A}_{i}^{\mathrm{t}}\right) .\end{array} $

(11)

其中：$ \boldsymbol{H}^{\mathrm{S}} \in \mathbf{R}^{l \times d}$和 $\boldsymbol{H}^{\mathrm{t}} \in \mathbf{R}^{l \times d^{\prime}} $分别是学生模型和教师模型中间层的前馈神经网络层的输出，l为输入文本的长度，d和d^′分别为学生模型和教师模型中间层的隐向量的维度，$ \boldsymbol{W}_{h} \in \mathbf{R}^{d \times d^{\prime}}$是在训练中学习线性变化参数，其用途是在学生模型和教师模型的隐向量的维度大小不同时，能够将H^S投影到H^t的表示空间进行计算。同时，$ \boldsymbol{A}_{i}^{\mathrm{S}} \in \mathbf{R}^{l \times l}$和 $ \boldsymbol{A}_{i}^{\mathrm{t}} \in \mathbf{R}^{l \times l} $分别是学生模型和教师模型中间层的第i个注意力头的注意力矩阵，h是多头自注意力机制中的头数，MSE()是均方误差损失函数。
3 FastBioNER3.1 NER的问题定义给定一个句子 $ X=\left\{x_{0}, x_{1}, x_{2}, \cdots, x_{l}\right\}, x$是句中的词，使用模型去预测对应的标签序列Y=(y₀, y₁, y₂, …, y_l)，NER任务可以定义为

$\begin{gathered}\left\{y_{0}, y_{1}, y_{2}, \cdots, y_{l}\right\}= \\\operatorname{model}\left(\left\{x_{0}, x_{1}, x_{2}, \cdots, x_{l}\right\}\right) .\end{gathered} $

(12)

在NER任务中，文本的标注通常采用BIO模式标记法：对于给定的实体，用“B-(类别名称)”标注该实体的第一个词，同时用“I-(类别名称)”标注该实体剩下的词，而句中其他无关的词则用“O”来进行标注。所以标签序列Y也被称作BIO序列(BIO sequences)。相较于通用领域的NER任务，生物医学领域的NER任务(BioNER)面向的是生物医学文献中的相关实体类别，比如：疾病名称、药物名称、基因名称等。
3.2 FastBioNER模型本文提出的基于知识蒸馏的高效生物医学命名实体识别的FastBioNER模型如图 2所示。第一步是使用任务微调的方法使BioBert针对BioNER任务进行训练；第二步是使用基于动态权重的知识蒸馏方法将训练好的BioBert模型作为教师模型压缩到一个规模更小的学生模型中去。

图 2 FastBioNER的结构图

图选项

首先，BioBert是基于Transformer结构的生物医学领域预训练语言模型，在Bert的基础上进行预训练而来，所以与Bert-base拥有相同的12层Transformer block结构。BioBert使用了从PubMed(https://www.ncbi.nih.gov/CBBresearch/Lu/Demo/Pubtator/)中获取的180亿词量级别的生物医学文献作为预训练语料库，在服务器上进行了为期10天的训练，在生物医学文献挖掘任务上拥有出色的表现。本文在BioBert的最后一层Transformer block输出的隐向量上加了一层线性全连接层作为BioNER的模型，使用任务微调的方式进行训练。
然后，将训练好的BioBert作为教师模型，利用本文提出的基于动态权重的知识蒸馏方法，将其压缩到一个规模更小的学生模型中。其中，教师模型的Transformer block层数为m，在BioBert中m为12，学生模型的Transformer block层数为n。本文按平均映射的策略将m层结构压缩到层结构中去，例如m=12、n=6的情况下，学生模型中每1层输出的A^S和H^S，按照顺序使用 $\mathcal{L} $_ML去拟合教师模型中每第2层输出的A^t和H^t。接着，利用softmaxT()分别计算教师模型和学生模型的软化标签，通过优化 $\mathcal{L} $_SL让学生模型去拟合教师模型的输出。同时，优化与真实标签的损失函数 $\mathcal{L} $_TL让学生模型去拟合真实标签。
FastBioNER的蒸馏过程同时考虑了模型的中间层和输出，并使用基于动态权重的知识蒸馏方法，通过在总损失函数 $\mathcal{L} $_total中引入动态权重函数φ(s)，使蒸馏过程从单一“模仿”教师模型行为过渡到到学生模型寻求“创新”的过程，即使用动态权重函数动态调整权重的分配，以此模拟真实的学习行为过程，提升学生模型预测的准确度。
4 实验及结果分析4.1 实验数据集及设置实验共采用了3组生物医学数据集，分别是NCBI疾病数据集^[23]、BC5CDR-chem药物化合物数据集^[24]以及BC4CHEMD药物化合物数据集^[25]，3组数据集的详细情况如表 1所示。此外，在实验过程中，本文将训练集和开发集合并为训练集。实验部分使用的教师模型均为BioBert，版本为BioBert-Base V1.1。使用的评估指标为实体级别的召回率(R)、精确率(P)和F1值。此外，在学生模型初始化参数方面，也按照平均映射的原则，将BioBert的参数用于学生模型对应层的初始化中。
表 1 实验数据集

数据集	实体类别	训练集	开发集	测试集
NCBI	疾病	5 423	922	939
BC5CDR-chem	药物/化合物	4 559	4 580	4 796
BC4CHEMD	药物/化合物	30 681	30 638	26 363

表选项

为了验证本文所提的基于动态权重的知识蒸馏方法以及FastBioNER模型的有效性，本文设置了3组不同的实验。
4.2 有效性对比实验及结果分析本实验同时对比了知识蒸馏方法的有效性和BioNER模型的有效性。为了公平对比知识蒸馏方法的有效性，实验将FastBioNER模型中的本文所提出的知识蒸馏方法分别替换为Bert-PKD和TinyBert所提的方法后，再与原本的FastBioNER作比较；同时，为了验证知识蒸馏的有效性，实验加入了不使用知识蒸馏而直接进行训练的方法Bert-OG进行比较。这4个方法在3组数据集上分别将12层的BioBert作为教师模型压缩到6层的学生模型后进行有效性对比。
为了验证FastBioNER在BioNER任务上的有效性，实验还加入了CollaboNet^[13]和Wang等^[12]的方法，以单任务学习的方式分别在3组数据上进行训练后与FastBioNER进行对比。实验的结果如表 2所示，最高和次高的结果分别以加粗和加下划线的方式表示出来了，相较于文[12-13]，FastBioNER的F1值明显更高，这得益于BioBert作为教师模型的指导；相较于不使用知识蒸馏的Bert-OG，其他3个基于知识蒸馏的方法的效果明显更好，这验证了知识蒸馏的有效性。但同时，3个基于知识蒸馏的方法的效果都不及教师模型BioBert，这说明知识蒸馏仅通过教师模型中的输出去模仿教师模型的行为还是具有一定局限性的。Bert-PKD是第一个提出对Transformer模型中间层的隐向量进行学习的方法；TinyBert同时考虑了中间层中前馈神经网络层的输出和多头自注意力机制的注意力矩阵输出，但这两个方法都使用了固定的权重。在3组数据集中，FastBioNER的效果都高于Bert-PKD和TinyBert，这是FastBioNER使用的基于动态权重的知识蒸馏方法不仅考虑了对Transformer模型中间层的学习，同时还使用动态权重模拟了真实的学习行为，这验证了基于动态权重的知识蒸馏方法的有效性。
表 2 有效性对比实验结果?

%
数据集	模型	R	P	F1
NCBI	BioBert^[14]	90.78	88.71	89.73
	Bert-OG	86.57	84.82	85.69
	Bert-PKD^[21]	89.11	86.85	87.96
	TinyBert^[22]	90.11	86.36	88.20
	CollaboNet^[13]	85.45	83.95	84.69
	Wang等^[12]	82.92	84.95	83.92
	FastBioNER	90.11	87.20	88.63
BC5CDR-chem	BioBert^[14]	95.02	92.38	93.68
	Bert-OG	90.35	92.77	91.55
	Bert-PKD^[21]	92.91	92.59	92.75
	TinyBert^[22]	93.76	91.30	92.51
	CollaboNet^[13]	91.50	94.02	92.74
	Wang等^[12]	86.87	93.05	89.05
	FastBioNER	93.57	92.08	92.82
BC4CHEMD	BioBert^[14]	92.27	93.22	92.75
	Bert-OG	89.78	91.60	90.68
	Bert-PKD^[21]	91.66	92.86	92.25
	TinyBert^[22]	92.30	92.41	92.36
	CollaboNet^[13]	85.96	90.50	88.19
	Wang等^[12]	87.04	90.53	88.75
	FastBioNER	92.18	93.03	92.60

表选项

4.3 模型压缩有效性和效率实验及结果分析实验使用FastBioNER将BioBert压缩到不同层数的学生模型后，对比模型的总参数量，并在单块英伟达2080ti显卡上，对BC4CHEMD的测试集进行计算，以对比模型的运算速度。实验结果如图 3所示。6层的FastBioNER能够有效的缩小BioBert模型39.26%的大小并缩短46.17%的运算时间；4层的FastBioNER能够有效的缩小BioBert模型52.35%的大小并缩短60.77%的运算时间；而3层的FastBioNER能够有效地缩小BioBert模型58.9%的大小并缩短68.78%的运算时间。值得注意的是，模型压缩的程度并没有与FastBioNER的层数成比例，该现象的出现主要是因为BioBert的嵌入层依然占有大量的参数，并没有随层数的减少而变少。同时，该实验结果也说明了FastBioNER能够有效地压缩BioBert的模型大小并能有效地提升运算速度。

图 3 模型压缩有效性和效率实验结果

图选项

4.4 层数对有效性影响实验及结果分析层数大小对FastBioNER在NCBI数据集上的有效性影响实验的结果如表 3所示，与直觉相符合，层数越高，FastBioNER的表现越接近BioBert的表现。其中，6层的FastBioNER与BioBert的F1值仅有1.10%的差距，而4层为3.72%，3层则为4.97%。虽然第2组实验证明了层数越小的FastBioNER能够更有效的压缩BioBert的大小，但也损失了深层Transformer结构带来的深层特征抽取能力。所以6层的FastBioNER是权衡表现后较为合理的选择，也更加贴近实际的使用。
表 3 层数对FastBioNER有效性的影响实验结果?

%
模型	R	P	F1
BioBert	90.78	88.71	89.73
FastBioNER (n=6)	90.11	87.20	88.63
FastBioNER (n=4)	86.84	85.19	86.01
FastBioNER (n=3)	85.44	84.09	84.76

表选项

5 结论在生物医学领域，BioBert在BioNER任务上具有优秀的实体识别性能，但存在模型规模过大、推理时间过慢的问题。针对网络模型如何高效压缩问题，本文提出了面向Transformer结构的一种基于动态权重的知识蒸馏方法；同时针对BioNER任务，本文提出了基于知识蒸馏的高效生物医学命名实体识别模型FastBioNER。提出的FastBioNER能在尽可能保留教师模型效果的基础上，有效压缩模型的大小，加速模型的推理速度。下一步工作中，会探索知识蒸馏与矩阵分解技术的结合，进一步提高压缩的效果。

参考文献

[1]	HANISCH D, FUNDEL K, MEVISSEN H T, et al. ProMiner: Rule-based protein and gene entity recognition[J]. BMC Bioinformatics, 2005, 6(S1): S14.
[2]	WANG X, ZHANG Y, LI Q, et al. PENNER: Pattern-enhanced nested named entity recognition in biomedical literature [C]//Proceedings of 2018 IEEE International Conference on Bioinformatics and Biomedicine. Madrid, Spain: IEEE, 2018: 540-547.
[3]	LI X, ROTH D. Learning question classifiers: The role of semantic information[J]. Natural Language Engineering, 2006, 12(3): 229-249. DOI:10.1017/S1351324905003955
[4]	LEE K J, HWANG Y S, KIM S, et al. Biomedical named entity recognition using two-phase model based on SVMs[J]. Journal of Biomedical Informatics, 2004, 37(6): 436-447. DOI:10.1016/j.jbi.2004.08.012
[5]	PONOMAREVA N, ROSSO P, PLA F, et al. Conditional random fields vs. hidden Markov models in a biomedical named entity recognition task [C]//Proceedings of 2007 International Conference on Recent Advances in Natural Language Processing. Borovets, Bulgaria: RANLP, 2007: 479-483.
[6]	LEAMAN R, WEI C H, ZOU C. Mining chemical patents with an ensemble of open systems[J]. Database: The Journal of Biological Databases and Curation, 2016, 2016.
[7]	KIM Y. Convolutional neural networks for sentence classification [C]//Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing. Doha, Qatar: Association for Computational Linguistics, 2014: 1746-1751.
[8]	ELMAN J L. Finding structure in time[J]. Cognitive Science, 1990, 14(2): 179-211. DOI:10.1207/s15516709cog1402_1
[9]	HOCHREITER S, SCHMIDHUBER J. Long short-term memory[J]. Neural Computation, 1997, 9(8): 1735-1780. DOI:10.1162/neco.1997.9.8.1735
[10]	LUO L, YANG Z H, YANG P, et al. An attention-based BiLSTM-CRF approach to document-level chemical named entity recognition[J]. Bioinformatics, 2018, 34(8): 1381-1388. DOI:10.1093/bioinformatics/btx761
[11]	CHO M, HA J, PARK C, et al. Combinatorial feature embedding based on CNN and LSTM for biomedical named entity recognition[J]. Journal of Biomedical Informatics, 2020, 103: 103381. DOI:10.1016/j.jbi.2020.103381
[12]	WANG X, ZHANG Y, REN X, et al. Cross-type biomedical named entity recognition with deep multi-task learning[J]. Bioinformatics, 2019, 35(10): 1745-1752. DOI:10.1093/bioinformatics/bty869
[13]	YOON W, SO C H, LEE J, et al. CollaboNet: Collaboration of deep neural networks for biomedical named entity recognition[J]. BMC Bioinformatics, 2019, 20(10): 249. DOI:10.1186/s12859-019-2813-6
[14]	LEE J, YOON W J, KIM S D, et al. BioBERT: A pre-trained biomedical language representation model for biomedical text mining[J]. Bioinformatics, 2020, 36(4): 1234-1240.
[15]	DEVLIN J, CHANG M W, LEE K, et al. BERT: Pre-training of deep bidirectional transformers for language understanding [C]//Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Minneapolis, USA: NAACL-HLT, 2019.
[16]	VASWANI A, SHAZEER N, PARMAR N, et al. Attention is all you need [C]//Proceedings of the 31st International Conference on Neural Information Processing Systems. Long Beach, USA: Curran Associates Inc., 2017: 5998-6008.
[17]	PETERS M, NEUMANN M, IYYER M, et al. Deep contextualized word representations [C]//Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. New Orleans, Louisiana, USA: NAACL-HLT, 2018: 2227-2237.
[18]	HINTON G, VINYALS O, DEAN J. Distilling the knowledge in a neural network[J]. arXiv preprint arXiv, 2015, 1503.02531.
[19]	SUN S, CHENG Y, GAN Z, et al. Patient knowledge distillation for bert model compression [C]//Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Hong Kong, China: EMNLP, 2019: 3-7.
[20]	CHEN H T, WANG Y H, XU C, et al. Data-free learning of student networks [C]//Proceedings of 2019 IEEE/CVF International Conference on Computer Vision (ICCV). Seoul, Korea (South): IEEE, 2019: 3513-3521.
[21]	JIAO X Q, YIN Y C, SHANG L F, et al. TinyBERT: Distilling BERT for natural language understanding[J]. arXiv preprint arXiv, 2019, 1909.10351.
[22]	DO AG?AN R I, LEAMAN R, LU Z Y. NCBI disease corpus: A resource for disease name recognition and concept normalization[J]. Journal of Biomedical Informatics, 2014, 47: 1-10. DOI:10.1016/j.jbi.2013.12.006
[23]	LI J, SUN Y P, JOHNSON R J, et al. BioCreative V CDR task corpus: A resource for chemical disease relation extraction[J]. Database: The Journal of Biological Databases and Curation, 2016, baw068.
[24]	KRALLINGER M, RABAL O, LEITNER F, et al. The CHEMDNER corpus of chemicals and drugs and its annotation principles[J]. Journal of Cheminformatics, 2015, 7(S1): S2. DOI:10.1186/1758-2946-7-S1-S2
[25]	MU X F, WANG W, XU A P. Incorporating token-level dictionary feature into neural model for named entity recognition[J]. Neurocomputing, 2020, 375: 43-50. DOI:10.1016/j.neucom.2019.09.005