LoRAより良いLLMモデルのFine-tuning方法はありませんか？

この記事は、Fusic Advent Calendar 2024 18日目の記事です。

前言

最近、「Bone微調整」という技術がLoRA微調整と比較して、より速く、そして効果的であることが分かりました。この技術について詳しく説明していきます。LoRAは、その資源節約効果と性能の良さから広く使用されてきました。その派生手法として、全パラメータ微調整との差を改善することを目指したOLoRAやLoRA-Gaなどがあります。これまで低ランク微調整手法の最適化は極限に達しているように見えました。

しかし、先月（2024年11月）、低ランク微調整とは完全に異なる新しい効率的な微調整方法が登場しました。この方法はGQAやMQAのアイデアを参考にし、重み行列を複数のブロックに分割し、各ブロックで更新可能な小さな重みを共有する仕組みを採用しています。この手法は「block-affine-adaptation（ブロックアフィン適応）」と呼ばれています。

Bone微調整は最新のPEFTライブラリに含まれており、transformers v4.46.3との組み合わせを推奨します。このバージョンでは累積勾配の問題が修正されています。一方で、v4.47.0は累積勾配において損失が倍増する問題が確認されているため、使用しないことをお勧めします。

なお、本文中の図で灰色のブロックはPAD（ゼロ埋め）を表しています。

原理

効率的なパラメータ微調整を実現するには、微調整可能なパラメータを削減する必要があります。この場合、問題の核心は「どのように縮小するか」に変わります。LoRAを代表とする低ランク微調整のアプローチでは、次のように操作します：

$$\begin{aligned} \bm{y} &= \left( \bm{W}_{out \times in} + \Delta \bm{W} \right) \bm{x}_{in} \\ &= \left( \bm{W}_{out \times in} + \bm{B}_{out \times r} \bm{A}_{r \times in} \right) \bm{x}_{in} \end{aligned}$$

ここで、 $r \ll \mathrm{Rank}[\bm{W}]$

BoneはGQAやMQAのアイデアを参考にし、次のように操作します：

$$\begin{aligned} \bm{y} &= \left( \bm{W}_{out \times in} + \Delta \bm{W} \right) \bm{x}_{in} \\ &= \left( \bm{W}_{out \times in} + \mathrm{Reshape} \left[ \bm{B}_{out \times r} \right]_{out \times in} \right) \bm{x}_{in} \end{aligned}$$

ここで、$\mathrm{Reshape}$操作は実際には$\bm{B}_{out \times r}$をコピーし、必要に応じてゼロ（0）で埋めて、$out \times in$の次元に調整するものです。そして、 $r \ll in$ が成立します。

$in$ は入力次元を、$out$ は出力次元を表します。本文の数式はBone論文の数式とは若干異なります。正直なところ、論文の数式と図表はやや簡略化されすぎている印象を受けます。

やってみよう

Bone微調整法の考え方は非常にシンプルです。実質的には、行列を分割して共有し、重みを更新することです。
その効果はどうでしょうか？元の論文の実験結果によると、収束が早く、loss（損失値）が低いことが示されています。

筆者は qwen2.5-0.5b において、多くの人が関心を寄せる自己認識タスクの微調整を行いました。その結果は以下の通りです：
論文で述べられている通り、収束が早く、lossが低いことが確認されました。また、必要な微調整パラメータの量もさらに少なくなりました。

具体実現

Bone

数学的には一般的に列ベクトルを使用しますが、機械学習では行ベクトルをよく使用します。次に、Bone（block-affine-adaptation） の具体的な実装について探ってみましょう：

入力 $𝑥$ の形状を調整した後、行列 $\mathrm{B}$を結合する代わりに合計を取ることで同じ目的を達成し、さらにメモリも節約できます。

import torch
import random
import numpy as np
from einops import rearrange
import torch.nn.functional as F

# シードを固定する
seed = 12
torch.manual_seed(seed)
if torch.mps.is_available():
    torch.mps.manual_seed(seed)
elif torch.cuda.is_available():
    torch.cuda.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)

