機械学習のための数学的最適化: Adamから凸最適化、ZeRO optimizerまで

目次

1. 最適化の基礎: 凸最適化とKKT条件
2. 勾配降下法系のoptimizer
3. 2次最適化法
4. 学習率スケジューリング
5. 正則化手法
6. 損失関数の設計
7. LLMトレーニング最適化
8. クイズ

最適化の基礎

凸最適化 (Convex Optimization)

関数 $f: \mathbb{R}^n \to \mathbb{R}$ が凸(convex)であるとは、任意の2点 $x, y$ と $\lambda \in [0,1]$ に対して次が成立することです。

$f(\lambda x + (1-\lambda)y) \leq \lambda f(x) + (1-\lambda)f(y)$

凸関数の重要な性質:

• すべての局所最小値が大域最小値と一致する
• 勾配降下法の収束が保証される
• 深層学習の損失関数はほとんど非凸だが、凸解析の手法が依然として有用

強凸性(Strongly Convex): $m > 0$ が存在して $f(y) \geq f(x) + \nabla f(x)^T(y-x) + \frac{m}{2}\|y-x\|^2$ が成立すれば、収束が線形(線形収束)になります。

ラグランジュ乗数法

等式制約付き最適化問題を扱います。

$\min_x f(x) \quad \text{subject to} \quad g_i(x) = 0, \; i = 1, \ldots, m$

ラグランジアン:

$\mathcal{L}(x, \lambda) = f(x) + \sum_{i=1}^{m} \lambda_i g_i(x)$

最適解では $\nabla_x \mathcal{L} = 0$、$\nabla_\lambda \mathcal{L} = 0$ が成立します。

KKT条件

不等式制約を含む一般的な最適化問題:

$\min_x f(x) \quad \text{s.t.} \quad g_i(x) \leq 0, \; h_j(x) = 0$

KKT条件（必要条件）:

1. 停留性(Stationarity): $\nabla f(x^*) + \sum_i \mu_i \nabla g_i(x^*) + \sum_j \lambda_j \nabla h_j(x^*) = 0$
2. 主実行可能性(Primal feasibility): $g_i(x^*) \leq 0$、$h_j(x^*) = 0$
3. 双対実行可能性(Dual feasibility): $\mu_i \geq 0$
4. 相補スラック(Complementary slackness): $\mu_i g_i(x^*) = 0$

凸問題ではKKT条件は十分条件にもなります。

鞍点 (Saddle Point)

深層学習の最適化では、局所最小値より鞍点の方が大きな問題です。鞍点では勾配がゼロになりますが、ある方向では関数値が増加し、別の方向では減少します。SGDの確率的ノイズが鞍点からの脱出を助けます。

勾配降下法系

SGDとその変形

基本SGD:

$\theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t; x_i, y_i)$

SGD with Momentum:

$v_{t+1} = \beta v_t + \nabla_\theta \mathcal{L}(\theta_t)$ $\theta_{t+1} = \theta_t - \eta v_{t+1}$

モメンタム $\beta = 0.9$ が標準的で、過去の勾配方向を保持してoscillationを低減します。

Nesterov加速勾配(NAG):

$v_{t+1} = \beta v_t + \nabla_\theta \mathcal{L}(\theta_t - \beta v_t)$ $\theta_{t+1} = \theta_t - \eta v_{t+1}$

現在位置ではなく「先読み」位置で勾配を計算します。

AdaGrad、RMSProp、Adam

AdaGrad: パラメータごとの適応学習率

$G_t = G_{t-1} + g_t^2$ $\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{G_t + \epsilon}} g_t$

頻出する特徴は学習率が小さく、稀な特徴は大きい学習率になります。欠点: 学習率が単調減少して学習が停滞します。

RMSProp: AdaGradの蓄積問題を解決

$v_t = \beta v_{t-1} + (1-\beta) g_t^2$ $\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{v_t + \epsilon}} g_t$

Adam (Adaptive Moment Estimation):

$m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \quad \text{(1次モーメント)}$ $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad \text{(2次モーメント)}$

バイアス補正:

$\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1 - \beta_2^t}$

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

デフォルトハイパーパラメータ: $\beta_1 = 0.9$、$\beta_2 = 0.999$、$\epsilon = 10^{-8}$

import torch
import torch.optim as optim

model = ...  # モデル定義

# 標準Adam
optimizer_adam = optim.Adam(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8
)

# AdamW (weight decayを分離)
optimizer_adamw = optim.AdamW(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01  # 勾配スケーリングとは独立して適用
)

AdamWとLion

AdamW: weight decayをL2ペナルティとしてではなく、パラメータ更新に直接適用します。

$\theta_{t+1} = \theta_t - \eta \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t\right)$

AdamにL2正則化を追加する場合と数学的に同等ではありません（クイズを参照）。

