SGD・Adam・AdamW・Lionの理論とPython実装

オプティマイザは深層学習モデルの学習効率と最終性能を大きく左右します。SGDから始まり、Adamが広く普及し、近年はAdamWやLionなど新しいオプティマイザが登場しています。

本記事では、主要なオプティマイザの数学的定式化から実装・比較まで詳しく解説します。

本記事の内容

  • 各オプティマイザの数学的定式化と更新則
  • Weight Decayとの関係
  • Pythonでの実装と比較実験

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

オプティマイザの役割

基本的な最適化問題

ニューラルネットワークの学習は、損失関数 $\mathcal{L}(\theta)$ を最小化するパラメータ $\theta$ を見つける問題です:

$$ \theta^* = \arg\min_\theta \mathcal{L}(\theta) $$

勾配降下法の基本更新則:

$$ \theta_{t+1} = \theta_t – \eta \nabla_\theta \mathcal{L}(\theta_t) $$

ここで $\eta$ は学習率です。

オプティマイザの発展

オプティマイザ 特徴
SGD 基本的な勾配降下法
2011 AdaGrad 適応的学習率
2012 RMSprop AdaGradの改良
2014 Adam モメンタム + 適応的学習率
2017 AdamW Adamの重み減衰修正
2023 Lion シンプルで高速

SGD with Momentum

数学的定式化

モメンタム付きSGDの更新則:

$$ \begin{align} \bm{v}_t &= \beta \bm{v}_{t-1} + \nabla_\theta \mathcal{L}(\theta_t) \\ \theta_{t+1} &= \theta_t – \eta \bm{v}_t \end{align} $$

$\beta$ はモメンタム係数(通常0.9)、$\bm{v}_t$ は速度(勾配の指数移動平均)です。

直感的な理解

モメンタムは「ボールが坂を転がる」ようなイメージです。過去の勾配方向を「慣性」として保持し、ノイズに対して頑健で、サドル点からの脱出を助けます。

Adam(Adaptive Moment Estimation)

数学的定式化

Adamは1次モーメント(平均)と2次モーメント(分散)の両方を推定します。

$$ \begin{align} \bm{m}_t &= \beta_1 \bm{m}_{t-1} + (1 – \beta_1) \bm{g}_t \\ \bm{v}_t &= \beta_2 \bm{v}_{t-1} + (1 – \beta_2) \bm{g}_t^2 \end{align} $$

ここで $\bm{g}_t = \nabla_\theta \mathcal{L}(\theta_t)$ は時刻 $t$ の勾配です。

バイアス補正

$\bm{m}_t, \bm{v}_t$ は初期値がゼロのため、学習初期にバイアスがあります。これを補正します:

$$ \begin{align} \hat{\bm{m}}_t &= \frac{\bm{m}_t}{1 – \beta_1^t} \\ \hat{\bm{v}}_t &= \frac{\bm{v}_t}{1 – \beta_2^t} \end{align} $$

バイアス補正の導出

$\bm{m}_t$ の期待値を計算します。$\bm{m}_0 = 0$ として:

$$ \bm{m}_t = (1 – \beta_1) \sum_{i=1}^{t} \beta_1^{t-i} \bm{g}_i $$

勾配が定常($\mathbb{E}[\bm{g}_i] = \bm{g}$ で一定)と仮定すると:

$$ \mathbb{E}[\bm{m}_t] = (1 – \beta_1) \bm{g} \sum_{i=1}^{t} \beta_1^{t-i} = \bm{g} (1 – \beta_1^t) $$

したがって、バイアスを補正するには $1 – \beta_1^t$ で割ります。

パラメータ更新

$$ \theta_{t+1} = \theta_t – \eta \frac{\hat{\bm{m}}_t}{\sqrt{\hat{\bm{v}}_t} + \epsilon} $$

$\epsilon$ は数値安定性のための小さな定数(通常 $10^{-8}$)です。

デフォルトパラメータ

パラメータ 推奨値 意味
$\eta$ 0.001 学習率
$\beta_1$ 0.9 1次モーメントの減衰率
$\beta_2$ 0.999 2次モーメントの減衰率
$\epsilon$ $10^{-8}$ 数値安定性

AdamW(Adam with Decoupled Weight Decay)

L2正則化とWeight Decayの違い

従来のL2正則化:

$$ \mathcal{L}_{\text{reg}}(\theta) = \mathcal{L}(\theta) + \frac{\lambda}{2} \|\theta\|^2 $$

この勾配は:

