第3章 自动微分 - 框架的核心引擎 (Automatic Differentiation: The Core Engine of Frameworks)
一、导读
1.1 本章学习目标
通过本章的学习，读者将能够：
理解自动微分的基本原理及其与符号微分、数值微分的本质区别
掌握前向模式自动微分和反向模式自动微分的算法流程及适用场景
深入分析PyTorch的tape-based自动微分实现机制
理解计算图的构建、执行和销毁过程
掌握高阶导数的计算方法及其在元学习和二阶优化中的应用
识别和处理自动微分中的常见问题（如原地操作、控制流、离散操作）
实现自定义操作的梯度函数
分析自动微分的性能和内存开销，并进行针对性优化
1.2 为什么本章如此重要
自动微分是现代深度学习框架的基石，它将用户从繁琐且易错的梯度推导中解放出来，使得研究人员能够快速实验新的模型架构。然而，自动微分并非魔法，其背后涉及精巧的数据结构和算法设计。
理解自动微分的工作原理对于多个方面至关重要。首先，当需要实现论文中的新颖操作时，必须正确定义其梯度行为。其次，调试训练问题时，经常需要检查梯度是否正确计算，这需要理解框架如何处理各种边界情况。第三，性能优化需要考虑自动微分的内存开销，特别是对于长序列或深层网络，中间激活值的存储可能成为显存瓶颈。最后，随着高阶优化方法和元学习的兴起，高阶导数的计算变得越来越重要。
本章将从理论到实践全面剖析自动微分系统，帮助读者建立对这一核心技术的深入理解。
1.3 核心问题引导
在学习本章内容时，请思考以下问题：
自动微分与符号微分、数值微分的根本区别是什么？为什么自动微分成为深度学习框架的首选？
前向模式和反向模式自动微分各有什么优缺点？在什么场景下应该选择其中一种？
PyTorch的动态计算图是如何构建的？grad_fn和next_functions的作用是什么？
为什么某些操作（如in-place操作）可能破坏自动微分的正确性？框架如何检测和处理这种情况？
如何为不可导操作（如argmax、采样）定义有意义的梯度？这些近似梯度的理论基础是什么？
高阶导数（如Hessian矩阵）的计算有哪些方法？它们在二阶优化和元学习中的应用场景是什么？
自动微分的内存开销主要来自哪里？有哪些技术可以减少这种开销？
JAX的函数式自动微分与PyTorch的命令式自动微分有何不同？各自的优劣是什么？
二、核心概念详解
自动微分的基本原理
自动微分（Automatic Differentiation, AD）是一种精确计算函数导数的技术，它既不同于符号微分，也不同于数值微分。理解这三者的区别是掌握自动微分的关键。
符号微分（Symbolic Differentiation）通过对数学表达式进行解析变换来推导导数的闭式表达。例如，对$f(x) = x^2$求导得到$f'(x) = 2x$。符号微分的优点是结果精确且可以复用，但缺点是表达式膨胀（Expression Swell）问题：复杂函数的导数表达式可能极其庞大，难以高效计算。此外，符号微分无法处理包含控制流（if-else、循环）的程序。
数值微分（Numerical Differentiation）通过有限差分近似导数：$f'(x) \approx \frac{f(x+h) - f(x)}{h}$。这种方法实现简单，可以应用于任何可计算的函数。但数值微分存在两个严重问题：一是截断误差，当$h$不够小时近似不准确；二是舍入误差，当$h$过小时浮点数精度问题导致结果失真。更重要的是，数值微分的计算复杂度与输入维度成正比，对于百万级参数的神经网络完全不可行。
自动微分通过将复杂函数分解为一系列基本操作的组合，利用链式法则逐层计算导数。自动微分的核心洞察是：任何计算机程序都可以分解为基本算术运算（加、减、乘、除）和基本函数（指数、对数、三角函数等）的组合，而这些基本操作的导数是已知的。通过记录程序的执行轨迹（称为trace或tape），自动微分系统能够精确计算任意复杂函数的导数，同时保持计算效率。
自动微分的关键特性：
- 精确性：与数值微分不同，自动微分计算的是精确的解析梯度，没有近似误差
- 高效性：计算梯度的时间复杂度与原函数的计算复杂度同阶，最多相差一个小的常数因子
- 通用性：可以处理包含控制流、递归等复杂结构的程序
前向模式自动微分
前向模式自动微分（Forward Mode AD）从输入变量开始，沿着计算的正向顺序同时计算函数值和导数值。
在前向模式中，每个变量都被扩展为一个二元组$(v, \dot{v})$，其中$v$是变量的值，$\dot{v} = \frac{\partial v}{\partial x}$是变量对某个选定输入$x$的导数。对于基本操作，我们可以直接定义其二元组的运算规则。
例如，对于加法$z = x + y$：
- 值：$z = x + y$
- 导数：$\dot{z} = \dot{x} + \dot{y}$
对于乘法$z = x \times y$：
- 值：$z = x \times y$
- 导数：$\dot{z} = \dot{x} \cdot y + x \cdot \dot{y}$（乘积法则）
对于复合函数$z = \sin(u)$，其中$u = g(x, y)$：
- 值：$z = \sin(u)$
- 导数：$\dot{z} = \cos(u) \cdot \dot{u}$
前向模式的计算流程如下：
初始化：设置输入变量的值为给定值，导数为1（对自身的导数）或其他输入变量的导数为0
按照程序的执行顺序，依次计算每个中间变量的值和导数
最终得到输出变量的值和对该输入的导数
前向模式的优点是直观、易于实现，内存开销小（不需要保存中间结果）。但缺点是，如果要计算函数对所有$n$个输入的梯度，需要执行$n$次前向传播，每次针对一个输入变量。对于深度学习中的典型场景（输入维度高达百万），这显然是不可行的。
因此，前向模式适用于输入维度低、输出维度高的场景，如雅可比矩阵的列计算。
反向模式自动微分
反向模式自动微分（Reverse Mode AD）先执行一次前向传播计算函数值，然后从输出反向传播计算所有输入的梯度。这正是深度学习中使用的反向传播算法的一般化形式。
反向模式分为两个阶段：
前向阶段（Forward Pass）：执行原始计算，记录所有中间变量和操作序列。这个记录称为tape或Wengert list。
反向阶段（Reverse Pass）：从输出开始，按照操作序列的逆序，应用链式法则计算梯度。
在反向模式中，我们维护一个伴随变量（adjoint）$\bar{v} = \frac{\partial y}{\partial v}$，表示最终输出$y$对变量$v$的梯度。反向传播的规则是：对于操作$z = f(x_1, x_2, ..., x_n)$，如果已知$\bar{z}$，则每个输入的伴随变量更新为：
$$\bar{x}_i += \bar{z} \cdot \frac{\partial z}{\partial x_i}$$
注意这里是累加（+=）而不是赋值，因为一个变量可能被多个操作使用，其梯度是所有使用路径贡献的总和。
反向模式的优势在于，无论输入维度多高，只需一次前向和一次反向传播即可计算所有输入的梯度。这正是深度学习场景所需要的：损失函数是标量（输出维度为1），参数数量巨大（输入维度高）。因此，反向模式成为深度学习框架的标准选择。
反向模式的代价是需要保存前向传播的所有中间结果，以便在反向传播时使用。这导致内存占用与前向传播的计算量成正比，是深度学习训练显存消耗的主要来源之一。
PyTorch的自动微分实现
PyTorch采用tape-based的反向模式自动微分实现。理解其内部机制有助于正确使用框架并诊断问题。
计算图节点：PyTorch中的每个张量（Tensor）都关联着一个grad_fn属性，它是一个指向Function对象的引用，该对象记录了创建此张量的操作。对于叶子节点（用户直接创建的张量），grad_fn为None；对于计算得到的张量，grad_fn指向产生它的操作。
Function对象：每个Function对象代表一个操作，它包含两个关键方法：forward()执行前向计算，backward()执行反向梯度计算。Function还保存了反向传播所需的上下文信息（通过ctx.save_for_backward()保存的张量）。
计算图的构建：当执行张量操作时，PyTorch动态地构建计算图。例如：
x = torch.tensor(2.0, requires_grad=True) # 叶子节点
y = x ** 2 # y.grad_fn指向Pow操作
z = y * 3 # z.grad_fn指向Mul操作在这个过程中，PyTorch创建了三个Function对象：AccumulateGrad（用于累积叶子节点的梯度）、Pow、Mul。这些对象通过next_functions链接形成图结构。
反向传播的执行：调用z.backward()时，PyTorch从z.grad_fn开始，按照拓扑排序的逆序遍历计算图。对于每个Function，调用其backward()方法，将上游梯度转换为下游梯度。对于叶子节点，梯度被累积到.grad属性中。
图的释放：默认情况下，反向传播完成后，计算图会被释放以节省内存。如果需要多次反向传播（如高阶导数计算），可以设置retain_graph=True。
高阶导数
高阶导数指对函数的一阶导数再次求导得到的结果，如二阶导数、Hessian矩阵等。高阶导数在多个领域有重要应用：
二阶优化：牛顿法利用Hessian矩阵信息进行参数更新，可以实现更快的收敛速度。虽然完整的牛顿法在深度学习中计算代价过高，但一些近似方法（如K-FAC、Shampoo）利用二阶信息取得了良好的效果。
元学习（Meta-Learning）：MAML（Model-Agnostic Meta-Learning）等元学习算法需要计算梯度对超参数的导数，这涉及二阶导数。
梯度惩罚：WGAN-GP等生成模型通过在真实数据和生成数据之间采样点上的梯度范数施加惩罚来提高训练稳定性，这需要计算梯度的梯度。
Hessian 向量积：完整 Hessian 矩阵的存储和计算代价过高（O(n²)），但 Hessian 与向量的乘积可以在 O(n) 时间内计算，这使得一些基于二阶信息的优化方法变得可行。
在PyTorch中，计算高阶导数的方法是嵌套调用autograd.grad()：
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3
# 一阶导数
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
# 二阶导数
d2y_dx2 = torch.autograd.grad(dy_dx, x)[0]
print(f"dy/dx = {dy_dx}") # 3x^2 = 12
print(f"d2y/dx2 = {d2y_dx2}") # 6x = 12注意第一次调用时需要设置create_graph=True，这样一阶导数的计算也会被记录在计算图中，从而可以对它再次求导。
处理不可导操作
实际程序中经常包含不可导的操作，如条件分支、离散采样、argmax等。自动微分系统需要通过特殊技巧来处理这些情况。
直通估计器（Straight-Through Estimator, STE）：对于离散操作，STE在前向传播时执行离散操作，在反向传播时将梯度直接传递过去，仿佛该操作是恒等映射。例如，对于量化操作：
class Quantize(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
return torch.round(x)
@staticmethod
def backward(ctx, grad_output):
return grad_output # 直通虽然这在数学上不正确，但在实践中往往有效，特别是在二值神经网络和量化感知训练中。
Gumbel-Softmax：用于从离散分布中采样的可微近似。标准的采样操作是不可导的，但Gumbel-Softmax通过添加Gumbel噪声并使用Softmax松弛，提供了一个可微的采样近似。当温度参数趋近于0时，Gumbel-Softmax趋近于真实的离散采样。
Surrogate Gradient：为不可导操作定义一个替代的梯度函数。例如，ReLU在零点不可导，但实践中通常将零点的导数设为0或1。类似地，阶跃函数的导数是狄拉克δ函数，无法直接计算，可以用Sigmoid的导数作为替代。
强化学习中的策略梯度：采样操作的梯度可以通过得分函数估计器（Score Function Estimator）或REINFORCE算法来计算，这利用了期望的梯度等于梯度乘以对数概率的期望这一性质。
控制流与自动微分
包含if-else分支和循环的程序给自动微分带来了挑战，因为计算图的结构依赖于运行时数据。
动态图框架（如PyTorch）天然支持控制流，因为计算图是在执行过程中逐步构建的。每次前向传播可能产生不同的计算图，反向传播只需根据实际的执行轨迹进行即可。这使得PyTorch非常适合研究原型开发。
静态图框架（如TensorFlow 1.x）需要在执行前确定完整的计算图，因此需要使用特殊的控制流操作（如tf.cond、tf.while_loop）来表示分支和循环。这些操作内部实现了相应的梯度逻辑。
JAX采用另一种方法：通过转换（transformation）来处理控制流。jax.lax.cond和jax.lax.while_loop是专门设计的控制流原语，它们知道如何正确地计算梯度。
控制流的梯度计算需要注意一些问题：
- 分支条件的变化不应影响梯度（即条件本身通常不接收梯度）
- 循环的迭代次数可能是动态的，反向传播需要展开相应次数的迭代
- 某些控制流可能导致梯度路径中断，需要小心处理
三、程序员的视角
3.1 在实际项目中的应用
自定义操作的梯度实现：当实现论文中的新颖操作时，可能需要手动定义梯度。这通过继承torch.autograd.Function实现，需要同时实现forward()和backward()静态方法。
调试梯度问题：当模型不收敛时，检查梯度是关键步骤。可以使用torch.autograd.gradcheck()验证自定义操作的梯度实现是否正确，通过与数值梯度的比较来发现潜在错误。
内存优化：对于长序列模型，前向传播保存的中间激活值可能耗尽显存。可以使用torch.utils.checkpoint（也称为gradient checkpointing或activation recomputation）来权衡计算和内存：只保存部分中间结果，在反向传播时重新计算缺失的部分。
3.2 常见陷阱与误区
原地操作破坏梯度计算：如果一个张量被多个操作使用，对其进行原地修改（如x.add_(1)）可能导致错误的梯度。PyTorch会检测这种情况并抛出错误，但有时可能漏检。
忘记设置requires_grad：只有requires_grad=True的张量才会被追踪并计算梯度。从文件加载的张量或某些操作的输出可能默认不要求梯度。
detach()的误用：detach()创建一个与原张量共享数据但不参与梯度计算的新张量。如果在需要梯度流动的地方意外使用了detach()，会导致梯度中断。
梯度累积未清零：PyTorch默认累积梯度而不是覆盖。如果不在每个训练步开始时调用optimizer.zero_grad()，梯度会不断累加，导致错误的更新。
create_graph的遗漏：计算高阶导数时，如果第一次求导没有设置create_graph=True，第二次求导将无法找到计算图。
3.3 与其他层级的联系
自动微分引擎的效率直接影响框架的整体性能。高效的实现需要：
- 最小化tape记录的开销
- 优化反向传播的调度顺序
- 合理管理中间结果的内存生命周期
- 并行化独立的梯度计算
底层硬件的特性也会影响自动微分的设计。例如，GPU的异步执行特性使得同步点的设计变得复杂，而分布式环境下的梯度同步需要额外的通信原语。
四、代码实战
示例1：实现自定义自动微分操作
import torch
import torch.nn as nn
class LogSumExp(torch.autograd.Function):
"""
数值稳定的LogSumExp操作
LogSumExp(x) = log(sum(exp(x)))
直接计算可能在x较大时溢出，因此使用以下技巧：
LogSumExp(x) = max(x) + log(sum(exp(x - max(x))))
"""
@staticmethod
def forward(ctx, x, dim=None):
# 保存输入和维度信息
ctx.dim = dim
x_max = torch.max(x, dim=dim, keepdim=True)[0]
# 数值稳定的计算
exp_x = torch.exp(x - x_max)
sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True)
result = x_max.squeeze(dim=dim) + torch.log(sum_exp_x.squeeze(dim=dim))
# 保存反向传播需要的中间结果
ctx.save_for_backward(x, x_max, sum_exp_x)
return result
@staticmethod
def backward(ctx, grad_output):
x, x_max, sum_exp_x = ctx.saved_tensors
dim = ctx.dim
# LogSumExp的梯度是softmax
exp_x = torch.exp(x - x_max)
softmax = exp_x / sum_exp_x
# 链式法则
grad_input = grad_output.unsqueeze(dim if dim is not None else -1) * softmax
return grad_input, None # 第二个None对应dim参数
def test_logsumexp():
"""测试自定义LogSumExp的正确性"""
x = torch.randn(5, 10, requires_grad=True)
# 自定义实现
result_custom = LogSumExp.apply(x, dim=1)
loss_custom = result_custom.sum()
loss_custom.backward()
grad_custom = x.grad.clone()
# 重置
x.grad = None
# PyTorch内置实现
result_builtin = torch.logsumexp(x, dim=1)
loss_builtin = result_builtin.sum()
loss_builtin.backward()
grad_builtin = x.grad.clone()
print("Custom result:", result_custom)
print("Built-in result:", result_builtin)
print("Result difference:", torch.abs(result_custom - result_builtin).max().item())
print("\nCustom gradient:", grad_custom)
print("Built-in gradient:", grad_builtin)
print("Gradient difference:", torch.abs(grad_custom - grad_builtin).max().item())
# 运行测试
test_logsumexp()示例2：梯度检查工具
import torch
def numerical_gradient_check(func, params, eps=1e-5):
"""
通过数值梯度验证解析梯度的正确性
Args:
func: 要检查的函数，接受参数并返回标量损失
params: 需要检查梯度的参数列表
eps: 有限差分的步长
Returns:
max_error: 最大相对误差
"""
max_error = 0.0
for param in params:
if param.grad is None:
continue
analytical_grad = param.grad.flatten()
numerical_grad = torch.zeros_like(analytical_grad)
param_flat = param.data.flatten()
for i in range(len(param_flat)):
original_value = param_flat[i].item()
# f(x + eps)
param_flat[i] = original_value + eps
loss_plus = func()
# f(x - eps)
param_flat[i] = original_value - eps
loss_minus = func()
# 中心差分
numerical_grad[i] = (loss_plus - loss_minus) / (2 * eps)
# 恢复原值
param_flat[i] = original_value
# 计算相对误差
denominator = torch.maximum(
torch.abs(analytical_grad) + torch.abs(numerical_grad),
torch.full_like(analytical_grad, 1e-8)
)
relative_error = torch.abs(analytical_grad - numerical_grad) / denominator
max_error = max(max_error, relative_error.max().item())
if relative_error.max().item() > 1e-4:
print(f"WARNING: Large gradient error detected!")
print(f" Max relative error: {relative_error.max().item():.6e}")
print(f" Analytical grad sample: {analytical_grad[:5]}")
print(f" Numerical grad sample: {numerical_grad[:5]}")
return max_error
# 使用示例
def test_gradient_check():
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x).sum()
model = SimpleModel()
x = torch.randn(3, 10)
# 前向传播
loss = model(x)
# 反向传播
loss.backward()
# 梯度检查
error = numerical_gradient_check(lambda: model(x).sum(), [model.fc.weight, model.fc.bias])
print(f"\nMax gradient error: {error:.6e}")
if error < 1e-5:
print("Gradient check PASSED!")
else:
print("Gradient check FAILED!")
test_gradient_check()示例3：使用checkpoint节省显存
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
class DeepNetworkWithCheckpoint(nn.Module):
"""
使用gradient checkpointing的深度网络
Checkpointing通过重新计算来节省显存，
适合训练非常深的网络或处理长序列
"""
def __init__(self, num_layers=100, hidden_size=256):
super().__init__()
self.layers = nn.ModuleList([
nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU()
)
for _ in range(num_layers)
])
self.input_proj = nn.Linear(128, hidden_size)
self.output_proj = nn.Linear(hidden_size, 10)
def forward(self, x):
x = self.input_proj(x)
for layer in self.layers:
# 使用checkpoint包装每层的计算
# 前向时不保存中间激活，反向时重新计算
x = checkpoint(layer, x, use_reentrant=False)
return self.output_proj(x)
def compare_memory_usage():
"""比较有/无checkpoint的显存占用"""
batch_size = 16
seq_length = 1000
input_size = 128
hidden_size = 256
num_layers = 50
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 带checkpoint的网络
model_checkpoint = DeepNetworkWithCheckpoint(num_layers, hidden_size).to(device)
# 不带checkpoint的网络
model_no_checkpoint = DeepNetworkWithCheckpoint(num_layers, hidden_size).to(device)
# 禁用checkpoint：直接调用layer
for name, module in model_no_checkpoint.named_modules():
if isinstance(module, type(model_no_checkpoint.layers[0])):
pass # 保持原样，只是不使用checkpoint
x = torch.randn(batch_size, seq_length, input_size).to(device)
def get_memory_mb():
if torch.cuda.is_available():
return torch.cuda.memory_allocated() / 1024 / 1024
return 0
# 测试带checkpoint的版本
torch.cuda.empty_cache()
output1 = model_checkpoint(x)
loss1 = output1.sum()
mem_with_checkpoint = get_memory_mb()
loss1.backward()
# 清空
del output1, loss1
for p in model_checkpoint.parameters():
p.grad = None
torch.cuda.empty_cache()
# 测试不带checkpoint的版本（需要修改forward不使用checkpoint）
# 这里简化演示，实际应创建另一个不使用checkpoint的类
print(f"Memory with checkpoint: {mem_with_checkpoint:.2f} MB")
print("(Full comparison requires separate model definitions)")
# 简单的checkpoint效果演示
def simple_checkpoint_demo():
"""演示checkpoint如何工作"""
def expensive_function(x):
"""模拟计算密集的操作"""
y = x ** 2
z = torch.sin(y)
w = torch.exp(z)
return w
x = torch.tensor(2.0, requires_grad=True)
# 不使用checkpoint：保存所有中间结果
result_no_ckpt = expensive_function(x)
result_no_ckpt.backward()
print(f"Without checkpoint - grad: {x.grad}")
x.grad = None
# 使用checkpoint：重新计算中间结果
result_ckpt = checkpoint(expensive_function, x, use_reentrant=False)
result_ckpt.backward()
print(f"With checkpoint - grad: {x.grad}")
simple_checkpoint_demo()五、深入思考
5.1 历史背景
自动微分的历史早于深度学习。早在20世纪60年代，就有人提出了类似的思想。Seppo Linnainmaa在1976年发表了关于反向模式自动微分的论文。Paul Werbos在1974年的博士论文中将反向传播应用于神经网络，但当时并未引起广泛关注。
20世纪90年代，自动微分在科学计算领域得到发展，出现了ADIFOR、ADOL-C等工具。但这些工具主要面向Fortran和C++，使用复杂，限制了普及。
2010年后，随着深度学习的兴起，自动微分成为研究热点。Theano（2010）是最早将自动微分与深度学习结合的框架之一。随后，TensorFlow（2015）和PyTorch（2016）将自动微分推向主流。近年来，JAX（2018）展示了函数式自动微分的优雅性和强大功能。
5.2 前沿方向
可逆神经网络：通过设计可逆的层结构，可以在反向传播时重新计算前向激活值，将内存复杂度从$O(n)$降低到$O(1)$（相对于深度）。RevNet、i-RevNet等工作在这方面取得了进展。
高阶自动微分：JAX等框架支持任意阶导数的计算，这使得高阶优化方法、物理模拟、概率推理等应用变得更加可行。
微分编程（Differentiable Programming）：将整个程序视为可微分的，包括控制流、数据结构、甚至编译器优化。这使得端到端的可微分系统设计成为可能。
自动微分与编译器集成：MLIR、Swift for TensorFlow等项目尝试将自动微分作为编译器的一等公民，在编译时进行梯度相关的优化。
5.3 延伸阅读
- "Evaluating Derivatives: Principles and Techniques" by Griewank and Walther - 自动微分的权威教材
- "Automatic Differentiation in Machine Learning: a Survey" by Baydin et al. - 综述文章
- PyTorch autograd文档和技术博客
- JAX autodiff cookbook - 展示高阶导数和高级用法
- "Reverse-Mode Automatic Differentiation: A Tutorial" - 详细的技术教程
六、实验笔记
实验1：验证自定义操作的梯度
目标：实现一个自定义的激活函数，并通过数值梯度检查验证其梯度实现的正确性
步骤：
选择一个非标准激活函数（如Swish: $x \cdot \sigma(x)$ 或 Mish: $x \cdot \tanh(\ln(1+e^x))$）
实现torch.autograd.Function，包括forward和backward
编写数值梯度检查函数
在不同输入范围内测试梯度误差
预期结果：相对误差应小于$10^{-5}$
分析要点：
- 哪些输入区域容易出现较大的梯度误差？
- 数值稳定性如何保证？
- 与内置实现的性能对比如何？
实验2：测量checkpoint的内存-时间权衡
目标：量化gradient checkpointing对显存占用和训练速度的影响
步骤：
构建一个可变深度的网络（如20-100层）
分别使用和不使用checkpoint训练
测量峰值显存占用和每轮训练时间
绘制深度vs显存、深度vs时间的曲线
预期结果：checkpoint显著减少显存占用，但增加训练时间（约20-30%）
分析要点：
- checkpoint的收益随深度如何变化？
- 哪些层适合使用checkpoint？
- 是否有最优的checkpoint策略？
七、本章小结
本章深入探讨了自动微分这一深度学习框架的核心技术。我们从基本原理出发，比较了前向模式和反向模式自动微分的优劣，分析了PyTorch的具体实现机制。通过讨论高阶导数、不可导操作的处理、控制流的支持等高级主题，我们展现了自动微分系统的丰富内涵。
自动微分不仅是工程实现的问题，更涉及深刻的数学和计算机科学原理。理解这些原理使我们能够更好地使用框架、实现新颖的操作、诊断训练问题，并在内存和计算之间做出明智的权衡。
随着深度学习的发展，自动微分技术也在不断演进。从简单的反向传播到支持任意阶导数、从命令式到函数式、从单一设备到分布式系统，自动微分将继续作为AI基础设施的核心组件发挥关键作用。
八、习题
手动执行前向模式自动微分，计算函数$f(x, y) = \sin(x) \cdot e^y$在点$(x=1, y=2)$处对$x$和$y$的偏导数。
证明反向模式自动微分计算标量函数$L$对所有输入$x_1, ..., x_n$的梯度的时间复杂度为$O(T_f)$，其中$T_f$是计算函数值的时间。
解释为什么PyTorch中叶子节点的grad_fn为None，而非叶子节点的grad_fn不为None。grad_fn.next_functions的结构是怎样的？
实现一个自定义的矩阵求逆操作的梯度函数。提示：利用公式$\frac{\partial A^{-1}}{\partial A_{ij}} = -A^{-1} E^{ij} A^{-1}$，其中$E^{ij}$是仅在$(i,j)$位置为1的矩阵。
对于一个使用Layer Normalization的Transformer块，画出其前向计算图，并标注哪些中间结果需要在反向传播时使用。估算一个batch大小为32、序列长度512、隐藏维度768的样本需要保存多少中间数据。
解释torch.no_grad()、torch.enable_grad()和torch.set_grad_enabled()的区别和使用场景。
为什么in-place操作可能破坏自动微分？PyTorch如何检测这种情况？给出一个触发检测错误的代码示例。
阅读JAX的autodiff文档，比较jax.grad、jax.value_and_grad、jax.jacfwd、jax.jacrev的功能和适用场景。
设计一个实验来测量不同长度的RNN在使用和不使用checkpoint时的显存占用和训练速度。总结checkpoint的效益如何随序列长度变化。
在元学习算法MAML中，需要计算"梯度的梯度"。详细描述MAML内层更新过程中的计算图结构，解释为什么需要create_graph=True，以及如果不设置会发生什么。