s08 优化器：从SGD到Adam — exercise.py 练习指南

练习目标

亲手实现 Momentum、RMSProp、NAG（Nesterov 加速梯度）三种优化器的更新规则，理解每种优化器"从 SGD 出发，一步步解决了什么问题"的设计演进脉络。

预备知识

建议先阅读 index.md 并运行 demo.py，确保理解：

概念	数学定义
SGD	$θ_{t + 1} = θ_{t} - α g_{t}$
指数滑动平均（EMA）	$m_{t} = β m_{t - 1} + (1 - β) g_{t}$
Momentum	$m_{t}$ 定方向， $θ_{t + 1} = θ_{t} - α m_{t}$
RMSProp	$v_{t}$ 缩步长， $θ_{t + 1} = θ_{t} - α g_{t} / \sqrt{v_{t} + ϵ}$
狭长峡谷	损失地形某些方向陡峭、某些平缓，条件数 $κ ≫ 1$

任务清单

任务1：实现 Momentum 更新规则

描述：补全 MomentumOptimizerExercise.step() 方法的三个 TODO——初始化速度向量、更新速度、更新参数。

公式回顾：

m_{t} = β \cdot m_{t - 1} + (1 - β) \cdot g_{t}

θ_{t + 1} = θ_{t} - α \cdot m_{t}

提示：

首次调用时 self.m is None，需要初始化为 np.zeros_like(theta)
self.m = self.beta * self.m + (1 - self.beta) * grad 更新速度
return theta - self.lr * self.m 更新参数
注意 (1 - self.beta) 不能省略——如果不乘，就变成了对梯度做纯滑动平均而没有衰减

期望行为：

从 $(3.0, 2.5)$ 出发，经过 5 步后动量向量应指向原点附近
每一步的动量 $m_{t}$ 是历史梯度的加权平均，不会像 SGD 那样每步方向剧烈变化

为什么 Momentum 比 SGD 好？ 在狭长峡谷中，SGD 在陡峭方向上来回震荡（一步正向一步反向），而 Momentum 通过平滑历史方向，抵消了这种震荡——正负梯度互相抵消，让 $m_{t}$ 在震荡方向上接近 0。

任务2：实现 RMSProp 自适应步长

描述：补全 RMSPropOptimizerExercise.step() 方法的三个 TODO——初始化 $v$ 、更新梯度平方的 EMA、自适应步长更新。

公式回顾：

v_{t} = β \cdot v_{t - 1} + (1 - β) \cdot g_{t} ⊙ g_{t}

θ_{t + 1} = θ_{t} - α \cdot \frac{g_{t}}{\sqrt{v_{t}} + ϵ}

提示：

grad ** 2 是逐元素平方，对应公式中的 $g_{t} ⊙ g_{t}$
np.sqrt(self.v) 计算 $\sqrt{v_{t}}$
self.eps（通常 $10^{- 8}$ ）加到分母上防止除零
有效学习率 $= α / (\sqrt{v_{t}} + ϵ)$ —— 可以打印出来观察

期望行为：

陡峭方向（ $θ_{1}$ ，系数 $a = 20$ ）的有效学习率应该小于平缓方向（ $θ_{2}$ ，系数 $b = 1$ ）
这是因为 $θ_{1}$ 方向梯度大 → $v_{t}$ 增长快 → 分母大 → 有效步长自动变小
最终输出应显示 θ₁ < θ₂: True

为什么 RMSProp 解决了步长不统一问题？ SGD 对所有参数用同一个学习率，但不同参数梯度的尺度可能差几个数量级。RMSProp 让每个参数拥有自己的"自适应学习率"——梯度历史大的参数步长自动变小，历史小的步长自动放大。

任务3：实现 Nesterov 加速梯度 (NAG)

描述：补全 NAGOptimizer.step() 方法的五个 TODO。NAG 是 Momentum 的进阶版——它"先沿动量方向看一步，再在那个位置计算梯度"。

与普通 Momentum 的区别：

