LoRAのための特異値分解 - 勾配降下党青年局

　特異値分解を解説する記事なんていくらでもありますが、LoRAに関連付けて話す記事なんてないと思うので、ここで書いてみます。まあ自分が特異値分解を理解するためでもあります。行列の右上カッコつき添え字に行数と列数を書きます。
参考記事：
yutomiyatake.github.io

LoRA

まずLoRAについて、重み $W^{(m,n)}$ の全結合層に対して、rank= $r$ のLoRAは重みの差分 $\Delta W^{(m,n)}=A^{(m,r)}B^{(r,n)}$ を学習します。学習後の重みは $W' = W+\Delta W$ になります。LoRAは $A$ の列、 $B$ の行を分解すると、 $\Delta W^{(m,n)}=(A_1 \cdots A_r) \begin{pmatrix} B_1 \\ \vdots \\ B_r\end{pmatrix}=\displaystyle{\sum_{i=1}^{r}A_i^{(m,1)}B_i^{(1,n)}}$
となって、rank= $1$ のLoRAをrank個分合計したような形になっています。

特異値分解

特異値分解は $\mathrm{rank}(W)=r$ の行列 $W^{(m,n)}$ を、

$W^{(m,n)}=U^{(m,m)}\Sigma^{(m,n)}V^{(n,n)}= U^{(m,m)} \begin{pmatrix} \sigma_{1} & & & & & \\ & \ddots & & & \huge{0} & \\ & & \sigma_{r} & & & \\ & & & 0 & & \\ & \huge{0} & & & \ddots & \\ & & & & & 0 \end{pmatrix} V^{(n,n)}$

とする操作です。ここで $U,V$ は直交行列（転置すると逆行列になる行列）で、 $\sigma_1,\cdots ,\sigma_r$ は特異値と呼ばれるものです。
図にすると以下のような形に変形できます。

どんどんLoRAと似たような形になっていくことが分かるでしょうか。

特異値分解の方法

$U,V$ が直交行列であることを考えると、
$\begin{eqnarray} W^TW &=&(U\Sigma V)^TU\Sigma V \\ &=&V^T\Sigma^TU^TU\Sigma V \\ &=&V^T \Sigma^T \Sigma V \\ &=&V^T\begin{pmatrix} \sigma_{1}^2 & & & & & \\ & \ddots & & & \huge{0} & \\ & & \sigma_{r}^2 & & & \\ & & & 0 & & \\ & \huge{0} & & & \ddots & \\ & & & & & 0 \end{pmatrix} V \end{eqnarray}$

$\begin{eqnarray} WW^T &=&U\Sigma V (U\Sigma V)^T \\ &=&U\Sigma V V^T\Sigma^T U^T \\ &=&U \Sigma \Sigma^T U^T \\ &=&U \begin{pmatrix} \sigma_{1}^2 & & & & & \\ & \ddots & & & \huge{0} & \\ & & \sigma_{r}^2 & & & \\ & & & 0 & & \\ & \huge{0} & & & \ddots & \\ & & & & & 0 \end{pmatrix} U^T \end{eqnarray}$

となります。これは対角化の形になっていますね。つまり $WW^T,W^TW$ の固有ベクトルを計算すれば、 $U,V$ を求めることができます。特異値は固有値の平方根になっています。実装上でどうやってるかは知らないですけど、アルゴリズム的には解けることはわかりますね。

特異値分解とrank

正則行列をかけても階数が変わらないという性質を考えると、 $\mathrm{rank}(A)=\mathrm{rank}(U\Sigma V)=\mathrm{rank}(\Sigma)$ になります。 $\Sigma$ は対角成分以外が0の行列ですので、非零の対角成分の個数がrankと一致します。つまり特異値の数がrankと一致します。

低ランク近似

特異値分解した結果、 $A=U\Sigma V=\displaystyle{\sum_{i=1}^{r}}\sigma_iU_iV_i$ で、特異値 $\sigma_1,\dots,\sigma_r$ は大きい順にソートされているとします。すると上位 $k$ 個の特異値のみを使った行列で近似することができます。 $A\fallingdotseq \tilde{A_k}=U\Sigma_{[\sigma_{k+1},\dots,\sigma_r:=0]}V=\displaystyle{\sum_{i=1}^{k}}\sigma_iU_iV_i$ です。0でない対角成分が $k$ 個となるわけですから、 $\mathrm{rank}(A_k)=k$ となります。この近似を低ランク近似と呼びます。

特異値分解によるLoRAの抽出

