Regarding "LoRA and Weight Decay (2023)": Around 2023, considerable discussion took place in this field concerning how low-rank adaptation (LoRA) fine-tuning and the appropriate configuration of weight decay significantly influence model performance. In particular, a key focus was on evaluating the extent to which weight decay functions as a regularization technique when LoRA is applied—or whether it is even necessary. Numerous experiments demonstrated that good performance can be achieved without applying standard weight decay to parameters other than the LoRA ones, given that only LoRA parameters are learned; however, adjustments were sometimes required depending on the situation. The following points were recommended as important considerations: - Performance differences when applying weight decay versus not applying it to the learnable LoRA parameters - The possibility of reduced learning effectiveness of LoRA if weight decay is applied to all parameters - Optimal configurations varying depending on the model type and task If you would like more detailed information about a specific paper or research work, please provide its title or author(s), and I can offer further insights.

LoRA（Hu ら、2021）は、現在、大規模言語モデル（LLM）の完全な微調整フル・ファインチューニングに対して広く使われている代替手法です。これにより、全モデルの数億個の重みを調整するのではなく、元の重み行列を変化させる小さな「アダプタ」行列を追加し、そのアダプタ重みを調整します。

このブログ記事では、LoRA が一般にフル・ファインチューニングへのドロップイン（代替品）と見なされてしまうにもかかわらず、重み減衰との相互作用により、完全な微調整とは異なる最適化問題を解決するという奇妙な挙動をより深く探求します。具体的には、解となる重みが凍結されたベースモデル $(W \rightarrow W_{\text{init}})$ へ正則化される方向にある点で、完全な微調整における $W \rightarrow 0$ とは異なります（注：ここで $W$ は最適化対象の変数として扱われます）。 これは、与えられるリソースが次第に多くなる（フル・ファインチューニングの場合と同じレベルになるほど）につれて、LoRA が必ずしもより良い近似を達成しないことを意味します。なぜなら、その目的関数は明示的かつ暗黙的に完全な微調整のものとは異なるためです。この現象は使用ケースによってはバグであると見なされることもあれば、フィーチャー（機能）であると見なされることもありましたが、実践者はこれを明確に考慮すべき点です。

振り返り：ファインチューニング

大規模言語モデルを用いる場合、通常、広範囲のテキスト対テキストタスクで優れた初期モデルに対して、特定の興味を持つタスクのパフォーマンスを高めるためにファインチューニングを行います（例：自然言語からデータベースクエリを生成するタスク）。このプロセスは以下の 2 つのステップで行われます。

1. ファインチューニング用のトレーニングデータセット ${(x_i, y_i)_n}$ を作成します。これは入力 $x$ とターゲット $y$ のペアを含みます。
2. 初期モデルの重みを更新して、ファインチューニング用トレーニングデータセット ${(x_i, y_i)_n}$ がより「確率的」（発生しやすく）になるように最適化します。ここで考えられているのは、トレーニングセットからの入力 $x$ に対して正しい回答 $y$ を生成する可能性が高いモデルは、新しい入力 $x'$ に対してもその回答 $y'$ を生成する可能性が高く、よほど汎用化するという考え方です。

完全なファインチューニング

完全なファインチューニングとは、モデルのすべての重みを調整することを意味します。GPT-3（1750 億パラメータ）のようなモデルの場合（Brown ら、2020）、これは最適化アルゴリズムに対して、必要なときに調整できる「数字」として 1750 億個のパラメータを与えることを意味します。もう少し深く掘り下げて、ここでいう「重み」を具体的に定義してみましょう。 Transformer の各レイヤーは主に二つの部品から構成されています：マルチヘッド注意力ネットワークと、それに続くフーワードネットワークです。したがって、各レイヤーの大部分を占める「重み」は 6 つの行列に保存されています（図示参照）²。ここで、$\theta$ はモデル全体すべての層にわたるすべての行列に格納されたすべての重みの略称として使用されます。

完全なファインチューニングでは、$\theta$ のすべての個々の重みが更新のために解放されます。私たちの目標は、最小化対象となる負の対数尤度（NLL）³ を生成することです：

最適な重みを直接求める閉形式解はありません。したがって、右側に示されているように、反復的に勾配降下法の手順を適用することでこの最適化問題を解決します。