Momentum：在当前点 $θ_{t}$ 计算梯度，然后沿 $m_{t}$ 方向更新
NAG：先沿 $m_{t - 1}$ 方向走一步到"前瞻位置"，在那个位置计算梯度，再更新

公式：

前瞻位置：

θ_{lookahead} = θ_{t} - α \cdot β \cdot m_{t - 1}

在"前瞻"位置计算梯度并更新动量：

m_{t} = β \cdot m_{t - 1} + (1 - β) \cdot \nabla L (θ_{lookahead})

参数更新：

θ_{t + 1} = θ_{t} - α \cdot m_{t}

提示：

注意：NAG 的 step() 接收的是 grad_fn（一个函数），而不是 grad（一个值）！因为需要在 lookahead 位置重新计算梯度
theta_lookahead = theta - self.lr * self.beta * self.m
grad_lookahead = grad_fn(theta_lookahead) —— 在"前瞻"位置调梯度函数
动量更新与 Momentum 相同，只是用的是 grad_lookahead 而非当前位置的梯度

直觉理解：NAG 像是一个"会思考的滚球"。普通 Momentum 是"盲人下山"——沿着惯性方向走，走到哪算哪。NAG 是"睁眼下山"——先往前看一眼，发现前面是悬崖（梯度变大），就可以提前调整方向。在数学上这对应更紧的收敛界。

期望行为：从 $(3.0, 2.5)$ 出发，经过 10 步后，NAG 的损失通常小于普通 Momentum——因为它更"聪明"地选择了更新方向。

优化器演进总览

优化器	改进点	要记住的关键行
SGD	基线	`theta - lr * grad`
Momentum	加惯性，平滑方向	`m = betam + (1-beta)grad`
RMSProp	自适应步长	`v = betav + (1-beta)(grad**2)`
Adam	方向 + 步长 + 修正	`m` + `v` + 偏差修正
NAG	"会思考的"Momentum	先在 lookahead 处算梯度

完整代码

# -*- coding: utf-8 -*-
"""
s08 优化器：从 SGD 到 Adam — 练习代码
======================================
请完成以下 TODO 任务，加深对优化器内部机制的理解。

每个 TODO 都有详细的中文指示和预期输出描述。
建议先阅读 README.md 并运行 demo.py，再尝试独立补全代码。
"""

import numpy as np
from typing import Tuple, List


# ============================================================================
# 辅助：损失地形（与 demo.py 一致）
# ============================================================================

class LossLandscape:
    """二维二次型损失函数 L(θ₁, θ₂) = 0.5·(a·θ₁² + b·θ₂²)"""

    def __init__(self, a: float = 20.0, b: float = 1.0):
        self.a = a
        self.b = b

    def __call__(self, theta: np.ndarray) -> float:
        return 0.5 * (self.a * theta[0] ** 2 + self.b * theta[1] ** 2)

    def gradient(self, theta: np.ndarray) -> np.ndarray:
        return np.array([self.a * theta[0], self.b * theta[1]])

    @property
    def optimum(self) -> np.ndarray:
        return np.array([0.0, 0.0])


# ============================================================================
# TODO 1: 实现 Momentum 更新规则
# ============================================================================

class MomentumOptimizerExercise:
    """
    Momentum 优化器（练习版本）。

    公式:
        m_t = β · m_{t-1} + (1-β) · g_t       ← 速度更新（指数滑动平均）
        θ_{t+1} = θ_t - α · m_t                ← 参数更新

    你需要实现 step() 方法。
    """

    def __init__(self, lr: float = 0.02, beta: float = 0.9):
        """
        初始化 Momentum 优化器。

        参数:
            lr: 学习率 α
            beta: 动量衰减系数 β
        """
        self.lr = lr
        self.beta = beta
        self.m = None  # 速度向量，首次 step 时初始化为零
        self.name = "Momentum"

    def step(self, theta: np.ndarray, grad: np.ndarray) -> np.ndarray:
        """
        执行一步 Momentum 更新。

        提示:
          1. 如果 self.m 是 None，初始化为 np.zeros_like(theta)
          2. 更新速度: self.m = β * self.m + (1 - β) * grad
          3. 更新参数: theta = theta - α * self.m
          4. 返回更新后的 theta

        参数:
            theta: 当前参数向量
            grad: 当前梯度向量

        返回:
            更新后的参数向量
        """
        # TODO: 如果 self.m 为 None，初始化为零向量
        if self.m is None:
            pass  # ← TODO: self.m = np.zeros_like(theta)

        # TODO: 更新速度 m_t
        pass  # ← TODO: self.m = self.beta * self.m + (1 - self.beta) * grad

        # TODO: 更新参数
        pass  # ← TODO: return theta - self.lr * self.m

        return theta  # 占位返回


