I-BERT: Integer-only BERT Quantization 论文学习

论文链接：
https://proceedings.mlr.press/v139/kim21d.html

摘要

基于转换器的模型，如BERT和RoBERTa，已经在许多自然语言处理任务中取得了最先进的结果。然而，它们的内存占用、推理延迟和功耗对于边缘甚至数据中心的有效推理都是禁止的。虽然量化是一个可行的解决方案，但以前基于Transformer的量化工作在推理中使用浮点算法，不能有效地利用仅整数的逻辑单元，如最近的Turing Tensor Cores，或传统的仅整数ARM处理器。在这项工作中，我们提出了I-BERT，一种新的基于Transformer的模型的量化方案，它量化了仅整数算法的整个推理。基于非线性操作的轻量级仅整数近似方法，如GELU、Softmax和层归一化，I-BERT执行端到端仅整数的BERT推理，而无需进行任何浮点计算。我们使用RoBERTaBase/Large对GLUE下游任务评估我们的方法。我们表明，对于这两种情况，I-BERT与全精度基线相比，实现了相似（且略高于此）的精度。此外，我们对I-BERT的初步实现显示，与FP32推理相比，在T4GPU系统上的INT8推理加速了2.4−4.0×。该框架已经在PyTorch中开发，并是开源的(Kim，2021)。

1. 引言

在这项工作中，我们建议使用I-BERT来解决这些挑战。I-BERT采用了一系列新的基于整数量化的方案。特别，我们的贡献是：

1. 我们提出了新的核来实现高效和精确的GELU和Softmax的仅整数计算。特别地，我们用轻量级的二阶多项式来近似GELU和Softmax，这可以用仅为整数的算术来计算。我们利用不同的技术来提高近似误差，GELU的最大误差为1.8×10−2，Softmax为1.9×10−3。详见3.4和3.5。
2. 对于LayerNorm，我们通过利用一个已知的平方根整数计算算法来执行仅整数计算(克兰德尔&Pomerance，2006)。详见3.6
3. 我们使用这些近似的GELU，Softmax，和LayerNorm来设计仅整数的量化基于Transformer的模型。具体来说，我们用INT8乘法和INT32积累来处理嵌入和矩阵乘法(MatMul)。然后根据INT32上的累积结果计算以下非线性操作(GELU、Softmax和LayerNorm)，然后将其返回回INT8。我们用整数表示整个计算图中的所有参数和激活，并且我们从不将它们转换为浮点数。示意图说明见图1（右）。
4. 我们将I-BERT应用于RoBERTa-Base/Large，并评估了它们在GLUE(Wangetal.，2018)下游任务上的准确性。与全精度基线相比，I-BERT基线取得了相似的结果。具体来说，在RoBERTa-Base和RoBERTaLarge的GLUE下游任务上，I-BERT分别比基线高出0.3和0.5。请参见”选项卡。2在4.1中的细节。
5. 我们使用TensorRT(NVIDIA，2018)部署INT8BERT模型，并只使用整数内核进行非线性操作。我们证明，与FP32推理相比，INT8推理实现了高达4×的加速。请参见”选项卡。3在4.2中的细节。
   

3. 方法

3.1.基本量化方法

量化

反量化：

3.2.具有仅限整数算法的非线性函数

纯整数量化的关键是用整数算法执行所有运算，而不使用任何浮点计算。与线性操作(例如MatMul)或分段线性操作(例如ReLU)不同，这对于非线性操作(例如GELU、Softmax和LayerNorm)来说并不简单。这是因为之前的工作中的仅整数量化算法(Jacobetal.，2018；Yao等人，2020)依赖于算子的线性性质。例如，对于线性MatMul操作的MatMul(Sq)等价于S·MatMul(q)。这个性质允许我们将整数MatMul应用于量化的输入q，然后将比例因子S相乘，得到与将浮点MatMul应用于去量化的输入Sq相同的结果。重要的是，这个属性不适用于非线性操作，例如，GELU(Sq)！=S·GELU(q)。一个简单的解决方案是计算这些操作的结果，并将它们存储在一个查找表中(Laietal.，2018)。然而，这种方法部署在芯片内存有限的芯片上可能会产生开销，并将产生与查找表执行速度成正比的瓶颈。另一种解决方案是将激活进行去量化，并将其转换为浮点运算，然后用单精度逻辑计算这些非线性操作(Bhandare等人，2019年；Zafrir等人，2019年)。然而，这种方法并不只是集成的，也不能在不支持浮点运算的专门高效硬件上使用，例如，ARMCortex-M(ARM，2020)。
为了解决这一挑战，我们近似的非线性激活函数，GELU和Softmax，与多项式，可以用仅整数算术计算。计算多项式只包括加法和乘法，这可以用整数算法来完成。因此，如果我们能找到对这些操作的良好的多项式近似，那么我们就可以用仅使用整数的算术来执行整个推理。例如，一个表示为a(x+b)2+c的二阶多项式可以用仅整数算法有效地计算出来，如Alg.1所示

3.3.非线性函数的多项式逼近

关于用多项式近似一个函数有大量的工作(Stewart，1996)。我们使用一类插值多项式，其中我们给出了一组n+1个不同数据点{(x0，f0)，…，…，(xn，fn)}，我们试图找到一个最多n的多项式，在这些点上的函数值完全匹配。众所周知，存在一个唯一的多项式的次数最多为n，它通过所有的数据点(Waring，1779)。我们用L来表示这个多项式，定义为：