特異値分解によって、LoRAの抽出が行えます。事前学習済みモデルの重み $W$ に対して、何らかのファインチューニングモデル $W'$ があるとすると、学習差分 $\Delta W = W'-W$ をLoRAの形 $AB$ に変換することができます。 $\Delta W$ を特異値分解すれば、 $\Delta W = \displaystyle{\sum_{i=1}^{r}}\sigma_iU_iV_i$ の形に変換できます。すると低ランク近似により、ランク $k$ のLoRAに近似できます。具体的には $\tilde{\Delta W_k}=\displaystyle{\sum_{i=1}^{k}}\sigma_iU_iV_i$ で、 $A=(\sigma_1U_1 \cdots \sigma_{k} U_{k})$ , $B=V_{1:k}$ のようにすればできます。 $\sigma$ は $A$ と $B$ どっちに入れてもいいですけどね。

LoRAのリサイズ

任意のLoRAをそれより小さいランクのLoRAに近似することもできます。といってもほぼ抽出と同じで、LoRAの $AB=\Delta W$ を計算して後は低ランク近似を用いて同様の操作をするだけです。ところでKohya氏によるLoRAのリサイズには $k$ を自動で決める実装があります。決め方に三通りあります。何らかの閾値 $r$ を決めて、

最大特異値に対して $r$ %の大きさまでの特異値を選ぶ。
特異値の累積和を求めて全体の $r$ %以上になるまで特異値を選ぶ。
特異値の累積二乗和を求めて全体の $r^2$ %以上になるまで特異値を選ぶ。

直感的には上の二つが分かりやすそうですが、3つ目には数学的な意味があります。

LoRAのマージ

LoRAのマージには重みをそのままマージする方法と、特異値分解を行う方法があります。
二つのLoRAをマージするとき、やりたいことは $\Delta W_1 + \Delta W_2 = A_1B_1 + A_2B_2$ となります(weightは省略)。しかし単に全重みをマージしたLoRAは、 $(A_1+A_2)(B_1+B_2)=A_1B_1+A_1B_2+A_2B_1+A_2+B_2$ となって目標とは違うものになってしまいます。特異値分解を行う方法では、 $A_1B_1 + A_2B_2$ の低ランク近似を行います。これも近似ではありますが、全く別のLoRAをマージするときは特異値分解を行う方法の方が精度が高くなりそうです。ちなみにランクが上がってしまいますが、 $A=(A_1 A_2), B = \begin{pmatrix} B_1 \\ B_2 \end{pmatrix}$ とすればマージ・・というかただの結合ですが精度を落とさず一つのLoRAにできます。

特異値とフロベニウスノルム

フロベニウスノルムとは、行列の各要素の二乗和の平方根です。 $\|A\|_F=\sqrt{\displaystyle{\sum_{i}\sum_{j}}a_{ij}^2}$ という感じです。ユークリッド距離をそのまま行列に適用したような形ですね。実は特異値の二乗和はフロベニウスノルムの二乗と一致します。まずそれを確認します。
補題1. $\|A\|_F^2=\mathrm{tr}(AA^T)=\mathrm{tr}(A^TA)$ ( $\mathrm{tr}$ は対角成分の和です。)
証明：
$A$ を $(m,n)$ 行列とすると、 $(AA^T)_{ij}=\displaystyle{\sum_{k=1}^{n}}a_{ik}a_{jk}$ ですので、対角成分は、 $(AA^T)_{ii}=\displaystyle{\sum_{k=1}^{n}}a_{ik}a_{ik}$ となり
$\mathrm{tr}(AA^T)= \displaystyle{\sum_{i=1}^m}(AA^T)_{ii}=\displaystyle{\sum_{i=1}^m}\displaystyle{\sum_{k=1}^{n}}a_{ik}a_{ik}=\|A\|_F^2$ です。 $A^TA$ の場合も多分同じようなもん。
補題2. 直交行列をかけてもフロベニウスノルムは変わらない。つまり $\|A\|_F=\|UA\|_F=\|AV\|_F$
証明：
$\|UA\|_F^2=\mathrm{tr}((UA)^TUA)=\mathrm{tr}(A^TU^TUA)=\mathrm{tr}(A^TA)=\|A\|_F^2$
$\|AV\|_F^2=\mathrm{tr}(AV(AV)^T)=\mathrm{tr}(AVV^TA^T)=\mathrm{tr}(AA^T)=\|A\|_F^2$
補題1を使ってます。正の数なので両辺のルートをとれば成立です。系として両側からかけても $\|UAV\|_F=\|UA\|_F=\|A\|_F$ となります。

定理：特異値の二乗和＝フロベニウスノルムの二乗
特異値分解された形を考えれば、補題2を使って、
$\|A\|_F^2=\|U\Sigma V\|_F^2 = \|\Sigma\|_F^2 = \displaystyle{\sum_{i=0}^{r}\sigma_i^2}$