# ---- 测试 TODO 1 ----
def test_momentum():
    """测试 Momentum 优化器的实现"""
    print("=" * 60)
    print("TODO 1 测试: Momentum 更新规则")
    print("=" * 60)

    landscape = LossLandscape(a=20.0, b=1.0)
    theta0 = np.array([3.0, 2.5])
    opt = MomentumOptimizerExercise(lr=0.02, beta=0.9)

    # 手动运行几步
    theta = theta0.copy()
    print(f"\n  初始: θ = {theta}, L = {landscape(theta):.4f}")

    for step in range(5):
        grad = landscape.gradient(theta)
        theta_new = opt.step(theta, grad)

        if theta_new is None or np.allclose(theta_new, theta):
            print("  TODO 未完成，请补全 MomentumOptimizerExercise.step()")
            return

        theta = theta_new
        if step < 5:
            print(f"  Step {step+1}: θ = [{theta[0]:.3f}, {theta[1]:.3f}], "
                  f"L = {landscape(theta):.4f}, m = [{opt.m[0]:.3f}, {opt.m[1]:.3f}]")

    # 验证动量向量的方向大致正确（应该指向原点附近）
    if opt.m is not None:
        m_direction = opt.m / (np.linalg.norm(opt.m) + 1e-10)
        target_direction = -theta0 / (np.linalg.norm(theta0) + 1e-10)
        similarity = np.dot(m_direction, target_direction)
        print(f"\n  动量方向与目标方向的相似度: {similarity:.3f} (>0 表示大致正确)")

    print()


# ============================================================================
# TODO 2: 实现 RMSProp 更新规则（v_t 计算）
# ============================================================================

class RMSPropOptimizerExercise:
    """
    RMSProp 优化器（练习版本）。

    公式:
        v_t = β · v_{t-1} + (1-β) · g_t²       ← 梯度平方的指数滑动平均
        θ_{t+1} = θ_t - α · g_t / (√v_t + ε)    ← 自适应步长更新

    你需要实现 step() 方法。
    特别注意 v_t 的计算——g_t 要先逐元素平方。
    """

    def __init__(self, lr: float = 0.05, beta: float = 0.9, eps: float = 1e-8):
        """
        初始化 RMSProp 优化器。

        参数:
            lr: 学习率 α
            beta: 衰减系数 β
            eps: 数值稳定常数 ε
        """
        self.lr = lr
        self.beta = beta
        self.eps = eps
        self.v = None
        self.name = "RMSProp"

    def step(self, theta: np.ndarray, grad: np.ndarray) -> np.ndarray:
        """
        执行一步 RMSProp 更新。

        提示:
          1. 如果 self.v 是 None，初始化为 np.zeros_like(theta)
          2. 更新 v_t: self.v = β * self.v + (1 - β) * (grad ** 2)
             注意：grad ** 2 是逐元素平方
          3. 更新参数: theta = theta - α * grad / (√v_t + ε)
             注意：np.sqrt(self.v) 也是逐元素操作
          4. 返回更新后的 theta

        参数:
            theta: 当前参数向量
            grad: 当前梯度向量

        返回:
            更新后的参数向量
        """
        # TODO: 如果 self.v 为 None，初始化为零向量
        if self.v is None:
            pass  # ← TODO: self.v = np.zeros_like(theta)

        # TODO: 更新梯度平方的滑动平均 v_t
        pass  # ← TODO: self.v = self.beta * self.v + (1 - self.beta) * (grad ** 2)

        # TODO: 自适应步长更新
        pass  # ← TODO: return theta - self.lr * grad / (np.sqrt(self.v) + self.eps)

        return theta  # 占位返回


