1. 梯度下降算法

先考虑一元情形。假设待更新的参数为 $\theta$，学习率为 $\eta$，目标函数为 $f(\theta )$，则更新策略表示为：

$$\theta \leftarrow \theta -\eta \cdot g$$

其中 $g=f^{\prime }(\theta )$ 是梯度，我们需要沿着梯度的反方向进行更新。

学习率 $\eta$ 是我们通常所说的步长。步长过小会导致更新缓慢，步长过大会导致振荡。以 $f(\theta )={\theta }^2$ 为例，$\theta =0$ 是它的极值点。我们从 $\theta =-5$ 开始进行梯度下降，并设置不同的学习率来观察更新的效果。

import numpy as np
import matplotlib.pyplot as plt


def f(x):
    return x**2

def grad_f(x):
    return 2*x

def gradient_descent(start_x, learning_rate, steps):
    x_vals = [start_x]
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        new_x = x_vals[-1] - learning_rate * grad
        x_vals.append(new_x)
    return np.array(x_vals)


start_x = -5
steps = 10
learning_rates = [0.01, 0.05, 0.2, 0.7, 0.9, 1.01]
x_range = np.linspace(-5.5, 5.5, 400)

fig, axes = plt.subplots(2, 3, figsize=(10, 6.6))

for idx, eta in enumerate(learning_rates):
    ax = axes[idx // 3, idx % 3]
    x_vals = gradient_descent(start_x, eta, steps)
    
    ax.plot(x_range, f(x_range), color='blue')
    ax.plot(x_vals, f(x_vals), 'o-', color='orange')
    
    ax.set_title(f'Learning Rate: {eta}')
    ax.set_xlabel('x')
    ax.set_ylabel('f(x)')
    ax.legend()


plt.tight_layout()
plt.show()

从图中可以观察到，当学习率较小时，参数更新的速度会明显减缓；当学习率 $>0.5$ 时，参数更新开始出现振荡现象；而当学习率 $>1$ 时，更新过程甚至会导致参数发散。

现考虑多元情形（以二元为例）。设 $\theta =({\theta }_1,{\theta }_2)$，$f(\theta )=f({\theta }_1,{\theta }_2)$，则更新策略为：

$$\begin{array}{rl}& \theta \leftarrow \theta -\eta \cdot \nabla f(\theta )\\ \iff & ({\theta }_1,{\theta }_2)\leftarrow ({\theta }_1,{\theta }_2)-\eta \cdot \left(\frac{\partial f}{\partial {\theta }_1},\frac{\partial f}{\partial {\theta }_2}\right)\\ \iff & {\theta }_i\leftarrow {\theta }_i-\eta \cdot \frac{\partial f}{\partial {\theta }_i},i=1,2\end{array}$$

即每个参数分别更新。以 $f({\theta }_1,{\theta }_2)={\theta }_1^2+6{\theta }_2^2$ 为例，显然 $(0,0)$ 是它的极值点，且 $\nabla f(\theta )=(2{\theta }_1,12{\theta }_2)$。我们从 $(-10,3)$ 开始进行梯度下降，并设置不同的学习率来观察更新的效果。

import numpy as np
import matplotlib.pyplot as plt

def f(x1, x2):
    return x1**2 + 6*x2**2

def grad_f(x1, x2):
    return np.array([2*x1, 12*x2])

def gradient_descent(start_x, learning_rate, steps):
    x_vals = [start_x]
    for _ in range(steps):
        grad = grad_f(*x_vals[-1])
        new_x = x_vals[-1] - learning_rate * grad
        x_vals.append(new_x)
    return np.array(x_vals)

start_x = np.array([-10, 3])
steps = 20
learning_rates = [0.01, 0.05, 0.1, 0.15, 0.2, 0.5]
x1_range = np.linspace(-12, 2, 200)
x2_range = np.linspace(-6, 6, 200)
x1_grid, x2_grid = np.meshgrid(x1_range, x2_range)
f_vals = f(x1_grid, x2_grid)

fig, axes = plt.subplots(2, 3, figsize=(10, 6.6))

for idx, eta in enumerate(learning_rates):
    ax = axes[idx // 3, idx % 3]
    x_vals = gradient_descent(start_x, eta, steps)
    
    ax.contour(x1_grid, x2_grid, f_vals, levels=15, cmap='Blues')
    ax.plot(x_vals[:, 0], x_vals[:, 1], 'o-', color='orange')
    
    ax.set_title(f'Learning Rate: {eta}')
    ax.set_xlabel(r'$x_1$')
    ax.set_ylabel(r'$x_2$')
    ax.legend()

    ax.set_xlim([-12, 2])
    ax.set_ylim([-6, 6])

plt.tight_layout()
plt.show()

从图中可以看出，多元情形依然有着相同的规律。学习率较小时，参数更新十分缓慢；当学习率增大到一定值时，开始出现振荡现象；继续增大学习率，则会直接导致发散。

据此可以总结出梯度下降法的优缺点。

优点：

• 计算简单，易于实现。
• 在凸优化问题中，梯度下降法有严格的收敛性保证。

缺点：

• 梯度下降法无法保证找到全局最优解，特别是对于非凸函数时，可能会陷入局部极小值或鞍点，而难以跳出。
• 对学习率敏感。过小的学习率会导致收敛速度过慢，过大的学习率则会使得参数更新出现振荡，甚至导致发散。
• 对初始点的选择敏感，若初始点距离最优点较远，可能需要更多的迭代次数才能收敛，甚至可能收敛到较差的局部极小值。

2. BGD、SGD、MBGD

在深度学习场景中，$\theta$ 通常指代模型的参数，而 $f(\theta )$ 则指代模型在训练集上的平均损失，即

$$f(\theta )=\frac{1}{n}\sum _{i=1}^n{f}_i(\theta )$$

其中 $n$ 是训练集的大小，$f_i(\theta )$ 是模型在第 $i$ 个样本上的损失。于是

$$\nabla f(\theta )=\frac{1}{n}\sum _{i=1}^n\nabla f_i(\theta )$$

如果采用梯度下降法，那么计算 $\nabla f(\theta )$ 的时间复杂度将是 $O(n)$ 的。特别是当训练集特别大时，这种方法往往是不可行的。

以上方法又叫Batch Gradient Descent（BGD，批量梯度下降），即每次采用所有样本进行更新。那我们能不能不用那么多样本呢？能不能每次随机挑一个样本呢？

随机挑选时，每个样本被选中的概率为 $\frac{1}{n}$，因此

$${\mathbb{E}}_i[\nabla f_i(\theta )]=\sum _{i=1}^n\frac{1}{n}\cdot \nabla f_i(\theta )=\nabla f(\theta )$$

这说明 $\nabla f_i(\theta )$ 是 $\nabla f(\theta )$ 的无偏估计，因此我们可以用 $\nabla f_i(\theta )$ 代替 $\nabla f(\theta )$ 进行梯度下降，这样一来，计算 $\nabla f(\theta )$ 的时间复杂度就降低至 $O(1)$ 了。这种方法又叫Stochastic Gradient Descent（SGD，随机梯度下降），即每次随机挑选一个样本进行更新。

只挑选一个样本或挑选所有样本未免有些太过极端了，有没有一种折中的办法呢？比如每次挑选 $b(1\le b\le n)$ 个样本来进行个更新，即用下式去代替 $\nabla f(\theta )$：

1 b ∑ i ∈ B ∇ f i ( θ ) , B ⊂ { 1 , 2 , ⋯   , n } , ∣ B ∣ = b \frac1b \sum_{i\in \mathcal{B}} \nabla f_i(\theta),\quad \mathcal{B}\subset \{1,2,\cdots,n\},\quad|\mathcal{B}|=b b1​i∈B∑​∇fi​(θ),B⊂{1,2,⋯,n},∣B∣=b

该式依然是 $\nabla f(\theta )$ 的无偏估计，且由于 $\nabla f_i(\theta )$ 是独立同分布的（具有相同的期望和方差），因此

$${\text{Var}}_i\left[\frac{1}{b}\sum _{i\in \mathcal{B}}\nabla f_i(\theta )\right]=\frac{1}{b^2}\sum _{i\in \mathcal{B}}{\text{Var}}_i\left[\nabla f_i(\theta )\right]=\frac{1}{b}{\text{Var}}_t\left[\nabla f_t(\theta )\right]$$

📝 诸 $\nabla f_i(\theta )$ 的方差均为 $\frac{1}{n}{\sum }_{i=1}^n\Vert \nabla f_i(\theta ){\Vert }^2-\Vert \nabla f(\theta ){\Vert }^2$，证明如下：
$$\begin{array}{rl}{\text{Var}}_i[\nabla f_i(\theta )]& ={\mathbb{E}}_i[\Vert \nabla f_i(\theta )-\nabla f(\theta ){\Vert }^2]={\mathbb{E}}_i[(\nabla f_i(\theta )-\nabla f(\theta ){)}^{\text{T}}(\nabla f_i(\theta )-\nabla f(\theta ))]\\ & ={\mathbb{E}}_i[\Vert \nabla f_i(\theta ){\Vert }^2-2\nabla f_i(\theta {)}^{\text{T}}\nabla f(\theta )+\Vert \nabla f(\theta ){\Vert }^2]\\ & =\frac{1}{n}\sum _{i=1}^n\Vert \nabla f_i(\theta ){\Vert }^2-\Vert \nabla f(\theta ){\Vert }^2\end{array}$$

这说明选取的 $b$ 越大，相应的方差就越小，我们对 $\nabla f(\theta )$ 的估计就更稳定。

每次选取 $b$ 个样本来进行更新的方法又叫Mini-Batch Gradient Descent（MBGD，小批量梯度下降）。$b$ 通常被称为batch size。

据此可以对三种方法进行总结：

• BGD：批量梯度下降，每次采用所有样本来更新参数。
• SGD：随机梯度下降，每次随机选择一个样本来更新参数。
• MBGD：小批量梯度下降，每次随机选择 $b$ 个样本来更新参数。

BGD、SGD可以看做是MBGD的特殊情形，一个是 $b$ 取 $n$，一个是 $b$ 取 $1$。

---

接下来对比BGD、SGD和MBGD的效果。由于BGD、SGD是MBGD的特殊情形，我们只需要实现MBGD即可。

假设训练集有 $100$ 个样本，MBGD中的batch size设为 $10$，采用MSE作为损失函数，训练 $5$ 个epoch。三种方法的更新步数如下：

• SGD：$500/1=500$ 步。
• BGD：$500/100=5$ 步。
• MBGD：$500/10=50$ 步。

📝 在深度学习中，完整遍历一次训练集称为一个 epoch，而处理一个 batch 数据称为一个 step。设训练集包含 $n$ 个样本，训练的 epoch 数为 $e$，batch size 为 $b$，则训练过程中总的参数更新步数为：$\frac{ne}{b}$。

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

n_samples = 100  # 样本数量
n_features = 1   # 特征数量

# 真实参数
w_true = np.array([2])
b_true = 5

# 生成特征和标签，添加一些噪声
X = np.random.randn(n_samples, n_features)
noise = 0.5 * np.random.randn(n_samples, 1)
y = X @ w_true.reshape(-1, 1) + b_true + noise


def compute_loss(X, y, w, b):
    n = X.shape[0]
    y_pred = X @ w + b
    loss = (1 / (2 * n)) * np.sum((y_pred - y) ** 2)
    return loss


def compute_gradient(X, y, w, b):
    n = X.shape[0]
    y_pred = X @ w + b
    error = y_pred - y
    dw = (1 / n) * X.T @ error
    db = (1 / n) * np.sum(error)
    return dw, db


def mini_batch_gradient_descent(X, y, w_init, b_init, lr, epochs, batch_size):
    w = w_init.copy()
    b = b_init
    n = X.shape[0]
    w_history = [w.copy()]
    b_history = [b]
    for _ in range(epochs):
        indices = np.arange(n)
        np.random.shuffle(indices)
        X_shuffled = X[indices]
        y_shuffled = y[indices]

        for i in range(0, n, batch_size):
            Xi = X_shuffled[i:i+batch_size]
            yi = y_shuffled[i:i+batch_size]
            dw, db = compute_gradient(Xi, yi, w, b)
            w -= lr * dw
            b -= lr * db
            w_history.append(w.copy())
            b_history.append(b)
    return w, b, w_history, b_history


w_init = np.array([[0.0]])
b_init = 0.0
lr = 0.1
epochs = 5

batch_sizes = [1, 10, n_samples]
labels = ['SGD (b=1)', 'MBGD (b=10)', 'BGD (b=n)']
colors = ['r', 'g', 'b']

# 创建等高线图的数据
w_range = np.linspace(-1, 5, 50)
b_range = np.linspace(0, 10, 50)
W, B = np.meshgrid(w_range, b_range)
Loss = np.zeros_like(W)

# 计算每个网格点的损失值
for i in range(W.shape[0]):
    for j in range(W.shape[1]):
        w_ij = np.array([[W[i, j]]])
        b_ij = B[i, j]
        Loss[i, j] = compute_loss(X, y, w_ij, b_ij)


plt.figure(figsize=(8, 6))
CS = plt.contour(W, B, Loss, levels=30, cmap='viridis')
plt.clabel(CS, inline=1, fontsize=8)
plt.xlabel('w')
plt.ylabel('b')
plt.title('Parameter Descent Paths on Loss Contour')

for batch_size, label, color in zip(batch_sizes, labels, colors):
    w_final, b_final, w_history, b_history = mini_batch_gradient_descent(
        X, y, w_init, b_init, lr, epochs, batch_size)
    w_vals = np.array(w_history).flatten()
    b_vals = np.array(b_history)
    plt.plot(w_vals, b_vals, marker='o', color=color, label=label, markersize=3, linewidth=1)

plt.legend()
plt.show()

可以看出：

• SGD方法下降的最快，但是下降的过程中波动较大，且无法及时收敛到最优解。
• BGD方法下降的最慢，5个epoch后离最优解依然还有很远的一段距离，但是它的每一步都很稳定。
• MBGD方法的速度介于SGD和BGD之间，且及时收敛到了最优解，并且更新过程也比较稳定。

3. momentum与dampening

在第一节的讨论中，我们提到了梯度下降容易陷入局部最优或鞍点。一个典型的例子是 $y=x^3$，显然 $x=0$ 是它的鞍点。从 $x=3$ 处出发，使用 $0.05$ 的学习率进行梯度下降：

import numpy as np
import matplotlib.pyplot as plt

def f(x):
    return x**3

def grad_f(x):
    return 3*x**2

def gradient_descent(start_x, learning_rate, steps):
    x_vals = [start_x]
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        new_x = x_vals[-1] - learning_rate * grad
        x_vals.append(new_x)
    return np.array(x_vals)

start_x = 3
learning_rate = 0.05
steps = 20
x_vals = gradient_descent(start_x, learning_rate, steps)

x_range = np.linspace(-5, 5, 400)
y_range = f(x_range)

plt.figure(figsize=(8, 6))
plt.plot(x_range, y_range, color='blue')
plt.plot(x_vals, f(x_vals), 'o-', color='orange', label='lr=0.05')
plt.title('Gradient Descent on $f(x) = x^3$')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.legend()
plt.grid(True)
plt.show()

📝 在高维空间中，鞍点的存在往往更加普遍。

除非调大学习率，否则最后总会收敛到 $x=0$ 处。这一结果显然不是我们想要的。如果把梯度下降比作小球在曲线上的滚动，我们更希望小球能够越过 $x=0$ 这一点继续向下滚动，这也更加符合真实的物理场景。

我们已经知道梯度下降的公式为 $\theta \leftarrow \theta -\eta g$，由于 $\eta$ 通常是固定的，因此可以将其视为一小段时间 $\Delta t$，这样一来，$\theta$ 便可以视作位置，$g$ 便可以视作速度，此时公式变为 $x\leftarrow x-v\Delta t$。在梯度下降的场景中，$v$ 取决于 $x$，即 $v=v(x)=3x^2$，当 $x$ 减小时，$v$ 也随之减小，说明小球在下落的过程中速度会越来越慢，这显然是不符合事实的。一个直观的想法是我们可以累积小球的历史速度，然后使用这个累积速度作为当前的速度进行下落：

$$\begin{gathered} m_t=v_1+v_2+\cdots +v_t \\ {x}_t=x_{t-1}-m_t\Delta t \end{gathered}$$

其中 $m_t$ 就是 $t$ 时刻的累积速度。

直接累加会有一个很明显的缺点，那就是 $m_t$ 可以记住很久以前的速度。换句话说，历史任一时刻的速度对当前累积速度的贡献都是相同的，而我们更希望累积速度 $m_t$ 具有遗忘性，即距离当前时刻越远的速度，相应的贡献就应当越少。

利用指数函数的衰减性，我们可以重新计算累积速度：

$$\begin{array}{rl}{m}_t& ={\beta }^0{v}_t+{\beta }^1{v}_{t-1}+{\beta }^2{v}_{t-2}+\cdots +{\beta }^{t-1}{v}_1,\beta <1\\ & =v_t+\beta (v_{t-1}+\beta v_{t-2}+\cdots +{\beta }^{t-2}{v}_1)\\ & =\beta m_{t-1}+v_t\end{array}$$

于是我们得到了带动量的梯度下降公式：

$$\begin{gathered} m_t=\beta m_{t-1}+g_t \\ {\theta }_t={\theta }_{t-1}-\eta \cdot m_t \end{gathered}$$

其中 $m_t$ 是动量，$\beta$ 是动量系数。

设置 $\beta =0.9$，更新步数为 $5$，我们来观察小球的运动轨迹：

def gradient_descent(start_x, learning_rate, steps):
    x_vals = [start_x]
    m = 0
    beta = 0.9
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        m = beta * m + grad
        new_x = x_vals[-1] - learning_rate * m
        x_vals.append(new_x)
    return np.array(x_vals)

可以发现此时小球成功越过了鞍点。

接下来以 $y=x^2$ 为例，比较带动量和不带动量的梯度下降法的效果：

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

def f(x):
    return x**2

def grad_f(x):
    return 2*x

# 带动量的梯度下降算法
def gradient_descent_with_momentum(start_x, learning_rate, steps):
    x_vals = [start_x]
    m = 0
    beta = 0.9
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        m = beta * m + grad
        new_x = x_vals[-1] - learning_rate * m
        x_vals.append(new_x)
    return np.array(x_vals)

# 不带动量的梯度下降算法
def gradient_descent_without_momentum(start_x, learning_rate, steps):
    x_vals = [start_x]
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        new_x = x_vals[-1] - learning_rate * grad
        x_vals.append(new_x)
    return np.array(x_vals)


start_x = 3
learning_rate = 0.05
steps = 50

x_vals_with_momentum = gradient_descent_with_momentum(start_x, learning_rate, steps)
x_vals_without_momentum = gradient_descent_without_momentum(start_x, learning_rate, steps)

x_range = np.linspace(-5, 5, 400)
y_range = f(x_range)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

ax1.plot(x_range, y_range, color='blue')
point1, = ax1.plot([], [], 'o', color='orange', markersize=10, label='With Momentum')
ax1.set_title('Gradient Descent with Momentum on $f(x) = x^2$', fontsize=16)
ax1.set_xlabel('x', fontsize=14)
ax1.set_ylabel('f(x)', fontsize=14)
ax1.legend(fontsize=12)
ax1.grid(True)

ax2.plot(x_range, y_range, color='blue')
point2, = ax2.plot([], [], 'o', color='green', markersize=10, label='Without Momentum')
ax2.set_title('Gradient Descent without Momentum on $f(x) = x^2$', fontsize=16)
ax2.set_xlabel('x', fontsize=14)
ax2.set_ylabel('f(x)', fontsize=14)
ax2.legend(fontsize=12)
ax2.grid(True)


def init():
    point1.set_data([], [])
    point2.set_data([], [])
    return point1, point2


def update(frame):
    point1.set_data(x_vals_with_momentum[frame], f(x_vals_with_momentum[frame]))
    point2.set_data(x_vals_without_momentum[frame], f(x_vals_without_momentum[frame]))
    return point1, point2

ani = FuncAnimation(fig, update, frames=len(x_vals_with_momentum), init_func=init, blit=True, interval=100)
ani.save('gradient_descent_comparison.gif', writer=PillowWriter(fps=10))

plt.show()

显然左图更贴近现实世界的情况，但其收敛速度不如右图快。原因是动量法会累积历史梯度，从而导致过冲和振荡。既然我们可以控制历史累积梯度的贡献程度（即 $\beta$），我们同样希望能控制当前梯度的贡献程度，由此可以引入阻尼系数 $\tau$（它相当于摩擦力），从而动量的更新公式变为：

$$m_t=\beta m_{t-1}+(1-\tau )g_t$$

有些教程会将 $\tau$ 固定为 $\beta$，使得 $g_t$ 的贡献程度恒为 $1-\beta$，这在一定程度上限制了灵活性。

为了充分验证动量系数和阻尼系数对优化过程的影响，我们选取函数 $f(x)=(x+2{)}^2(x-1{)}^2+1.5x^2$ 进行实验。该函数具有两个极小值点，其中一个为局部极小值，另一个为全局极小值。实验从初始点 $x=-2.5$ 开始，学习率设定为 $0.01$，并通过以下六组设置进行比较：

1. 不使用动量和阻尼。
2. 动量系数设置为 $0.9$，阻尼系数为 $0$。
3. 动量系数设置为 $0.9$，阻尼系数为 $0.5$。
4. 动量系数设置为 $0.9$，阻尼系数为 $0.9$。
5. 动量系数设置为 $0.99$，阻尼系数为 $0$。
6. 动量系数设置为 $0.999$，阻尼系数为 $0$。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter

def f(x):
    return (x + 2)**2 * (x - 1)**2 + 1.5 * x**2

def grad_f(x):
    return 2 * (x + 2) * (x - 1)**2 + 2 * (x - 1) * (x + 2)**2 + 3 * x

def gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, tau=0.0):
    x_vals = [start_x]
    m = 0
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        m = beta * m + (1 - tau) * grad
        new_x = x_vals[-1] - learning_rate * m
        x_vals.append(new_x)
    return np.array(x_vals)

start_x = -2.5
learning_rate = 0.01
steps = 100

x_vals_no_momentum = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.0, tau=0.0)
x_vals_momentum_no_damping = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, tau=0.0)
x_vals_momentum_medium_damping = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, tau=0.5)
x_vals_momentum_high_damping = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, tau=0.9)
x_vals_momentum_0_99 = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.99, tau=0.0)
x_vals_momentum_0_999 = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.999, tau=0.0)

