Paper Reading||Training With Quantization Noise for Extreme Model Compression

论文(ICLR 2021)：https://arxiv.org/abs/2004.07320
源码(Pytorch)：https://github.com/pytorch/fairseq/tree/master/examples/quant_noise

Motivation

1. STE的作用会在梯度中引入偏差，该偏差取决于量化电平，而这样的梯度不利于参数更新。
   
   

Prior Knowledge

Description of Quantification

图 (1). 定义实矩阵及其分块矩阵

 　如图 (1) 左侧所示，定义一个 $n$ 行 $p$ 列的实矩阵 $W$ ，其元素由 $e_{kl}，(1\le k\le n,1\le l\le p)$ 表示。那么通过将 $W$ 分块成 $m\times q$ 的分块矩阵，则可以有如右箭头所示。然后可以根据这个分块矩阵得到一个对应的索引矩阵 $I$ ，如图 (2) 所示。

图 (2). 由分块矩阵对应出索引矩阵

图 (3). 索引与Codebook的关系

 　如图 (3) 左侧所示，接着定义一个 码本集 (Codebook Set) $C$ ，包含 $K$ 个码字 (Codeword) $c[i],(1\le i\le K)$ 。而索引矩阵 $I$ 的元素 $I_{kl}$ 即能够索引到对应的Codeword，具体是哪个Codeword和 $b_{kl}$ 的最近邻四舍五入有关。需要注意的是，标量量化时 $c[i]$ 和 $b_{kl}$ 为标量 (即仅代表一个元素) ，而矢量量化时 $c[i]$ 和 $b_{kl}$ 为向量 (即可代表多个元素) 。

 　通常来说，量化就是为了得到码本和索引矩阵。一个真量化后的模型在实际部署时不再储存 $W$ ，而储存索引矩阵 $I$ 和码本集 $C$ ，在实际前向推理时通过索引映射操作得到近似结果 $\hat{W}$ ，如图 (4) 所示。

图 (4). 通过码本和索引矩阵恢复出近似结果

Scalar Quantization

 　Scalar Quantization (标量量化) 即是Element-wise的量化，对应于图 (1) 中有 $e_{kl}=b_{kl}$，即每个块只表示一个元素。而对于码本中码字的个数，有 $K=2^N-1$ ，这个 $N$ 即为位宽 (bit-width)，减1是因为避免运算溢出。我们最常讨论的均匀量化、线性量化就是标量量化。

 　如图 (5) 所示， $e$ 经过缩放变换再四舍五入后得到的 $\hat{e}$ 为整数，其实就是索引；而索引再通过反变换得到码字 $c$ 。当然这样的一个近似过程是个特例，是线性量化的特性，没有通过索引直接映射，没有如图 (4) 那样的过程。

图 (5). 标量量化的过程

 　需要注意的是，标量量化的作用对象可以是Weight和Activation。

Vector Quantization

 　Vector Quantization (矢量量化) 即是Group-wise的量化，对应于图 (1) 中有 $e_{kl}\ne b_{kl}$，即每个块可表示多个元素。而使用聚类算法比如K-means来Optimize出质心 (Centroid) 作为Codeword。矢量量化中Codeword的个数即为质心的个数。

 　举个例子，若设图 (1) 中分块矩阵的列数与原矩阵中的列数相等，即 $q=p$ ，并要求提供8个Codeword (即8个近似值) 时，可得到如图 (6) 所示的可视化效果。此时有 $2\times p$ 个 vectors。

图 (6). 矢量量化过程
 　因为8个Codeword需要8个质心，那么质心对应的索引值就要 Int3 来表征。每个Codeword向量能表示 $\frac{n}{2}$ 个元素，而Codeword的每个元素由Float32表示。也就是说采用矢量量化后占用储存空间 (bits) 为：

 　需要注意的是，矢量量化的作用对象只可以是Weight。 因为压缩程度太大，不适用于推理过程中会变化的Activation。

Product Quantization

 　Product Quantization (乘积量化) 属于矢量量化的一种。同样使用聚类算法来Optimize出质心，不过质心通过排列组合的方式才表示成码字。

 　还是举个例子，若设图 (1) 中分块矩阵的列数与原矩阵中的列数相等，即 $q=p$ ，并要求提供8个Codeword (即8个近似值) 时，可得到如图 (7) 所示的可视化效果。此时有 $2\times 3\times p$ 个 vectors。

图 (6). 乘积量化过程

 　“乘积”，顾名思义，采用“笛卡尔乘积”来完成质心 $⟶$ 码字 的映射。8个码字仅需要2个质心，质心通过排列组合有 $2^3=8$。那么质心对应的索引值只要 Int1 来表征。

 　采用乘积量化后占用储存空间 (bits) 为：

 　可以看到乘积量化比矢量量化还要更进一步压缩空间。
 　需要注意的是，乘积量化的作用对象只可以是Weight。 因为压缩程度太大，不适用于推理过程中会变化的Activation。

