1bit量化

paper(OneBit: Towards Extremely Low-bit Large Language Models)

rank-1 approximation

奇异值分解
$A=U\Sigma V^T$
取最大的奇异值(${\sigma }_1$)，和最大的奇异值对应的左奇异向量和右奇异向量
$\hat{A}=u_1{\sigma }_1{v}_1^T$

(《Foundations of Data Science》的3.5节证明了这样分解是Frobenius范数下最好的秩1分解，证明我还没看)

$W=W_{sign}\odot W_{value}$，其中 $W_{value}=|W|$，$W_{sign}=\mathrm{Sign}(W)$（符号和绝对值的 Hadamard 积）

$W_{value}\approx a{b}^T$（把绝对值矩阵做秩 1 分解）

$W\approx W_{sign}\odot (a{b}^T)$

proposition 1

$\mathbf{X}{\mathbf{W}}^{\mathrm{T}}\approx \left[\left(\mathbf{X}\odot {\mathbf{b}}^{\mathrm{T}}\right){\mathbf{W}}_{\text{sign }}^{\mathrm{T}}\right]\odot {\mathbf{a}}^{\mathrm{T}}$
这个命题的是为了得到想到的网络结构。其中$b^T$赋给模型架构里的$g$，$a^T$赋给$h$。$X\odot b^T$操作表示对$b^T$进行广播（$b^T$是一个行向量，广播的意思就是把$b^T$一行一行重复地排列下去，直到和X的形状一样。），然后和X作hadamard积。

proof

$w_{ij}\approx s_{ij}\cdot a_i{b}_j$, where $s_{ij}$ is the element of ${\mathbf{W}}_{\text{sign }}$. Hence we have

$$\begin{array}{rl}{\left(\mathbf{X}{\mathbf{W}}^{\mathrm{T}}\right)}_{ij}& =\sum _k{x}_{ik}{w}_{kj}^{\mathrm{T}}=\sum _k{x}_{ik}{w}_{jk}\\ & \approx \sum _k{x}_{ik}{s}_{jk}{a}_j{b}_k\\ & =\sum _k{x}_{ik}{b}_k{s}_{jk}{a}_j\\ & =\sum _k{\left(\mathbf{X}\odot {\mathbf{b}}^{\mathrm{T}}\right)}_{ik}{s}_{kj}^{\mathrm{T}}{a}_j\\ & ={\left[\left(\mathbf{X}\odot {\mathbf{b}}^{\mathrm{T}}\right){\mathbf{W}}_{\text{sign }}^{\mathrm{T}}\right]}_{ij}{a}_j\\ & ={\left\{\left[\left(\mathbf{X}\odot {\mathbf{b}}^{\mathrm{T}}\right){\mathbf{W}}_{\text{sign }}^{\mathrm{T}}\right]\odot {\mathbf{a}}^{\mathrm{T}}\right\}}_{ij}.\end{array}$$

lemma 1

Let ${\sigma }_i(\mathbf{W})$ denote the $i$-th biggest singular value of matrix $\mathbf{W}$. The following inequality holds:

$${\sigma }_1(|\mathbf{W}|)\ge {\sigma }_1(\mathbf{W}).$$

Proof

According to the definition of induced norm, there are

$$\begin{array}{c}{\sigma }_1(\mathbf{W})=\Vert \mathbf{W}{\Vert }_2=\underset{\mathbf{x},\Vert \mathbf{x}{\Vert }_2=1}{\max }\Vert \mathbf{Wx}{\Vert }_2,\\ {\sigma }_1(|\mathbf{W}|)=\Vert |\mathbf{W}|{\Vert }_2=\underset{\mathbf{y},\Vert \mathbf{y}{\Vert }_2=1}{\max }\Vert |\mathbf{W}|\mathbf{y}{\Vert }_2.\end{array}$$

Note that for $\forall \mathbf{x},\Vert \mathbf{x}{\Vert }_2=1$ and we have

$$\begin{array}{rl}\Vert |\mathbf{W}||\mathbf{x}|{\Vert }_2^2& =\sum _i{\left(\sum _j\left|w_{ij}\right|\left|x_j\right|\right)}^2\\ & \ge \sum _i{\left(\left|\sum _j{w}_{ij}{x}_j\right|\right)}^2\\ & =\sum _i{\left(\sum _j{w}_{ij}{x}_j\right)}^2=\Vert \mathbf{Wx}{\Vert }_2^2.\end{array}$$