x_range = np.linspace(-3, 2, 400)
y_range = f(x_range)

fig, axes = plt.subplots(2, 3, figsize=(10, 6))

titles = [
    'No Momentum, No Damping', 
    'Momentum=0.9, No Damping',
    'Momentum=0.9, Damping=0.5',
    'Momentum=0.9, Damping=0.9',
    'Momentum=0.99, No Damping',
    'Momentum=0.999, No Damping'
]
colors = ['orange', 'green', 'blue', 'red', 'purple', 'brown']
x_vals_list = [
    x_vals_no_momentum,
    x_vals_momentum_no_damping,
    x_vals_momentum_medium_damping,
    x_vals_momentum_high_damping,
    x_vals_momentum_0_99,
    x_vals_momentum_0_999
]

plots = []
for i, ax in enumerate(axes.flat):
    ax.plot(x_range, y_range, color='blue')
    point, = ax.plot([], [], 'o', color=colors[i], markersize=8)
    ax.set_title(titles[i])
    ax.set_xlim([-3, 2])
    ax.set_ylim([0, 30])
    ax.set_xlabel('x')
    ax.set_ylabel('f(x)')
    ax.grid(True)
    plots.append(point)

def init():
    for plot in plots:
        plot.set_data([], [])
    return plots

def update(frame):
    for i in range(6):
        plots[i].set_data(x_vals_list[i][frame], f(x_vals_list[i][frame]))
    return plots