# ---- 测试 TODO 2 ----
def test_rmsprop():
    """测试 RMSProp 优化器的实现"""
    print("=" * 60)
    print("TODO 2 测试: RMSProp 自适应步长")
    print("=" * 60)

    landscape = LossLandscape(a=20.0, b=1.0)
    theta0 = np.array([3.0, 2.5])
    opt = RMSPropOptimizerExercise(lr=0.05, beta=0.9)

    theta = theta0.copy()
    print(f"\n  初始: θ = {theta}, L = {landscape(theta):.4f}")

    for step in range(10):
        grad = landscape.gradient(theta)
        theta_new = opt.step(theta, grad)

        if theta_new is None or np.allclose(theta_new, theta):
            print("  TODO 未完成，请补全 RMSPropOptimizerExercise.step()")
            return

        theta = theta_new
        if step < 5 or step == 9:
            effective_lr = opt.lr / (np.sqrt(opt.v) + opt.eps)
            print(f"  Step {step+1:2d}: θ = [{theta[0]:.3f}, {theta[1]:.3f}], "
                  f"L = {landscape(theta):.4f}")
            print(f"         有效学习率: [{effective_lr[0]:.4f}, {effective_lr[1]:.4f}]")

    # 验证：陡峭方向(θ₁)的有效学习率应该小于平缓方向(θ₂)
    if opt.v is not None:
        effective_lr = opt.lr / (np.sqrt(opt.v) + opt.eps)
        print(f"\n  最终有效学习率: θ₁={effective_lr[0]:.5f}, θ₂={effective_lr[1]:.5f}")
        print(f"  θ₁ < θ₂ 吗？{effective_lr[0] < effective_lr[1]} (RMSProp 应该压小陡峭方向的步长)")

    print()


# ============================================================================
# TODO 3: 实现 Nesterov 加速梯度 (NAG)
# ============================================================================

class NAGOptimizer:
    """
    Nesterov 加速梯度 (Nesterov Accelerated Gradient, NAG)。

    NAG 和普通 Momentum 的关键区别：
      - Momentum: 在当前点计算梯度，然后沿动量方向更新
      - NAG:     先沿动量方向"前瞻"一步，在"前瞻位置"计算梯度，再更新

    公式:
        θ_lookahead = θ_t - α · β · m_{t-1}            ← 沿当前动量方向看一步
        m_t = β · m_{t-1} + (1-β) · ∇L(θ_lookahead)   ← 在"前瞻"位置计算梯度
        θ_{t+1} = θ_t - α · m_t                        ← 用前瞻梯度更新

    直观上，NAG 像是"先试着往前走一步，发现不对再调整方向"。

    你需要实现 step() 方法。

    参数:
        lr: 学习率 α
        beta: 动量衰减系数 β
    """

    def __init__(self, lr: float = 0.02, beta: float = 0.9):
        self.lr = lr
        self.beta = beta
        self.m = None
        self.name = "NAG"

    def step(self, theta: np.ndarray, grad_fn) -> np.ndarray:
        """
        执行一步 NAG 更新。

        注意：与普通优化器不同，NAG 需要 grad_fn（梯度计算函数），
        因为需要在"前瞻"位置重新计算梯度。

        提示:
          1. 如果 self.m 为 None，初始化为零向量
          2. 计算前瞻位置: theta_lookahead = theta - lr * beta * m
          3. 在前瞻位置计算梯度: grad_lookahead = grad_fn(theta_lookahead)
          4. 更新动量: m = beta * m + (1 - beta) * grad_lookahead
          5. 更新参数: theta = theta - lr * m

        参数:
            theta: 当前参数向量
            grad_fn: 梯度计算函数，签名为 grad_fn(theta) -> gradient

        返回:
            更新后的参数向量
        """
        # TODO: 初始化动量
        if self.m is None:
            pass  # ← TODO: self.m = np.zeros_like(theta)

        # TODO: 计算前瞻位置
        theta_lookahead = None  # ← TODO: theta - self.lr * self.beta * self.m

        # TODO: 在前瞻位置计算梯度
        grad_lookahead = None  # ← TODO: grad_fn(theta_lookahead)

        # TODO: 更新动量
        pass  # ← TODO: self.m = self.beta * self.m + (1 - self.beta) * grad_lookahead

        # TODO: 更新参数
        pass  # ← TODO: return theta - self.lr * self.m

        return theta  # 占位返回


