ICLR 2022｜唯快不破！面向极限压缩的全二值化BiBERT 作者秦浩桐量子位转载|公众号

作者秦浩桐
量子位转载|公众号QbitAI
近年来，预训练语言模型在自然语言处理上表现出色，但其庞大的参数量阻碍了它在真实世界的硬件设备上的部署。
近日，机器学习顶会ICLR2022接收论文结果已经正式公布，至少有9项工作展示了神经网络量化方向的相关进展。
本文将介绍首个用于自然语言任务的全二值量化BERT模型——BiBERT ，具有高达56.3倍和31.2倍的FLOPs和模型尺寸的节省。
这项研究工作由北京航空航天大学刘祥龙教授团队、南洋理工大学和百度公司共同完成。

文章图片
预训练语言模型在自然语言处理上表现出色，但其庞大的参数量阻碍了它在真实世界的硬件设备上的部署。
现有的模型压缩方法包括参数量化、蒸馏、剪枝、参数共享等等。
其中，参数量化方法高效地通过将浮点参数转换为定点数表示，使模型变得紧凑。
研究者们提出了许多方案例如Q-BERT[1]、Q8BERT[2]、GOBO[3]等，但量化模型仍旧面临严重的表达能力有限和优化困难的问题。
幸运的是，知识蒸馏作为一种惯用的辅助优化的手段，令量化模型模仿全精度教师模型的特征表达，从而较好地解决精度损失问题。
在本文中，来自北航、NTU、百度的研究人员提出了BiBERT ，将权重、激活和嵌入均量化到1比特（而不仅仅是将权重量化到1比特，而激活维持在4比特或更高）。
这样能使模型在推理时使用逐位运算操作，大大加快了模型部署到真实硬件时的推理速度。
我们研究了BERT模型在二值化过程中的性能损失，作者在信息理论的基础上引入了一个高效的Bi-Attention（二值注意力）机制，解决前向传播中二值化后的注意力机制的信息退化问题；提出方向匹配蒸馏（Direction-MatchingDistillation）方法，解决后向传播中蒸馏的优化方向不匹配问题。

文章图片
△图1BiBERT的整体框架图
BiBERT首次证明了BERT模型全二值化的可行性，在GLUE数据集上的准确性极大地超越了现有的BERT模型二值化算法，甚至超过了更高比特表示的模型。
在模型计算量和体积上， BiBERT理论上能够带来56.3倍和31.2倍的FLOPs和模型尺寸的减少。
方法
Bi-Attention：二值化注意力机制
我们的研究表明，在BERT模型的注意力机制中， softmax函数得到的归一化注意力权重被视为遵循一个概率分布，而直接对其进行二值化会导致完全的信息丧失，其信息熵退化为0（见图2）。

文章图片
△图2直接对softmax函数应用二值化导致完全的信息丧失
一个缓解这种信息退化的常用措施是，在应用sign函数之前对输入张量的分布进行移位，可以表示为：

文章图片
其中，移位参数也被认为是二值化的阈值，希望能使二值化后的熵达到最大。
我们注意到， softmax函数是保序的，这意味着存在一个固定的阈值使二值化表示的信息熵最大化。
受到HardAttention的启发[4] ，作者用bool函数来二值化注意力权重A：

文章图片
通过应用bool函数，注意权重中值较低的元素被二值化为0 ，因此得到的熵值最大的注意权重可以过滤出关键部分的元素。