ani = FuncAnimation(fig, update, frames=len(x_vals_no_momentum), init_func=init, blit=True, interval=100)
ani.save('momentum_damping_experiments.gif', writer=PillowWriter(fps=10))

plt.show()

从图中可以得出以下结论：

• 在未使用动量法（即无动量和阻尼）的情况下，或当阻尼系数较大时，优化过程更容易陷入局部极小值。
• 动量系数过大时，对历史梯度的记忆会更加明显，可能导致优化过程中的振荡现象。

一般而言，建议设置 $\beta =0.9,\tau =0$，但更多时候还是要视具体情况而定。

⚠️ 学术界关于阻尼系数的研究目前还不是很充分，大多数情况下建议直接设置为 $0$，即不考虑阻尼。

回顾第一节中的二元梯度下降，当学习率设置为 $0.15$ 时，优化过程呈振荡式下降。设置合适的动量系数可以有效缓解振荡，并加速下降。

import numpy as np
import matplotlib.pyplot as plt

def f(x1, x2):
    return x1**2 + 6*x2**2

def grad_f(x1, x2):
    return np.array([2*x1, 12*x2])

def gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0):
    x_vals = [start_x]
    m = 0
    for _ in range(steps):
        grad = grad_f(*x_vals[-1])
        m = beta * m + grad
        new_x = x_vals[-1] - m * learning_rate
        x_vals.append(new_x)
    return np.array(x_vals)

