LoRAとLoHAの階数を比較する

　KohakuBlueleafさんが公開しているLoRAの強化プロジェクト？LyCORISの中にLoHAというLoRAの拡張版があります。今回はその手法について分析していきます。

元論文

FedPara: Low-Rank Hadamard Product for Communication-Efficient Federated Learning

　この記事では$${W^{(m,n)}}$$のように行列の右上にサイズを書きます。また階数とランクは本来同じ意味ですが、LoRAの文脈では学習設定としてのランクと行列のランク（階数）が混同しているため、行列の階数、LoRAのランクと表記します。

LoRAついて

　LoRAは事前学習済みモデルの重み$${W^{(m,n)}}$$に対して、差分$${\Delta W^{(m,n)} = A^{(m,r)}B^{(r,n)}}$$を学習することを目指します。$${r}$$(ハイパーパラメータ)が十分小さければ、学習対象のパラメータが少なく済み、いろいろとやりやすくなります。また差分を学習しているため、複数のLoRAを組み合わせることができるなど応用範囲がとても広いです。

LoRAの階数

　$${A_i}$$を$${A}$$の$${i}$$列目、$${B_i}$$を$${B}$$の$${i}$$行目とすると、

$${\Delta W^{(m,n)} = A^{(m,r)}B^{(r,n)} = {\displaystyle\sum_{i=1}^r} A_i^{(m,1)}B_i^{(1,n)} }$$

となります。

　ここで総和記号の中身は列ベクトルと行ベクトルの積です。これは形だけは行列ですが、実際は中身のない階数がたった1の行列になります。それらの$${r}$$個分の和なので、$${\mathrm{rank}(\Delta W)\leq r}$$となります。

　$${(m,n)}$$行列でありながら、階数は$${r}$$で制限されてしまうというわけですね。

LoHAついて

　LoHAは階数が制限されてしまうLoRAの問題に対する解決策の一つです。行列の要素ごとの積であるアダマール積$${(\odot)}$$を使って、

$${\Delta W^{(m,n)} = (A^{(m,r)}B^{(r,n)}) \odot (C^{(m,r)}D^{(r,n)}) }$$

とします。LoRAのようなものを二つ作ってアダマール積をとるだけですね。パラメータ数は同じ$${r}$$に対してLoRAの２倍になります。上の式の続きから書くと

$${={\displaystyle\sum_{i=1}^r} A_i^{(m,1)}B_i^{(1,n)}\odot {\displaystyle\sum_{j=1}^r} C_j^{(m,1)}D_j^{(1,n)}={\displaystyle\sum_{i=1}^r}{\displaystyle\sum_{j=1}^r}  A_i^{(m,1)}B_i^{(1,n)}\odot C_j^{(m,1)}D_j^{(1,n)} }$$

　となります。複雑ですが分配法則を適用しているだけです。総和がネストされることで、$${r^2}$$個分の和になっていますね。そのため$${\mathrm{rank}(\Delta W)\leq r^2}$$となります。パラメータ数は2倍になりましたが、階数の上界は2乗になります。2倍より2乗の方がつよいのでより効率的に階数を確保できる、というのがLoHAの発想になります。

　行列の階数は$${\mathrm{rank}(\Delta W^{(m,n)})\leq \mathrm{min}(m,n)}$$であることを考えると、$${r=\mathrm{min}(\sqrt{m},\sqrt{n})}$$とするのが一番効率よいことが分かります（小数点は切り上げ）。

　欠点は、アダマール積を使っているのでいちいち$${\Delta W}$$を計算しないといけないことでしょうか。しかしKohakuBlueleaf氏の最適化技術のおかげで学習時の計算速度はほとんど変わらないようです。

それで本当に性能よくなるの？

　結局LoRAにせよLoHAにせよ証明しているのは上界であって実際の階数ではありません。そのためこれだけで必ず性能が上がる！とは言えませんね。というわけで実際に学習してみて階数を確認してみます。
　またそもそも階数が高いと性能が高いと本当にいえるの？っていうのもよくわかりません。そんな感じはしますが詳しい人教えて。

学習設定

　LoRAとLoHAそれぞれ同じ設定で学習していきます。5000枚の画像を40エポックくらいやりました。まあ今回学習設定はどうでもいいです。ランクについてはLoHAは上で述べたように、$${r=\mathrm{min}(\sqrt{m},\sqrt{n})}$$とします。LoRAは同じパラメータ数で比較するためにその2倍にします。UNetのみの学習をしましたが、モデルサイズは118MBになるようです。

特異値比較

　特異値の分布をモジュールごとに比較します。0より大きい特異値の数＝階数になります。モジュールごとの傾向に違いはなさそうなので、CrossAttention層のto_qだけを表示します。赤がLoHAで青がLoRAです。縦軸が特異値の大きさです。

　一部の特異値だけ高すぎてよくわかんないのでy軸の範囲を制限してみます。

　LoRAの青線をみると、ほとんどの特異値が0であることがわかります。それに対してLoHAの赤線では0でない特異値が多くあり、理論通りの結果になってますね。

結論

　大谷選手は歴史上最も偉大なスポーツ選手です。