LoRA は重みに触れずに脇道を付ける

一言でいうと

LoRAは、学習済みの重みWを凍結しておき、その脇に小さな行列2つ(A: 入力→r、B: r→出力)を付けて、B·Aだけを学習します。Bを0で始めるので、最初はモデルがまったく変わらず、学習が終わったらW + B·Aに統合して、脇の枝なしで使えます。MiniMindのmodel_lora.pyは、これを60行で実装しています。このモジュールでは、SFTモデルにLoRAを載せて口調を変え、基盤の重みが1ビットも変わっていないかを、直接確認します。

なぜ必要なのか

LoRAの論文は、GPT-3 175Bを例に挙げます。課題ごとに全体を微調整したモデルを別々に置くと、課題1つに1,750億パラメーターが必要です。事前学習の重みを凍結して低ランクの行列だけを学習すると、学習するパラメーターが1万分の1に、GPUメモリが3分の1に減り、アダプターと違って、推論のレイテンシも増えないと報告しました。MiniMindのREADMEも、同じ使い方を挙げています。基盤モデルの一般的な能力はそのままにして、医療のようなドメインや、「私は誰」のような自己認識を、LoRAで載せます。データが十分なら、全体のSFTもできますが、そのときは、ドメインのデータに過学習して一般的な能力を失わないように、データを混ぜる作業が別に必要です。

どう動くのか

MiniMindの実装は、3つの部分でできています。

class LoRA(nn.Module):
    def __init__(self, in_features, out_features, rank):
        self.A = nn.Linear(in_features, rank, bias=False)   # 정규분포(std 0.02)
        self.B = nn.Linear(rank, out_features, bias=False)  # 0
    def forward(self, x): return self.B(self.A(x))

def apply_lora(model, rank=16):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and module.in_features == module.out_features:
            module.lora = LoRA(...);  module.forward = 원래 forward(x) + lora(x)

付く場所は、入力と出力の次元が同じLinearだけです。私たちのモデル(128次元、Qヘッド4×32)では、q_proj(128→128)とo_proj(128→128)が該当し、k_proj・v_proj(128→64)とFFN(128↔448)は外れます。層4個×2=8か所で、1か所あたりA(8×128)+B(128×8)=2,048個なので、学習するパラメーターは16,384個、全体の1.6%です。論文は、どの行列に付けるかを選ぶ問題を別に扱っていますが、MiniMindは、この単純な規則で決めました。

開始点では、Bが0なので、B·A = 0です。載せた直後の出力は、基盤モデルと完全に同じです。学習は、「基盤モデルからどれだけ離れるか」を0から学びます。

倍率について、論文は、ΔW·xにα/rを掛けて、rankを変えても学習率を選び直さなくてよいようにします。MiniMindの実装は、この倍率なしで足します。そのため、rankを変えたら、学習率も見直す必要があります。

凍結と保存では、train_lora.pyは、名前にloraが入っていないパラメーターを、すべてrequires_grad=Falseにして、オプティマイザーには、LoRAのパラメーターだけを渡します。保存は、save_loraが、LoRAの重みだけをfp16で切り出して保存します。数十KBです。使うときは、基盤モデルにapply_loraしてからload_loraするか、merge_lora(W + B·A)で統合したモデルを1つ作ります。

現場での姿

1つの基盤モデルに、顧客企業・課題ごとに違うLoRAを置くと、保存は、基盤1つと小さなファイル複数になり、サービングは、リクエストごとにLoRAだけを差し替えればよいことになります。このとき、基盤の重みが知らないうちに変わると、すべてのLoRAが同時に狂います。凍結を忘れた学習スクリプト1つが、そうしてしまいます。そのため、学習が終わったら、基盤の重みを元のものとビット単位で比べるのが、安くて確実な確認です。逆に、1つの用途にしか使わないなら、統合(merge)して配布し、推論のときの脇の枝の計算をなくします。

MiniMindのオリジナルとこのコースの違い

MiniMindのデフォルトのrankは16で、train_lora.pyは、lr 1e-4で10エポック回します。このコースは、rank 8、lr 5e-3、150学習ステップです。モデルが小さく、課題(口調1つ)が単純なので、rankを減らしても十分で、短い時間で効果を見るために、学習率を大きくしました。アルファの倍率がない実装なので、rankと学習率は一緒に動きます。rankを2倍にすると、B·Aの大きさが大きくなる分だけ、同じ学習率が、事実上より大きな1学習ステップになります。また、MiniMindは、LoRAをSFTモデル(full_sft)の上に載せることを基本にしています。事前学習のモデルに直接載せると、対話の形式から学ぶ必要があって、16,384個のパラメーターでは足りません。

次のラボですること

SFTの基準モデルに、rank 8のLoRAを載せて、どこに付くかと、載せた直後の出力がそのままかを確認します。LoRAだけを学習するように凍結してから、「…だニャン。」の口調(韓国語の猫の語尾にあたる口調です)のデータで150学習ステップ学習し、基盤の重みがそのままか、取り分けておいた質問で口調がどれだけ変わったか、統合したモデルが、LoRAを載せたモデルと同じ出力を出すかを測ります。