LoRA(低ランク適応)

実装: llm/lora/ / 実行: go test ./llm/lora/

大規模モデルを丸ごと微調整するには全パラメータ分の勾配とオプティマイザ状態が要り、メモリが巨大になる。LoRA は元の重みを凍結したまま、その隣に小さな低ランク行列の積を足す。学習するのはこの小行列だけで、全体の 1% 未満にすぎない。学習後は元の重みに足し込めば 1 枚に戻り、推論の追加コストはゼロだ。この章では低ランク補正・初期恒等性・マージを実装し、パラメータ削減を数字で確かめる。

この章で作るもの

学習パイプライン の SFT や RLHF は、モデルの全パラメータを更新するフル微調整だった。だが 70B モデルをフル微調整するには、重みだけでなく勾配とオプティマイザの状態(Adam なら重みの 2 倍)も持つ必要があり、必要メモリは推論の数倍に膨れる。手元の GPU で公開モデルを自分用に調整したくても、この壁に阻まれる。

LoRA(low-rank adaptation)はこれを崩す。元の重み行列 W を凍結し、更新しない。かわりに W の隣に、小さな 2 枚の行列 A と B の積を補正として足す。学習するのは A と B だけだ。W が d×k なら A は d×r、B は r×k で、r(ランク)を小さく取れば学習パラメータは W 全体のごく一部で済む。

        ┌─────── W (d×k、凍結) ───────┐
入力 x ─┤                              ├→ 足す → 出力
        └── A(d×r) → B(r×k) 学習する ─┘

r ≪ d,k なので:
  フル微調整   d×k 個(例 4096×4096 = 1600万)
  LoRA         (d+k)×r 個(例 4096×8×2 = 6.5万 = 0.4%)

LoRA の構造。元の重み W は凍結し、低ランクの A·B を補正として並列に足す。学習するのは細い A と B だけで、d×k の W に対して (d+k)×r と激減する

順に見ていく。

1. 低ランク補正で更新を代理する: W を凍結し、A·B だけ学習する。微調整に必要な変化は低次元に収まるという観察が根拠
2. 初期は base と恒等: B を 0 で始めるので学習前の補正は 0。元のモデルの挙動を壊さない
3. マージで追加コストなし: 学習後 A·B を W に足せば 1 枚の行列に戻る。推論時は通常と同じ

① 低ランク補正: 学習パラメータを削る

実装の中心は、base の出力に低ランクの補正を足すことだ。B を 0 初期化するのが重要で、これで学習開始時の補正がゼロになる:

// LoRA は凍結した base 行列 (d×k) に、低ランク補正 A(d×r)·B(r×k) を足す層。
//
//	y = x·W + (alpha/r)·x·A·B     W は凍結、A・B だけ学習
//
// A は小さな乱数、B は 0 で初期化するので、学習開始時の補正は 0
// (= モデルの挙動を変えない)。
type LoRA struct {
	base  *tensor.Tensor // 凍結された元の重み (d×k)
	a     *tensor.Tensor // (d×r)
	b     *tensor.Tensor // (r×k)
	rank  int
	scale float32 // alpha/rank
}

// New は base の隣に rank 次元の LoRA を作る。alpha は補正の強さ。
// A は決定的な擬似乱数、B は 0 初期化(初期補正ゼロ)。
func New(base *tensor.Tensor, rank int, alpha float32) (*LoRA, error) {
	d, k := base.Rows, base.Cols
	if rank < 1 {
		return nil, errors.New("lora: rank must be >= 1")
	}
	if rank > d || rank > k {
		return nil, errors.New("lora: rank must not exceed matrix dimensions")
	}
	return &LoRA{
		base:  base,
		a:     randMatrix(d, rank, 42),
		b:     tensor.New(rank, k), // ゼロ初期化
		rank:  rank,
		scale: alpha / float32(rank),
	}, nil
}