因为我们总是可以让y取x的绝对值，所以：

$$\underset{\mathbf{y},\Vert \mathbf{y}{\Vert }_2=1}{\max }\Vert |\mathbf{W}|\mathbf{y}{\Vert }_2\ge \underset{\mathbf{x},\Vert \mathbf{x}{\Vert }_2=1}{\max }\Vert \mathbf{Wx}{\Vert }_2.$$

This lemma is proved.

proposition 2

Given matrices $\mathbf{W}$ and $|\mathbf{W}|$, $\mathbf{W}={\mathbf{W}}_{\text{sign }}\odot |\mathbf{W}|$. We decompose these matrices in the way $\mathbf{W}=\mathbf{a}{\mathbf{b}}^{\mathrm{T}}+{\mathbf{E}}_1$ and $|\mathbf{W}|=\tilde{\mathbf{a}}{\tilde{\mathbf{b}}}^{\mathrm{T}}+{\mathbf{E}}_2$, where ${\mathbf{E}}_i$ denotes the error matrices. In terms of the Frobenius-norm, the SVID is closer to the original matrix $\mathbf{W}$ :

$${\Vert \mathbf{W}-{\mathbf{W}}_{\text{sign }}\odot \tilde{\mathbf{a}}{\tilde{\mathbf{b}}}^T\Vert }_F^2\le {\Vert \mathbf{W}-\mathbf{a}{\mathbf{b}}^T\Vert }_F^2.$$

这个命题主要想说，在Frobenius范数意义下，先取绝对值再分解比直接分解的误差小。

proof

Here we consider SVD to prove it. For SVD, the norm of the error matrix $\mathbf{E}$ in the rank1 approximation is the sum of the squares of all singular values except for the largest one. We have
$E=A-\hat{A}=A-u_1{\sigma }_1{v}_1^T$

$\Vert E{\Vert }_F^2={\sum }_{i,j}{\left(A_{ij}-{\sigma }_1{u}_{1i}{v}_{1j}\right)}^2$

$\Vert E{\Vert }_F^2={\sum }_{i,j}{A}_{ij}^2-2{\sigma }_1{\sum }_{i,j}{A}_{ij}{u}_{1i}{v}_{1j}+{\sigma }_1^2{\sum }_{i,j}{u}_{1i}^2{v}_{1j}^2$

u和v是U和V里的奇异向量，$u_1\ast {\sigma }_1=A{v}_1$, 且u和v都是单位向量。${\sum }_{i,j}{u}_{1i}^2{v}_{1j}^2=\Vert u_1{\Vert }_2^2\Vert v_1{\Vert }_2^2=1$,所以

$\Vert E{\Vert }_F^2={\sum }_{i,j}{A}_{ij}^2-2{\sigma }_1{\sum }_{i,j}{A}_{ij}{u}_{1i}{v}_{1j}+{\sigma }_1^2$

然后使用迹技巧

$\begin{array}{rl}& \sum _{i,j}{A}_{ij}{u}_{1i}{v}_{1j}=\mathrm{tr}\left(A^T{u}_1{v}_1^T\right)\\ & =\mathrm{tr}\left(v_1^T{A}^T{u}_1\right)\\ & =\mathrm{tr}\left(u_1^T{\sigma }_1{u}_1\right)\\ & ={\sigma }_1\mathrm{tr}\left(u_1^T{u}_1\right)\\ & ={\sigma }_1\left(u_1^T{u}_1\right)\\ & ={\sigma }_1\end{array}$

所以

$\begin{array}{rl}& \Vert E{\Vert }_F^2=\sum _{i,j}{A}_{ij}^2-2{\sigma }_1^2+{\sigma }_1^2\\ & \Vert E{\Vert }_F^2=\sum _{i,j}{A}_{ij}^2-{\sigma }_1^2\end{array}$

因为矩阵的Frobenius范数的平方等于矩阵的所有奇异值的平方和，即

$\Vert \mathbf{A}{\Vert }_F=\sqrt{{\sigma }_1^2+{\sigma }_2^2+\dots +{\sigma }_n^2}$

所以误差矩阵的Frobenius范数的平方等于除了最大奇异值的其他所有奇异值的平方和。

