【PyTorch】同じ層内のパラメータに別々の学習率を設定する方法

問題

1つのPyTorchモデルにて，パラメータ別に異なる学習率で学習したい層があるとします．具体的には

• 更新しない層
• 更新する層
  • 学習率を高めにして更新するパラメータ
  • 学習率をちょっと低めにして更新するパラメータ
  • ...などなど

とわかれている状況です．

この問題で難しいのは，同じ層内に異なる学習率で学習したいパラメータがあることです．下の記事では，層別に異なる学習率で学習する方法について紹介されています．

しかし，この方法では，同じ層内のパラメータに対して，別々の学習率を指定することができません．

同じ層内のパラメータに対して，別々の学習率を指定するにはどうすればいいでしょうか？

問題例

この問題例としては，LoRA+の実装があります．

詳しくは上などの記事を参照されたいのですが，簡単に言うと

• 更新しない層：元のモデル
• 更新する層:LoRA
  • 学習率を高めにして更新するパラメータ：LoRAのB
  • 学習率を低めにして更新するパラメータ：LoRAのA

というような学習を行います．LoRA層のそれぞれのパラメータA，Bに対して，別々の学習率を指定するにはどうすればいいでしょうか？

実装方法を例を交えて説明する

状況設定

PyTorchで作業するため，まずは必要なパッケージをインポートします．

コード：必要なパッケージをインポート

import torch
from torch import nn, optim, Tensor, gradient, tensor

次に示すような構造のモデル，MyModelを考えます．

MyModel(
  (t0): Target()
  (d0): Dummy()
  (t1): Target()
  (m0): Mother(
    (t): Target()
    (d): Dummy()
  )
  (t2): Target()
  (d1): Dummy()
)

Target層内にはパラメータが2つあります．それぞれa，bとします．Dummy内にもまた別名のパラメータがあります．Motherは内部にTargetとDummyを入れ子として持つ層です．

コード：モデルの定義

class Target(nn.Module):
    def __init__(self, a0:float, b0:float, *args, **kwargs ) -> None:
        super().__init__(*args, **kwargs)
        self.a = nn.Parameter(Tensor([a0]))
        self.b = nn.Parameter(Tensor([b0]))
        
        self.a.requires_grad = True
        self.b.requires_grad = True
    def forward(self, x:Tensor)->Tensor:
        return self.a*x + self.b

class Mother(nn.Module):
    def __init__(self, *args, **kwargs ) -> None:
        super().__init__(*args, **kwargs)
        self.t = Target(10., 20.)
        self.d = Dummy()
    def forward(self, x:Tensor)->Tensor:
        x = self.t(x)
        x = self.d(x)
        return x

class Dummy(nn.Module):
    def __init__(self, *args, **kwargs ) -> None:
        super().__init__(*args, **kwargs)
        self.c = nn.Parameter(Tensor([-1.]))
        self.d = nn.Parameter(Tensor([-2.]))
    def forward(self, x:Tensor)->Tensor:
        return self.c*x + self.d

class MyModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.t0 = Target(1.,2.)
        self.d0 = Dummy()
        self.t1 = Target(3.,4.)
        self.m0 = Mother()
        self.t2 = Target(5.,6.)
        self.d1 = Dummy()
    def forward(self, x:Tensor)->Tensor:
        x = self.t0(x)
        x = self.d0(x)
        x = self.t1(x)
        x = self.m0(x)
        x = self.t2(x)
        x = self.d1(x)
        return x

パラメータの更新方法を次のように定めます．

• 更新しない層：Dummy
• 更新する層:Target
  • 学習率を高めにして更新するパラメータ：a
  • 学習率を低めにして更新するパラメータ：b

実装

❶層の取り出し

パラメータa，bを取り出すためには，まずはそれぞれが属する層，つまりTarget層を取り出してくる必要があります．
次のコードでは，Targetのみを取り出し，Targetだけの配列を作ります．

target_list = []
for module in model.modules(): 
    if isinstance(module, Target):
        print("✅", end="")
        target_list.append(module)
    else:
        print("🟥", end="")
    print(module)
target_list

model.modulesメソッドは，自分のモジュール内の子モジュールを再帰的に挙げるものです．

結果

🟥MyModel(
  (t0): Target()
  (d0): Dummy()
  (t1): Target()
  (m0): TargetMother(
    (t): Target()
    (d): Dummy()
  )
  (t2): Target()
  (d1): Dummy()
)
✅Target()
🟥Dummy()
✅Target()
🟥TargetMother(
  (t): Target()
  (d): Dummy()
)
✅Target()
🟥Dummy()
✅Target()
🟥Dummy()

[Target(), Target(), Target(), Target()]

Targetだけうまく選べていることが分かります．

❷パラメータの取り出し

次に，Target層だけの配列から，aだけの配列とbだけの配列を作ります．

a_param_list = []
b_param_list = []
for target in target_list:
    a_param_list.append(target.a)
    b_param_list.append(target.b)
a_param_list, b_param_list

結果

([Parameter containing:
  tensor([1.], requires_grad=True),
  Parameter containing:
  tensor([3.], requires_grad=True),
  Parameter containing:
  tensor([10.], requires_grad=True),
  Parameter containing:
  tensor([5.], requires_grad=True)],
 [Parameter containing:
  tensor([2.], requires_grad=True),
  Parameter containing:
  tensor([4.], requires_grad=True),
  Parameter containing:
  tensor([20.], requires_grad=True),
  Parameter containing:
  tensor([6.], requires_grad=True)])

結果を見ると，`a`と`b`だけが取り出されていることが分かります．

❸最適化アルゴリズムの起動

最後に，PyTorchの機能を使って，Optimizerを起動します．

optimizer = optim.SGD([
    {"params":a_param_list, "lr":1.},
    {"params":b_param_list, "lr":0.0001}
])

このように書くと，a_param_list内のパラメータに対しては学習率を1に，b_param_list内のパラメータに対しては学習率を0.0001にして学習が行われます．また，ここに出てこなかったDummy層などのパラメータは更新されません．

❹学習過程

試しに簡単な例で学習を行い，パラメータ更新前と後でどのパラメータがどのくらい変わったかを見てみます．

コードと結果

print("Before")
for name, param in model.named_parameters():
    print(f"{name:<10s}", param.detach().numpy()[0])

x = tensor([3.0]) # 入力データ
y = model(x)
y # 損失関数

y.backward() # dy/dθ，それぞれの微分指定されたパラメータθ を計算する
optimizer.step() # パラメータを更新

print("After")
for name, param in model.named_parameters():
    print(f"{name:<10s}", param.detach().numpy()[0])

Before
t0.a       1.0
t0.b       2.0
d0.c       -1.0
d0.d       -2.0
t1.a       3.0
t1.b       4.0
m0.t.a     10.0
m0.t.b     20.0
m0.d.c     -1.0
m0.d.d     -2.0
t2.a       5.0
t2.b       6.0
d1.c       -1.0
d1.d       -2.0
After
t0.a       451.0
t0.b       2.015
d0.c       -1.0
d0.d       -2.0
t1.a       353.0
t1.b       3.995
m0.t.a     95.0
m0.t.b     19.9995
m0.d.c     -1.0
m0.d.d     -2.0
t2.a       153.0
t2.b       6.0001
d1.c       -1.0
d1.d       -2.0

結果を見ると，$a$が大きく更新されていて，$b$が小さく更新されていることが分かります．ちなみに，`Dummy`層にあるパラメータ$c,d$に関しては変化していないことが分かります．

Target層内にあるパラメータa，bに対して，別々の学習率を設定して学習を進めることができました．