祝大家中秋快乐,文末送书啦~~~
随着GPT、BART等大型深度语言模型的问世,语言模型的运行效率、内存开销成为了备受关注的研究点。在实际应用中,可能没有那么多机器,或者需要提升运行效率,这个时候就需要对语言模型进行压缩,让模型的运行性能能够达到应用标准。基于量化的模型压缩在BERT以及CV领域的一些模型中实现了模型效果微小下降前提下的压缩,华为和港大提出了适用于GPT等自回归语言模型的量化压缩技术,并被评选为ACL 2022的杰出论文。接下来先给大家介绍基于量化的模型压缩方法的基本思路,然后介绍这篇论文提出的改进方法。
1
量化的模型压缩
BinaryConnect: Training Deep Neural Networks with binary weights during propagations(NIPS 2015)中提出了基于量化的模型压缩方法,核心是将模型中的参数转换成1或-1这两种数值,从而将乘法运算转换为加法运算,进而实现减小运行开销和压缩模型尺寸。
BinaryConnect的核心主要在于两点,一个是如何将参数实数值转换为1或-1,另一个是二值化后的参数如何参与网络训练和参数更新。针对第一点,文中提出两种二值化方法,第一种是大于0的参数映射到1,其他参数映射到-1;第二种是随机二值化,根据权重大小生成一个概率,映射到1或-1,公式如下:
第二个点是如何应用BinaryConnect,在前向传播和反向传播过程中,使用二值化的参数参与计算,在梯度更新参数时,使用原始精度更高的参数更新。因为在随机梯度下降中,参数的更新需要保证较高的精度才能使更新过程稳定。
二值化的一个好处是,相当于给权重和激活函数增加了噪声,起到了网络参数正则化的作用,和dropout等有类似的功效。
BinaryConnect方法存在的问题是,量化的过程没有可学习的参数,模型不知道如何调整量化的尺度适应模型训练过程中的损失。在量化权重参数时,由于权重的分布比较集中,并且可以加正则化约束,因此问题不大。但是在量化激活函数时,relu的输出没有上界,因此如何选择量化的阈值对于量化效果至关重要。PACT: PARAMETERIZED CLIPPING ACTIVATION FOR QUANTIZED NEURAL NETWORKS(2018)提出了PACT方法,让模型动态学习量化的clipping阈值,截断部分离群点,让参数的量化结果更精确。LEARNED STEP SIZE QUANTIZATION(ICLR 2020)中则对反量化进行优化,通过网络学习反量化的step size实现更精准的参数还原。
2
针对BERT模型的量化压缩
在NLP领域的BERT方面的模型压缩上,量化也是主流方法。Ternarybert: Distillation-aware ultra-low bit ber(ACL 2020)将BERT中的权重参数三值化为{-1, 0, 1}中的一个值,以及一个缩放参数,并且对于每层网络、每列参数使用不同的缩放参数。对于激活函数则用min-max的方法量化到8bit。由于量化后的BERT表征能力变弱,本文还是用了知识蒸馏的方式,让量化BERT作为student网络,使用一个teacher网络指导student网络的学习。整体的网络结构如下图所示:
Ternarybert对于权重的量化公式如下,目标是让缩放函数*量化后的值尽可能接近原始参数:
Binarybert: Pushing the limit of bert quantization(ACL 2021)对BERT实现了进一步的1-bit压缩。一般的方法压缩到1-bit效果损失很严重(如下图所示)。
造成binary压缩效果损失很大的原因在于,1-bit量化后的优化超平面非常复杂,导致模型的训练过程很困难,下面是一个可视化的优化超平面图,可以看到binary model由于量化过细导致超平面被分割成很多不平滑的区间:
为了解决这个问题,本文采用的思路是首先训练一个最终模型尺寸一半的TernayBERT,利用其更加平滑的超平面实现初始参数的学习。然后利用TWS Operator将参数进行转换,作为BinaryBERT的初始化参数,再基于此参数进行模型finetune。
3
量化压缩在PLM中的问题
将以前在BERT上取得成功的量化压缩方法应用到预训练语言模型(PLM)上,会存在比较明显的效果下降。下图是之前的PACT、LSQ等量化压缩方法随着压缩后的bit减少困惑度的变化。可以看到压缩到2-bit时出现非常明显的效果折损。
为什么之前的量化压缩方法应用到预训练语言模型上效果会这么差的?作者分析原因表示,量化压缩后,由于模型的容量减小,导致每个token的embedding趋于同化,即难以学到不同token的差异性embedding。此外,不同的Transformer层的权重参数分布差异也很大。这些问题在预训练语言模型这种单向结构中更加严重,因为单向的学习会导致量化误差累计问题。
下面两个图分别展示了使用不同量化压缩方法生成的embedding表示分布图,以及句子中两两token之间的cosine相似度矩阵。可以看到,PACT、LSQ等量化方法在应用到预训练语言模型后,embedding容易形成簇,也就是embedding之间差异性减弱。从下面的cosine相似度也能看出同样的问题。
下图则显示了不同层的权重分布,可以看到不同层的差异很大,并且每层的权重分布都是偏度非常高的,即有很多outliner。这对于确定合适的clipping阈值非常困难,即使自动学习clipping阈值的PACT模型也很难优化。
4
本文提出的方法
本文提出的方法主要包括两个核心,第一个是token粒度的对比学习,让每个token生成的表示更具备差异性,对应于上面提到的第一个难点;第二个是动态scaling学习,对应于上面提到的不同层权重分布差异大的难点。模型的整体结构如下图所示,包括一个量化student网络和全精度的teacher网络,这和TernaryBert比较类似。
在token粒度的对比学习中,作者使用一个token的量化student网络表示和全精度teachder网络表示互为正样本,以句子中的其他token为负样本。对比学习的架构使用了Momentum Encoder,即动量更新的方式,可以参考之前的文章对比学习中的4种经典训练模式,有详细介绍。对比学习部分的loss如下,其中s代表量化网络,t代表全精度网络,q代表表示在memory bank中的平滑版本:
此外,模型中也增加了直接优化student网络和teachder网络对应token表示距离的损失函数,其中z网模型输出的logits:
最终的损失函数是上述两个损失函数的加权求和。
第二个模块动态scaling学习,主要是针对模型不同层、不同module学习自适应的scaling,用来进行设置clipping的阈值。和PACT的自动学习scaling的差别主要体现在两点,一个是本方法不是直接学scale,而是先学习一个gamma,然后用下面公式计算clipping阈值:
通过这种方式,让学到的alpha能够更接近每层权重的整体scaling不偏离太远,其实相比让模型直接学习alpha更容易一些。其次,在梯度反传上,原来的PACT只回传在clipping区域外的权重,而本文提出的方法也会回传clipping区域内的权重,让clipping区域内和区域外的权重共同决定alpha,这样其实是更加合理的。
5
实验结果
下面是不同的量化压缩方法在多个数据集上的困惑度效果、模型尺寸等方面的对比,基础模型为GPT2。可以看到本文提出的方法在效果上均要显著优于之前的量化方法。
此外,作者也对比模型不同量化模型在Next Utterance Prediction和Abstractive Summarization等两个具体任务上的效果,也都优于其他量化方法。
6
总结
本文以ACL 2022杰出论文——量化压缩预训练语言模型为出发点,首先介绍了经典的量化压缩方法,包括BinaryConnect、PACT等,又进一步介绍了在BERT上的量化压缩工作。通过以上工作在预训语言模型上应用的不足,带大家理解了本文面临的难点和核心解法。