Lion (EvoLved Sign Momentum):

$u_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$ $\theta_{t+1} = \theta_t - \eta \cdot \text{sign}(u_t)$ $m_t = \beta_2 m_{t-1} + (1-\beta_2) g_t$

符号(sign)のみを使用するため、更新量が均一でメモリ効率が良いです。

2次最適化法

Newton法

2次微分(Hessian)を活用します。

$\theta_{t+1} = \theta_t - H_t^{-1} \nabla f(\theta_t)$

ここで $H_t = \nabla^2 f(\theta_t)$ はHessian行列です。2次収束しますが、$n \times n$ 行列の逆行列計算が $O(n^3)$ となり深層学習では非現実的です。

L-BFGS (Limited-memory BFGS)

Hessianを直接保存せず、最近 $m$ 個の勾配差分で近似します。

$H_t^{-1} \approx \text{ベクトル列 } \{s_k\}, \{y_k\}\text{ による近似}$

ここで $s_k = \theta_{k+1} - \theta_k$、$y_k = \nabla f_{k+1} - \nabla f_k$ です。

import torch
import torch.optim as optim

# L-BFGSはクロージャ(closure)関数が必要
optimizer = optim.LBFGS(
    model.parameters(),
    lr=1.0,
    max_iter=20,
    history_size=10,
    line_search_fn='strong_wolfe'
)

def closure():
    optimizer.zero_grad()
    output = model(input_data)
    loss = criterion(output, target)
    loss.backward()
    return loss

optimizer.step(closure)

自然勾配降下法 (Natural Gradient Descent)

Fisher情報行列を使用してパラメータ空間の曲率を考慮します。

$\theta_{t+1} = \theta_t - \eta F(\theta_t)^{-1} \nabla \mathcal{L}(\theta_t)$

Fisher行列: $F(\theta) = \mathbb{E}\left[\nabla \log p(y|x;\theta) \nabla \log p(y|x;\theta)^T\right]$

K-FAC (Kronecker-factored Approximate Curvature) は自然勾配法を層ごとに分解して実用的に実装します。

学習率スケジューリング

Warmup

初期に学習率を徐々に増加させてトレーニングを安定化します。

$\eta_t = \eta_{\max} \cdot \frac{t}{T_{\text{warmup}}} \quad (t \leq T_{\text{warmup}})$

Cosine Annealing

$\eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\frac{\pi t}{T}\right)$

from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau

# Cosine Annealing
scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# OneCycleLR: Warmup + コサイン減衰
scheduler = OneCycleLR(
    optimizer,
    max_lr=1e-3,
    total_steps=1000,
    pct_start=0.3,         # 30% warmupフェーズ
    anneal_strategy='cos'
)

# ReduceLROnPlateau: 検証損失が改善しない場合に減少
scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.5,
    patience=10,
    min_lr=1e-6
)

Cyclical Learning Rate (CLR)

学習率を周期的に変動させて鞍点からの脱出を助けます。

$\eta_t = \eta_{\min} + (\eta_{\max} - \eta_{\min}) \cdot \max\left(0, 1 - \left|\frac{t}{\text{step\_size}} - 2k + 1\right|\right)$

正則化手法

L1 / L2 正則化

L2正則化 (Ridge):

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

勾配: $\nabla_\theta \mathcal{L}_{\text{reg}} = \nabla_\theta \mathcal{L} + \lambda \theta$

L1正則化 (Lasso):

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

L1はスパース(疎)な解を誘導して多くの重みを正確にゼロにします。

Batch NormalizationとLayer Normalization

Batch Normalization (BN):

$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}} \cdot \gamma + \beta$

ここで $\mu_B$、$\sigma_B^2$ はミニバッチ内の統計量です。バッチ方向に正規化します。

Layer Normalization (LN):

$\hat{x} = \frac{x - \mu_L}{\sqrt{\sigma_L^2 + \epsilon}} \cdot \gamma + \beta$

ここで統計量は各サンプルの特徴次元を使って計算されます。

Weight DecayとL2正則化

SGDでは:

$\theta_{t+1} = \theta_t - \eta(\nabla \mathcal{L} + \lambda \theta_t) = (1 - \eta\lambda)\theta_t - \eta \nabla \mathcal{L}$

この場合、weight decayとL2正則化は同一です。しかしAdamでは:

• L2 Adam: 勾配に $\lambda\theta$ を加えた後、適応スケーリング係数で除算 → 正則化効果が弱まる
• AdamW: パラメータ更新後に $\lambda\theta$ を直接引く → すべてのパラメータに均等なweight decay

import torch.nn as nn

class RegularizedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(512, 256)
        self.bn1 = nn.BatchNorm1d(256)
        self.ln1 = nn.LayerNorm(256)
        self.dropout = nn.Dropout(p=0.3)

    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)   # Transformer向けはself.ln1(x)
        x = torch.relu(x)
        x = self.dropout(x)
        return x

損失関数の設計

Cross-Entropy Loss

$\mathcal{L}_{CE} = -\sum_{c=1}^{C} y_c \log \hat{p}_c$

二値分類: $\mathcal{L}_{BCE} = -[y \log p + (1-y)\log(1-p)]$

Focal Loss

クラス不均衡問題を解決します。簡単なサンプルの寄与を減らします。

$\mathcal{L}_{FL} = -(1-p_t)^\gamma \log(p_t)$

ここで $p_t$ は正解クラスの予測確率、$\gamma \geq 0$ はfocusing parameterです。$\gamma = 0$ の場合は通常のCross-Entropyと同じです。

import torch
import torch.nn.functional as F

class FocalLoss(torch.nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, logits, targets):
        bce_loss = F.binary_cross_entropy_with_logits(
            logits, targets.float(), reduction='none'
        )
        p = torch.sigmoid(logits)
        p_t = p * targets + (1 - p) * (1 - targets)
        focal_weight = (1 - p_t) ** self.gamma
        alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
        loss = alpha_t * focal_weight * bce_loss
        return loss.mean()

Contrastive LossとTriplet Loss

Contrastive Loss (Siameseネットワーク):

$\mathcal{L} = (1-y)\frac{d^2}{2} + y \cdot \max(0, m - d)^2$

ここで $d = |f(x_1) - f(x_2)|_2$、$y=0$ が類似ペア、$y=1$ が非類似ペアです。

Triplet Loss:

$\mathcal{L}_{trip} = \max(0, \|f(a) - f(p)\|_2^2 - \|f(a) - f(n)\|_2^2 + m)$

アンカー(a)、ポジティブ(p)、ネガティブ(n)サンプルを使用します。

InfoNCE Loss (NT-Xent)

対照学習(Contrastive Learning)の核心損失関数です。

$\mathcal{L}_{InfoNCE} = -\log \frac{\exp(\text{sim}(z_i, z_j)/\tau)}{\sum_{k=1}^{2N} \mathbf{1}_{k \neq i} \exp(\text{sim}(z_i, z_k)/\tau)}$

ここで $\tau$ はtemperatureパラメータ、$\text{sim}$ はコサイン類似度です。

import torch
import torch.nn.functional as F

def info_nce_loss(features, temperature=0.07):
    """
    features: (2N, D) - 各画像の2つのaugmentationビュー
    """
    N = features.shape[0] // 2
    features = F.normalize(features, dim=1)

    # 類似度行列を計算
    similarity = torch.matmul(features, features.T) / temperature

    # 自己類似度を除去（対角を-infに）
    mask = torch.eye(2 * N, dtype=torch.bool, device=features.device)
    similarity.masked_fill_(mask, float('-inf'))

    # ポジティブペア: iとi+N、i+Nとi
    labels = torch.cat([
        torch.arange(N, 2 * N),
        torch.arange(N)
    ]).to(features.device)

    loss = F.cross_entropy(similarity, labels)
    return loss

LLMトレーニング最適化

勾配クリッピング (Gradient Clipping)

勾配爆発(exploding gradient)を防ぎます。

$g \leftarrow g \cdot \min\left(1, \frac{\text{clip\_norm}}{\|g\|_2}\right)$

import torch

def train_with_clipping(model, optimizer, loss, max_norm=1.0):
    optimizer.zero_grad()
    loss.backward()

    # クリッピング前の勾配ノルムを監視
    total_norm = 0
    for p in model.parameters():
        if p.grad is not None:
            param_norm = p.grad.data.norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm ** 0.5

    # クリッピング適用
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)
    optimizer.step()
    return total_norm

ZeRO Optimizer (Zero Redundancy Optimizer)

モデル学習時のメモリを3段階で最適化します。

混合精度(FP16/BF16) + ZeRO-3で数十億パラメータのモデルを単一ノードで学習可能です。

8-bit Adam

量子化を通じてoptimizer状態をFP32ではなくINT8で保存します。

• Optimizer状態のメモリをFP32比で75%削減
• ブロックごとの量子化で精度損失を最小化
• bitsandbytes ライブラリで実装可能

# bitsandbytes 8-bit Adam
import bitsandbytes as bnb

optimizer = bnb.optim.Adam8bit(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.999)
)

Adafactor

Adamの2次モーメントを行列分解で近似します。

$V_t \approx \hat{r}_t \hat{v}_t^T \quad \text{(ランク1近似)}$

パラメータサイズに比例したメモリのみ使用（行ベクトル + 列ベクトル）。T5、PaLMなどの超大規模モデルの学習に使用されます。

クイズ