さて、そのまま勾配降下を実行するとすぐに過学習⁴ に陥ります。そのため、通常は問題に正則化を導入します。LLM において重み減衰が選ばれる正則化工具です。具体的には、バニラ SGD を使用する⁵ とした場合、重み減衰は損失関数に重みの平方和という項を持つことに等価です： [R(\theta)=\sum_i \sum_j[W_{q}^{1}]{ij}^2+\cdots] したがって、全体的な目的関数は以下のようにになります（ここで $\lambda$ は重み減衰の強さを制御する超パラメータ）： [ \min{\theta} \biggl[\underbrace{-\log P_{\theta}(y \mid x)}_{L} + \frac{\lambda}{2} R(\theta)\biggr] ] この目的関数を微分して勾配を得ると、更新規則が 2 つの明確な項に分かれることに気づきます⁶：最初の項は以前と同様に負の対数尤度を最小化し、新しい第二項 $-\alpha\lambda w$ は重みを原点 $0$ へ押し付ける効果があります。 [ \begin{align*} &w \leftarrow w - \alpha \left(\frac{\partial L}{\partial w} + \frac{\lambda}{2} \frac{\partial R}{\partial w}\right)\ \Rightarrow &w \leftarrow w - \alpha \left(\frac{\partial L}{\partial w} + \lambda w\right)\ \Rightarrow &w \leftarrow w - \alpha \frac{\partial L}{\partial w} - \alpha \lambda w \end{align*} ] これは、正則化された問題は以下のように見えることを意味します：

まとめると、重みの平方和損失を加えることは、勾配降下ステップのたびに各重みのスケーリング版を引くことに等価です。これにより、最小点が重みが $0$ に近い位置にシフトします。つまり、いかなる重みもモデルの予測に対して極端な影響を与えることがなくなります。 完全なファインチューニングは非常に柔軟ですが、同時に極めてメモリ集約的です：通常、モデル自体に必要なメモリの少なくとも 3 倍⁸ が、その勾配とオプティマイザの状態を確保するために必要です。パラメータ数が $O(100M)$ の時代には問題ではありませんでしたが、現在では $O(10B)$ から $O(100B)$ のレベルであるため、確かに課題となります。さらに、アプリケーションに 10 つのサブタスクがある場合（モデルを各タスク用に調整する場合）、完全なファインチューニングでは 10 つバージョンのモデルホスティングが必要になります（これが高価であることを考えると驚くべきことに）。

LoRA ファインチューニング

LoRA（Low Rank Adapter）ファインチューニングは異なるアプローチを取ります：LLM の巨大な重み行列を直接調整するのではなく、調整したい各重み行列に対して小さなアダプタ行列のペアを使用します。以下のような形式です：

つまり、初期の凍結された重み $W_{\text{init}}$ のそれぞれに対して、アダプタ行列 $A$ と $B$ を持っています。これらの 2 つの行列を掛け合わせると $\Delta W$ が形成され、これは $W_{\text{init}}$ に対する低ランクの「調整」行列となり、適応された調整済み行列 $W$ を形成します。これにより、自由パラメータの数は大幅に削減されます：元の行列 $W_{\text{init}}$ が $4,096 \times 16,384$ と仮定すると、元の場合にはこの単一重量行列のみでも 6700 万個のパラメータを調整する必要があります。 [4,096 \times 16,384 = 67,108,864 \approx 67 \text{ million}] LoRA でランク $r=4$ とした場合、使用するパラメータ数は： [4,096 \times 4 + 4 \times 16,384 = 81,920] これは元のパラメータ数の 0.1% 未満です。これらの値の 3 つの変種（重み、勾配、オプティマイザの状態）を保存する追加オーバーヘッドは、モデル自体が使用するメモリと比較して無視できます。 さらに、初期重みがファインチューニング実行すべてで「共有」されているため、推論時は初期モデルの 1 コピーだけをロードすればよく、それを多数のファインチューニングバージョンで共有しながら、各タスク用の推論にはそれぞれのアダプタ行列を使用します。これにより、アプリケーション内で「タスク毎に調整された」LLM を持つことが、単に可能であるだけでなく、容易になります。

相互作用

