深入理解torch.optim.Adam的用法及优化流程

深入理解 torch.optim.Adam 的用法及优化流程

在深度学习中，优化器是训练模型的重要组件，用于更新模型参数以最小化损失函数。torch.optim.Adam 是 PyTorch 中常用的优化器之一，它结合了动量法和自适应学习率的优点，被广泛应用于各种神经网络的训练。

本文将结合以下代码，详细讲解 torch.optim.Adam 的使用方法及其背后的逻辑。

---

代码背景

在代码中，我们实现了一个强化学习 PPO（Proximal Policy Optimization）的优化过程，涉及 Actor 模型（策略）和 Critic 模型（值函数）的更新。以下是完整代码片段：具体可参考笔者的另一篇博客：RLHF (PPO) 流程详解: Proximal Policy Optimization

# Actor Loss: 策略梯度损失（PPO 损失函数）
actor_loss = -torch.min(ratios * advantages, clipped_ratios * advantages).mean()

# Critic Loss: 值函数损失
critic_loss = nn.MSELoss()(critic_values, returns)

# 总损失
total_loss = actor_loss + critic_loss

# 5. 反向传播和更新
optimizer = torch.optim.Adam(list(actor_model.parameters()) + list(critic_model.parameters()))
optimizer.zero_grad()
total_loss.backward()
optimizer.step()

---

1. torch.optim.Adam 的基础知识

Adam 是一种自适应优化算法，结合了两种方法的优点：

1. 动量法（Momentum）：利用梯度的指数移动平均加速收敛。
2. RMSProp：利用历史梯度平方的移动平均动态调整学习率。

Adam 的核心更新公式为：
$${\theta }_{t+1}={\theta }_t-\frac{\alpha }{\sqrt{{\hat{v}}_t}+ϵ}{\hat{m}}_t$$
其中：

• ( ${\hat{m}}_t$ ) 是梯度的一阶动量估计（类似于动量）。
• ( ${\hat{v}}_t$ ) 是梯度的二阶动量估计（类似于 RMSProp 的平方梯度）。
• ( $\alpha$ ) 是学习率。
• ( $ϵ$ ) 是防止除零的小数。

PyTorch 中 torch.optim.Adam 的关键参数：

• params: 要优化的参数（如模型的 parameters()）。
• lr: 学习率（默认值为 0.001）。
• betas: 一阶和二阶动量的衰减系数（默认值为 (0.9, 0.999)）。
• eps: 防止除零的微小值（默认值为 1e-8）。
• weight_decay: 权重衰减系数，用于 L2 正则化。

---

2. torch.optim.Adam 的使用流程

以下是使用 torch.optim.Adam 的典型步骤：

(1) 初始化优化器

optimizer = torch.optim.Adam(list(actor_model.parameters()) + list(critic_model.parameters()), lr=3e-4)

• 将需要优化的参数列表传入优化器，如 actor_model.parameters() 和 critic_model.parameters()。
• 可以为不同模块设置不同的学习率。

(2) 清零梯度

optimizer.zero_grad()

• 每次反向传播前需要清零梯度，避免梯度累加导致错误更新。

(3) 反向传播

total_loss.backward()

• 对 total_loss 求梯度，梯度会存储在每个参数的 .grad 属性中。

(4) 更新参数

optimizer.step()

• 根据梯度和优化公式更新参数。

---

3. PPO 损失函数的优化

在 PPO 中，损失函数包含两部分：

1. Actor Loss（策略损失）：用于更新 Actor 模型，使其生成更优的策略。

   actor_loss = -torch.min(ratios * advantages, clipped_ratios * advantages).mean()
   

   • ratios 是当前策略与参考策略的概率比值。
   • clipped_ratios 是裁剪后的概率比值，避免策略更新过大。

2. Critic Loss（值函数损失）：用于更新 Critic 模型，使其对状态值的估计更准确。

   critic_loss = nn.MSELoss()(critic_values, returns)
   

   • 使用均方误差（MSE）衡量 Critic 模型预测值（critic_values）与目标值（returns）之间的差距。

3. 总损失：两者相加，用于同时优化 Actor 和 Critic。

   total_loss = actor_loss + critic_loss
   

---

4. 示例代码分析

以下代码展示了使用 torch.optim.Adam 完成 Actor 和 Critic 模型更新的过程：

# 初始化优化器
optimizer = torch.optim.Adam(
    list(actor_model.parameters()) + list(critic_model.parameters()), lr=3e-4
)

# 清零梯度
optimizer.zero_grad()

# 计算总损失
total_loss = actor_loss + critic_loss

# 反向传播
total_loss.backward()

# 更新参数
optimizer.step()

优化流程

1. zero_grad 清除上一轮梯度，确保每次更新的独立性。
2. backward 计算损失函数对所有模型参数的梯度。
3. step 根据梯度更新模型参数，完成一次优化。

---

5. 代码运行结果示例

假设：

• ratios = torch.tensor([1.2, 0.8])，clipped_ratios = torch.tensor([1.0, 0.8])。
• advantages = torch.tensor([2.0, 1.5])。
• critic_values = torch.tensor([1.0, 2.0])，returns = torch.tensor([1.5, 2.5])。

我们可以打印每个步骤的值：

print("Actor Loss:", actor_loss.item())
print("Critic Loss:", critic_loss.item())
print("Total Loss:", total_loss.item())

输出示例：

Actor Loss: -1.35
Critic Loss: 0.25
Total Loss: -1.10

---

6. 总结

通过以上分析可以看到，torch.optim.Adam 的使用主要包括以下几点：

