APGD梯度下降算法实战：解决深度学习模型训练中的收敛难题

1次阅读

共计 2252 个字符，预计需要花费 6 分钟才能阅读完成。

在训练深度学习模型时，传统梯度下降（如 SGD）经常遇到两个典型问题：

非光滑目标函数的优化困难：当损失函数包含 L1 正则化等不可导部分时，标准梯度下降无法直接应用
高维稀疏数据的低效收敛：在推荐系统等场景下，参数稀疏性会导致优化路径出现 ” 之字形 ” 震荡

以推荐系统的矩阵分解为例，其目标函数通常形如：
$$\min_W \frac{1}{2}||X-WH^T||_F^2 + \lambda||W||_1$$
其中 L1 项使得目标函数在零点不可导。这时就需要能够处理非光滑函数的优化方法。

算法	更新公式	特点
SGD	$w_{t+1} = w_t – \eta \nabla f(w_t)$	基础版本，震荡严重
Momentum	$v_{t+1} = \gamma v_t + \eta \nabla f(w_t)$	加入惯性缓解震荡
Adam	自适应矩估计	适合稀疏梯度
APGD	$y_t = w_t + \beta_t(w_t – w_{t-1})$	Nesterov 加速 + 近端算子

APGD 的核心创新在于将 Nesterov 加速与近端梯度结合：
$$w_{t+1} = \text{prox}_{\eta h}(y_t – \eta \nabla g(y_t))$$
其中 $g$ 是光滑部分，$h$ 是非光滑部分（如 L1 正则）。

import torch

class APGDOptimizer(torch.optim.Optimizer):
    def __init__(self, params, lr=0.01, beta=0.9, lambda_l1=0.01):
        defaults = dict(lr=lr, beta=beta, lambda_l1=lambda_l1)
        super().__init__(params, defaults)

    def step(self):
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None: continue

                state = self.state[p]
                if len(state) == 0:  # 初始化状态
                    state['momentum'] = torch.zeros_like(p.data)
                    state['previous'] = p.data.clone()

                # Nesterov 加速
                y_t = p.data + group['beta'] * state['momentum']

                # 计算光滑部分梯度
                p.data = y_t.clone()  # 临时替换参数
                g = torch.autograd.grad(self._loss(p), p)[0]

                # 近端映射 (L1 正则的闭式解)
                with torch.no_grad():
                    update = y_t - group['lr'] * g
                    p.data = torch.sign(update) * torch.clamp(torch.abs(update) - group['lr']*group['lambda_l1'], 
                        min=0)

                    # 更新动量项
                    state['momentum'] = p.data - state['previous']
                    state['previous'] = p.data.clone()

    def _loss(self, params):
        """假设损失函数已绑定到优化器实例"""
        return self.loss_fn(params)

关键实现细节：

近端算子处理 L1 正则时，采用软阈值函数：
$$\text{prox}_{\lambda|\cdot|}(v) = \text{sign}(v)\max(|v|-\lambda, 0)$$
Nesterov 加速通过在计算梯度时使用 ” 前瞻点 ”y_t 实现
梯度裁剪可通过在计算 g 后添加 g = torch.clamp(g, -clip, clip) 实现

使用 LeNet- 5 在 MNIST 上的对比实验（batch_size=128）：