​EMNLP 2022 | 预训练语言模型的时空效率优化

©PaperWeekly 原创 · 作者 | 申博文

单位 | 中国科学院信息工程研究所

研究方向 | 自然语言处理

论文标题： 

COST-EFF: Collaborative Optimization of Spatial and Temporal Efficiency with Slenderized Multi-exit Language Models

收录会议： 

EMNLP 2022

论文链接：

https://arxiv.org/abs/2210.15523

代码链接：

https://github.com/sbwww/cost-eff

动机

众多基于 Transformer 的预训练语言模型（Pre-trained Language Models，PLMs）不断刷新着各项任务的性能，却存在体积大、推断慢等效率问题。对于资源受限的设备和应用场景，需要一种在空间和时间上高效，且在任务推断上准确的模型。

为了得到这样的模型，现有一些研究对 PLM 进行静态压缩 [1]。然而，单纯地进行静态压缩难以得到一个合适的模型，因为压缩后的模型很可能对简单样本而言仍有冗余，对复杂样本而言能力不足。为了使模型意识到输入样本的复杂性差异，Xin et al [2]，Liu et al [3] 等将 PLM 修改为多出口模型（即模型的多个部位都具有输出分类器），并使用动态提前退出方法进行推断加速。

我们发现，使用动态提前退出方法来加速小容量的压缩模型推断会造成较大的性能损失，其原因在于，多出口模型的浅层与深层在目标上存在不一致性。具体来说，浅层模块需兼顾做出预测和提取更深层所需的信息两个目标，而深层更多关注做出预测。这种不一致性在多出口模型中普遍存在，大容量模型有较好的能力缓解该问题，但小容量的压缩模型难以做出权衡。

为了解决上述问题，我们提出了 COST-EFF 来整合静态模型压缩和动态推断加速，实现空间和时间上的协同优化。具体来说，我们

1. 将 PLM 的宽度细化，而深度保持不变，保留模型提取深层知识的能力 [4]。同时，使用逐层的动态提前退出来减小模型深度带来的推断开销，根据样本复杂性动态地控制模型计算量，加速推理。

2.提出了一种联合训练方法，能够校正 PLM 的细化过程和压缩后的恢复训练过程，缓解在压缩的多出口模型上，浅层和深层目标不一致带来的权衡问题，提升协同优化模型的性能。

▲ 图1 COST-EFF示意图，Emb为嵌入层，Tfm为Transformer层，Clf指输出分类器

方法

COST-EFF 主要针对 Transformer 结构的预训练语言模型进行效率优化，优化方法包括：

1. 静态模型细化

• 词嵌入（Word embedding）矩阵的分解

• 多头自注意力（Multi-Head Attention，MHA）和前馈网络（Feed-Forward Network，FFN）的结构化剪枝

2. 动态推断加速

• 多出口网络的推断

• 多出口网络的训练

3. 联合训练流程

具体的，COST-EFF 结构和流程如下图所示，接下来将逐步介绍其细节。

▲ 图2 COST-EFF 结构（上半部分）和各部分对应的优化方法（下半部分）

2.1 静态模型细化

2.1.1 词嵌入矩