自动微分

自动微分#

在本节中，您将学习 JAX 中自动微分（autodiff）的基础应用。JAX 拥有一个非常通用的自动微分系统。计算梯度是现代机器学习方法中的关键部分，本教程将引导您了解一些自动微分的入门主题，例如：

1. 使用 jax.grad 计算梯度
2. 在线性逻辑回归中计算梯度
3. 对嵌套列表、元组和字典进行微分
4. 使用 jax.value_and_grad 同时评估函数值及其梯度
5. 通过数值差分进行校验

请务必查看 “高级自动微分”指南以了解更深入的主题。

虽然理解自动微分的“底层原理”对于在大多数场景下使用 JAX 并非必不可少，但仍鼓励您观看这个非常易懂的视频，以更深入地了解其运行机制。

1. 使用 `jax.grad` 计算梯度#

在 JAX 中，您可以使用 jax.grad() 变换对标量函数进行微分。

import jax
import jax.numpy as jnp
from jax import grad

grad_tanh = grad(jnp.tanh)
print(grad_tanh(2.0))

0.070650816

jax.grad() 接收一个函数并返回一个函数。如果您有一个计算数学函数 \(f\) 的 Python 函数 f，那么 jax.grad(f) 就是一个计算数学函数 \(\nabla f\) 的 Python 函数。这意味着 grad(f)(x) 表示 \(\nabla f(x)\) 的值。

由于 jax.grad() 作用于函数，您可以将其应用于自身的输出，从而进行任意次数的微分。

print(grad(grad(jnp.tanh))(2.0))
print(grad(grad(grad(jnp.tanh)))(2.0))

-0.13621868
0.25265405

JAX 的自动微分使得计算高阶导数变得简单，因为计算导数的函数本身也是可微的。因此，高阶导数就像堆叠变换一样简单。这可以通过单变量情况来说明：

函数 \(f(x) = x^3 + 2x^2 - 3x + 1\) 的导数可以计算为：

f = lambda x: x**3 + 2*x**2 - 3*x + 1

dfdx = jax.grad(f)

\(f\) 的高阶导数为：

\[\begin{split} \begin{array}{l} f'(x) = 3x^2 + 4x -3\\ f''(x) = 6x + 4\\ f'''(x) = 6\\ f^{iv}(x) = 0 \end{array} \end{split}\]

在 JAX 中计算这些导数，只需链式调用 jax.grad() 函数即可。

d2fdx = jax.grad(dfdx)
d3fdx = jax.grad(d2fdx)
d4fdx = jax.grad(d3fdx)

在 \(x=1\) 处评估上述导数将得到：

\[\begin{split} \begin{array}{l} f'(1) = 4\\ f''(1) = 10\\ f'''(1) = 6\\ f^{iv}(1) = 0 \end{array} \end{split}\]

使用 JAX

print(dfdx(1.))
print(d2fdx(1.))
print(d3fdx(1.))
print(d4fdx(1.))

2. 在线性逻辑回归中计算梯度#

下一个示例展示了如何在线性逻辑回归模型中使用 jax.grad() 计算梯度。首先是设置：

key = jax.random.key(0)

def sigmoid(x):
  return 0.5 * (jnp.tanh(x / 2) + 1)

# Outputs probability of a label being true.
def predict(W, b, inputs):
  return sigmoid(jnp.dot(inputs, W) + b)

# Build a toy dataset.
inputs = jnp.array([[0.52, 1.12,  0.77],
                    [0.88, -1.08, 0.15],
                    [0.52, 0.06, -1.30],
                    [0.74, -2.49, 1.39]])
targets = jnp.array([True, True, False, True])

# Training loss is the negative log-likelihood of the training examples.
def loss(W, b):
  preds = predict(W, b, inputs)
  label_probs = preds * targets + (1 - preds) * (1 - targets)
  return -jnp.sum(jnp.log(label_probs))

# Initialize random model coefficients
key, W_key, b_key = jax.random.split(key, 3)
W = jax.random.normal(W_key, (3,))
b = jax.random.normal(b_key, ())

使用 jax.grad() 函数及其 argnums 参数，可以针对位置参数对函数进行微分。

# Differentiate `loss` with respect to the first positional argument:
W_grad = grad(loss, argnums=0)(W, b)
print(f'{W_grad=}')