# 重み行列 (in_feature x out_feature = 7 x 3)
w = torch.randn(7, 3)
# 微調整行列 (rank x out_feature = 2 x 3)
delta_w = torch.randn(2, 3)
# Boneの行列更新 (rank = 2)
r = 2
# 入力 x の次元 (batch x in_feature = 1 x 7)
x = torch.randn(1, 7)

# 補完が必要なサイズを計算する
padding_size = (r - x.size(-1) % r) % r
# x にパディングを追加する
x_padded = F.pad(x, (0, padding_size))
# x をリシェイプする
x_reshaped = rearrange(x_padded, '... (d r) -> ... d r', r=r)
# 出力を計算する
y = x @ w + torch.sum(x_reshaped, dim=-2) @ delta_w
print(y.numpy())
# 出力次元 y: (batch x out_feature)
# [[-0.5105,  4.1959, -0.2296]]

最後にPEFTのソースコード実装を貼り付ける.

# peft/tuners/bone/layer.py 323-332
result = self.base_layer(x, *args, **kwargs)
for active_adapter in self.active_adapters:
    if active_adapter not in self.bone_block.keys():
        continue
    bone = self.bone_block[active_adapter]
    r = bone.size(0)
    if x.size(-1) % r != 0:
        padding_size = (r - x.size(-1) % r) % r
        x = F.pad(x, (0, padding_size))
    result = result + torch.sum(x.reshape(*x.shape[:-1], x.size(-1) // r, r), dim=-2) @ bone

BAT

Bone微調整では、分割された行列が同じ微調整行列を共有していることがわかります。しかし、本来であれば分割された行列は異なるものであるべきであり、それらの勾配情報は同じではないはずです。さらに、分割された行列の情報も十分に活用されていません。

この点に着目し、分割行列を有効に活用する方法が模索され、論文中で提案されたのがBlock-Affine-Transformation (BAT) です。

import torch
from einops import rearrange

# 重み行列 (in_feature x out_feature = 6 x 3)
w = torch.randn(6, 3)
delta_w = torch.randn(3, 3)

# 次元データを計算
in_feature = w.size(0)  # 入力特徴量の次元
out_feature = w.size(1)  # 出力特徴量の次元
r = 3  # 分割サイズ

# 重み行列をリシェイプする
w_reshaped = w.reshape(in_feature // r, r, out_feature // r, r)
# リシェイプ後の次元: out/r x in/r x r x r

# 行列の順番を入れ替える
w_reshaped = rearrange(w_reshaped, 'in_r a out_r b -> out_r in_r a b')

# 重みガイドの計算: 分割された重み行列と delta_w の行列積
weight_guide = w_reshaped @ delta_w

# ガイドを追加して更新された delta_w を生成
delta_w_with_guide = weight_guide + delta_w

# 次元を元に戻す
delta_w_with_guide = rearrange(delta_w_with_guide, 'out_r in_r a b -> in_r a out_r b')
delta_w_with_guide = rearrange(delta_w_with_guide, 'in_r a out_r b -> (in_r a) (out_r b)')

# 更新された重み行列を計算
new_w = w + delta_w_with_guide
print(new_w.numpy())

# 出力結果例:
# [[-0.10902774  2.5378582  -0.77570546]
#  [-1.0661874  -0.7394609  -0.07972464]
#  [-4.799402   -0.55822074 -0.9818473 ]
#  [ 2.281693    1.301961   -0.33605066]
#  [ 3.692684    0.92067444 -0.47727013]
#  [-2.1276853   2.8093538  -1.9311792 ]]

PEFT

# peft/tuners/bone/layer.py 225-230
w = (orig_weight.reshape(orig_weight.size(0) // r, r, orig_weight.size(1) // r, r).permute(2, 0, 1, 3)
    @ weight_bone
    + weight_bone)
output_tensor = w.permute(1, 2, 0, 3).reshape(*orig_weight.shape)

論文

Kang, J. (2024, November 28). Bone: Block-Affine Adaptation of Large Language Models. arXiv. https://doi.org/10.48550/arXiv.