Adamで勾配を伝えていないはずなのにパラメータ更新がされる場合

はじめに

Adamを使ったニューラルネットワークの訓練をしていた際に，ぶつかった問題とその原因についてせつめいする．

事例

ざっくり書くと下記のような二又構造のモデルを使っていた．場面によって2つの経路のどちらかのモデルを選択して使うといった感じ．このとき，片方のモデルを使わずに誤差逆伝播をしてパラメータ更新を行ったにも関わらず，使っていない方のパタメータも更新されているという事象に遭遇した．

Class Model(nn.Module):
    def __init__(self):
        self.fc = Linear()
        self.fc1 = Linear()
        self.fc2 = Linear()
    def forward(self, x, right=False):
        h = self.fc(x)
        if right:
            return self.fc1(h)
        else:
            return self.fc2(h)

原因・解決方法

Adam内に含まれているMomentumが原因だった．Adamの仕様上，直近で行ったパラメータ更新時の勾配が保存されているため，たとえ勾配を逆伝播させておらず勾配が0であっても，残っているMomentumの項のみを使ってパラメータの更新が行われてしまう．そのため，使っていないモデルのパラメータも変化する．SGDなどのMomentum項を持たないoptimizerを使うと，上記の問題が発生しない．どうしてもAdamを使いたい場合はモデルの構造をいじる必要がある．

簡単な再現例

import torch
import torch.nn as nn
import torch.optim as optim


class Model(nn.Module):
	def __init__(self):
		super().__init__()
		self.fc = nn.Linear(1, 1)

	def forward(self, x):
		return self.fc(x)


model = Model()
optimizer = optim.Adam(model.parameters())

x = torch.tensor([2.])
y = torch.tensor([0.])

p = model(x)
loss = nn.functional.mse_loss(p, y)
print(model.fc.weight) # tensor([[0.5239]], requires_grad=True)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(model.fc.weight) # tensor([[0.5229]], requires_grad=True)

optimizer.zero_grad()
optimizer.step()
print(model.fc.weight) # tensor([[0.5222]], requires_grad=True)

実際に出力される値は異なるかもしれないが，この例でもbackwardしていないのにも関わらず，パラメータの値が変化する．

おわりに

細かい仕様を意識して書かないと面倒くさいことになりますね．