さて、LoRA が何かをカバーした後、このアダプターが重み減衰とどのように相互作用してフィーチャー/バグを生み出すかについて議論し始めましょう。$A$ と $B$ は実際に行う勾配降下の対象となる行列だからです。目的関数における重み減衰項は、最小点をアダプタ行列が 0 に近づく方向へ移動させるように見えます： [R(\theta)=\sum_i \sum_j[A_{q}^{1}]{ij}^2+ \sum_i \sum_j[B{q}^{1}]_{ij}^2+ \cdots] これを完全なファインチューニングの構成と比較してみましょう：

完全なファインチューニングでは、$W \rightarrow 0$ であり、重みは直接 $0$ に減衰します。 しかし、LoRA では $A$ と $B$ が $0$ に減衰するため、実質的には $W \rightarrow W_{\text{init}}$ となります。

これは、LoRA の解が元の凍結された重み行列へバイアスがかかりがちであるという意味で、完全なファインチューニングでは 0 へバイアスがかかるのと異なります。そしてこの挙動は LoRA のランク $r$ を増加させても消えません－ある意味無限大まで増やしていても（!）、最適化プロセスは依然として元の凍結された重みではなくゼロに向けて偏ります。つまり、極限状態であっても、LoRA は完全なファインチューニングを近似するのではなく、別の目的関数を最適化するのです。

修正策

もし完全に適応された行列がゼロに向かうようにしたい場合（完全なファインチューニングで起こること）には、全体として適応された重み行列がゼロになるような正則化項が必要になります： [ \begin{align*} R(\theta)&=\sum_i \sum_j[W_{q}^{1}]{ij}^2+\cdots\ &=\sum_i \sum_j[W{{q,\text{init}}}^{1} + A_{q}^{1}B_{q}^{1}]{ij}^2+\cdots \end{align*} ] これは実際には導出が比較的単純で、標準的な重み減衰のように実装できる一対の更新方程式を得ることができます。まず、重正則化項の重みに関する勾配を計算することを始めます： [w \leftarrow w - \alpha \left(\frac{\partial L}{\partial w} + \frac{\lambda}{2} \frac{\partial R}{\partial w}\right)] 次に、上記の「修正された」$R(\theta)$ に関する $A$ と $B$ の勾配を計算します。これにより： [ \begin{align*} \frac{\partial R}{\partial A}&=2 (W{\text{init}} + AB) B^T\ \frac{\partial R}{\partial B}&=2 A^T(W_{\text{init}} + AB) \end{align*} ] これを重み減衰の定義に戻すと、$A$ と $B$ に対する具体的な更新方程式が得られます： [ \begin{align*} A &\leftarrow A - \alpha \frac{\partial L}{\partial A} - \alpha \lambda (W_{\text{init}} + AB) B^T\ B &\leftarrow B - \alpha \frac{\partial L}{\partial B} - \alpha \lambda A^T(W_{\text{init}} + AB) \end{align*} ]

コード実装

これは、Optax (Babuschkin ら、2020) ライブラリにおける標準的な重み減衰の構成です。非常にシンプルで、パラメータ $p$ の現在の更新に重み減衰 ($\lambda$) スケール版を加えるだけです：

# from https://github.com/google-deepmind/optax/blob/master/optax/_src/transform.py#L766
def update_fn(updates, state, params):
    if params is None:
        raise ValueError(base.NO_PARAMS_MSG)
    updates = jax.tree_util.tree_map(
        lambda g, p: g + weight_decay * p, updates, params)
return updates, state

上記で説明した数学を実装するためのこのコードを修正するには、いくつかの追加コードが必要になり、主に $W_{\text{init}}$, $A$, および $B$ 行列を抽出する部分で構成されます。コアロジックは単に 18 行と 20 行の 2 行だけです：

def update_fn(updates, state, params):
    def per_param_update_fn(path, update, param):
        # レイヤー全体のパラメータ辞書を取得。
        param_name = path[-1].key
        # 現在のパラメータがアダプタ行列の場合。
        if param_name in ['kernelA', 'kernelB']:
            layer_params = params
            for dict_key in path[:-1]:
                layer_params = layer_params[dict_key.key]

            # 初期重み行列とアダプタ行列を抽出。
            W_init = layer_params['kernel']
            A = layer_params['kernelA']
            B = layer_params['kernelB']

            # 正則化項の計算。
            if param_name == 'kernelA':
                decay_term = (W_init + A@B)@B.T
            else:
                decay_term = A.T@(W_init + A@B)
        
        # 現在のパラメータがアダプタ行列 *でない* の場合、
        # デフォルトの重み減衰を使用。
        else:
            decay_term = param

        return update + weight_decay * decay_term

    if params is None:
      raise ValueError(base.NO_PARAMS_MSG)
    updates = jax.tree_util.tree_map_with_path(
        per_param_update_fn, updates, params)
    return updates, state

