如果只能做整数Integer运算还能用BERT吗？

想当年，其实估摸着也就大半年前，多多同学还在实验室瞪大眼睛盯着一种叫做xilinx系列的板子，调试着一种叫做VHDL的语言，还记得那个写代码的工具叫做Vivado，不知道大家听说过没有？那个时候，我想实现一个复杂的公式，涉及的计算稍微复杂点（比如来个开方）就要写一大串代码（虽然常用的复杂函数是有IP核可以调的），同时调试过程十分麻烦，甚至要具体到clock对齐。总而言之，十分难忘。那个时候业余时间写下一行Python代码解决一个问题，简直可以直呼“爽啊”。当然，硬件代码虽然难写，但毕竟计算速度、能耗比、并行优势一直很好，所以即便不好写，还是依旧使用广泛。

那么回到今天的主题，在这样只能进行整型数计算的设备上，我们的无敌的BERT如何部署上去呢？当多多看到这篇文章的时候，立刻就来分享给大家了。以下内容主要参考文章：I-BERT: Integer-only BERT Quantization。代码已开源，HuggingFace Git 仓库也有。

相关背景

1. 最近两年，以BERT为代表的预训练模型大火，在自然语言处理/最近计算机视觉超越各种达到新SOTA，当然预训练模型虽然效果一个比一个好，参数大小也是越来越大。
2. 想要把这些大模型高效（存储空间小，运算速度快）跑起来大致有这么几个办法： A. Pruning，去掉不太重要的参数； B. 知识蒸馏，找个小一点的模型仿大模型； C. 更快更好的计算框架设计； D. 结合计算硬件对模型进行针对性的优化设计； E. Quantization，量化，比如32bit浮点数变成16bit浮点数甚至8bit整型数。
3. 目前已有很多压缩Transformer的研究，其中Quantization方面的也不少，但这些Quantization的方法大多不是完完全全的整型数量化（GELU、Softmax、LayerNorm部分都还是用浮点数计算的）。
4. 新的Integer-only优秀硬件不断更新，比如ARM处理器。

I-BERT主要贡献

1. 针对激活函数GELU，Softmax函数，I-BERT提出了一种高效并且正确的integer-only的估算方法，该方法为基于整型计算实现的轻量级二阶多项式。
2. 针对Transformer里的LayerNorm，使用integer-only迭代算法对平方根进行估算。
3. 基于GELU、Softmax、LayerNorm的估算，将Transformer进行了量化。 具体的Transformer量化为：用INT8乘法和INT32累加器处理Embedding和Matrix multiplication；然后基于INT32的累加结果计算GELU、Softmax和LayerNorm，然后再将INT32量化回INT8。整个计算图中的参数和激活函数计算都是用Intergers。图1展示了I-BERT和其他工作的区别。
4. I-BERT在基于Roberta-base/large进行了实现，并在GLUE上取得了和全精度模型相近的效果； 和32位浮点计算相比，将I-BERT基于TensorRT部署到T4 GPU上，有4倍的加速效果。

图1 对比I-BERT和之前的工作

I-BERT实现

I-BERT全文都在传递一个思想，其实和我上一篇写的贝叶斯优化很相似：如果一个函数的计算解决不了，就找个近似函数，如果近似函数不好找，就限定自变量范围来找！

1

基本量化方法

这里简单介绍以基本的量化概念和方法，在uniform symmetric quantization框架下（对称、均匀的量化），一个实数x会被映射成一个整数q，q的范围是[-2^(b-1), 2^(b-1) -1]，b为量化的bit数，正式定义为：

简单来讲为：

1. 如果一个数x不在[-a,a]范围内则直接用x=a或者x=-a，如果在范围内则保持x。
2. 得到x除以2^(b-1)-1得到的整数商作为量化结果。

把量化的结果转化为原始值也就是用q乘以上式中的分母，可以看到量化是损失的，一般来说，量化之后再反量化过程无法得到和原来一摸一样的数字。除了均匀量化，还有很多非均匀、不对称的量化方法，这里不再深入介绍。

1

Integer-Only的非线性函数GELU计算

我们先放出非线性GELU的表达式，然后思考下这么复杂的函数如何用整数计算进行近似估计，orz我感觉挺难的反正：

I-BERT是这么解决这个问题的？

由于整型计算可以计算2阶多项式，比如a(x+b)^2 + c。那么先用2阶多项式估计GELU函数，再用整型计算2阶多项式。为什么是2阶不是更高阶呢？因为阶数约高，估算函数也越复杂，作者发现2阶多项式足够估计GELU了。

2阶多项式估算GLUE又可以转化为优化如下目标函数，主要是对erf函数进行估计：找到一个2阶多项式，系数为a,b,c，能让GELU函数与2阶多项式误差最小。

erf函数有这么两个特点：对于输入很大的值而言，erf结果基本趋近于1或者-1，并且erf函数还是一个奇函数，因此只考虑优化正实数域的优化。在这两个特点下，作者找到一个相对较优的2阶多项式：

其中a=-0.2888，b=-1.769。sgn为符号函数。最终GELU的近视表达为：

整个GELU的估算和效果如图所示：

图2 GELU算法和效果

1

Integer-Only的softmax计算

个人觉得这几个估算十分精妙！

softmax函数的表达式是和特点：

然后可以发现x=x-x_max之后都会变成非正数，然后任何一个非正数可以表达为x=(-ln2)z + p，其中z是一个非负整数，p是一个在[-ln2,0]之间的实数，因此x的指数可以表示为：

>>是移位运算，在整型逻辑计算硬件中十分常见。最终将估计问题变化为简单的：估计[-ln2,0]之间的p的指数值。同样使用2阶多项式进行估计，最终得到：

到这里，整个integer-only的softmax计算逻辑如图所示：

图3 Softmax估计算法

1

Integer-Only的Layer Norm

Layer Norm在Transformer中大量使用，并且使用了非线性函数来计算方差和均值。

而且，注意LayerNorm的计算是在模型运行的是时候动态计算的，所以这些开方、平方运算也需要跟着一起算。不过这些平方、开方运算在整型运算里面还算常见，常见的方法就是迭代求解。如下图所示，便不再详细介绍啦：

图4 Layer Norm中的平方根估算方法

1

I-BERT的效果

虽然说I-BERT能完成用整型计算得到Transformer的结果了，但也不能损失太多的效果。总体上I-BERT的效果是不错的。如下图所示：

图5 GLUE任务上的效果

同时I-BERT在硬件上的运算速度也很好，相比于浮点数计算有2-4倍的加速。

图6 硬件运行加速对比

总结

本文对I-BERT进行了介绍，整体上I-BERT能完全基于整型计算进行，并且对于GELU、Softmax、LayerNorm的估算非常有借鉴意义。笔者能从这几个估算过程感受懂啊数学之美啊！！！！！