LoRA の原理を手書きで理解する

前の課で、本当にファインチューニングが必要なときは LoRA が第一候補だと言いました。この課ではそれが何なのかを理解し、モジュール 09 で学習させた小さな GPT で実際に実装してみます。

python lora_from_scratch.py

フルファインチューニングの面倒

ファインチューニングとは、学習済みのモデルを新しいデータでさらに学習させることです。最も直接的なやり方はすべてのパラメータを学習に参加させるもので、フルファインチューニングと呼ばれます。

LLM にとって、これは高くつきます。学習のとき、各パラメータはそれ自体に加えて勾配も保存する必要があり、AdamW オプティマイザはさらにパラメータごとに数を二つ保存します（モジュール 08 第 5 課で触れた、各パラメータの「モーメント」を記録するものです）。70 億パラメータのモデルなら、これだけで 100 GB を超える GPU メモリが必要になります。

しかもファインチューニングするたびに、完全な新しいモデルが一つできます。顧客 10 社のためにそれぞれ一つずつファインチューニングすれば、数十 GB のファイルを 10 個保存しなければなりません。

LoRA の考え方

LoRA（Low-Rank Adaptation、低ランク適応）の考え方はこうです。ファインチューニングで重みに必要な変更量は、実はそれほど「複雑」でなくてよい。

線形層の重みは行列 W で、たとえば 128×384 です。フルファインチューニングは同じ大きさの変更量 ΔW を学んで W に足すことです。LoRA は ΔW を直接学ぶのではなく、とても細長い二つの行列の積に分解します。

      ΔW       =       A       ×     B
  (128 × 384)      (128 × 8)     (8 × 384)

  49152 个数       1024 个数   +  3072 个数  =  4096 个数

真ん中の 8 をランク（rank、r）と呼びます。A と B を合わせても 4096 個の数しかなく、ΔW の 12 分の 1 です。r が小さいほど学ぶパラメータは少なくなりますが、表せる変化も限られます。

学習のとき、元の W は凍結して動かさず、A と B だけを学習させます。

输出 = x × W  +  x × A × B × 缩放系数
       ─────     ─────────────────────
      原来的路        新加的支路

実装

学習済みの nn.Linear を包むのに、十数行で足ります。

