MNN量化—ADMM优化算法

作者：糖心他爸
链接：https://zhuanlan.zhihu.com/p/81243626
来源：知乎

目录

量化的模型建立

ADMM算法

---

量化的模型建立

现在我们知道了如何做量化推断，下一步是如何去建立量化模型，或者说我们应该用一个什么样的方式才能求得量化权重和量化输入呢？现在已知我们的输入为fp32，我们想用int8来对原fp32的数据进行表示，其中的转化关系假设为：

fp32和int8值之间的转化关系

对于每个数据块，我们都希望用一个尺度因子来对其进行int8的刻画，这个尺度因子就是上面式子中的s。经过上面式子之后，我们就能找到int8值和fp32值的一个对应关系。下面我们看看具体的从fp32到int8的值求解过程，这里假设我们已经知道了尺度因子s和阶段关系，那么int8的值我们可以通过如下式子获得：

                                                                    

int8值的求解过程

其中E是一个分段函数，表示如下。在数轴上看这个函数，是一个分段线性函数，且每个分段的导数为0，在阶梯跳跃处的整数点不可导。这个性质是一个很好用的性质，文中后续将会用到这个性质。

                                                        

截断函数

在经过了这一系列过程后，我们就由fp32值得到了int8的值，我们进一步抽象这个过程，为了后面的数学推导做准备。为此，我们再定义一个函数f，我们称之为编码函数；反过来f的逆，我们称之为反编码函数，由上式我们很容易得到：

                                              

所以，在进行int8值的量化编码时候，我们执行了如下操作，

                                                                         

int8的编码过程，由于进行了数据截断，该过程不可逆，所以这个是一个信息损失的过程，为了度量损失的程度，我们下面将会介绍两种思路，也是MNN中解决该问题的两种算法。其目的都是为了，尽量的在编码过程中减少信息的损失。一种是从概率分布的角度出发的，目的是让编码后的数据分布和编码前的数据分布尽可能的一致，采用的概率分布度量为KL-散度；另一种是从最优化优化角度出发的，目的是让编码后的数据在进行反编码后的结果，尽可能和原数据接近，采用的度量可以是L2、L1或是广义范数，但由于L2度量的优化算法求解方便，我们后面也主要基于L2度量进行讨论。

 

ADMM算法

MNN中的ADMM算法是从优化的角度出发，来保证编码前后数据尽可能相似的方法。

首先定义一个目标：也就是希望编码后的数据（int8）经过反编码后，跟原数据的”距离“尽可能的接近。该描述可以通过如下数学公式进行刻画：

                                                                             

进一步的，假设我们这里选定的度量D为L2度量，f和E的选取跟文章一开始时候介绍的一致，我们可以将上面式子转化如下目标函数：

                                                                       

所以，求解尺度因子s的过程，就变成了极小化上述目标函数的过程。而且，上面我们也提到过，E的函数性质很有特点，其在阶段域内整点处处不可导，其余点处处导数为0。利用该阶段函数的特点，我们很容对上述目标函数进行求解。不过先别急，我们先算算上面目标函数相对于s的导数，如下：

                                                           

由于E函数的特点，我们知道E关于s的隐函数求导，在可导处的函数值处处为0，不可导的地方我们先忽略掉（工程实现只需要避免奇点即可）。我们可以将上述式子再进一步简化如下：

                                                               

再得到目标函数关于s的导数公式以后，后续的求解就非常简单了，我们给出几种求解思路。

第一种，比较直接的方法，可以采用梯度下降法，流程简单如下：

## 给定s的初值（可以统计数据的min、max进行强映射初始化）
1、 计算L关于s的梯度
2、 更新s的值
3、 判断终止条件
4、 更新步长

第二种，稍微麻烦点，由第一种算法我们可以知道，步长的选取是启发性的，收敛性虽然基本都能 得到保证，但如果我们想自适应的计算最优步长，可以采用golden-search的求解方式，求解过程跟梯度下降法很接近，如下：

## 给定s的初值（可以统计数据的min、max进行强映射初始化）
1、 计算L关于s的梯度
2、 计算最优步长
3、 更新s的值
4、 判断终止条件

第三种方式，类似一阶收敛（梯度下降、golden-search）和二阶收敛（牛顿等）我就不再拓展了；我们这里介绍一下MNN中采用的优化算法，ADMM（交替方向法）。在得到L关于s的梯度公式后，如下图：

                                                             

其进行了一个概念转化，求解L的极小化问题，转化成求解L关于s的导数为0的数学问题。但由于函数E中还包含了未知数s，MNN在求解这个数据问题的时候，采用了ADMM的思想，先freezeE中的s变量，进行s值的求解；然后再通过求解后的s去估计E的值，如此反复直到s值收敛到最优。于是上述的式子要弱化为如下公式：

                                                         

所以迭代的过程就变成了交替迭代求解如下两个公式的过程，

                                        

ADMM的迭代过程

于是算法流程可构造如下：

## 给定s的初值（可以统计数据的min、max进行强映射初始化）
1、 计算E的值
2、 更新s的值
3、 判断终止条件

关于三种算法的终止条件判断，可以采用l2 loss的方式即||L(k+1) - L(L(k)) || / N。也可以采用相对变化率的方式，即当前迭代变化和上次迭代变化的比值。不过其实不判断问题也不大，如MNN中的源码实现中是按照最大固定迭代次数来实现的。