// Forward は y = x·base + scale·(x·A·B)。base は変えず補正だけ足す。
func (l *LoRA) Forward(x *tensor.Tensor) *tensor.Tensor {
	out := tensor.MatMul(x, l.base)
	delta := tensor.MatMul(tensor.MatMul(x, l.a), l.b) // (x·A)·B
	for i := range out.Data {
		out.Data[i] += l.scale * delta.Data[i]
	}
	return out
}

// SetB は学習相当の更新をテストするため B を差し替える(通常は勾配で更新される)。
func (l *LoRA) SetB(b *tensor.Tensor) { l.b = b }

テストで固定したのは初期の恒等性だ。作った直後の LoRA は、出力が base 単独(x·W)と完全に一致する。B が 0 だからで、これは「学習を始めた瞬間にモデルの振る舞いが変わらない」という要件を満たす。フル微調整のように初期状態が不安定にならず、既存の能力を保ったまま少しずつ調整に入れる。

なぜ低ランクで足りるのか。Aghajanyan らの観察(2020)と LoRA 論文(2021)は、事前学習済みモデルを新しいタスクに適応させるとき、必要な重みの変化 ΔW が本質的に低ランク(少数の方向にしか動かない)であることを示した。W 全体は高次元でも、タスク適応で動く部分は細い。だから ΔW を r 本の方向の積 A·B で近似すれば足りる。

② パラメータ削減を数える

削減量を会計する。学習対象は A と B の 2 枚だけだ:

// TrainableParams は学習対象(A + B)のパラメータ数。base は含まない。
func (l *LoRA) TrainableParams() int {
	return l.a.Rows*l.a.Cols + l.b.Rows*l.b.Cols
}

// BaseParams は凍結された base のパラメータ数(参考: フル微調整で学習する量)。
func BaseParams(base *tensor.Tensor) int { return base.Rows * base.Cols }

4096×4096 の重みをフル微調整すると 1600 万パラメータを学習するが、rank 8 の LoRA なら A・B 合わせて 6.5 万、全体の 0.4% で済む。テストでは学習パラメータが base 全体の 0.5% 未満に収まることを固定した。

この削減が効くのは勾配とオプティマイザ状態だ。凍結した W には勾配が要らないので、メモリを食うのは小さな A・B の分だけになる。結果として、フル微調整では複数 GPU が要ったモデルが、単一 GPU で調整できるようになる。量子化章の 4bit 凍結重みと組み合わせた QLoRA は、この 2 つを重ねて、消費者向け GPU 1 枚で 65B モデルの微調整を実現した。

③ マージ: 推論時のコストをゼロにする

LoRA のもう 1 つの利点は、学習後に補正を base に畳み込めることだ:

// Merge は base に scale·(A·B) を足し込んだ 1 枚の行列を返す。
// 学習後にこれを使えば、推論時は通常の行列 1 枚と同じで追加コストが無い。
func (l *LoRA) Merge() *tensor.Tensor {
	ab := tensor.MatMul(l.a, l.b) // (d×k)
	merged := tensor.New(l.base.Rows, l.base.Cols)
	for i := range merged.Data {
		merged.Data[i] = l.base.Data[i] + l.scale*ab.Data[i]
	}
	return merged
}

学習が終われば A·B は固定なので、W + (alpha/r)·A·B を計算して 1 枚の行列にしてしまえる。テストでは、マージした行列での出力が Forward(base + 補正を別々に計算)と一致し、サイズが base と同じ(追加パラメータなし)であることを固定した。

これが adapter(層の間に小さなモジュールを挿す別の微調整法)に対する LoRA の優位だ。adapter は推論時も追加の層を通るのでレイテンシが増えるが、LoRA はマージすれば元と同じ 1 枚の行列積になり、推論コストが増えない。さらに、base を共有したまま複数の LoRA(タスクごと、顧客ごと)を切り替えられる。マージせず持っておけば、同じ base モデルに別々の小さな LoRA を差し替えて別タスクをこなせる。

動かす