$$ \nabla \mathcal{L}_{\text{reg}} = \nabla \mathcal{L} + \lambda \theta $$

SGDではこれをそのまま使えますが、Adamでは適応的学習率により、Weight Decayの効果が減衰してしまいます。

AdamWの更新則

AdamWは、Weight Decayを勾配更新から分離します:

$$ \begin{align} \bm{m}_t &= \beta_1 \bm{m}_{t-1} + (1 – \beta_1) \bm{g}_t \\ \bm{v}_t &= \beta_2 \bm{v}_{t-1} + (1 – \beta_2) \bm{g}_t^2 \\ \hat{\bm{m}}_t &= \frac{\bm{m}_t}{1 – \beta_1^t} \\ \hat{\bm{v}}_t &= \frac{\bm{v}_t}{1 – \beta_2^t} \\ \theta_{t+1} &= \theta_t – \eta \left( \frac{\hat{\bm{m}}_t}{\sqrt{\hat{\bm{v}}_t} + \epsilon} + \lambda \theta_t \right) \end{align} $$

最後の行で、Weight Decay $\lambda \theta_t$ が適応的学習率の影響を受けずに適用されます。

なぜAdamWが効くのか

Adamでは、大きな勾配を持つパラメータは $\sqrt{\hat{\bm{v}}_t}$ も大きくなり、更新が小さくなります。L2正則化の項 $\lambda \theta$ も同様にスケールダウンされてしまいます。

AdamWでは、Weight Decayが独立して適用されるため、正則化の効果が一貫して働きます。

Lion(Evolved Sign Momentum)

数学的定式化

Lion(2023年、Googleが発表)は、AutoMLで発見されたシンプルなオプティマイザです。

$$ \begin{align} \bm{u}_t &= \text{sign}(\beta_1 \bm{m}_{t-1} + (1 – \beta_1) \bm{g}_t) \\ \theta_{t+1} &= \theta_t – \eta (\bm{u}_t + \lambda \theta_t) \\ \bm{m}_t &= \beta_2 \bm{m}_{t-1} + (1 – \beta_2) \bm{g}_t \end{align} $$

Lionの特徴

  1. sign関数: 更新量は常に $\pm \eta$(大きさが一定)
  2. 2つのモメンタム: 更新用 $(\beta_1)$ と状態用 $(\beta_2)$
  3. メモリ効率: 2次モーメント不要でメモリ半減
  4. Weight Decay内蔵: $\lambda$ が更新式に含まれる

デフォルトパラメータ

パラメータ 推奨値 意味
$\eta$ 0.0001 (Adamの1/10) 学習率
$\beta_1$ 0.9 更新のモメンタム
$\beta_2$ 0.99 状態のモメンタム
$\lambda$ 0.01 Weight Decay

Lionは学習率をAdamより小さく設定するのが一般的です。

Pythonでの実装

各オプティマイザのスクラッチ実装

import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