1. 初始化优化器：传入需要优化的参数及超参数。
2. 清零梯度：避免梯度累加。
3. 反向传播：通过 .backward() 计算梯度。
4. 更新参数：通过 .step() 完成参数更新。

在 PPO 的优化中，Adam 的动量机制和自适应学习率帮助我们更高效地优化策略和值函数，为模型生成更优质的响应提供了强大支持。

补充内容：为什么传入的是 list(actor_model.parameters())？

在 PyTorch 中，model.parameters() 返回的是一个 生成器（generator）对象，它惰性地迭代模型的所有参数。如果需要对参数进行进一步操作（如合并多个模型的参数列表），就需要将其显式地转换为 list 类型。

例如：

# 获取参数生成器
params_gen = actor_model.parameters()
print(type(params_gen))  # 输出：<class 'generator'>

# 转换为列表
params_list = list(params_gen)
print(type(params_list))  # 输出：<class 'list'>

在以下代码中，我们将 actor_model 和 critic_model 的参数合并传入优化器：

optimizer = torch.optim.Adam(
    list(actor_model.parameters()) + list(critic_model.parameters()), lr=3e-4
)

• actor_model.parameters() 和 critic_model.parameters() 分别返回两个生成器。
• 使用 list() 将生成器转为列表后，可以通过 + 运算符合并两个模型的参数。

这样，Adam 优化器就可以同时更新 actor_model 和 critic_model 的参数。

---

具体的数值模拟：Adam 的计算过程

模拟参数初始化

假设 actor_model 中有两个可训练参数：

1. ( ${\theta }_1=0.5$ )
2. ( ${\theta }_2=-0.3$ )

目标是最小化以下损失函数：
$$\mathcal{L}={\theta }_1^2+2{\theta }_2^2$$
对应的梯度为：
$$\frac{\partial \mathcal{L}}{\partial {\theta }_1}=2{\theta }_1,\frac{\partial \mathcal{L}}{\partial {\theta }_2}=4{\theta }_2$$

---

Adam 优化过程

Adam 的核心更新公式为：
$${\theta }_{t+1}={\theta }_t-\frac{\alpha }{\sqrt{{\hat{v}}_t}+ϵ}\cdot {\hat{m}}_t$$

其中：

1. 梯度 ( $g_t$ ) 是损失函数对参数的偏导数。
2. ( ${\hat{m}}_t$)：梯度的一阶动量估计。
   $$m_t={\beta }_1\cdot m_{t-1}+(1-{\beta }_1)\cdot g_t$$
   $${\hat{m}}_t=\frac{m_t}{1-{\beta }_1^t}$$
3. ( ${\hat{v}}_t$ )：梯度平方的二阶动量估计。
   $$v_t={\beta }_2\cdot v_{t-1}+(1-{\beta }_2)\cdot g_t^2$$
   $${\hat{v}}_t=\frac{v_t}{1-{\beta }_2^t}$$

超参数取值：

• 学习率 ( $\alpha =0.1$ )
• 动量系数 ( ${\beta }_1=0.9$ )
• 二阶动量系数 ( ${\beta }_2=0.999$ )
• 防止除零项 ( $ϵ=10^{-8}$ )

---

初始状态

• 参数：( ${\theta }_1=0.5$ ), ( ${\theta }_2=-0.3$ )
• 梯度：
  $$g_{{\theta }_1}=2\cdot 0.5=1.0,g_{{\theta }_2}=4\cdot (-0.3)=-1.2$$
• 初始一阶动量：( $m_0=0$ )
• 初始二阶动量：( $v_0=0$ )

---

第 1 步迭代

1. 计算一阶动量：
   $$m_{{\theta }_1}^{(1)}=0.9\cdot 0+0.1\cdot 1.0=0.1,m_{{\theta }_2}^{(1)}=0.9\cdot 0+0.1\cdot (-1.2)=-0.12$$
   偏差修正：
   $${\hat{m}}_{{\theta }_1}^{(1)}=\frac{0.1}{1-0.9^1}=1.0,{\hat{m}}_{{\theta }_2}^{(1)}=\frac{-0.12}{1-0.9^1}=-1.2$$

2. 计算二阶动量：
   $$v_{{\theta }_1}^{(1)}=0.999\cdot 0+0.001\cdot (1.0^2)=0.001,v_{{\theta }_2}^{(1)}=0.999\cdot 0+0.001\cdot (-1.2^2)=0.00144$$
   偏差修正：
   $${\hat{v}}_{{\theta }_1}^{(1)}=\frac{0.001}{1-0.999^1}=1.0,{\hat{v}}_{{\theta }_2}^{(1)}=\frac{0.00144}{1-0.999^1}=1.44$$

3. 更新参数：
   $${\theta }_1^{(1)}=0.5-\frac{0.1}{\sqrt{1.0}+10^{-8}}\cdot 1.0=0.4$$
   $${\theta }_2^{(1)}=-0.3-\frac{0.1}{\sqrt{1.44}+10^{-8}}\cdot (-1.2)=-0.2$$

---

第 2 步迭代

重复上述步骤，计算新的梯度、一阶动量、二阶动量和参数更新。

---

总结：Adam 的优点

1. 动态学习率：针对每个参数自动调整学习率，收敛更稳定。
2. 动量机制：缓解梯度更新的噪声，提高收敛速度。
3. 广泛适用性：在稀疏数据、高维数据中表现优异。

通过数值模拟可以更直观地理解 Adam 的核心更新机制，也能帮助更好地理解为什么要使用 list(actor_model.parameters()) 传入优化器以更新多个模型的参数。

后记

2024年12月13日13点00分于上海，在GPT4o大模型辅助下完成。