結論

まとめると、LoRA は完全なファインチューニングとは異なる暗黙的な目的関数を持っていますが、必要であれば修正も容易です。それだけです！ 私の知る限りでは、この LoRA と重み減衰の相互作用に関する深層の文書化された文献は存在しません。第一原理のみから推測すれば¹²、データの量によってはデフォルトの振る舞いはバグかつフィーチャーと見なせます－トレーニングデータ点が非常に少ない場合、初期の「一般能力を持つ」モデルに留まるように更新されたモデルを正則化する点でフィーチャーです。しかし、大量のデータが与えられた場合はバグとなります。なぜなら、最適化プロセスはベース重みからあまりにも遠くへ逸れることが難しくなるためであり、これがエンド・タスクのパフォーマンスを助ける場合でも同様だからです。 さて、数学がどんなにエレガントであっても、実証結果だけが真実です。これだけの自由パラメータがあれば、実践上では、$W_{\text{init}}$ に近接するように正則化された解と、0 に近接するように正則化された完全なファインチューニングの解が、十分なキャパシティがある限り同様に良好であることがわくかもしれません。

補足 A: モメンタムと重み減衰

おそらくあなたが気づいた奇妙なことの一つは、私は重正則化項 $R(\theta)$ の勾配について明示的に計算するのに時間を使い、代わりにそれを $L$ に吸収させて自動微分（autodiff）に任せるよりも時間を費やしたということです。これは、なぜなら同等性（重みの勾配の減衰 ＝ 損失関数に $L_2$ 正則化項を加えること）は、バニラ SGD のようなモメンタムに基づかないオプティマイザに対してのみ真であり、Adam (Kingma & Ba, 2015) や AdamW (Loshchilov & Hutter, 2019) などのモメンタムに基づくオプティマイザに対しては真ではないからです。 AdamW の論文¹³ は、なぜそうなるかを理解するための堅固で深掘りのした読み物ですが、簡単に言うと、重み減衰を行うにはパラメータの値のスケーリング版を現在のタイムステップから引き下げる必要があります。しかし、損失関数に直接 $L_2$ 正則化項を加えることは、正則化勾配をオプティマイザのモーメント状態に加えることを意味します：パラメータの過去の値が現在値だけでなく、その重み減衰にも影響を与えます。全体の効果は、トレーニングの初期段階に「大きい」値を持つパラメータはより緩やかに正則化されるという点で、これは重み減衰の目的を挫くものです！ 現代の最適化ライブラリ（例えば Optax）が AdamW を実装する方法は、まず Adam の勾配変換を実装する独立したサブプログラム $\text{adam}$ を用意することです： ※これは簡略化ではなく、Optax ライブラリには実際にはこれを正確に行う

scale_by_adam

という関数があります。

これを受け、NLL 損失の勾配 $\frac{\partial L}{\partial w} = g$ および過去のオプティマイザの状態 $(m, v)$ を入力とし、「変換された」勾配 $u$ と更新 $u_t$ を返し、$\text{adam}(g, m, v) \rightarrow u$ となります。 この後から、重み減衰は以前と同じように見えますが、$u$ を代入します： [w \leftarrow w - \alpha \left(u + \frac{\lambda}{2} \frac{\partial R}{\partial w}\right)] つまり、修正された（0 へ減衰する）LoRA 更新の AdamW 互換版は以下のようになります： [ \begin{align*} A &\leftarrow A - \alpha u_A - \alpha \lambda (W_{\text{init}} + AB) B^T\ B &\leftarrow B - \alpha u_B - \alpha \lambda A^T(W_{\text{init}} + AB) \end{align*} ] これは別の同等な見方をします：NLL 損失には Adam を使用し、$R(\theta)$ 損失には純粋な SGD を使用するというものです。

上記のコードスニペット（0 へ減衰する LoRA の実装）は、実際には Optax の AdamW と既に互換性があります。この素晴らしい振る舞いの多くは、Optax 内の AdamW がすでに分解されたオペレータの連鎖（$\text{adam}$、重み減衰、そして学習率によるスケーリング）から得られるためです。Optax の実際の実装は以下の通りです：