class SGDMomentum:
    """モメンタム付きSGD"""

    def __init__(self, params, lr=0.01, momentum=0.9):
        self.params = list(params)
        self.lr = lr
        self.momentum = momentum
        self.velocities = [torch.zeros_like(p) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                self.velocities[i] = self.momentum * self.velocities[i] + p.grad
                p -= self.lr * self.velocities[i]

class AdamOptimizer:
    """Adam"""

    def __init__(self, params, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        self.params = list(params)
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.t = 0
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        self.t += 1
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                g = p.grad

                # 1次・2次モーメントの更新
                self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * g
                self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * g**2

                # バイアス補正
                m_hat = self.m[i] / (1 - self.beta1**self.t)
                v_hat = self.v[i] / (1 - self.beta2**self.t)

                # パラメータ更新
                p -= self.lr * m_hat / (torch.sqrt(v_hat) + self.eps)

class AdamWOptimizer:
    """AdamW"""

    def __init__(self, params, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.01):
        self.params = list(params)
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.weight_decay = weight_decay
        self.t = 0
        self.m = [torch.zeros_like(p) for p in self.params]
        self.v = [torch.zeros_like(p) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        self.t += 1
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                g = p.grad

                # 1次・2次モーメントの更新
                self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * g
                self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * g**2

                # バイアス補正
                m_hat = self.m[i] / (1 - self.beta1**self.t)
                v_hat = self.v[i] / (1 - self.beta2**self.t)

                # パラメータ更新(Weight Decayを分離)
                p -= self.lr * (m_hat / (torch.sqrt(v_hat) + self.eps) + self.weight_decay * p)

class LionOptimizer:
    """Lion"""

    def __init__(self, params, lr=0.0001, beta1=0.9, beta2=0.99, weight_decay=0.01):
        self.params = list(params)
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.weight_decay = weight_decay
        self.m = [torch.zeros_like(p) for p in self.params]

    def zero_grad(self):
        for p in self.params:
            if p.grad is not None:
                p.grad.zero_()

    def step(self):
        with torch.no_grad():
            for i, p in enumerate(self.params):
                if p.grad is None:
                    continue
                g = p.grad

                # 更新方向の計算(sign関数)
                update = torch.sign(self.beta1 * self.m[i] + (1 - self.beta1) * g)

                # パラメータ更新
                p -= self.lr * (update + self.weight_decay * p)

                # モメンタムの更新
                self.m[i] = self.beta2 * self.m[i] + (1 - self.beta2) * g

# Rosenbrock関数での可視化
def rosenbrock(x, y):
    """Rosenbrock関数(最適化のベンチマーク)"""
    return (1 - x)**2 + 100 * (y - x**2)**2

def optimize_rosenbrock(optimizer_class, lr, n_steps=500, **kwargs):
    """Rosenbrock関数を最適化"""
    # 初期点
    x = torch.tensor([-1.5], requires_grad=True)
    y = torch.tensor([1.5], requires_grad=True)

    optimizer = optimizer_class([x, y], lr=lr, **kwargs)

    trajectory = [(x.item(), y.item())]
    losses = []

    for _ in range(n_steps):
        optimizer.zero_grad()
        loss = rosenbrock(x, y)
        loss.backward()
        optimizer.step()

        trajectory.append((x.item(), y.item()))
        losses.append(loss.item())

    return trajectory, losses

# 各オプティマイザで最適化
trajectories = {}
losses = {}

trajectories['SGD'], losses['SGD'] = optimize_rosenbrock(SGDMomentum, lr=0.0001)
trajectories['Adam'], losses['Adam'] = optimize_rosenbrock(AdamOptimizer, lr=0.01)
trajectories['AdamW'], losses['AdamW'] = optimize_rosenbrock(AdamWOptimizer, lr=0.01)
trajectories['Lion'], losses['Lion'] = optimize_rosenbrock(LionOptimizer, lr=0.001)

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# 軌跡のプロット
ax1 = axes[0]
x_range = np.linspace(-2, 2, 100)
y_range = np.linspace(-1, 3, 100)
X, Y = np.meshgrid(x_range, y_range)
Z = rosenbrock(X, Y)

ax1.contour(X, Y, Z, levels=np.logspace(0, 3, 20), cmap='viridis', alpha=0.5)
ax1.contourf(X, Y, np.log(Z + 1), levels=50, cmap='viridis', alpha=0.3)

colors = {'SGD': 'blue', 'Adam': 'red', 'AdamW': 'green', 'Lion': 'orange'}
for name, traj in trajectories.items():
    traj = np.array(traj)
    ax1.plot(traj[:, 0], traj[:, 1], '-', label=name, color=colors[name], linewidth=1.5)
    ax1.scatter(traj[0, 0], traj[0, 1], color=colors[name], s=50, marker='o')
    ax1.scatter(traj[-1, 0], traj[-1, 1], color=colors[name], s=100, marker='*')

ax1.scatter([1], [1], color='black', s=200, marker='x', label='Optimum')
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_title('Optimization Trajectories on Rosenbrock Function')
ax1.legend()
ax1.set_xlim(-2, 2)
ax1.set_ylim(-1, 3)

# 損失曲線
ax2 = axes[1]
for name, loss in losses.items():
    ax2.plot(loss, label=name, color=colors[name], linewidth=1.5)
ax2.set_xlabel('Step')
ax2.set_ylabel('Loss')
ax2.set_title('Loss Curves')
ax2.set_yscale('log')
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

ニューラルネットワークでの比較実験

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib.pyplot as plt

class MLP(nn.Module):
    def __init__(self, input_dim=20, hidden_dims=[128, 64], num_classes=10):
        super().__init__()
        layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
                nn.BatchNorm1d(hidden_dim)
            ])
            prev_dim = hidden_dim
        layers.append(nn.Linear(prev_dim, num_classes))
        self.net = nn.Sequential(*layers)

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

def create_dataset(n_samples=5000, input_dim=20, num_classes=10):
    np.random.seed(42)

    X = np.random.randn(n_samples, input_dim).astype(np.float32)
    W = np.random.randn(input_dim, num_classes).astype(np.float32)
    logits = X @ W + np.random.randn(n_samples, num_classes).astype(np.float32) * 0.3
    y = np.argmax(logits, axis=1)

    split = int(0.8 * n_samples)
    return (X[:split], y[:split]), (X[split:], y[split:])

def train_with_optimizer(optimizer_name, train_loader, test_loader,
                         input_dim, num_classes, n_epochs=50):
    """指定されたオプティマイザで学習"""
    torch.manual_seed(42)
    model = MLP(input_dim=input_dim, num_classes=num_classes)
    criterion = nn.CrossEntropyLoss()

    if optimizer_name == 'SGD':
        optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    elif optimizer_name == 'Adam':
        optimizer = optim.Adam(model.parameters(), lr=0.001)
    elif optimizer_name == 'AdamW':
        optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)
    elif optimizer_name == 'Lion':
        # PyTorchにはLionが標準でないため、手動実装版を使用
        optimizer = LionOptimizer(model.parameters(), lr=0.0001, weight_decay=0.01)
    else:
        raise ValueError(f"Unknown optimizer: {optimizer_name}")

    train_losses = []
    test_accuracies = []

    for epoch in range(n_epochs):
        # Training
        model.train()
        epoch_loss = 0
        for X_batch, y_batch in train_loader:
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = criterion(outputs, y_batch)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        train_losses.append(epoch_loss / len(train_loader))

        # Evaluation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for X_batch, y_batch in test_loader:
                outputs = model(X_batch)
                _, predicted = torch.max(outputs, 1)
                total += y_batch.size(0)
                correct += (predicted == y_batch).sum().item()
        test_accuracies.append(correct / total)

    return train_losses, test_accuracies