class LoRALinear(nn.Module):
    """包住一个已经训练好的线性层：原来的权重冻结不动，旁边加一条 A、B 两个小矩阵的"支路"。"""

    def __init__(self, base: nn.Linear, rank=8, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False  # 原来的参数一个都不训练
        self.A = nn.Parameter(torch.randn(base.in_features, rank) * 0.01)  # 输入维度 × r
        self.B = nn.Parameter(torch.zeros(rank, base.out_features))  # r × 输出维度，初始为 0
        self.scale = alpha / rank

    def forward(self, x):
        # 原来的输出 + 支路的输出。B 一开始是 0，所以刚加上时模型的行为和原来完全一样
        return self.base(x) + (x @ self.A @ self.B) * self.scale

細かい点が二つあります。

B を 0 で初期化する。こうすると最初は A×B が 0 で、脇道はまったく作用せず、モデルは元とまったく同じです。学習は元のモデルから出発し、ランダムな脇道でかき乱されたモデルから出発するのではありません。A まで 0 にしてはいけません。そうすると A と B の勾配がどちらも 0 になり、いつまでも学習が進みません。

スケーリング係数 alpha / r。脇道の影響の大きさを制御します。こうしておけば r を調整するときに学習率を調整し直す必要がありません。r=8、alpha=16 はとてもよく使われる設定です。

それからモデルの各アテンション層の qkv と proj をすべて LoRA に置き換えます。

for p in model.parameters():
    p.requires_grad = False
for block in model.blocks:
    block.attn.qkv = LoRALinear(block.attn.qkv)
    block.attn.proj = LoRALinear(block.attn.proj)

タスク：五言絶句だけを書かせる

モジュール 09 で学習させた小さな GPT は、どんな形式の詩でも書きます。今度は五言絶句だけでファインチューニングして、五言絶句だけを書くように変えられるか見ます。

原来的训练数据里，五言绝句占 10.4%（3538 首）
微调前：生成 200 首，其中五言绝句 67 首

模型共 1,639,296 个参数，要训练的只有 24,576 个（1.50%）
刚装上 LoRA 时，输出和原模型一样吗？True

ファインチューニングの前、200 首のうち 67 首が五言絶句です。これは学習データでの 10.4% よりかなり高く、五言絶句が最も短く、最後まで書き切りやすいからかもしれません。

LoRA を付けた後、学習させるのは 24576 個のパラメータだけで、モデル全体の 1.5% です。4 層で、各層に LoRA が二つあります。qkv の A と B は 128×8 と 8×384、proj のものは 128×8 と 8×128 で、1 層で 6144 個、4 層で 24576 個です。

「LoRA を付けたばかりのとき出力は同じか？ True」で、B を 0 で初期化した効果が確かめられました。

学習

五言絶句だけをデータにして、300 ステップ学習させます。

  第 1 步  损失 4.356
  第 100 步  损失 4.294
  第 200 步  损失 4.335
  第 300 步  损失 4.213
训练 300 步用了 23 秒

微调后：生成 200 首，其中五言绝句 191 首。前 5 首：
  白头还作雨，一鬓自如霜。何时之饮在，试死势悠扬。
  白发勤鸿急，秋贫梦不闲。夜寒休绕郡，秋杀杜陵陂。
  才子俭高人，何言亦重重。言知不见人，辄空一一言。
  一丛生碧霄，危棹度清湍。沙渚穿花发，苍苍隔鹤闲。
  九陌两金阁，一枝归草间。何当重枕簟，归复独裴回。

200 首のうち 191 首が五言絶句になり、33.5% から 95.5% に上がりました。動かしたのはパラメータの 1.5% だけで、学習は 23 秒です。

面白いことに、損失はほとんど変わっておらず、4.36 から 4.21 です。これは、モデルが詩を書く「腕前」（各文字をどれだけ正確に予測できるか）は上がっておらず、変わったのは「習慣」だということを示しています。五言の句を 4 つ書き終えたら、続けて改行を出力してその詩を終える、ということを学んだのです。LoRA が得意なのはまさにこの種の変化で、形式、スタイル、ふるまいを調整することであって、大量の新しい知識をモデルに教えることではありません。前の課で言った「ファインチューニングが変えるのはふるまいで、知識は RAG に頼る」ことが、ここではっきり見て取れます。

保存とマージ

LoRA 的权重单独存下来只有 96 KB，整个模型是 6.3 MB
把 LoRA 合并回原来的权重之后，模型又变回 1,614,720 个参数，输出和合并前最大差别 5.7e-06

LoRA の一つ目の利点：A と B だけを保存すればよいことです。ここでは 96 KB で、モデル自体は 6.3 MB です。LLM ではこの差はもっと極端になり、数十 GB のモデルでも、LoRA の重みは数十 MB しかないかもしれません。顧客 10 社のためにそれぞれ一度ずつファインチューニングしても、ベースモデル一つと小さなファイル 10 個で済みます。

二つ目の利点：マージできることです。学習が終わったら、A×B×スケーリング係数を元の W に直接足せば、普通の線形層になります。

def merged(self):
    """把支路合并回原来的权重，得到一个普通的线性层：推理时没有任何额外开销。"""
    layer = nn.Linear(self.base.in_features, self.base.out_features)
    with torch.no_grad():
        # nn.Linear 的权重形状是 (输出, 输入)，所以要转置
        layer.weight.copy_(self.base.weight + (self.A @ self.B).T * self.scale)
        layer.bias.copy_(self.base.bias)
    return layer

マージした後、モデルの構造もパラメータ数も元とまったく同じで、推論が遅くなることもありません。出力とマージ前との差は 5.7e-06 で、浮動小数点の誤差にすぎません。

LoRA はなぜ効くのか

直感的な説明はこうです。事前学習済みのモデルはすでにほとんどのことを学んでいて、ファインチューニングはその上で「小さな調整」をしているだけ。2021 年に LoRA を提案した論文（Hu ら、『LoRA: Low-Rank Adaptation of Large Language Models』）の出発点はまさにここで、この種の調整に必要な変更量はそれ自体が「低ランク」で、とても小さな r で表せる、というものです。

私たちの実験もこの説明に合っています。モデルに「4 句書いたら終える」ようにさせるのはとても単純なふるまいの変化で、パラメータの 1.5% で十分です。しかし大量のまったく新しい知識、たとえばある新しい分野まるごとの内容をモデルに教えたいなら、LoRA の効果はフルファインチューニングに及ばないかもしれません。

LoRA にはよく使われる変種もあり、QLoRA と呼ばれます。凍結した元のモデルを 4 ビットに量子化して（量子化は第 4 課で扱います）GPU メモリを節約し、LoRA の A と B は引き続き普通の精度で学習させるものです。これで、一般向けの GPU 一枚で数十億パラメータのモデルをファインチューニングできるようになりました。

練習問題

1. rank を 1、2、32 に変えると、ファインチューニング後の五言絶句の割合はそれぞれいくつになりますか。パラメータ数はそれぞれいくつですか。
2. LoRA を使わず、そのままフルファインチューニングして（すべてのパラメータを学習させ、学習率は 1e-4 に変える）、同じく 300 ステップ学習させ、五言絶句の割合と学習時間を比べてください。
3. フィードフォワードネットワーク（block.mlp[0] と block.mlp[2]）にだけ LoRA を付け、アテンションには付けないと、効果にどんな違いがありますか。

確認テスト

1. LoRA が学習させるのはどのパラメータですか？元のパラメータはどう扱いますか？

元の重みはすべて凍結し、学習には参加させません。LoRA は選んだ線形層のそばに小さな行列 A と B を二つ加え、それだけを学習させます。出力は、元の出力に x×A×B×スケーリング係数を足したものになります。

2. B を 0 で初期化するのはなぜですか？A も 0 で初期化してよいですか？

B が 0 なら A×B は 0 になり、LoRA を付けたばかりのモデルは元のモデルとまったく同じで、学習は元のモデルから出発します。A も 0 にしてはいけません。A と B がどちらも 0 だと、両方の勾配も 0 になり、パラメータはいつまでも更新されません。

3. LoRA の学習を終えた後、推論は遅くなりますか？

遅くならないようにできます。学習後に A×B×スケーリング係数を元の重みに足してマージすれば、得られるモデルは構造もパラメータ数も元のモデルと同じで、推論の速度も同じです。マージしない場合は脇道の計算が一本増えるので少し遅くなりますが、いつでも別の LoRA に切り替えられるという利点があります。