Stochastic Gradient Descent (SGD)

確率的勾配降下法（SGD）は、データセット全体ではなく、学習データの小さなランダムサンプルに対する損失の勾配を使って重みを繰り返し更新し、モデルを学習させる最適化アルゴリズムです。各サンプルの勾配はノイズを含みますが、データセット全体の勾配に対して不偏な推定値です。そのため、SGDは、全データを使う勾配降下法が1回の高コストな更新を行う間に、多数の低コストな更新を実行できます。この手法はRobbinsとMonroが1951年に提唱した確率近似法に由来し、モメンタムと組み合わせたものは畳み込みニューラルネットワーク（CNN）の標準的なオプティマイザーとして今も使われています。

SGDの仕組み#

SGDの各ステップでは、ミニバッチBを取り出し、バックプロパゲーションでそのデータに対する損失の勾配を計算し、各重みを勾配と逆方向へ少しだけ移動させます。w ← w − η·∇L_B(w)です。ηは学習率です。厳密には、SGDは1ステップにつき1つのサンプルを使いますが、ディープラーニングのフレームワークでは数十から数千サンプルのミニバッチに適用するため、実際には「SGD」と「ミニバッチ勾配降下法」は同義です。バッチサイズを大きくすると勾配のノイズが減り、ハードウェアの利用効率も高まりますが、エポックあたりの更新回数は少なくなります。

学習では各エポックの開始時にデータをシャッフルし、ミニバッチ単位で処理して、各バッチで1回更新します。ノイズを含む1,000個の点に直線を当てはめる例で、手順全体を確認できます。

import torch

torch.manual_seed(0)

# 1,000 noisy samples of y = 3x + 2
x = torch.randn(1000, 1)
y = 3 * x + 2 + 0.1 * torch.randn(1000, 1)

model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)

for epoch in range(5):
    # Reshuffle every epoch, then update the weights once per 32-sample mini-batch
    for idx in torch.randperm(1000).split(32):
        loss = torch.nn.functional.mse_loss(model(x[idx]), y[idx])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch}: w = {model.weight.item():.3f}, b = {model.bias.item():.3f}, loss = {loss.item():.4f}")

合計64回の更新を経て、2エポック後には重みが3.081、バイアスが2.012になります。その後の3エポックでは、各ミニバッチの勾配にサンプリングノイズが含まれるため、重みは収束せず2.967から3.003の間を動き続けます。最後のバッチの損失は約0.01で推移します。これはyに加えられたノイズの分散です。

学習率、ノイズ、バッチサイズ#

学習率を一定に保つと、SGDは最小値の周辺で変動します。その領域の大きさは、学習率と勾配ノイズが大きくなるほど広がります。確率近似理論で収束が保証されるのは学習率がゼロに近づく場合に限られるため、実際の学習では短いウォームアップの後、ステップ減衰やコサインスケジュールを使って時間とともに学習率を下げます。小さなバッチサイズによる勾配ノイズは、SGDを大規模バッチ学習で見つかりやすい鋭い最小値よりも汎化性能の高い、平坦な最小値へ導く働きもします。これは正則化の暗黙的な形態です。

バッチサイズと学習率は相互に作用します。学習率をバッチサイズに比例して引き上げ、5エポックのウォームアップを行うことで、ResNet-50は8,192枚の画像からなるミニバッチを使用してImageNetでの学習を1時間で完了し、精度を損なわずに済みました。

モメンタムとNesterovモメンタム#

損失曲面に細長い谷がある場合、勾配は主に谷を横切る方向を向くため、通常のSGDは谷底に沿ってゆっくり進みながら、両側の壁の間をジグザグに動きます。モメンタムは速度の累積値 v ← μ·v + g を保持し、それに沿って w ← w − η·v と更新します。そのため、一貫した勾配方向では速度が増し、振動する方向は相殺されます。係数 μ = 0.9 が一般的なデフォルト値です。Nesterovモメンタムは、速度に沿って先読みした位置で勾配を評価するため、行き過ぎをさらに抑えます。PyTorchのSGDは両方を実装しており、勾配に λ·w を加算して weight_decay を適用します。これはSGDの場合、重みに対するL2ペナルティと等価です。

SGDと他のオプティマイザーの比較#

• バッチ勾配降下法: 更新のたびにデータセット全体の正確な勾配を計算するため、各エポックでの更新は1回だけで、必要なメモリはデータセットのサイズに応じて増加します。
• Adam: 勾配と勾配の二乗の移動平均から、各パラメーターに固有のステップサイズを設定します。また、モメンタム付きSGDがパラメーターごとに状態値を1つ保持するのに対し、Adamは2つ保持します。Adamは学習初期の進みが速く、学習率への感度も低い一方、適切に調整したモメンタム付きSGDは、CNNで最終的な精度が同等またはそれ以上になることがよくあります。重み減衰を分離した派生手法のAdamWは、Transformerで標準的なオプティマイザーです。
• Muon: 各重み行列のモメンタム更新を行列全体として直交化し、大規模言語モデルの学習にスケールアップされています。

Ultralytics YOLOのトレーニングでは、optimizer="SGD"を設定すると、Nesterovモメンタム付きのPyTorch SGDを構築し、デフォルトでlr0=0.01とmomentum=0.937を使用します。デフォルトのoptimizer="auto"はlr0を無視し、10,000反復を超える実行にはSGDとMuonの更新を融合したハイブリッド手法MuSGDを選び、それより短い実行にはAdamWを選択します。

SGDの利用例#

モメンタム付きSGDは、代表的なImageNet CNNの大半の学習に使われました。初期のResNetモデルでは、モメンタム0.9、重み減衰0.0001、画像256枚のミニバッチが使用され、誤差が停滞するたびに0.1から始まる学習率を10分の1に下げていました。ディープラーニング以外では、scikit-learnのSGDClassifierとSGDRegressorが、ロジスティック回帰や線形サポートベクターマシンなどの線形モデルをSGDで学習します。また、これらの partial_fit メソッドは、メモリに収まらない大規模なデータセットを一度に1チャンクずつ学習します。