# ---- 测试 TODO 3 ----
def test_nag():
    """测试 NAG 优化器"""
    print("=" * 60)
    print("TODO 3 测试: Nesterov 加速梯度 (NAG)")
    print("=" * 60)

    landscape = LossLandscape(a=20.0, b=1.0)
    theta0 = np.array([3.0, 2.5])

    # NAG
    nag = NAGOptimizer(lr=0.02, beta=0.9)
    theta_nag = theta0.copy()

    # 普通 Momentum 用于对比
    from types import SimpleNamespace
    mom = SimpleNamespace()
    mom.lr = 0.02
    mom.beta = 0.9
    mom.m = np.zeros_like(theta0)
    theta_mom = theta0.copy()

    print(f"\n  对比 NAG vs Momentum (前10步)")
    print(f"  {'Step':<6} {'NAG L':<14} {'NAG θ₁':<12} {'Mom L':<14} {'Mom θ₁'}")
    print(f"  {'-'*60}")

    for step in range(10):
        # NAG 更新
        grad_fn = landscape.gradient
        theta_nag_new = nag.step(theta_nag, grad_fn)

        if theta_nag_new is None or np.allclose(theta_nag_new, theta_nag):
            print("  TODO 未完成，请补全 NAGOptimizer.step()")
            return

        theta_nag = theta_nag_new

        # 普通 Momentum 更新
        grad_mom = landscape.gradient(theta_mom)
        mom.m = mom.beta * mom.m + (1 - mom.beta) * grad_mom
        theta_mom = theta_mom - mom.lr * mom.m

        print(f"  {step+1:<6} {landscape(theta_nag):<14.6f} {theta_nag[0]:<12.4f} "
              f"{landscape(theta_mom):<14.6f} {theta_mom[0]:.4f}")

    # NAG 通常比 Momentum 收敛更快
    loss_nag = landscape(theta_nag)
    loss_mom = landscape(theta_mom)
    print(f"\n  最终: NAG loss={loss_nag:.6f}, Momentum loss={loss_mom:.6f}")
    print(f"  NAG 是否更快: {'是' if loss_nag < loss_mom else '否'} (预期: 是)")

    print()


# ============================================================================
# 主程序
# ============================================================================
if __name__ == "__main__":
    print("\n╔══════════════════════════════════════════════════════════════╗")
    print("║   s08 优化器：从 SGD 到 Adam — 动手练习                    ║")
    print("║   请依次完成 TODO 1, 2, 3                                   ║")
    print("╚══════════════════════════════════════════════════════════════╝\n")

    test_momentum()
    test_rmsprop()
    test_nag()

    print("=" * 60)
    print("所有测试完成！请检查输出结果。")
    print("如有未通过的测试，请回到对应的 TODO 部分补全代码。")
    print()
    print("优化器设计演进:")
    print("  SGD → 只看当前梯度，简单但有震荡、慢缩等问题")
    print("  Momentum → 加惯性，方向更平滑")
    print("  RMSProp → 自适应步长，每个参数有自己的学习率")
    print("  Adam → Momentum + RMSProp，集两者之长")
    print("  NAG → 比 Momentum 更进一步：先"看一步"再调整方向")
    print("=" * 60)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383

s08 优化器：从SGD到Adam — exercise.py 练习指南 ​

练习目标 ​

预备知识 ​

任务清单 ​

任务1：实现 Momentum 更新规则 ​

任务2：实现 RMSProp 自适应步长 ​

任务3：实现 Nesterov 加速梯度 (NAG) ​

优化器演进总览 ​

完整代码 ​