start_x = np.array([-10, 3])
steps = 100
learning_rate = 0.15
beta_momentum = 0.3
x1_range = np.linspace(-12, 2, 200)
x2_range = np.linspace(-6, 6, 200)
x1_grid, x2_grid = np.meshgrid(x1_range, x2_range)
f_vals = f(x1_grid, x2_grid)

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

x_vals = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0)
ax = axes[0]
ax.contour(x1_grid, x2_grid, f_vals, levels=15, cmap='Blues')
ax.plot(x_vals[:, 0], x_vals[:, 1], 'o-', color='orange')
ax.set_title(f'Gradient Descent (LR: {learning_rate})')
ax.set_xlabel(r'$x_1$')
ax.set_ylabel(r'$x_2$')
ax.set_xlim([-12, 2])
ax.set_ylim([-6, 6])

x_vals_momentum = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=beta_momentum)
ax = axes[1]
ax.contour(x1_grid, x2_grid, f_vals, levels=15, cmap='Blues')
ax.plot(x_vals_momentum[:, 0], x_vals_momentum[:, 1], 'o-', color='orange')
ax.set_title(f'Momentum (LR: {learning_rate}, Beta: {beta_momentum})')
ax.set_xlabel(r'$x_1$')
ax.set_ylabel(r'$x_2$')
ax.set_xlim([-12, 2])
ax.set_ylim([-6, 6])