$$\begin{array}{c}{\Vert {\mathbf{E}}_1\Vert }_F^2=\sum _{i=2}^n{\sigma }_i^2(\mathbf{W}),\\ {\Vert {\mathbf{E}}_2\Vert }_F^2=\sum _{i=2}^n{\sigma }_i^2(|\mathbf{W}|).\end{array}$$

Based on $\Vert \mathbf{W}{\Vert }_F^2=\Vert |\mathbf{W}|{\Vert }_F^2$, we have

$$\sum _{i=1}^n{\sigma }_i^2(\mathbf{W})=\sum _{i=1}^n{\sigma }_i^2(|\mathbf{W}|).$$

According to Lemma 1, we can conclude

$${\Vert {\mathbf{E}}_2\Vert }_F^2\le {\Vert {\mathbf{E}}_1\Vert }_F^2.$$

From the equation in this proposition, we can formulate

$${\mathbf{W}}_{\text{sign }}\odot |\mathbf{W}|={\mathbf{W}}_{\text{sign }}\odot \tilde{\mathbf{a}}{\tilde{\mathbf{b}}}^{\mathrm{T}}+{\mathbf{W}}_{\text{sign }}\odot {\mathbf{E}}_2$$

Hence we have

$$\mathbf{W}-{\mathbf{W}}_{\text{sign }}\odot \tilde{\mathbf{a}}{\tilde{\mathbf{b}}}^{\mathrm{T}}={\mathbf{W}}_{\text{sign }}\odot {\mathbf{E}}_2.$$

Therefore

$$\begin{array}{rl}{\Vert {\mathbf{W}}_{\text{sign }}\odot {\mathbf{E}}_2\Vert }_F^2& =\sum _{i,j}{s}_{ij}^2{e}_{ij}^2=\sum _{i,j}{e}_{ij}^2\\ & ={\Vert {\mathbf{E}}_2\Vert }_F^2\le {\Vert {\mathbf{E}}_1\Vert }_F^2,\end{array}$$

where $s_{ij}=\pm 1$ is the element of ${\mathbf{W}}_{\text{sign }}$. Hence the inequation in this proposition is proved.

知识蒸馏

这里就是用命题1得到的Y的表达式。
$\begin{array}{r}{\mathbf{W}}_{\pm 1}=\mathrm{Sign}(\mathbf{W}),\\ \mathbf{Y}=\left[(\mathbf{X}\odot \mathbf{g}){\mathbf{W}}_{\pm 1}^{\mathrm{T}}\right]\odot \mathbf{h},\\ \mathbf{Z}=\text{ LayerNorm }(\mathbf{Y}),\end{array}$

$${\mathcal{L}}_{\mathrm{CE}}=-\frac{1}{n_s}\sum _{i=1}^{n_s}\sum _c{P}_c^{\mathcal{T}}\left({\mathbf{o}}_i\right)\log P_c^{\mathcal{S}}\left({\mathbf{o}}_i\right),$$

where $c$ denotes the number of classes and $n_s$ denotes the number of training samples in the current batch. $\mathcal{T}$ and $\mathcal{S}$ are the teacher model and student model, respectively. The error of hidden states is defined as

$${\mathcal{L}}_{\text{MSE }}=\sum _{i=1}^{n_s}\sum _{j=1}^{n_l}{\Vert \frac{{\mathbf{q}}_{i,j}^{\mathcal{T}}}{{\Vert {\mathbf{q}}_{i,j}^{\mathcal{T}}\Vert }_2}-\frac{{\mathbf{q}}_{i,j}^{\mathcal{S}}}{{\Vert {\mathbf{q}}_{i,j}^{\mathcal{S}}\Vert }_2}\Vert }_2^2,$$

where $n_l$ denotes the number of layers and $\mathbf{q}$ denotes the hidden state. Hence the final objective function can be formulated as

$${\mathcal{L}}_{\mathrm{KD}}={\mathcal{L}}_{\mathrm{CE}}+\alpha {\mathcal{L}}_{\mathrm{MSE}},$$

${\mathcal{L}}_{\mathrm{CE}}$使得学生模型在输出上接近教师模型。
${\mathcal{L}}_{\mathrm{MSE}}$使得学生模型在隐藏状态上接近教师模型。