# Since argnums=0 is the default, this does the same thing:
W_grad = grad(loss)(W, b)
print(f'{W_grad=}')

# But you can choose different values too, and drop the keyword:
b_grad = grad(loss, 1)(W, b)
print(f'{b_grad=}')

# Including tuple values
W_grad, b_grad = grad(loss, (0, 1))(W, b)
print(f'{W_grad=}')
print(f'{b_grad=}')

W_grad=Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32)
W_grad=Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32)
b_grad=Array(-0.69001776, dtype=float32)
W_grad=Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32)
b_grad=Array(-0.69001776, dtype=float32)

jax.grad() API 与 Spivak 的经典著作《流形上的微积分》（1965）中的优秀符号体系直接对应，该体系也用于 Sussman 和 Wisdom 的《经典力学的结构与解释》（2015）及其《函数微分几何》（2013）。这两本书均为开放获取。特别参阅《函数微分几何》的“序言”部分，其中对该符号体系进行了辩护。

本质上，当使用 argnums 参数时，如果 f 是用于评估数学函数 \(f\) 的 Python 函数，则 Python 表达式 jax.grad(f, i) 将评估为用于计算 \(\partial_i f\) 的 Python 函数。

3. 对嵌套列表、元组和字典进行微分#

由于 JAX 的 PyTree 抽象（参见 Pytrees），针对标准 Python 容器的微分可以直接工作，因此您可以随心所欲地使用元组、列表和字典（以及任意嵌套结构）。

接续上一个示例：

def loss2(params_dict):
    preds = predict(params_dict['W'], params_dict['b'], inputs)
    label_probs = preds * targets + (1 - preds) * (1 - targets)
    return -jnp.sum(jnp.log(label_probs))

print(grad(loss2)({'W': W, 'b': b}))

{'W': Array([-0.433146 , -0.7354605, -1.2598922], dtype=float32), 'b': Array(-0.69001776, dtype=float32)}

您可以创建自定义 pytree 节点，使其不仅能与 jax.grad() 一起使用，还能与其他 JAX 变换（jax.jit(), jax.vmap() 等）配合使用。

4. 使用 `jax.value_and_grad` 同时评估函数值及其梯度#

另一个便利函数是 jax.value_and_grad()，用于高效地在一次计算中同时获取函数值及其梯度值。

接续之前的示例：

loss_value, Wb_grad = jax.value_and_grad(loss, (0, 1))(W, b)
print('loss value', loss_value)
print('loss value', loss(W, b))

loss value 2.9729187
loss value 2.9729187

5. 通过数值差分进行校验#

导数的一大优点是可以通过有限差分轻松进行校验。

接续之前的示例：

# Set a step size for finite differences calculations
eps = 1e-4

# Check b_grad with scalar finite differences
b_grad_numerical = (loss(W, b + eps / 2.) - loss(W, b - eps / 2.)) / eps
print('b_grad_numerical', b_grad_numerical)
print('b_grad_autodiff', grad(loss, 1)(W, b))

# Check W_grad with finite differences in a random direction
key, subkey = jax.random.split(key)
vec = jax.random.normal(subkey, W.shape)
unitvec = vec / jnp.sqrt(jnp.vdot(vec, vec))
W_grad_numerical = (loss(W + eps / 2. * unitvec, b) - loss(W - eps / 2. * unitvec, b)) / eps
print('W_dirderiv_numerical', W_grad_numerical)
print('W_dirderiv_autodiff', jnp.vdot(grad(loss)(W, b), unitvec))

b_grad_numerical -0.6890297
b_grad_autodiff -0.69001776
W_dirderiv_numerical 1.3041496
W_dirderiv_autodiff 1.3006744

JAX 提供了一个简单的便利函数，其功能基本相同，但可以校验任意阶数的导数：

from jax.test_util import check_grads

check_grads(loss, (W, b), order=2)  # check up to 2nd order derivatives

下一步#

“高级自动微分”指南提供了关于本文档中涵盖的概念如何在 JAX 后端实现的高级详细说明。某些功能（例如针对 JAX 可变换 Python 函数的自定义导数规则）依赖于对高级自动微分的理解，因此如果您感兴趣，请务必查看该部分。