有趣的是，对于我们的问题，我们有两个旋钮来改变，以找到最好的多项式近似。由于我们知道实际的目标函数，并且可以查询任何输入的精确值，所以我们可以选择插值点(xi，fi)作为函数上的任何点。第二个旋钮是选择多项式的次数。当选择一个高阶多项式会导致更小的误差(见附录B)时，在这方面有两个问题。首先，高阶多项式具有更高的计算量和内存开销。其次，用低精度的仅整数算法来计算它们是具有挑战性的，因为当将整数值相乘时可能会发生溢出。对于每一个乘法，我们都需要使用双精度，以避免溢出。因此，挑战是找到一个好的低阶多项式，可以非常接近于在Transformers中使用的非线性函数。这就是我们接下来讨论的，对于GELU和Softmax，分别在3.4和3.5中，我们证明了只使用一个二阶多项式可以得到一个接近的近似。

3.4.仅整数的GELU

GELU可以看作 dropout的思想和relu的结合，(在keras和torch的transformer实现的代码里都是直接用relu+dropout而没有使用gelu)　，主要是为激活函数引入了随机性使得模型训练过程更加鲁棒

本人对此激活函数不感兴趣，感兴趣的可参考原文，此处只贴一个算法：

3.5.仅整数softmax

用整数算法来近似的Softmax层是相当具有挑战性的，因为在Softmax中使用的指数函数是无界的，并且变化迅速。因此，先前的Transformer量化技术(Bhandare等人，2019年；Zafrir等人，2019年)使用浮点算法处理这一层。之前的一些工作提出了用插值来查找表（施劳道夫，1999），但与之前一样，我们避免查找表，并努力实现一个基于纯算术的近似。此外，虽然(Hauser&Purdy，2001)提出了指数函数的多项式逼近方法，但它使用了显著的高度多项式，并且只适用于有限的有限域。
与GELU类似，我们不能使用高阶多项式，但即使使用这样的多项式也不能近似Softmax中的指数函数。然而，我们可以通过限制Softmax的近似范围来解决这个问题。首先，为了稳定我们减去输入数值最大值：

其中，xmax=maxi(xi)。请注意，现在对指数函数的所有输入，即˜xi=xi−xmax，都变成了非正的。我们可以将任何非正实数˜x分解为˜x=(−ln2)z+p，其中商z是一个非负整数，其余的p是在(−ln2,0]中的一个实数。那么，˜x的指数可以写成：

其中>>是移位操作。因此，我们只需要在p∈(−ln2,0]的紧致区间中近似出指数函数。与所有实数的域相比，这是一个要小得多的范围。有趣的是，这种方法的一种变体在e Itanium 2 machine
from HP使用了(Detrey&deDinechin，2005；Thomas等人，2004)，但用一个查找表来评估exp§。
我们使用一个二阶多项式来近似在这个范围内的指数函数。为了找到多项式的系数，我们在(−ln2,0]的区间内最小化到指数函数的l2距离。这就导致了以下近似值：

用等式中的指数项代替Eq.12用这个多项式得到的结果是i-exp：

其中，z=[−˜x/ln2]和p=˜x+zln2。这可以用整数算法来计算。算法3描述了使用i-exp的软最大功能的仅整数计算。图2（右）绘制了i-exp的结果，它与指数函数几乎相同。我们发现这两个函数之间的最大差距只有1.9×10⁻³。考虑到单位区间的8位量化引入了1/256=3.9×10⁻³的量化误差，我们的近似误差相对可以忽略不计，可以纳入量化误差中。

3.6.仅整数LayerNorm

LayerNorm通常用于Transformers中，涉及一些非线性操作，如除法、平方和平方根。此操作用于跨通道维度规范化输入激活。标准化过程描述如下：

这里，µ和σ是输入跨通道维度的平均值和标准差。这里的一个微妙挑战是，NLP任务的输入统计信息（即µ和σ）变化迅速，这些值需要在运行期间动态计算。虽然计算µ很简单，但计算σ需要平方根函数。
通过(Crandall，2006)中提出的迭代算法，可以用仅整数算法有效地计算平方根函数，如Alg所述。 4.给定任意非负整数输入n，该算法基于牛顿法迭代搜索[n^(1/2)]的精确值，只需要整数运算。该算法在计算上是轻量级的，因为对于任何INT32输入，它最多可以在4次迭代中收敛，并且每次迭代只包含一个整数除法、一个整数加法和一个位移操作。而LayerNorm中的其他非线性运算，如除法和平方运算，则直接用整数算法计算。

4.结果

在本节中，我们首先使用通用语言理解评估(Wangetal.，2018)(GLUE)基准测试（4.1）来衡量I-BERT的准确性。然后，我们再讨论使用直接硬件部署的I-BERT的延迟加速，并将其与纯FP32模型（4.2）进行比较。最后，我们进行消融研究以展示我们的整数近似方法（4.3）的有效性。

4.1.用GLUE基准准确性评价

4.2.延迟评估

4.3.消融研究

5.结论

我们提出了一种新的只有整数的变压器I-BERT量化方案，其中整个推理都是用纯整数算法进行的。I-BERT的关键元素是非线性操作的近似方法，如GELU、Softmax和LayerNorm，它们使它们能够通过整数计算进行近似。我们在RoBERTa-Base/大型模型上对I-BERT进行了实证评估，其中我们的量化方法将平均GLUE评分提高了0.3/0.5分。此外，我们直接部署了量化模型，并测量了端到端推断延迟，表明与浮点基线相比，I-BERT在 TeslaT4 GPU上可以实现高达4.00×的加速。作为未来工作的一部分，我们也可以考虑使用我们的近似值来提高训练速度。例如，人们可以考虑在训练期间用i-GELU替换GELU。此外，还需要进一步的研究来评估i-GELU与GELU相比的性能效益。