iterative Product Quantization

 　iterative Product Quantization (迭代乘积量化) 是乘积量化+迭代式策略，针对于模型压缩问题而提出的。

 　该方法可总结出以下流程：

1. 计算出所有各层的质心。
2. 逐层进行量化 (这才使用质心来近似、替代Weight)。被量化的层保留量化效果，然后Fine-tune其他没有被量化的层。
3. 最后再对所有层进行一次全局Fine-tune。

 　注意Fine-tune时会对对应的质心也做Fine-tune，更新方式如下：

 　其中 J c = { ( k , l ) ∣ c [ I k l ] = c } J_{c}=\{(k,l) | c[I_{kl}]=c\} Jc​={(k,l)∣c[Ikl​]=c} 是一个集合，是和归属于该质心的矢量的脚标集合。

Method

Proposal：Combining Scalar Quantization and Product Quantization

 　作者提出，可以将标量量化和乘积量化/迭代乘积量化结合起来。具体就是：

1. 对Weight使用迭代乘积量化的方法。
2. 对各层生成的质心进行标量量化、对Activation使用标量量化。

 　因为质心原本是以FP32格式储存的，对其进行量化可以进一步提高压缩率。

Proposal：Quant-Noise

 　这是作者这篇论文提出的最核心的创新。作者把Quant-Noise看作是一种观念，一种可以在各种量化方式中通用的观念。个人认为可以将其理解为一种trick。可视化如Fig 1 所示。

 　作者首先提出畸变操作的定义：

 　其中 $J$ 也是一个下脚标集合，而该集合是分块矩阵中脚标全集的子集，即 J ∈ { ( k , l ) } J \in \{(k,l)\} J∈{(k,l)} for $1\le k\le m，1\le l\le q$ 。$\phi$ 定义为Noise function。畸变操作的作用对象是一个分块矩阵，而Noise function的作用对象是某个Block。需要注意的是， $J$ 是根据给定的比例随机生成的，比如random_ratio=0.5时，会随机选取、抽样分块矩阵中50%的块来生成 $J$ 。当random_ratio=1时，等价于所有Block进行量化并使用STE来传播梯度，等价于纯QAT。

 　接着作者称该畸变操作是可以叠加的，总畸变操作可以等于子畸变操作的叠加，且顺序可调换：

 　畸变操作可以在各种量化方式中通用，具体表现为Noise function的定义。

• 在标量量化中，Noise function定义为：

• 在乘积量化中，Noise function定义为：
  
  
  
   　并且作者同时定义另一个可微可传递梯度的Noise function “proxy”，并通过叠加的方式使用proxy和PQ：

• 作者引申概念，在把剪枝操作也看作是一个产生噪声的原因。而由于作者的初衷并不是真正删除掉某个卷积核，作者只是想做类似LayerDrop的操作，也就是说此时剪枝等价于LayerDrop，且不可微不可传递梯度，那么Noise function定义为：

 　总的来说，作者把量化\剪枝当成一种噪声源。在Training中每次量化\剪枝时随机选择一定比例的Block来执行量化，反向传播时通过有偏梯度来更新参数；而剩下的Block不执行量化\剪枝，反向传播时通过无偏梯度来更新参数。这么做就可以达到类似Dropout、LayerDrop的效果，使得Weight的各个元素之间的相关性减少、整体的鲁棒性增强。

Experiment

 　由于该论文做的实验太多了，包括补充的实验，本人就摘选一些比较有代表性的实验来描述。

 　Fig 2 (a) 是作者在Wikitext-103任务中对16层的Transformer进行量化后与其他方法的对比。Fig 2 (b) 是作者在MNLI任务中对RoBERTa进行量化后与其他模型的对比。Fig 2（c）是作者在Imagenet任务中对EfficientNet-B3进行量化后与其他模型的对比。可以看到作者的方法还是靠近理想的。

 　Table 1 是Quant-Noise与其他Trick在不同量化方法上的对比试验。需要注意的是PPL (困惑度) 越低模型性能越好。

 　Table 2 第一大列是作者在Wikitext-103任务中对16层的Transformer进行量化后与其他压缩方法的对比。第二大列是作者在MNLI任务中对RoBERTa进行量化后与其他压缩方法的对比。第三大列是作者在Imagenet任务中对EfficientNet-B3进行量化后与其他压缩方法的对比。

 　Table 3 是作者在MNLI任务中对RoBERTa进行迭代乘积量化的消融试验。

个人思考

 　这篇论文开头就讲了很多的定义和概念，对量化的理解还是起到一定作用的。作者的核心创新就是Quant-Noise观念，这种观念可以应用到量化、剪枝中。Quant-Noise思想与Dropout、LayerDrop有异曲同工之妙，Quant-Noise减小了元素间相关性、增强了鲁棒性。