plt.tight_layout()
plt.show()

在不使用动量的情况下，每次参数更新的梯度可以分解为沿 $x_1$ 和 $x_2$ 方向的分量。其中，沿 $x_1$ 方向的梯度始终为正，因此参数在该方向上会持续减小；而沿 $x_2$ 方向的梯度则时而为正、时而为负，导致参数在该方向上的更新呈现出上下振荡的趋势。

动量法通过累积梯度，有效缓解了这种振荡现象。由于沿 $x_2$ 方向的梯度正负交替变化，动量累积过程中这些相反的梯度部分抵消，从而减小了沿 $x_2$ 方向的振荡幅度。与此同时，沿 $x_1$ 方向的梯度始终为正，因此动量会不断累积，导致参数沿 $x_1$ 方向下降得越来越快。

当然，在此场景中，若设置 $\beta =0.9$，参数更新就会出现过冲现象。

3.1 另一种形式的momentum

3.1.1 学习率固定

我们之前提到的带动量的更新公式为：

$$\begin{gathered} m_t=\beta m_{t-1}+g_t \\ {\theta }_t={\theta }_{t-1}-\eta \cdot m_t \end{gathered}$$

在不少教程中，还存在另外一种形式：

$$\begin{gathered} m_t=\beta m_{t-1}-\eta \cdot g_t \\ {\theta }_t={\theta }_{t-1}+m_t \end{gathered}$$

