Motivation
矩陣相乘 $y=Wx$ 是神經網路的基本單元, 其中 $W\in\mathbb{R}^{m\times n}$, $x\in\mathbb{R}^{n\times1}$, $y\in\mathbb{R}^{m\times 1}$.
試想一下如果 $W$ 裡面的值存在 outliers, 則對整個 matrix 量化時, 為了能覆蓋 outlier 則會犧牲精度, 這正是現今 LLM 極致量化時的魔王關卡.
考慮旋轉矩陣 $U\in\mathbb{R}^{m\times m}$, $V\in\mathbb{R}^{n\times n}$, 如果先對 $W$ 和 $x$ 作如下座標旋轉操作:
$$\begin{align}
\tilde{W}\leftarrow UWV^T \\
\tilde{x}\leftarrow Vx
\end{align}$$ 則我們觀察一下 $\tilde{W}\tilde{x}$ 與原來的 $Wx$ 差異在哪
$$\begin{align}
\tilde{W}\tilde{x} &= (UWV^T)(Vx) =UWx
\end{align}$$ 所以如果再多乘上 $U^T$ 旋轉一下, 則就能還原出原本的 output activation $y=Wx$ 了.
那為什麼要套這些旋轉矩陣? 反正最後又要還原回去, 不是多此一舉嗎?
其實不是白做, 旋轉有很大的好處
這是因為透過旋轉矩陣, 我們能夠把 outlier 消除, 進而能夠很好的量化
$\tilde{W}$ 比 $W$, $\tilde{x}$ 比 $x$ 好量化很多! 👍🏻
Toy Example 實驗
用下面這個簡單的 3D 例子舉例
最極端的情況就是一個 one-hot vector, $[1, 0, 0]^T$, 如下圖紅色箭頭
透過旋轉矩陣:
$$\begin{aligned}
R = \begin{bmatrix} 0.577 & 0.707 & 0.408 \\ 0.577 & -0.707 & 0.408 \\ 0.577 & 0.000 & -0.816 \end{bmatrix}
\end{aligned}$$ 轉到藍色箭頭的向量 $[0.577, 0.577, 0.577]^T$, 此時所有數值都一樣, 達到完美量化的情況.
在用一個例子, 維度 1000 的 one hot vector 做實驗
|
|
畫出 v_rotated 的 histogram 如下圖:
大部分的值已經落在 ±0.032 之間了
因此, 我們期望透過 $U,V$ 去旋轉, 使得 $\tilde{W}$ 變得非常容易量化
難點: 高維度旋轉代價太高以及如何選擇旋轉矩陣
但要這麼做, 有幾個關鍵問題得回答, 這裡先點出兩個:
高維度的矩陣相乘將會是無法忽視的負擔, 需要輕量化的旋轉矩陣乘法
QuIP 採用 Kronecker 方法拆解高維度 $n$ 成兩個低微度相乘 $n=p\times q$, 因此只需要兩個低微度的隨機旋轉矩陣, 更多請參考論文.
而 QuIP# (QuIP續作) 更進一步改進, 採用隨機阿達馬轉換 (Randomized Hadamard Transform, RHT), 因為矩陣的元素只有 $\{-1, +1\}$, 所以計算過程甚至不需要進行浮點數的乘法運算. 注意到 RHT 也是正交矩陣, 論文更證明 RHT 能得到的量化誤差更小.該選擇哪一個旋轉(正交)矩陣?
隨機旋轉就可以! 聽起來很玄, 但 QuIP 這篇論文證明了”隨機”也能保證有高機率消除 outlier. 本文後面將重點補充這個重要的數學證明.
另外, SpinQuant 採用 Cayley SGD 來學習最佳的旋轉 (上一篇筆記), 並且這些旋轉矩陣可以直接合併到原來的參數矩陣 $W$ 裡, 不會有額外開銷. 這種學習更直接, 也不隨機了!
PS: SpinQuant 在推論時, 可以套上一個預先設定好的 RHT 旋轉矩陣, 以便讓 activation 可以動態量化, 進一步壓縮. 這裡有個細節, 由於 $W$ 的旋轉矩陣是學出來的, 因此這個學習可以預先考慮到動態量化的 RHT 的誤差.
隨機旋轉的數學保證
接著, 我們來看看最重要的前提: ”隨機”也能保證有高機率消除 outlier
QuIP 給出了數學保證
機率理論中有個不等式: Concentration inequality
這是在探討什麼樣的情況, random variable 偏離平均值太遠的機率, 不會太高
首先我們看 QuIP 論文的 [Lemma 9]:
設 $x\in\mathbb{R}^n$, $\|x\|_2=1$ (unit sphere 表面的一個向量), 對任意 1-Lipschitz 函數 $F:x\mapsto\mathbb{R}$, 存在 $A$, $C$ 與維度 $n$ 無關, 滿足:
$$\begin{align}
\mathbf{P}_x(F(x)-\mathbb{E}[F(x)]\geq t)\leq C\exp\left(-\frac{nt^2}{A}\right)
\end{align}$$ 白話來說, $F(x)$ 偏離期望值超過 $t$ 的機率不會太大 (不會超過 $C\exp(-nt^2/A)$)
接著 [Lemma 10]:
設 $B\in\mathbb{R}^{m\times n}$, 且 $x$ 是 $\mathbb{R}^n$ unit sphere 表面的一個均勻採樣向量
同樣存在 $A$, $C$ 與維度 $n$ 無關, 滿足:
$$\begin{align}
\mathbf{P}\left(\left(\frac{\|Bx\|}{\|B\|_F}\right)^2\geq \frac{A'}{n}\log\left(C'\over\delta\right)\right)\leq \delta
\end{align}$$ 其中定義 $\delta$ 為 $A$, $C$ 的函數
$$\begin{aligned}
\delta:=C\exp\left(-\frac{nt^2}{A}\right)
\end{aligned}$$ 讓我們用 $B\in\mathbb{R}^{1\times n}$ 為一個 row vector 來看會更容易明白
定義 $b^T:=B$ 此時相當於將球面的隨機向量 $x$ 投影到 $b$.
注意到如果是 vector 的話, Frobenius 等於 spectral norm, i.e. $\|b\|_F=\|b\|$, 所以 $F(x)=\|b^Tx\|/\|b\|_F$ 相當對 $x$ 在 $b$ 向量上投影
得到的結論就是: 這個隨機投影的長度平方, 要異常大的機率是極度微小的!
注意到如果維度 $n$ 愈大, 因為 (5) 機率裡的不等式右邊反比於 $n$, 所以理論上愈不容易發生投影到 $b$ 後有 outliers 出現
最後, 若把 $x$ 和 $b$ 的角色對調, 給定一個 $x$, 令 $U=\{b_i\}_{i=1}^n$ 為隨機 orthogonal 矩陣, 則 $Ux$ 表示旋轉後投影在各 basis 的結果
根據上述討論, 我們知道其能量都不會太大! 這就是隨機旋轉能消除極值的理論奧秘和保證!
Appendix: 補充 Lemma10 的詳細推導證明
QuIP 給的 Lemma 10 證明中間有幾個步驟對我來說省略太多, 因此這裡補充起來
定義 $F(x)=\|Bx\|/\|B\|_F$, 則
$$\begin{aligned}
\nabla F(x)=\frac{1}{\|B\|_F}\nabla (x^TB^TBx)^{1/2}=\frac{1}{\|B\|_F}\cdot\frac{B^TBx}{(x^TB^TBx)^{1/2}}=\frac{B^TBx}{\|Bx\|\cdot\|B\|_F} \\
\Longrightarrow \|\nabla F(x)\|=\frac{\|B^TBx\|}{\|Bx\|\cdot\|B\|_F}\leq \frac{\|B\|^T\cdot\|Bx\|}{\|Bx\|\cdot\|B\|_F}=\frac{\|B\|}{\|B\|_F}\leq1
\end{aligned}$$
最後的不等式 $\|B\|\leq\|B\|_F$ 可以參考 “量量矩陣的大小” 這篇筆記
因此 $F$ 符合 1-Lipschitz 條件
因為 $\mathrm{Var}(x)=\mathbb{E}(x^2)-(\mathbb{E}x)^2\geq0$, 所以 $\mathbb{E}x\leq\sqrt{\mathbb{E}(x^2)}$, 所以 $\mathbb{E}[F(x)]$:
$$\begin{aligned}
\mathbb{E}[F(x)] &\leq \sqrt{\mathbb{E}[F(x)^2]}=\sqrt{\mathbb{E}\left[\left(
\frac{\|Bx\|}{\|B\|_F}
\right)^2\right]} \\
&= \frac{1}{\|B\|_F}\cdot \sqrt{\mathbb{E}(x^TB^TBx)}
\end{aligned}$$ 因為 $x^TB^TBx$ 是一個 scalar 所以等於 $\mathrm{Tr}(x^TB^TBx)$, 加上 $\mathrm{Tr}(AB)=\mathrm{Tr}(BA)$ 特性
$$\begin{aligned}
\mathbb{E}(x^TB^TBx) &= \mathbb{E}\left[\mathrm{Tr}\left(x^TB^TBx\right)\right]=\mathbb{E}\left[\mathrm{Tr}\left(Bxx^TB^T\right)\right] \\
&= \mathrm{Tr}\left(B\mathbb{E}[xx^T]B^T\right)
\end{aligned}$$ 注意到 $\mathbb{E}[xx^T]$ 是一個 $n\times n$ 的矩陣, 對角項為 $\mathbb{E}[x_i^2]$ 表示維度 $i$ 的能量大小, 而非對角項 $\mathbb{E}[x_ix_j]$ 可以想成是 維度 $i$ 和 $j$ 的相關性
由於 $x$ 來自於球面上的隨機向量, 理論上不同維度之間是獨立, 因此 $\mathbb{E}[x_ix_j]=0$.
也因為維度之間是獨立, 所以每個維度分配到的能量理論上期望值也相同, 相當於從總能量 $1$ 去均分, 因此 $\mathbb{E}[x_i^2]=1/n$.
得到結論 $\mathbb{E}[xx^T]=\frac{1}{n}\cdot I$. 所以:
$$\begin{aligned}
\mathbb{E}(x^TB^TBx) &= \mathrm{Tr}\left(B\mathbb{E}[xx^T]B^T\right) = \mathrm{Tr}\left(B\frac{1}{n}IB^T\right) \\
&=\frac{1}{n}\mathrm{Tr}(BB^T)=\frac{1}{n}\mathrm{Tr}(B^TB)=\frac{1}{n}\|B\|_F^2
\end{aligned}$$ 代回去
$$\begin{aligned}
\mathbb{E}[F(x)] &\leq \frac{1}{\|B\|_F}\cdot \sqrt{\mathbb{E}(x^TB^TBx)}=\frac{1}{\|B\|_F}\cdot\frac{1}{\sqrt{n}}\|B\|_F=\frac{1}{\sqrt{n}}
\end{aligned}$$ 套用 Lemma 9, 存在 $A,C$ 獨立於維度 $n$, s.t.
$$\begin{aligned}
\mathbf{P}\left(\frac{\|Bx\|}{\|B\|_F}\geq t + \mathbb{E}[F(x)]\right)\leq C\exp\left(-\frac{nt^2}{A}\right)
\end{aligned}$$ 因為 $\mathbb{E}[F(x)]\leq\frac{1}{\sqrt{n}}$, 所以同樣成立
$$\begin{aligned}
\Longrightarrow\mathbf{P}\left(\frac{\|Bx\|}{\|B\|_F}\geq t+\frac{1}{\sqrt{n}}\right)\leq C\exp\left(-\frac{nt^2}{A}\right)
\end{aligned}$$ 舉個簡單例子, 由於 $0.8>0.5$, 如果 $\mathbf{P}(x\geq 0.5)\leq \delta$, 則 $\mathbf{P}(x\geq 0.8)\leq \delta$ 同樣成立, 因為條件更嚴苛
定義
$$\begin{aligned}
\delta:=C\exp\left(-\frac{nt^2}{A}\right) \\
\Longrightarrow t^2=\frac{A}{n}\log\left(C\over\delta\right)
\end{aligned}$$ 繼續改寫, 由於機率裡面不等式條件兩邊都是正
$$\begin{aligned}
\Longrightarrow \mathbf{P}\left(\left(\frac{\|Bx\|}{\|B\|_F}\right)^2\geq \left(t+\frac{1}{\sqrt{n}}\right)^2\right)\leq \delta
\end{aligned}$$ 利用 $(a+b)^2\leq 2a^2+2b^2$, 得到 $(t+1/\sqrt{n})^2\leq 2t^2+2/n$. 並帶入 $t^2$ 與 $\delta$ 的關係式
$$\begin{aligned}
\left(t+\frac{1}{\sqrt{n}}\right)^2 &\leq 2\frac{A}{n}\log\left(C\over\delta\right)+\frac{2}{n} \\
&=\frac{2A}{n}\left(\log\left(C\over\delta\right)+\frac{1}{A}\right) \\
&=\frac{2A}{n}\left(\log\left(C\over\delta\right)+\log(e^{1/A})\right) \\
&=\frac{2A}{n}\log\left(\frac{Ce^{1/A}}{\delta}\right) \\
&=\frac{A'}{n}\log\left(C'\over\delta\right)
\end{aligned}$$ 只要令 $A'=2A$, $C'=Ce^{1/A}$ 即可. 因此
$$\begin{aligned}
\mathbf{P}\left(\left(\frac{\|Bx\|}{\|B\|_F}\right)^2\geq \left(t+\frac{1}{\sqrt{n}}\right)^2\right)\leq \delta \\
\Longrightarrow \mathbf{P}\left(\left(\frac{\|Bx\|}{\|B\|_F}\right)^2\geq \frac{A'}{n}\log\left(C'\over\delta\right)\right)\leq \delta
\end{aligned}$$ Q.E.D.