近期回頭讀 GPTQ (針對 LLM 的 PTQ) 論文, 同時想認真理解其背後數學精隨, 才發現以前理解的太淺薄了.
當我理解 Schur 互補 (Complement) 和 Cholesky 分解之間的關聯裡, 突然有種直擊核心的感覺, 因此決定做比較詳細的數學筆記.
最後會用一個簡單的情境來說明 GPTQ 怎麼應用 Cholesky 分解.
分塊矩陣逆定理 (Block Matrix Inversion Theorem)
定義分塊矩陣 $M=\begin{pmatrix} A & B \\ C & D \end{pmatrix}$, 且 $S_D:=A-BD^{-1}C$ 定義為 $D$ 在 $M$ 中的 Schur 互補.
若 $D$ 和 $S_D$ 可逆, 則
$$\begin{align}
M^{-1}=
\begin{pmatrix}
A & B \\ C & D
\end{pmatrix}^{-1}
= \begin{pmatrix}
S_D^{-1} & -S_D^{-1}BD^{-1} \\
-D^{-1}CS_D^{-1} & D^{-1}+D^{-1}CS_D^{-1}BD^{-1}
\end{pmatrix}
\end{align}$$ 證明可參考此 YouTube [Schur complement in applications] 前 6 分半, 我覺得滿清楚的.
這定理有一個很大的應用: 快速計算分塊的反矩陣.
若矩陣 $H$ 的 inverse $M:=H^{-1}$ 已知, 注意到 $H$ 不用已知, 定義他們對應的 submatrix 如下:
$$\begin{aligned} H=\begin{pmatrix} H_a & H_b \\ H_c & H_d \end{pmatrix},\quad M:=H^{-1}=\begin{pmatrix} A & B \\ C & D \end{pmatrix} \end{aligned}$$ 則 $H_a$ 這個 submatrix 的 inverse, i.e. $(H_a)^{-1}$, 可由已知的 $M$ 做 Schur 互補快速求得:
$$\begin{align}
(H_a)^{-1}=S_D:=A-BD^{-1}C
\end{align}$$ 這是因為 $M^{-1}=H$, 由分塊矩陣的逆矩陣定理 (1) 知 $H_a$ 剛好就是 $S_D^{-1}$, 則有 $(H_a)^{-1}=S_D$ 的結論.
另外當 $D$ 是 scalar 時, 因為 $C$ 變 row vector、$B$ 變 col vector, 則等同於做 Rank-1 update.
基本矩陣 和 基本列運算
為了簡潔, 請參考: Week 2 Systems of Linear Equations
高斯消去法的一步 = 計算一次 Schur Complement
令矩陣 $M=\begin{pmatrix} a & \mathbf{b} \\ \mathbf{c} & D \end{pmatrix}_{n\times n}$, 其中 $a\in\mathbb{R}$、$\mathbf{b}\in\mathbb{R}^{1\times(n-1)}$、$\mathbf{c}\in\mathbb{R}^{(n-1)\times1}$、$D\in\mathbb{R}^{(n-1)\times(n-1)}$.
為了要把 column 1 的 $\mathbf{c}$ 消掉變成 $\mathbf{0}$, 我們會將 row 1 乘上 $\mathbf{c}/{a}$, 然後從 row 2 (也就是下方所有 rows) 減去它, 相當於左乘一個 “取代” 的基本列運算矩陣 $E$:
$$\begin{align} EM= \begin{pmatrix} 1 & \mathbf{0}^T \\ -\frac{\mathbf{c}}{a} & I \end{pmatrix} \begin{pmatrix} a & \mathbf{b} \\ \mathbf{c} & D \end{pmatrix} = \begin{pmatrix} a & \mathbf{b} \\ \mathbf{0} & D - \frac{\mathbf{c} \mathbf{b}}{a} \end{pmatrix} \end{align}$$ 右下角的 $(n-1)\times(n-1)$ 矩陣, 正好就是 $a$ 在 $M$ 中的 Schur 互補.
$$\begin{aligned} D_{new} = D - \frac{\mathbf{c} \mathbf{b}}{a} \end{aligned}$$ 這也正好也是做 Rank-1 Update.
連續的高斯消去 = 逐步建構 Cholesky 分解
Cholesky 分解在說明, 對於一正定矩陣 $M$ 來說, 存在唯一下三角矩陣 $L$ (對角項皆為正) 滿足 $M=LL^T$.
因為 $M$ 正定, 所以主對角項 $>0$, 則可重複高斯消去法 (3) 使得最後變成上三角矩陣 $U$:
$$\begin{aligned} E_{n-1} \dots E_2 E_1 M = U \\ \end{aligned}$$ 注意到 $E_i$ 全部都是下三角矩陣且皆可逆 (因為基本列運算可逆), 而下三角矩陣的連乘或反矩陣仍然是下三角.
全部移到右邊後令 $L_{unit}:=(E_{n-1} \dots E_2 E_1)^{-1}$ 為一下三角矩陣, 得到:
$$\begin{aligned} M = L_{unit} U \end{aligned}$$ 在我們的問題裡, 因為 $M$ 是 (對稱) 正定矩陣:
$$\begin{aligned} M=M^T \\ \Longrightarrow L_{unit}U=U^TL_{unit}^T \\ \Longrightarrow L_{unit}^{-1}(L_{unit}U)(L_{unit}^T)^{-1} = L_{unit}^{-1}(U^TL_{unit}^T)(L_{unit}^T)^{-1} \\ \Longrightarrow U(L_{unit}^T)^{-1}=L_{unit}^{-1}U^T \end{aligned}$$ 最後一行, L.H.S. 是兩個上三角矩陣相乘其結果仍為上三角, 同理 R.H.S. 結果為下三角.
上三角等於下三角, 所以結論為對角矩陣, 令為 $\Lambda:=U(L_{unit}^T)^{-1}$, 得到 $U=\Lambda L_{unit}^T$.
所以我們可推導出 Cholesky 分解:
$$\begin{aligned} M = L_{unit} U = L_{unit} \Lambda L_{unit}^T \\ \Longrightarrow M = (L_{unit} \Lambda^{1/2}) (L_{unit} \Lambda^{1/2})^T \\ \Longrightarrow M = LL^T \end{aligned}$$ 從代數角度來看, 連續對矩陣進行 Rank-1 Update 的 Schur 互補操作, 數學上完全等價於在對該矩陣執行高斯消去法, 也就是在一步步計算它的 Cholesky 分解.
Cholesky 分解與分塊反矩陣的關聯
文章開頭我們介紹如何利用 Schur 互補 (2) 快速計算分塊的反矩陣.
而上一節介紹對矩陣連續進行 Schur 互補等於計算它的 Cholesky 分解.
我們現在把這兩件事關連起來!
令正定矩陣 $H$ 和其反矩陣 $M$ 的分塊矩陣定義如下:
$$\begin{aligned} H=\begin{pmatrix} h_a & \mathbf{h}_c^T \\ \mathbf{h}_c & {\color{green}{H_d}} \end{pmatrix},\quad M:=H^{-1}=\begin{pmatrix} a & \mathbf{c}^T \\ \mathbf{c} & D \end{pmatrix} \end{aligned}$$ 注意到 $h_a$ 和 $a$ 都是 scalar, 所以這裡的分塊矩陣概念上是希望從一個大矩陣降一個維度.
因為 $M$ 也是正定, 所以存在唯一 Cholesky 分解 $M = L L^T$, 並將 $L$ 分塊:
$M = \begin{pmatrix} l_a & \mathbf{0}^T \\ \mathbf{l}_c & {\color{orange}{L_d}} \end{pmatrix} \begin{pmatrix} l_a & \mathbf{l}_c^T \\ \mathbf{0} & {\color{orange}{L_d}}^T \end{pmatrix} = \begin{pmatrix} l_a^2 & l_a\mathbf{l}_c^T \\ l_a\mathbf{l}_c & \mathbf{l}_c\mathbf{l}_c^T + {\color{orange}{L_dL_d^T}} \end{pmatrix}$
與 $M$ 的分塊矩陣對照馬上得到: $l_a=\sqrt{a}$ 和 $\mathbf{l}_c=\mathbf{c}/\sqrt{a}$.
觀察右下角區塊的等式:
$$\begin{aligned} D = \mathbf{l}_c\mathbf{l}_c^T + {\color{orange}{L_dL_d^T}} \\ \Longrightarrow {\color{orange}{L_dL_d^T}} = D - \mathbf{l}_c\mathbf{l}_c^T \end{aligned}$$ 代入 $\mathbf{l}_c=\mathbf{c}/\sqrt{a}$ 進去得到:
$$\begin{align} \Longrightarrow {\color{orange}{L_dL_d^T}} = {\color{green}{ D - \frac{\mathbf{c}\mathbf{c}^T}{a} }} \end{align}$$ R.H.S. 是 $S_a$ ($a$ 在 $M$ 中的 Schur 互補), 根據之前的討論 (2) 知道它就是 ${\color{green}{H_d^{-1}}}$.
L.H.S. 表明 ${\color{green}{H_d^{-1}}}$ 的 Cholesky 分解, 剛好就是原 $H^{-1}$ 的 Cholesky 矩陣 $L$ 的子區塊 ${\color{orange}{L_d}}$.
這在代數上給出了一個強大的結論: 對矩陣執行一次完整的 Cholesky 分解, 等同於預先計算了所有可能遞迴的 Submatrix Inverse. 矩陣 $L$ 內的每一個元素, 都已經是各階段 Rank-1 Update 所需的現成係數.
我們用個簡單的例子來說明:
令 $H_7$ 是維度 $7$ 的正定矩陣, 且我們已經花第一次的代價算出其反矩陣 $M_7:=H_7^{-1}$.
此時對 $M_7$ 計算 Cholesky 分解得到 $M_7=L_7 L_7^T$.
令 $M_5$ 表示 $M_7$ 的右下角維度為 $5$ 的 submatrix, 同理 $H_5$ 對應 $H_7$ 的右下角子矩陣, $L_5$ 對應 $L_7$ 的右下角子矩陣
則 $H_5^{-1}$ 直接就用已經有的 $L_5$ 計算 $L_5 L_5^T$ 就可得到!
這樣做的好處是如果矩陣維度一旦很大 (GPTQ 裡這個矩陣的維度就是該層的 input channel 數, LLM 動輒上千甚至上萬), 每個子矩陣的反運算代價就很高.
不如一開始算一次高效的 Cholesky 分解 (數值又穩定), 這樣之後每個子矩陣的反矩陣一次全部都能得到!
這就是 GPTQ 「神來之筆」的數學核心!