⚠️ 這部份屬於我的知識盲區, 如有不嚴謹/錯誤之處還請見諒/指正
又再一次體會到數學的美
正交矩陣 $X\in\mathbb{R}^{n\times p}$ 形成的集合稱為 Stiefel manifold
$$\begin{align}
\mathcal{M}=\{X\in\mathbb{R}^{n\times p}:X^TX=I\}
\end{align}$$ 以下用方陣方便解釋, i.e. $X\in\mathbb{R}^{n\times n}$.
讓我們發揮點想像力, 把 Stiefel manifold 想像成在 $n^2$ 維度歐式空間中的一個集合形成的 “流形”.
閱讀本文你將了解到:
- 正交矩陣的幾何結構: 了解 Stiefel Manifold 的定義, 以及為什麼我們將正交矩陣視為高維空間中一塊「彎曲」的流形
- 瞬時旋轉的本質: 透過泰勒展開式推導出切空間 (Tangent Space) 的約束條件, 並揭示為什麼 反對稱矩陣 (Skew-symmetric matrix) 本質上就是單位矩陣 $I$ 處的「瞬時旋轉方向」
- Cayley 轉換的橋樑作用: 掌握如何利用 Cayley Transform 在反對稱矩陣與正交矩陣之間進行無損轉換, 並理解其作為「微小步長旋轉」的物理意義
- Cayley SGD 的運作邏輯: 學習如何將普通的歐幾里得梯度投影至流形切空間, 並透過迭代式更新 (Iterative Update) 在不計算矩陣 inverse 的情況下, 維持矩陣的正交性
Stiefel Manifold (正交矩陣) 的 Tangent Space (切空間)
在這“流形”上面我們隨便抓出一條連續平滑曲線, 並用 $t$ 參數化表示: $X(t)\in\mathcal{M},\forall t\in\mathbb{R}$.
因為正交矩陣物理意義等同於旋轉矩陣, 所以我們把 $X(t)$ 想像成隨著時間變化的旋轉矩陣, 其中時間 $t=0$ 時定義 $X(0):=X$.
我們令 Stiefel manifold 上的 $X$ 這點, 其 tangent (切線, 即一階導數) 為 $Z$.
對 $X(t)$ 寫出泰勒展開式, 在 $t=0$ 點展開:
$$\begin{aligned}
X(dt)=X(0)+X'(0)dt+\mathcal{O}(dt^2) \\
\Longrightarrow X(dt)=X+Zdt+\mathcal{O}(dt^2)
\end{aligned}$$ 當 $dt$ 極小時 $X(dt)$ 仍要待在”流形”中, 即 $X(dt)\in\mathcal{M}$, 因此 $dt\rightarrow0$ 時需滿足
$$\begin{aligned}
I &= X(dt)^TX(dt) \\
&= (X+Zdt+\mathcal{O}(dt^2))^T(X+Zdt+\mathcal{O}(dt^2)) \\
&=X^TX + (X^TZ+Z^TX)dt+\mathcal{O}(dt^2) \\
&= I + (X^TZ+Z^TX)dt+\mathcal{O}(dt^2) \\
&\Longrightarrow (X^TZ+Z^TX)dt+\mathcal{O}(dt^2) = 0 \\
&\Longrightarrow X^TZ+Z^TX=\frac{\mathcal{O}(dt^2)}{dt}\rightarrow0,\quad\text{as }dt\rightarrow0
\end{aligned}$$ 因此 Stiefel manifold 上, $X$ 這點其 tangent 矩陣 $Z$ 必須滿足條件 $X^TZ+Z^TX=0$.
上面我們只是利用通過 $X$ 這點的一條曲線, 可以想像不同曲線的 tangent 矩陣都不同, 但不管如何也都必須滿足條件.
因此定義 $T_X\mathcal{M}$ 表示通過 $X$ 這點的 tangent 矩陣 (如同論文 [arxiv] 的定義)
$$\begin{align}
T_X\mathcal{M}=\{Z:X^TZ+Z^TX=0\}
\end{align}$$ 用更具體化的描述, $T_X\mathcal{M}$ 表示在 $X$ 這點合法的 “瞬時旋轉矩陣”
考慮特殊情況, 當 $X=I$ 時, tangent $Z$ 條件蛻化為 $Z=-Z^T$
說明在 $I$ 處的瞬時旋轉(正交)矩陣必定為反對稱矩陣 (Skew-symmetric matrix).
反之$Z$是反對稱矩陣, 則 $Z\in T_I\mathcal{M}$.
因此我們對反對稱矩陣有了更深一層的理解, 他們就是 Stielfel manifold 位於 $I$ 點的 “瞬時旋轉矩陣”!
有意思的是, 某一特殊情況的正交矩陣會與反對稱矩陣一對一對應.
Cayley 轉換描述了這樣的關係
Cayley Transform
任何一個不含特徵值 $-1$ 的正交矩陣 $Q$ 與反對稱矩陣 $A$ 有 1 對 1 的對應關係
首先 $A,B\in\mathbb{R}^{n\times n}$ 且都可逆, 如果滿足 $AB=BA$ 的話, 則
- $A^{-1}B=BA^{-1}$, 即 $A^{-1}$ 與 $B$ 矩陣乘法可交換
- $A^{-1}B^{-1}=B^{-1}A^{-1}$, 即 $A^{-1}$ 與 $B^{-1}$ 矩陣乘法可交換
很容易證明:
$$\begin{aligned}
AB=BA\Longrightarrow B^{-1}AB=A\Longrightarrow B^{-1}A=AB^{-1} \\
\Longrightarrow B^{-1}=AB^{-1}A^{-1}\Longrightarrow A^{-1}B^{-1}=B^{-1}A^{-1}
\end{aligned}$$ 因為 $(I-A)(I+A)=(I+A)(I-A)$ (都乘開來驗證即可)
所以 $(I+A)^{-1}(I-A)=(I-A)(I+A)^{-1}$.
定義 Cayley transform $T(X)$ 為:
$$\begin{align}
T(X)=(I-X)^{-1}(I+X)=(I+X)(I-X)^{-1}
\end{align}$$ $\circ$ 當 $A$ 為反對稱矩陣時, Cayley transform $T(A)$ 為正交矩陣
$\circ$ 當 $Q$ 為正交矩陣時, Cayley transform $T(Q)$ 為反對稱矩陣
證明請參考周老師的文章: Cayley 變換
設定 $A$ 為反對稱矩陣 我們來觀察一下 Cayley transform, 先定義 $Q(\alpha,A)$:
$$\begin{align}
Q(\alpha,A):=T\left(\frac{\alpha}{2} A\right)=\left(I-\frac{\alpha}{2}A\right)^{-1}\left(I+\frac{\alpha}{2}A\right)
\end{align}$$ 利用 Neumann series $(I-A)^{-1}=\sum_{k=0}^\infty A^k$, 分析一下 $Q(\alpha,A)$ 這個正交旋轉矩陣的長相:
$$\begin{aligned}
Q(\alpha,A) &=T\left(\frac{\alpha}{2} A\right)=\left(I-\frac{\alpha}{2}A\right)^{-1}\left(I+\frac{\alpha}{2}A\right) \\
&=\left(I+\frac{\alpha}{2}A+\left(\frac{\alpha}{2}A\right)^2+\cdots\right)\left(I+\frac{\alpha}{2}A\right) \\
&= I + \alpha A + \mathcal{O}(\alpha^2) \\
&\approx I + \alpha A
\end{aligned}$$ 得到這樣的結果 $Q(\alpha,A)\approx I+\alpha A$.
結合之前對反對稱矩陣 $A$ 的理解: 在位置 $I$ 的順時旋轉方向
則 Cayley transform $Q(\alpha,A)$ 物理意義就相當清晰了, 本質上就是「從 $I$ 出發, 沿著 $A$ 方向前進步長 $\alpha$ 的微小旋轉矩陣」
如果想要套用在已經有的舊旋轉矩陣 $R_{\text{old}}$ 後, 再旋轉的話, 則可以這麼用
$$\begin{align}
R_{\text{new}}=Q(\alpha,A)R_{\text{old}}
\end{align}$$ 其中 $A$ 為反對稱矩陣, 表示在位置 $I$ 的順時旋轉方向; $\alpha$ 則理解為一個微小的步長
根據 Cayley transform, $Q(\alpha,A)$ 本身是正交矩陣, 所以 $R_{\text{new}}$ 仍是正交矩陣 (仍在 Stiefel manifold 上)
延伸思考: 線性變換都是”拉伸”和”旋轉”的組合
任何矩陣都 $M$ 可拆成對稱和反對稱矩陣的組合:
$$\begin{align}
M= \underbrace{\frac{1}{2}\left(M+M^T\right)}_{\text{symmetric}} + \underbrace{\frac{1}{2}\left(M-M^T\right)}_{\text{skew symmetric}}
\end{align}$$ $\circ$ 反對稱矩陣我們已經知道其意義為瞬時旋轉矩陣, 即 ”旋轉” 的成分
$\circ$ 而對稱矩陣則代表著線性轉換中 “拉伸” 的成分
這是由於對稱矩陣必定可對角化, 將 eigenvectors 設定為 basis 座標來看的話, eigenvalues 表示對應方向的 scaling, 即拉伸
關於對稱和反對稱戶為補空間, 更多可以參考周老師的文章, 特殊矩陣 (13)：反對稱矩陣
投影到 $I$ 的切空間 $T_{I}\mathcal{M}$
對於任意一個線性變換如果我們只想保留其旋轉部分, 則取其反對稱成分即可
換句話說, 對任意矩陣 $A$ 投影到 $T_{\color{orange}{I}}\mathcal{M}$ 空間的結果就是 $(A-A^T)/2$ 這個反對稱矩陣.
$$\begin{align}
A\xrightarrow[]{P}(A-A^T)/2\in T_{\color{orange}{I}}\mathcal{M}
\end{align}$$ 很容易可以驗證 $P$ 為投影算子, i.e. 驗證 $P^2=P$, 因為如果 $A$ 已經是反對稱矩陣了, 投影的計算結果仍是自己 $A$
💡 不厭其煩多解釋一下, 注意到是 $T_{\color{orange}{I}}\mathcal{M}$ 而非 $T_{\color{orange}{X}}\mathcal{M}$. 本文上面已經做過許多解釋, 這是因為在 Stiefel manifold 的 $I$ 點的切空間, $T_{\color{orange}{I}}\mathcal{M}$, 等同於反對稱矩陣.
投影到 $X$ (屬於正交矩陣) 的切空間 $T_{X}\mathcal{M}$
對任意矩陣 $A$ 投影到 $I$ 的切空間, $T_{\color{orange}{I}}\mathcal{M}$, 結果是 $(A-A^T)/2$.
如果要投影到任意 Stiefel manifold 位置 $X$ 的切空間呢?
此問題等價於尋找一個切線向量 $V^\ast \in T_X \mathcal{M}$, 使得它與 $A$ 的距離最短, 即 Frobenius norm 最小.
$$\begin{aligned}
V^\ast=\arg\min_{V\in T_X \mathcal{M}}\|V - A\|_F^2
\end{aligned}$$ 由於正交(旋轉)矩陣 “保長” 特性, 左乘一個正交矩陣不會改變 norm, 因此我們是不是可以先 “旋轉” 回去 $I$ 位置, 然後套用 $I$ 的切空間投影公式, 最後再 “旋轉” 回去?
讓我們抱著這個想法推導看看
對任意矩陣 $A$, 我們求其投影到 Stiefel manifold 位置 $X$ 的切空間, i.e. 投影到 $T_X{\mathcal{X}}$.
[步驟一] 拉回原點 $I$:
讓我們把整個空間作旋轉, 把 $X$ 轉到 $I$, 所以所有矩陣都要左乘 $X^T$
這意味著我們要投影的矩陣 $A$ 現在應該當作 ${\color{orange}{X^T}}A$ 來看待, 然後變成投影到 $I$ 的切空間
[步驟二] $I$ 的切空間投影:
套用公式 (7) 對矩陣 $X^TA$ 投影到 $T_I{\mathcal{X}}$ 得到
$$\begin{aligned}
X^TA\xrightarrow[]{P}(X^TA-A^TX)/2\in T_I\mathcal{M}
\end{aligned}$$ [步驟三] 回復原來的位置到 $X$:
對投影結果左乘 $X$ 回復得到目標 $V^\ast$
$$\begin{aligned}
V^\ast &= \frac{1}{2}{\color{orange}{X}}\left(X^TA-A^TX\right) \\
&= \frac{1}{2}\left(A-XA^TX\right)
\end{aligned}$$ [結論] 對任意矩陣 $A$ 投影到 $T_{\color{orange}{X}}\mathcal{M}$ 空間的結果 ($X$ 為正交矩陣):
$$\begin{align}
A\xrightarrow[]{P}(A-XA^TX)/2\in T_{\color{orange}{X}}\mathcal{M}
\end{align}$$ 注意到, 將論文 [arxiv] 裡的式 (2) 化簡之後跟上式 (8) 是一樣的結果
Cayley SGD
試想我們從一個正交矩陣 $R_{\text{old}}$ 出發, 透過計算 loss 得到一個梯度方向 $G$, 如果做 SGD 更新, 則得到的 $R_{\text{old}}+\alpha G$ 無法保證仍屬於 Stiefel manifold, 即無法保持正交性.
Cayley SGD 就是為了解決這個問題而誕生的流形優化器 (Manifold Optimizer)
設 $R_{\text{old}}$ 為正交矩陣, 梯度方向 $G$ 為任意矩陣 (不必是反對稱或正交)
Cayley SGD 更新公式為:
$$\begin{align}
R_{new} = \underbrace{\left[ (I - \frac{\alpha}{2}Y)^{-1}(I + \frac{\alpha}{2}Y) \right]}_{Q(\alpha,Y)} R_{old} \\
Y=GR_{\text{old}}^T-R_{\text{old}}G^T
\end{align}$$ 注意到 $Y$ (10) 的計算是在把矩陣 $GR_{\text{old}}^T$ 投影到 $T_{\color{orange}{I}}\mathcal{M}$ 空間上, 即取出反對稱矩陣成分 (見式 (7)), 因此 $Y$ 必定為反對稱矩陣, 物理意義為要旋轉的方向
接著把 $Y$ 透過 Cayley transform $Q(\alpha,Y)$ (式 (4)) 得到對應的旋轉矩陣後, 套在 (5) 的結果最終就得到 Cayley SGD 的公式 (9) 了
Cayley SGD 公式原理解釋
這裡要特別解釋為什麼 $Y$ 的計算是使用 $GR_{\text{old}}^T$ 做投影, 而不是直接把梯度方向 $G$ 做投影就好, i.e. $G-G^T$? 還有, 為什麼不是用 $R_{\text{old}}^TG$ 做投影? 🤔
先解釋為什麼不是用 $G$ 做投影
從幾何上來說, 梯度 $G$ 是在位置 $R_\text{old}$ 上計算的, 因此應該要投影到的是 $T_{\color{orange}{R_\text{old}}}\mathcal{M}$ 這個切空間, 式 (8)
但我們遇到一個困難, 我們無法算出對應的旋轉矩陣
這是因為 Cayley transform (4) 只能作用在 $T_{\color{orange}{I}}\mathcal{M}$ 這個切空間上!
所以如同之前找非 $I$ 的切空間投影想法, 先將位置 $R_\text{old}$ 校正回 $I$ 位置 (先乘上 $R_\text{old}^T$ 做反向旋轉), 接著再對 $G$ 套用投影公式 (7), 這樣才能用 Cayley transform 找對應的旋轉矩陣
那麼該怎麼解釋為什麼不是用 $R_{\text{old}}^TG$ 做投影, 而是用 $GR_{\text{old}}^T$?
意思是, 反向旋轉 $R_\text{old}^T$ 為什麼是乘在 $G$ 的右邊而非左邊?
這個問題困擾我許久, 特別感謝偉大的 Gemini 😆
沒有現在的 AI 我真不知道何年何月才能理解通, 分析思路好精彩!
設 $Y$ 表示位於 $I$ 點的 “瞬時旋轉矩陣” 所以是反對稱矩陣, 其對應的旋轉矩陣透過 Cayley transform 得到為 $Q$
之前我們分析過 $Q\approx I+\alpha Y$, 所以
$$\begin{aligned}
R_\text{new}=QR\approx(I+\alpha Y)R=R+\alpha YR
\end{aligned}$$ 對照 SGD 更新公式我們知道更新的方向為 $YR$.
由於 $G$ 是梯度, 其必定是對損失函數來說最陡峭的方向
那如果朝著 $YR$ 這個方向移動的話, 損失函數仍然夠陡峭媽? 下降了多少?
在數學上, 這就是 “原始歐幾里得梯度 $G$” 與 “實際移動方向 $YR$” 的內積 (參考 Frobenius norm)
$$\begin{aligned}
\text{變化量} = \langle G, YR \rangle=\mathrm{Tr}(G^TYR)
\end{aligned}$$ 藉由 trace 內的矩陣乘法可循環, 我們繼續推導
$$\begin{aligned}
\text{變化量} = \mathrm{Tr}(RG^TY) = \langle GR^T,Y\rangle
\end{aligned}$$ 看出來了嗎?
原本我們是在位置 $R$ 計算 $G$ 和 $YR$ 的內積 $\langle G, YR \rangle$.
經過 Trace 轉換後, 它在數學上完全等價於在單位矩陣 $I$ 的位置, 計算 $GR^T$ 和 $Y$ 的內積: $\langle GR^T, Y \rangle$.
為了讓 $Y$ (定義在 $I$ 處的反對稱矩陣) 能直接和某個東西做內積運算, 我們 “被迫” 提煉出了 $GR^T$.
這就是為什麼代表拉回原點的梯度是 $GR^T$!
Iterative Cayley SGD
原來的 Cayley SGD 必須依賴 inverse 矩陣運算, 高維度將會是個災難:
重新對公式做個整理:
$$\begin{aligned}
R' = (I - \frac{\alpha}{2}Y)^{-1}(I + \frac{\alpha}{2}Y) R \\
\Longrightarrow (I - \frac{\alpha}{2}Y)R'=(I + \frac{\alpha}{2}Y) R \\
\Longrightarrow R'=R+\frac{\alpha}{2}Y(R+R')
\end{aligned}$$ 定義 $F(R')=R+\frac{\alpha}{2}Y(R+R')$, 如果 $F$ 是 contraction mapping function, 則存在收斂點 fixed point $R^\ast$ 滿足 $R^\ast=F(R^\ast)$.
因此 iterative Cayley SGD 更新公式為:
$$\begin{align}
R^{t+1}=R^{t}+\frac{\alpha}{2}Y(R^t+R^{t+1})
\end{align}$$ 論文 [arxiv] 的 Theorem 1&2 即在證明, 當 step size $\alpha$ 滿足某個條件夠小, 則 $F$ 是 contraction mapping, (11) 收斂且實務上迭代個 2, 3 次就足夠. (論文 [arxiv] Table 5)
同時該論文還有一個重要貢獻是結合 Momentum 的 iterative Cayley SGD
核心想法就是 momentum 更新在 Stiefel manifold 原點 $I$ 的切空間上, i.e. $T_{\color{orange}{I}}\mathcal{M}$, 而不是在原來的歐式空間上. 詳細請參考論文.
總結 Takeaway Messages
- 「旋轉」位於反對稱性中: 正交矩陣在單位元素 $I$ 處的切向量必定是反對稱矩陣 ($A = -A^T$), 這意味著在進行流形優化時, 我們處理的不是權重的「增量」, 而是權重的「旋轉」
- 投影路徑是「旋轉 -> 投影 -> 轉回」: 由於我們只擅長在 $I$ 點處理反對稱矩陣, 因此對於任何位置 $X$ 的梯度, 標準做法是先左乘 $X^T$ 將其拉回原點, 投影成反對稱矩陣後, 再映射回流形
- Cayley SGD 是一種「乘法更新」: 不同於傳統 SGD 的 $W_{new} = W + \Delta W$ (會破壞正交性), Cayley SGD 採用 $R_{new} = Q R_{old}$. 透過 Cayley 轉換生成的 $Q$ 本身就是正交矩陣, 因此更新過程會像是在流形表面「滾動」, 自然地維持了正交約束
References
- Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform [arxiv]
- Cayley 變換
- 特殊矩陣 (13)：反對稱矩陣