つまりLoRAのリサイズに特異値の二乗和を利用するのは、フロベニウスノルムを使った行列間の距離に基づく方法になります。閾値を $r^2$ %としたのは特異値の累積二乗和の平方根をとるより閾値を二乗したほうが計算も実装が楽だからだと思います。

低ランク近似の根拠

LoRAの抽出において、特異値を大きい順にソートして、 $k$ 個で打ち切ることで近似しましたが、これには根拠があります。以下の定理が成り立ちます。
エッカート・ヤングの定理.
$\displaystyle{\mathrm{argmin}_{X,\mathrm{rank}(X)\leq k}} \| A-X\|_F = \tilde{A_k}$
ここで $\tilde{A_k}$ は低ランク近似を行った行列です。つまり低ランク近似はフロベニウスノルムを距離としたときに最適な近似ということです、
証明：
$A=U\Sigma V$ と特異値分解できるとすると、
$\begin{eqnarray} \| A-X\|_F^2 &=& \| U^T(A-X)V^T\|_F^2 \\ & = & \mathrm{tr}(U^T(A-X)V^T(U^T(A-X)V^T)^T) \\ & = & \mathrm{tr}((U^TAV^T-U^TXV^T)(U^TAV^T-U^TXV^T)^T ) \\ & = & \mathrm{tr}((\Sigma-U^TXV^T)(\Sigma-U^TXV^T)^T ) \\ & = & \|\Sigma - U^TXV^T \|_F^2 \end{eqnarray}$
となります。ここで $U^TXV^T=G$ として、 $(i,j)$ 要素を $g_{ij}$ とすると、
$\begin{eqnarray} \| A-X\|_F^2 &=& \|\Sigma - U^TXV^T \|_F^2 \\ & = & \displaystyle{\sum_{i=1}^{r}(\sigma_i - g_{ii})^2} + \displaystyle{\sum_{i=k+1}^{m}g_{ii}^2} + \displaystyle{\sum_{i\neq j}g_{ij}^2} \end{eqnarray}$

ここで $\mathrm{rank}(X)=\mathrm{rank}(G)$ となるので、上の値と $G$ のrankをなるべく小さくしたいですね。このとき $i=k+1,\dots,r$ までの $g_{ii}$ と $i\neq j$ で $g_{ij}$ は0にすればrankも $\| A-X\|_F$ もあがりません。あとは $i=1,\dots,k$ での、 $g_{ii}$ です。これはそれぞれ、 $g_{ii}=\sigma_i$ とすれば0になりますが、rankの制限のことを考えると、0より大きくしていいのは最大 $k$ 個です。そこで大きい順から $k$ 個選んで、それ以外を0とすれば最適解になります。これはまさに $\Sigma_{[\sigma_{k+1},\dots,\sigma_r:=0]}$ そのものであり、 $X=UGV=A_k$ になります。

LoConの特異値分解

いままでは全結合層の重みについて話していきましたが、畳み込みニューラルネットワークでも同様の操作が行えます。LoConですね。畳み込みニューラルネットワークは、フィルター1回分の計算を考えてみると、入力チャンネル×フィルターサイズ→出力チャンネル×1の全結合層になります。この全結合層を位置を動かしながら適用するだけです。

フィルターサイズ3×3の畳み込み層の重みを $W^{(m\times 9,n)}$ とすれば、LoConは $\Delta W^{(m\times 9,n)} = A^{(m\times 9,r)}B^{(r,n)}$ であり、要素の順番さえ気をつかえば特異値分解によるLoRA抽出も行えます。
CP分解についてはLyCORISで実装されていますが特異値分解とはまた別の分解でありよく分かっていません。
upの方を3×3にする方式は多分特異値分解による抽出は出来ないと思います。down層はin×9次元からrank×9次元への変換ですが、実際には9マスそれぞれ同じ重みを共有しているからです。つまり $\Delta W = A^{(m\times 9,r \times 9)}B^{(r \times 9,n)}$ としたいのですが、実際にはdownは $A^{(m\times 9,r \times 9)}$ ではなく、 $A^{(m,r)}$ になります。ただし特異値分解による抽出ができないだけで、マージはできます。2×2を2回適用する方法も特異値分解できないと思います。こっちはもはや行列としても表現しにくいです。ただしこちらもマージはできます。まあこの辺多分そうでしょうという感じで語ってます。

まとめ

LoRAは特異値分解のような形状になっている。そのため特異値分解によってリサイズや抽出ができる。
特異値分解による低ランク近似はフロベニウスノルムに基づいており、特異値の二乗和によって計算できる。
畳み込みのLoRAで特異値分解の形をしているのは、downを3×3の畳み込み、upを1×1にする方法のみっぽい。