从手动推导到自动求导：一个简单线性回归的JAX实现，带你吃透自动微分的数学本质-编程阁

从手动推导到自动求导：一个简单线性回归的JAX实现，带你吃透自动微分的数学本质

在机器学习的实践中，我们常常会听到"自动微分"这个术语。它像一位隐形的助手，默默地在背后计算着梯度，驱动着模型的参数更新。但你是否曾好奇过，这位助手究竟是如何工作的？本文将从一个简单的线性回归模型出发，先手动推导其梯度公式，再借助JAX这一现代工具实现自动微分，通过对比两者结果，揭示自动微分背后的数学本质。

1. 线性回归模型与手动梯度推导

线性回归是机器学习中最基础的模型之一，其数学表达式为：

y_pred = w * x + b

其中，w是权重，b是偏置，x是输入特征，y_pred是预测值。我们的目标是找到最优的w和b，使得预测值尽可能接近真实值y。

1.1 损失函数的定义

常用的损失函数是均方误差(MSE)：

def loss_fn(w, b, x, y): y_pred = w * x + b return ((y_pred - y) ** 2).mean()

1.2 手动计算梯度

为了最小化损失函数，我们需要计算其对参数w和b的梯度。根据微积分知识：

对w的偏导数： $$\frac{\partial L}{\partial w} = \frac{2}{N}\sum_{i=1}^N (w x_i + b - y_i) x_i$$
对b的偏导数： $$\frac{\partial L}{\partial b} = \frac{2}{N}\sum_{i=1}^N (w x_i + b - y_i)$$

注意：这里的N是样本数量，求和是对所有样本进行的。

2. 引入JAX实现自动微分

JAX是一个结合了NumPy风格接口和自动微分功能的Python库。它提供了grad函数，可以自动计算任意函数的导数。

2.1 基本使用

import jax import jax.numpy as jnp # 定义损失函数 def loss_fn(params, x, y): w, b = params y_pred = w * x + b return jnp.mean((y_pred - y) ** 2) # 获取梯度函数 grad_fn = jax.grad(loss_fn)

2.2 梯度计算对比

让我们用具体数据来验证手动推导和自动微分的结果是否一致：

# 生成数据 x = jnp.array([1.0, 2.0, 3.0]) y = jnp.array([2.0, 4.0, 6.0]) params = (1.0, 0.0) # w=1.0, b=0.0 # 自动微分计算梯度 auto_grad = grad_fn(params, x, y) # 手动计算梯度 def manual_grad(params, x, y): w, b = params N = len(x) dw = 2/N * jnp.sum((w * x + b - y) * x) db = 2/N * jnp.sum(w * x + b - y) return (dw, db) manual_grad_val = manual_grad(params, x, y)

比较结果会发现auto_grad和manual_grad_val完全一致，验证了自动微分的正确性。

3. 自动微分的数学原理

自动微分既不是符号微分，也不是数值微分，而是一种基于计算图和链式法则的精确微分方法。

3.1 计算图的概念

任何计算都可以表示为计算图。以我们的线性回归为例：

输入x → 乘法(w) → 加法(b) → 减法(y) → 平方 → 平均 → 输出L

3.2 前向模式与反向模式

自动微分有两种主要模式：

前向模式：沿着计算图正向传播，同时计算函数值和导数
反向模式：先正向计算函数值，再反向传播导数（深度学习框架常用）

JAX主要使用反向模式自动微分，这也是为什么我们调用jax.grad就能得到梯度。

3.3 向量-雅可比积(VJP)

反向模式自动微分的核心是向量-雅可比积。对于函数$f: ℝ^n → ℝ^m$，其雅可比矩阵$J$是一个$m×n$矩阵。反向模式计算的是：

$$ v^T J $$

其中$v$通常是标量函数对输出的梯度（在我们的例子中就是1）。

4. JAX自动微分的高级特性

JAX提供了比传统深度学习框架更灵活的自动微分功能。

4.1 高阶导数

JAX可以轻松计算高阶导数：

# 计算二阶导数 hessian_fn = jax.grad(jax.grad(loss_fn)) hessian = hessian_fn(params, x, y)

4.2 自定义导数规则

可以定义自定义函数的导数规则：

@jax.custom_jvp def custom_fn(x): return x * x @custom_fn.defjvp def custom_fn_jvp(primals, tangents): x, = primals dx, = tangents primal_out = custom_fn(x) tangent_out = 2 * x * dx return primal_out, tangent_out

4.3 批处理与向量化

JAX的vmap可以自动向量化函数，处理批量数据：

batch_loss_fn = jax.vmap(loss_fn, in_axes=(None, 0, 0))

5. 实际应用中的注意事项

虽然自动微分强大，但在实际应用中仍需注意以下几点：

数值稳定性：某些数学表达式可能导致数值不稳定，即使数学上正确
内存消耗：反向模式需要存储中间结果，可能消耗大量内存
控制流处理：循环和条件语句需要特殊处理

提示：在JAX中，使用jax.lax.cond和jax.lax.while_loop等函数来处理控制流，而不是Python原生控制结构。

6. 性能优化技巧

为了充分发挥JAX自动微分的性能，可以考虑以下优化：

JIT编译：使用jax.jit加速计算

@jax.jit def jitted_loss_fn(params, x, y): return loss_fn(params, x, y)

设备放置：明确指定计算设备

with jax.default_device(jax.devices('gpu')[0]): # GPU计算

并行计算：利用pmap进行多设备并行

from jax import pmap parallel_grad = pmap(grad_fn)

7. 扩展应用：超越简单线性回归

理解了自动微分的原理后，我们可以将其应用到更复杂的模型中：

神经网络：自动计算各层参数的梯度
物理模拟：求解微分方程
概率模型：变分推断中的梯度估计
优化问题：约束优化的梯度计算

在实际项目中，我发现自动微分特别适合原型开发阶段。它让我们能够快速尝试不同的模型结构，而无需手动推导复杂的梯度公式。特别是在研究新型神经网络架构时，自动微分大大提高了实验效率。

从手动推导到自动求导：一个简单线性回归的JAX实现，带你吃透自动微分的数学本质