# データ準備
(X_train, y_train), (X_test, y_test) = create_dataset()

train_dataset = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))
test_dataset = TensorDataset(torch.tensor(X_test), torch.tensor(y_test))

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 各オプティマイザで学習
optimizers = ['SGD', 'Adam', 'AdamW', 'Lion']
results = {}

for opt_name in optimizers:
    print(f"Training with {opt_name}...")
    train_losses, test_accs = train_with_optimizer(
        opt_name, train_loader, test_loader,
        input_dim=20, num_classes=10, n_epochs=50
    )
    results[opt_name] = {'train_loss': train_losses, 'test_acc': test_accs}
    print(f"  Final Test Accuracy: {test_accs[-1]:.4f}")

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

colors = {'SGD': 'blue', 'Adam': 'red', 'AdamW': 'green', 'Lion': 'orange'}

ax1 = axes[0]
for opt_name in optimizers:
    ax1.plot(results[opt_name]['train_loss'], label=opt_name, color=colors[opt_name], linewidth=2)
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Training Loss')
ax1.set_title('Training Loss Comparison')
ax1.legend()
ax1.grid(True, alpha=0.3)

ax2 = axes[1]
for opt_name in optimizers:
    ax2.plot(results[opt_name]['test_acc'], label=opt_name, color=colors[opt_name], linewidth=2)
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Test Accuracy')
ax2.set_title('Test Accuracy Comparison')
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

オプティマイザの選び方

一般的な推奨

タスク 推奨オプティマイザ
画像分類(CNN) SGD with Momentum
自然言語処理(Transformer) AdamW
大規模モデル AdamW or Lion
ファインチューニング AdamW
GAN Adam

使い分けの指針

  1. 安定性重視: SGD with Momentum(収束は遅いが安定)
  2. 速度重視: Adam(収束が速い)
  3. 正則化が重要: AdamW(Weight Decayが正しく機能)
  4. メモリ制約: Lion(メモリ使用量が少ない)
  5. 大規模学習: AdamW or Lion(実績あり)

ハイパーパラメータの調整

オプティマイザ 学習率の目安 調整ポイント
SGD 0.01-0.1 学習率スケジュールが重要
Adam 0.0001-0.001 $\beta_2$ を大きくすると安定
AdamW 0.0001-0.001 Weight Decay(0.01-0.1)
Lion Adamの1/10 学習率を小さめに

まとめ

本記事では、主要なオプティマイザ(SGD, Adam, AdamW, Lion)について解説しました。

  • SGDはシンプルだが、モメンタムで加速可能
  • Adamは適応的学習率を持ち、広く使われる
  • AdamWはWeight Decayを分離し、正則化が正しく機能
  • Lionはsign関数を使い、メモリ効率が高い

次のステップとして、以下の記事も参考にしてください。