先来推导第一个形式。不难得出

$$m_t=\sum _{i=1}^t{\beta }^{t-i}{g}_i,{\theta }_t={\theta }_0-\eta \sum _{i=1}^t{m}_i$$

合并两式即可得到

$${\theta }_t={\theta }_0-\eta \sum _{i=1}^t\sum _{j=1}^i{\beta }^{i-j}{g}_j$$

再来推导第二个形式。不难得出

$$m_t=-\eta \sum _{i=1}^t{\beta }^{t-i}{g}_i,{\theta }_t={\theta }_0+\sum _{i=1}^t{m}_i$$

合并两式即可得到

$${\theta }_t={\theta }_0-\eta \sum _{i=1}^t\sum _{j=1}^i{\beta }^{i-j}{g}_j$$

由此可知，学习率固定的情况下，这两种形式是完全等价的。

3.1.2 学习率不固定

此时每一步的学习率都不一样，第一种形式的更新公式为：

$$\begin{gathered} m_t=\beta m_{t-1}+g_t \\ {\theta }_t={\theta }_{t-1}-{\eta }_t\cdot m_t \end{gathered}$$

不难得出

$$m_t=\sum _{i=1}^t{\beta }^{t-i}{g}_i,{\theta }_t={\theta }_0-\sum _{i=1}^t{\eta }_i{m}_i$$

合并两式即可得到

$${\theta }_t={\theta }_0-\sum _{i=1}^t{\eta }_i\sum _{j=1}^i{\beta }^{i-j}{g}_j$$

第二种形式的更新公式为：

$$\begin{gathered} m_t=\beta m_{t-1}-{\eta }_t\cdot g_t \\ {\theta }_t={\theta }_{t-1}+m_t \end{gathered}$$

不难得出

$$m_t=-\sum _{i=1}^t{\beta }^{t-i}{\eta }_i{g}_i,{\theta }_t={\theta }_0+\sum _{i=1}^t{m}_i$$

合并两式即可得到

$${\theta }_t={\theta }_0-\sum _{i=1}^t\sum _{j=1}^i{\beta }^{i-j}{\eta }_j{g}_j$$

可以看出学习率不固定的情况下，两者并不等价。

4. nesterov