下のデモは 2 つの見方を用意した。「パラメータ削減」は行列サイズとランクを変え、フル微調整と LoRA の学習パラメータ数の差を見る。「学習と恒等性」は B が 0 の初期状態から補正が育ち、base の出力から少しずつズレていく様子と、マージしても出力が変わらないことを追う。

デモLoRA(低ランク適応)0.39%

パラメータ削減学習と恒等性r=4r=8r=16r=64

重み行列: 1024×10242048×20484096×4096

フル微調整16.8M

LoRA (A·B)66k

学習するパラメータの比。LoRA は全体の 0.39% だけ

4096×4096 の重みを rank 8 で適応。フル微調整は 16.8M 個を学習するが、LoRA は A·B の 66k 個(0.39%)だけ。凍結した W には勾配が要らないので、省くのはこの差分のメモリ

LoRA は元の重みを凍結し、隣に低ランクの A·B を足して学習する。学習量は全体の 1% 未満で、 B=0 初期化により始めは base と恒等。学習後は A·B を W に畳み込めば 1 枚に戻り、 推論の追加コストはゼロ。量子化と組めば(QLoRA)単一 GPU で大規模微調整ができる。

設計の観点

• ランクの選び方: r が小さいほど省メモリだが表現力が下がる。タスクが base から遠いほど大きい r が要る。実務では 4〜64 の範囲で、多くは 8〜16。r を上げても頭打ちになりやすい
• alpha の役割: 補正の強さを alpha/rank でスケールする。r を変えても補正の大きさが揃うよう alpha も連動させるのが通例(alpha=2r など)
• どの層に挿すか: 全層でなく attention の Q/V 射影だけに挿すのが効率的、という LoRA 論文の知見。挿す場所の選択も設計変数
• マージの功罪: マージすれば推論コストゼロだが、複数タスクの切り替えはできなくなる。マージせず持てば切り替え可能だが推論に少し足される。運用の要求で選ぶ
• QLoRA との重なり: 量子化で base を 4bit に縮め、LoRA で微調整する。凍結重みは量子化で圧縮でき、学習する A・B だけ高精度で持つ。2 つの技術が直交して効く好例

対照と実例

手法	学習パラメータ	推論コスト	特徴
フル微調整	全部(100%)	変わらず	最高精度、最大メモリ
Adapter	挿入層のみ	増える	層を挿す。切り替え容易だが遅い
LoRA	A·B のみ(<1%)	マージで増えない	低ランク補正。実質標準
QLoRA	A·B のみ + 4bit base	マージで増えない	単一 GPU で大規模微調整
Prefix/Prompt tuning	少数トークン	わずかに増える	入力側に学習ベクトルを足す

裏どり:

• LoRA(2021): Hu et al.。低ランク補正の原典。GPT-3 175B を 0.01% のパラメータで微調整できることを示した
• Intrinsic Dimensionality(2020): Aghajanyan et al.。微調整に必要な次元が驚くほど低い、という LoRA の理論的下地
• QLoRA(2023): Dettmers et al.。4bit 量子化 + LoRA で単一 GPU 微調整
• DoRA(2024): Liu et al.。重みを大きさと向きに分解して LoRA を改良

簡略化したこと

• 勾配・学習ループなし: B の更新は差し替えで代用し、低ランク補正の性質だけを検証した
• 1 層のみ: 実物は複数の射影に LoRA を挿す。ここは 1 枚の重みで仕組みを示した
• 量子化併用なし: QLoRA は 量子化章と本章の組み合わせで、ここでは別々に扱った
• 変種なし: DoRA・LoRA+ などの改良は設計の観点で言及のみ

参考資料

• Hu et al., LoRA: Low-Rank Adaptation of Large Language Models(2021)
• Aghajanyan et al., Intrinsic Dimensionality Explains the Effectiveness of Fine-Tuning(2020)
• Dettmers et al., QLoRA(2023)
• 実装: llm/lora