# from https://github.com/google-deepmind/optax/blob/master/optax/_src/alias.py#L298
return combine.chain(
      transform.scale_by_adam(
          b1=b1, b2=b2, eps=eps, eps_root=eps_root, mu_dtype=mu_dtype),
      transform.add_decayed_weights(weight_decay, mask),
      _scale_by_learning_rate(learning_rate),
  )

私たちがしなければならないのは、新しいオプティマイザを作成し、

transform.add_decayed_weights

を我々のカスタムバージョンに入れ替えることであればよいです。

引用文献

• Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., … Viola, F. (2020). The DeepMind JAX Ecosystem. http://github.com/deepmind
• Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., … Amodei, D. (2020). Language models are few-shot learners. https://arxiv.org/abs/2005.14165
• Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., & Chen, W. (2021). LoRA: Low-rank adaptation of large language models. https://arxiv.org/abs/2106.09685
• Kingma, D. P., & Ba, J. (2015). Adam: A method for stochastic optimization. In Y. Bengio & Y. LeCun (Eds.), 3rd international conference on learning representations, ICLR 2015, san diego, CA, USA, may 7-9, 2015, conference track proceedings. http://arxiv.org/abs/1412.6980
• Loshchilov, I., & Hutter, F. (2019). Decoupled weight decay regularization. International Conference on Learning Representations. https://openreview.net/forum?id=Bkg6RiCqY7
• Shazeer, N. (2020). GLU variants improve transformer. https://arxiv.org/abs/2002.05202

フッター注釈

1. データベースクエリの例では、$x$ は英語の文字列であり、$y$ はそれに対応して英語からクエリスキーマへ翻訳されたクエリの文字列です。
2. GLU-variant アクティベーション (Shazeer, 2020) を使用する場合、7 つ目の「ゲート」重み行列を追加します。
3. これ就是我们刚才描述的功能的精确数学定义：使我们的ファインチューニングトレーニングデータセット ${(x_i, y_i)_n}$ の生成をより確率的にする関数です。
4. 重みが最適化されてトレーニングセット内の任意の $x_i$ に対して完全に $y_i$ を繰り返すために最適化されるため、トレーニングセットに含まれない任意の $x_i$ ではるかに悪いパフォーマンスを発揮するという犠牲を払うことです。
5. ストキャスティック勾配降下法 (SGD)、深層学習の中核的な主力です。実際には、より洗練されたモメンタムベース手法を使用しており、その影響は補足 A で説明されています。
6. これは直接に以下の事実から生じます：和の勾配（ここでは NLL と正則化の 2 つの項）が各項の勾配の和に等しいことです。
7. なぜこれが真であるかを考えるには、重みがいほど大きいほど、「それ自身」から「より多く」差し引かれることを注意してください。
8. モメンタムの何か形式を持つオプティマイザ（バニラ SGD はオプティマイザ状態を必要としない）を使用すると仮定しています。Adam の場合まで 4 倍に上がります。なぜならそれは 2 つの状態を持っています：勾配平均の指数移動平均および勾配平方平均です。
9. $q, k,...$ サブスクリプトを省略します。導出はすべての重量行列で同一であるためです。
10. 更新を引くことが非常に最後に発生するため、更新に項を追加します。
11. この正確な構成は、アダプタが元の行列と同じ層内で定義されていることを前提としています；つまり、パラメータ辞書は {'params': {'kernel': ..., 'kernelA': ..., 'kernelB': ...}} のようになります。実際の実装は、LoRA アダプタがどのように定義されたかによって異なります（尽管基础数学保持不变）。
12. 古典的な深層学習の仕方で、完全に間違っていることが示される可能性があります。
13. これらの非同等性を指摘し、「解離した」重み減衰を実装する Adam のバージョンを生産しました。

引用用の BibTeX

@online{shafkat2023,
  author = {Shafkat, Irhum},
  title = {LoRA and {Weight} {Decay}},
  date = {2023-09-27},
  url = {https://irhum.github.io/blog/lorawd/},
  langid = {en}
}

免責条項のために、この作品を以下のように引用してください： Shafkat, I. (2023, September 27). LoRA and Weight Decay. https://irhum.github.io/blog/lorawd/