필사 모드: 机器学习的数学优化：从 Adam 到凸优化、ZeRO optimizer

目录

1. 优化基础：凸优化与 KKT 条件
2. 梯度下降系列 optimizer
3. 二阶优化方法
4. 学习率调度
5. 正则化技巧
6. 损失函数设计
7. LLM 训练优化
8. 测验

优化基础

凸优化 (Convex Optimization)

函数 $f: \mathbb{R}^n \to \mathbb{R}$ 若为凸(convex)函数，则对任意两点 $x, y$ 及 $\lambda \in [0,1]$，以下不等式成立。

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

凸函数的核心性质：

• 任何局部最小值(local minimum)都是全局最小值(global minimum)
• 梯度下降法的收敛得到保证
• 深度学习的损失函数大多是非凸(non-convex)的，但凸分析技巧依然有用

强凸(Strongly Convex)：若存在 $m > 0$，使得 $f(y) \geq f(x) + \nabla f(x)^T(y-x) + \frac{m}{2}\|y-x\|^2$ 成立，则收敛速度会加快为线性收敛(linear convergence)。

拉格朗日乘子法

用于处理等式约束的优化问题。

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

拉格朗日函数(Lagrangian)：

$\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 (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 加速梯度 (Nesterov Accelerated Gradient, 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{(一阶矩)}$ $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad \text{(二阶矩)}$

偏差校正(Bias correction)：

$\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
optimizer_adam = optim.Adam(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8
)

# AdamW optimizer (分离 weight decay)
optimizer_adamw = optim.AdamW(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01  # 独立于 L2 应用
)

AdamW 与 Lion

AdamW：在 Adam 的基础上，不把 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$

Lion 只使用符号(sign)，因此内存效率高，更新幅度也保持均匀。

二阶优化

Newton 方法

利用二阶导数(Hessian)。

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

这里 $H_t = \nabla^2 f(\theta_t)$ 是 Hessian 矩阵。虽然具有二次收敛(quadratic convergence)，但 $n \times n$ Hessian 矩阵求逆的计算量为 $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)

利用 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 + Cosine Decay
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 会诱导出稀疏(sparse)解。

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$ 是小批量(mini-batch)内的统计量。沿批量方向做归一化。

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

# Dropout
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)   # 或使用 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 Network)：

$\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)$

使用锚点(anchor, a)、正样本(positive, p)、负样本(negative, 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 parameter，$\text{sim}$ 是余弦相似度。

import torch
import torch.nn.functional as F

def info_nce_loss(features, temperature=0.07):
    """
    features: (2N, D) - 每张图像的两个 augmentation view
    """
    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)

模型训练时分三个阶段优化内存。

混合精度(mixed precision) + ZeRO-3 可以让数十亿参数量级的模型在单节点上完成训练。

8-bit Adam

通过量化(quantization)，把 optimizer 状态以 INT8 而非 FP32 存储。

• optimizer 状态内存节省 75%(相较于 FP32)
• 通过 block-wise quantization 将精度损失降到最低
• 可通过 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 的二阶矩。

$V_t \approx \hat{r}_t \hat{v}_t^T \quad (\text{rank-1 近似})$

只使用与参数量成比例的内存(行向量 + 列向量)。被用于 T5、PaLM 等超大规模模型的训练。

测验