Nesterov Accelerated Gradient (NAG) 是一种加速梯度下降的优化算法，最早由苏联数学家 Yurii Nesterov 提出。它是对标准的带动量梯度下降方法的改进，能够在收敛性和效率上表现得更好。与标准动量方法不同，NAG 在计算梯度时并不是根据当前的参数位置，而是根据向前预测的参数位置。这个改进让NAG可以在每一步中更具前瞻性，避免在接近局部最优点或鞍点时的惯性过冲现象。

NAG的主要思想是：利用当前的动量预测出未来参数的位置，然后在该预测位置计算梯度，再利用这个梯度来调整动量和更新参数。

这样解释或许有些抽象，为方便接下来的讲解，我们假定 $\beta =\eta =1$，从而更新公式变为：

$$\begin{gathered} m_t=m_{t-1}+g_t \\ {\theta }_t={\theta }_{t-1}-m_t \end{gathered}$$

在高维空间中，上述公式其实就是向量的加减法。下图展示了从 ${\theta }_{t-1}$ 到 ${\theta }_t$ 的更新过程：

⚠️ 图中和 $g,m$ 相关的均在前面加上了负号，这是因为它们原本表示的是上升的方向，加上负号才能变成下降的方向。

模型初始参数为 ${\theta }_0$，动量初始为 $m_0=0$（因为还没开始累积）。

在第 $1$ 时刻，我们会在 ${\theta }_0$ 处计算梯度，得到 $g_1=\nabla f({\theta }_0)$，然后我们会用当前时刻的梯度 $g_1$ 和上一时刻的动量 $m_0$（历史累积梯度）合成当前时刻的动量：$m_1=m_0+g_1$。当前时刻的动量代表了 $\theta$ 将要更新的方向。

以此类推，在第 $t$ 时刻，我们会在 ${\theta }_{t-1}$ 处计算梯度，得到 $g_t=\nabla f({\theta }_{t-1})$。然后用当前时刻的梯度 $g_t$ 和上一时刻的动量 $m_{t-1}$（历史累积梯度）合成当前时刻的动量 $m_t$。

仔细观察，$m_t$ 决定了第 $t$ 时刻的更新方向，即当前时刻的动量决定了当前时刻的参数更新，并且参数更新均是使用动量来完成的。既然 $m_t$ 的计算依赖于 $m_{t-1}$ 和 $g_t$，那在计算之前，我们能不能用上一时刻的动量来提前预判当前时刻的参数更新呢？即用 $m_{t-1}$ 去预判 ${\theta }_t$ 的位置。

我们在第 $t$ 时刻不去计算 $m_t$，而是先用 $m_{t-1}$ 去预判 ${\theta }_t$ 的位置，假设预判的结果为 ${\tilde{\theta }}_t$，那么有 ${\tilde{\theta }}_t={\theta }_{t-1}-m_{t-1}$，然后我们在预判的位置上重新计算梯度 ${\tilde{g}}_t=\nabla f({\tilde{\theta }}_t)$，于是可以得到第 $t$ 时刻的更新结果：

$${\theta }_t={\tilde{\theta }}_t-{\tilde{g}}_t={\theta }_{t-1}-m_{t-1}-\nabla f({\theta }_{t-1}-m_{t-1})={\theta }_{t-1}-(m_{t-1}+\nabla f({\theta }_{t-1}-m_{t-1}))$$

完整的更新公式如下：

$$\begin{gathered} m_t=m_{t-1}+\nabla f({\theta }_{t-1}-m_{t-1}) \\ {\theta }_t={\theta }_{t-1}-m_t \end{gathered}$$

Nesterov相比原始的动量法，无非是计算梯度的位置发生了变化，原先是在 ${\theta }_{t-1}$ 处计算梯度，而现在是在 ${\tilde{\theta }}_t={\theta }_{t-1}-m_{t-1}$ 处计算梯度，更具有前瞻性。

Nesterov原论文中的更新公式如下：

$$\begin{gathered} m_t=\beta m_{t-1}-\eta \nabla f({\theta }_{t-1}+\beta m_{t-1}) \\ {\theta }_t={\theta }_{t-1}+m_t \end{gathered}$$

这对应了我们之前所说的第二种形式的momentum。那第一种形式对应的Nesterov公式又是什么样的呢？首先可以知道

$$\begin{array}{rl}{\theta }_t& ={\theta }_{t-1}-\eta m_t\\ & ={\theta }_{t-1}-\eta (\beta m_{t-1}+g_t)\\ & =({\theta }_{t-1}-\eta \beta m_{t-1})-\eta g_{t-1}\end{array}$$

于是可以推测 $g_{t-1}=\nabla f({\theta }_{t-1}-\eta \beta m_{t-1})$，从而该形式下的Nesterov公式为：

$$\begin{gathered} m_t=\beta m_{t-1}+\nabla f({\theta }_{t-1}-\eta \beta m_{t-1}) \\ {\theta }_t={\theta }_{t-1}-\eta m_t\text{\hspace{1em}(∗)} \end{gathered}$$

4.1 PyTorch中的Nesterov

PyTorch中的Nesterov公式为

$$\begin{gathered} m_t=\beta m_{t-1}+\nabla f({\theta }_{t-1}) \\ {\theta }_t={\theta }_{t-1}-\eta (\beta m_t+\nabla f({\theta }_{t-1})) \end{gathered}$$

