PyTorch 核心机制：自动微分与雅可比向量积详解

1. 自动微分概述

自动微分（Automatic Differentiation, AD）是现代深度学习框架的基石。与数值微分和符号微分不同，自动微分利用链式法则，通过记录计算过程中的中间变量，精确地计算导数。在 PyTorch 中，这一功能由 autograd 模块提供。

当创建一个张量并设置 requires_grad=True 时，PyTorch 会为该张量构建一个动态计算图（Dynamic Computation Graph）。每一个操作都会被记录，形成一个有向无环图（DAG）。这使得我们可以方便地进行反向传播，计算梯度。

1.1 基本示例

以下是一个简单的标量函数求导示例：

import torch

# 创建张量，启用梯度追踪
x = torch.tensor([2.0], requires_grad=True)

# 定义函数 y = x^2
y = x ** 2

# 执行反向传播
y.backward()

# 获取梯度
print(x.grad)  # 输出：tensor([4.])

在此例中，系统自动计算了 $y$ 对 $x$ 的偏导数 $rac{ ext{d}y}{ ext{d}x} = 2x$。当 $x=2$ 时，梯度为 4。

2. 雅可比矩阵与向量积

在深度学习中，我们处理的通常是多输入多输出的函数。假设有一个函数 $F: ext{R}^n o ext{R}^m$，其雅可比矩阵（Jacobian Matrix）$J$ 定义为：

$$ J_{ij} = rac{ ext{d}F_i}{ ext{d}x_j} $$

其中 $i$ 表示输出维度，$j$ 表示输入维度。完整的雅可比矩阵大小为 $m imes n$。对于复杂的神经网络，这个矩阵可能极其庞大，直接计算和存储往往是不现实的。

因此，我们通常关注两种乘积形式：

雅可比向量积（JVP）：$J imes v$，其中 $v$ 是长度为 $n$ 的向量。结果是一个长度为 $m$ 的向量。这代表了函数在某一点沿方向 $v$ 的变化率。
向量雅可比积（VJP）：$v^T imes J$，其中 $v$ 是长度为 $m$ 的向量。结果是一个长度为 $n$ 的向量。这正是反向传播算法所计算的梯度形式。

PyTorch 的 backward() 方法默认计算的是 VJP（即梯度）。而 torch.autograd.grad 允许我们显式地控制计算 JVP 或其他形式的导数。

3. PyTorch 中的雅可比向量积实现

虽然 backward() 是最常用的接口，但在某些高级优化算法（如二阶优化、Hessian 矩阵近似）中，我们需要直接访问 JVP。

3.1 使用 `torch.autograd.grad`

torch.autograd.grad 函数允许我们指定输出张量和输入张量，以及可选的 grad_outputs（即向量 $v$）。

import torch

# 定义输入
x = torch.tensor([2.0], requires_grad=True)

# 定义函数 y = x^2
y = x ** 


v = torch.tensor([])



Jv = torch.autograd.grad(y, x, v, retain_graph=)[]

(Jv)

PyTorch 核心机制：自动微分与雅可比向量积详解