论文标题:
OMNIQUANT: OMNIDIRECTIONALLY CALIBRATED QUANTIZATION FOR LARGE LANGUAGE MODELS
研究动机:
最新的训练后量化(PTQ)方法在减少内存占用和提高LLM的计算效率方面非常有效,但它们手动设计量化参数,导致性能较低,并且无法处理极低位量化。
为了解决这个问题,作者引入了一种名为(OmniQuant)的技术,这个技术既可以进行仅权重量化,又可以进行权重-激活量化。通过有效地优化各种量化参数,在不同的量化设置下实现良好的性能,同时保持PTQ的计算效率。
背景知识:
量化方法:
接报告开头所说,目前主要的量化方法有QAT和PTQ。QAT能在训练过程中进行量化,但是由于成本过高难以应用于LLM。PTQ在模型训练结束后使用梯度优化,可惜调整大模型权重的过程非常耗时,但不管怎么说,成本是比较低的。于是大多数LLM优先考虑PTQ,但是PTQ在低比特情况下量化效果不佳。考虑中庸之道,作者希望在LLM量化中集成梯度更新,在保持PTQ的效率的同时反映一些QAT的方法。
AWQ:
这篇论文的一个主要思路来自AWQ这篇文章,而AWQ的思路来自于SmoothQuant(因为这俩本来就是同一个作者),本身也是相当于SmoothQuant的升级版。AWQ的一个核心思想就是,权重的对模型的重要程度,存在极强的不均衡性。1%的权重,可能主导的量化过程中损失的性能。假如我们在量化中保护一下这1%的参数,就能极大程度保护模型的性能不受影响。
首先我们要想办法找这1%的重要权重。AWQ和SmoothQuant一样,也是放眼channel。但是实验发现,保护那些权重值单纯比较大的数对性能改进没什么帮助,反倒是根据激活幅度选择的权重可以显著提升性能。虽然是仅权重量化,但是AWQ的一个关键思想是,让大的激活值对应的权重保持精度可以更有效的保持这些特征的精度。通过这种方式可以找到1%的权重。
找到之后就得考虑保护。其实最好的方法是保持这些值不量化,但是混合精度数据类型会使系统实现变得困难,我们需要一种平替。AWQ的保护方式类似SmoothQuant的思想,是寻找一个缩放比例s,在权重量化之前乘以这个比例,这样量化对该值的影响不会那么强烈,激活矩阵计算时再除以这个比例,这样数学上保持不变。通过上述方法可减小量化时产生的误差。
那什么样的缩放比例是好的呢?作者描述了下述公式:
这个是原权重乘激活矩阵。这个是保护并量化后的权重乘激活矩阵。那么只需要找一个数,使该loss函数值最小,那这个数就是我们要的s。这个比例呢,作者最后是用穷举法找出来的(作者在23年6月份的文章提出了一个看起来很巧妙的方法,然后10月份的第二版中放弃之)。
以上就是AWQ的主要原理。
文章的另一个主要思路来自OS+这篇文章。
Outlier Suppression+:
该文在作者前文Outlier Suppression针对激活异常值的基础上,提出了(OS+)框架。研究的动机是作者发现激活异常值的新特征:它们在通道之间呈非对称形状且比较集中。有问题的通道有的全是负值,有的全是正值,而且堆一块。
因此,OS+提出了包含了针对不对称性的通道级偏移和针对集中性的通道级缩放。
大写字母表示矩阵(例如,X)),小写字母表示向量(例如,x)。符号⊙和⊘分别表示矩阵或向量的逐元素乘法和除法。使用WX表示矩阵与矩阵的乘法。
其中 z是一个行向量 (z∈Rn),作用是对每个通道的激活进行偏移。通过精心设计的 z,新的张量X′可以消除异常的不对称属性。行向量s∈Rn对每个通道的偏移值进行缩放,并生成最终适合量化的激活值X。
为了保持数学上的等效性,OS+还提出了一种统一的迁移模式,将翻转后的偏移向量和缩放向量转移到后续模块中。
考虑一个常见情况,即该层后是个线性层紧随其后。原始的浮点数激活权重矩阵输出公式为:
将公式2反向后带入,得到
展开可得:
观察公式(3)结果,权重和偏置可以分别吸收s和z,从而得到:
这样等价于在下一层中更新W和b,完成了在后续模块中的更新。
那么,我们怎么计算偏移向量z和缩放向量s呢?
首先是z。偏移向量的设计应该消除通道间的不对称影响。因此,OS+设计将每个通道的中心z对齐到0,以便异常通道不会仅占据正面或负面。详细地说,z被定义为每个通道中最小和最大信号的平均值,即这个公式:
然后是s。
由于公式(4) 同样也会对权重进行缩放,可能导致较差的权重量化结果,这要求我们计算精细的缩放s值以实现激活和权重之间的量化平衡。
OS+认为关键点在于最小化由于迁移和量化导致的交互输出变化。因此提出了一个新的优化目标。
首先研究问题激活作为一个线性层的输入的简单情况(大图中的b)。与单独最小化激活和单独最小化权重的量化误差不同,OS+采用了任务损失的观点,关注它们的矩阵乘法输出。致力于最小化权重和激活进行缩放和量化后的变化,并使用均方误差计量变化。最小化目标公式如下:
之后研究多个线性层的情况。有点类似于自注意力(大图中的a),其中三个权重将与相同的缩放向量相乘,并使用相同的抑制激活进行计算。在这种情况下,通过缩放和量化矩阵产生的矩阵乘法输出被标记为Qq,Kq,Vq(原始输出记为Q,K,V)。
与上面的方法类似,多层最小化目标公式为:
为了实现上述的两个优化目标,OS+提出仅对异常通道进行缩放;
同时引入一个异常值阈值 t,用于将只有激活范围超过 t 的通道压缩到 (-t, t) 范围内,同时保持其他通道不变;这种技术将具有大量变量的复杂问题s简化为单变量t。然后我们采用简单的网格搜索来实现最小化目标。
在获得有效的t之后,缩放向量计算如下:
这样就可以得到缩放向量s。综上可实现偏移和缩放操作。
以上为OS+的介绍。
主要方法:
1、BLOCK-WISE QUANTIZATION ERROR MINIMIZATION
OmniQuant在AWQ关于loss的公式的基础上进行了改良,采用分块量化误差最小化的方式,并且可以通过可微分的方式优化其他量化参数,而无需转换整个模型。公式如下:
F表示LLM中的transformer块的映射函数,W和X是浮点数的权重和激活值,Qw()和 Qa()分别表示权重和激活值的量化器,Θ1是可学习的权重剪裁(LWC)的量化参数,Θ2可学习的等效转换(LET)中的量化参数,这两个会在后面的模块里讨论。
该方法具有两个优点。首先,OmniQuant可以同时优化在LWC和LET中的量化参数,使其能够包含仅权重和权重-激活两种量化,LWC只和权重相关,LET和权重激活都相关。其次,该方法对资源的需求极低。OmniQuant只使用一些最优量化参数,较之以前基于PTQ的方法更容易优化所有权重。
那么,这个LWC和LET模块如何发挥作用呢?
2、LEARNABLE WEIGHT CLIPPING(LWC)
首先是LWC,可学习的权重剪裁。
OmniQuant采用可学习的权重剪裁(LWC)模块,以降低在LLM中量化权重的难度。LWC通过优化剪裁阈值来确定权重的最佳动态范围。
之前那些旧的方法直接取固定剪切阈值,比如minmax量化,只固定取最大最小,再分成2^n个点,把每个值round到这个这些点上。
与此不同,OmniQuant通过优化剪切强度来实现,公式如下:
前面这部分基本就是上一篇一直在说的minmax量化,但h这里不一样。OmniQuant在这里用了两个改良的参数γ∈[0,1] 和 β∈[0,1] 来调整学习的上界和下界裁剪强度。这两个是用经典的sigmoid函数来实例化的。因此,在1的方程中,Θ1={γ,β}。
当γ和β都取1时,就是经典款MinMax量化。
通过用可学习的变量γ和β代替固定的参数,可以提高性能。
3、LEARNABLE EQUIVALENT TRANSFORMATION(LET)
最后是可学习等效转换(LET),这个东西可以进一步降低权重-激活量化的难度。这一块基本就是借鉴的刚才讲的Outlier Suppression+,也是采用逐通道缩放和逐通道平移来操纵激活的分布,公式都很像。
Linear layer:
结合上文,数学上的线性层表达式为:
其中,Y表示输出,δ和s分别是按通道偏移向量和缩放向量,X˜,W˜和B˜分别是等效激活函数,权重和偏差参数,‘⊘’和‘⊙’表示逐元素除法和乘法。
这里先是把量化难度从激活转到权重,再结合LWC降低权重量化难度的机制,可得下述公式:
其中,Qa是针对激活的基础的MinMax量化,而Qw是针对权重的(LWC),刚才说的可学习权重裁剪。
X˜中的缩放和偏移参数可以合并到归一化或线性层中,而W˜中的缩放因子可以融入到原始的线性权重W中。因此,上述方程中的等效变换可以有效地减少量化误差,而不会引入额外的参数或成本。
Attention operation:
在权重-激活量化设置中,文章也是搬OS+,将Q/K/V矩阵量化为低位。具体而言,自注意力亲和矩阵的可学习等效变换可以表示为:
其中,sa是缩放向量。量化关联矩阵计算可以表示为
这两个Qa都是MinMax量化。同时原文称,V的分布已经通过与输出投影线性层相关的逆变换进行了通道尺度的改变,故略去。
从这一小节可以得知,公式(1)中的Θ2={δ、s、sa}。
以上为OmniQuant的主要方法。