可以发现它并没有使用偏移后的参数进行梯度计算，这是为什么呢？

在 $(\ast )$ 式中令 ${\tilde{\theta }}_{t-1}={\theta }_{t-1}-\eta \beta m_{t-1}$，然后利用 $(\ast )$ 式中的结果去推导 ${\tilde{\theta }}_t$，

$$\begin{array}{rl}{\tilde{\theta }}_t& ={\theta }_t-\eta \beta m_t\\ & ={\theta }_{t-1}-\eta m_t-\eta \beta m_t\\ & ={\theta }_{t-1}-\eta (\beta m_{t-1}+\nabla f({\tilde{\theta }}_{t-1}))-\eta \beta m_t\\ & =({\theta }_{t-1}-\eta \beta m_{t-1})-\eta \nabla f({\tilde{\theta }}_{t-1})-\eta \beta m_t\\ & ={\tilde{\theta }}_{t-1}-\eta (\beta m_t+\nabla f({\tilde{\theta }}_{t-1}))\end{array}$$

由此可以得知，PyTorch在实际进行计算的时候，是把参数 ${\theta }_t$ 当作是已经偏移过后的参数 ${\tilde{\theta }}_t$，所以其实一直更新的是 ${\tilde{\theta }}_t$。按理来说，在更新结束后，我们应当修正回原来的参数：

$${\theta }_t={\tilde{\theta }}_t+\eta \beta m_t$$

但PyTorch的源码中并未做这一修正，所以PyTorch关于Nesterov的实现实际上是一种近似。

4.2 Polyak与Nesterov的比较

📝 Polyak动量法就是传统的动量法，也叫Classical Momentum（CM）。

我们使用PyTorch中的更新方法，初始时先进行偏移，即 ${\tilde{\theta }}_0={\theta }_0-\eta \beta m_0={\theta }_0$，之后的每次更新按照如下流程计算：

$$\begin{gathered} m_t=\beta m_{t-1}+\nabla f({\tilde{\theta }}_{t-1}) \\ \nabla f({\tilde{\theta }}_{t-1})=\nabla f({\tilde{\theta }}_{t-1})+\beta m_t \\ {\theta }_t={\theta }_{t-1}-\eta \nabla f({\tilde{\theta }}_{t-1}) \end{gathered}$$

计算完成后修正参数即可。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter


def f(x):
    return (x + 2)**2 * (x - 1)**2 + 1.5 * x**2

def grad_f(x):
    return 2 * (x + 2) * (x - 1)**2 + 2 * (x - 1) * (x + 2)**2 + 3 * x


def gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, nesterov=False):
    x_vals = [start_x]
    m = 0
    for _ in range(steps):
        grad = grad_f(x_vals[-1])
        m = beta * m + grad
        if nesterov:
            grad += beta * m
        else:
            grad = m
        new_x = x_vals[-1] - learning_rate * grad
        x_vals.append(new_x)
    if nesterov:
        x_vals[-1] += beta * learning_rate * m
    return np.array(x_vals)


start_x = -3
learning_rate = 0.01
steps = 100


x_vals_momentum = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, nesterov=False)
x_vals_nesterov = gradient_descent_with_momentum(start_x, learning_rate, steps, beta=0.9, nesterov=True)


x_range = np.linspace(-3.5, 2, 400)
y_range = f(x_range)

fig, axes = plt.subplots(1, 2, figsize=(12, 6))

titles = ['Momentum=0.9, No Nesterov', 'Momentum=0.9, Nesterov']
x_vals_list = [x_vals_momentum, x_vals_nesterov]
colors = ['green', 'red']

plots = []
for i, ax in enumerate(axes):
    ax.plot(x_range, y_range, color='blue')
    point, = ax.plot([], [], 'o', color=colors[i], markersize=8)
    ax.set_title(titles[i])
    ax.set_xlim([-3, 2])
    ax.set_ylim([0, 30])
    ax.set_xlabel('x')
    ax.set_ylabel('f(x)')
    ax.grid(True)
    plots.append(point)

def init():
    for plot in plots:
        plot.set_data([], [])
    return plots

def update(frame):
    for i in range(2):
        plots[i].set_data(x_vals_list[i][frame], f(x_vals_list[i][frame]))
    return plots

ani = FuncAnimation(fig, update, frames=len(x_vals_momentum), init_func=init, blit=True, interval=100)
ani.save('momentum_vs_nesterov.gif', writer=PillowWriter(fps=10))

plt.show()

可以看到Nesterov明显收敛的更快。

---

Ref

[1] https://pytorch.org/docs/stable/generated/torch.optim.SGD.html
[2] https://d2l.ai/chapter_optimization/momentum.html
[3] https://www.bilibili.com/video/BV1jh4y1q7ua
[4] https://www.bilibili.com/video/BV1YF411n7Dr
[5] https://optimi.benjaminwarner.dev/optimizers/sgd/
[6] https://towardsdatascience.com/is-pytorchs-nesterov-momentum-implementation-wrong-5dc5f5032008
[7] https://www.bilibili.com/video/BV1kL411E7P5