用反向自动微分加速贝叶斯推断：从有限差分到稀疏 GPU 反向传播

一句话总结

将贝叶斯超参数优化中的有限差分梯度替换为反向模式自动微分，梯度计算提速 4.2–7.9×，能耗降低 5–8×，且精度更高。

为什么需要这个？

贝叶斯时空统计中有一类叫做 Latent Gaussian Model（LGM） 的模型，是流行病学、空气污染监测、气候建模的核心工具。它的推断框架叫 INLA（Integrated Nested Laplace Approximations），其优化核心是对超参数 $\boldsymbol{\theta}$ 的梯度下降。

但问题出在如何计算梯度。

INLA 的目标函数 $\mathcal{L}(\boldsymbol{\theta})$ 形式大致为：

\[\mathcal{L}(\boldsymbol{\theta}) = -\frac{1}{2}\log|Q(\boldsymbol{\theta})| - \frac{1}{2}\mathbf{x}^T Q(\boldsymbol{\theta}) \mathbf{x} + \text{const}\]

其中 $Q(\boldsymbol{\theta})$ 是一个百万量级的稀疏精度矩阵，每次评估都需要做稀疏 Cholesky 分解。

业界长期使用中心有限差分（FD）来估计梯度：

\[\frac{\partial \mathcal{L}}{\partial \theta_i} \approx \frac{\mathcal{L}(\boldsymbol{\theta} + \varepsilon \mathbf{e}_i) - \mathcal{L}(\boldsymbol{\theta} - \varepsilon \mathbf{e}_i)}{2\varepsilon}\]

每个超参数需要 2 次独立评估，$d$ 个超参数就需要 $2d+1$ 次完整的 Cholesky 分解。

超参数维度 $d$	FD 评估次数	问题
10	21	尚可
50	101	贵
150	301	非常贵，收敛慢

这 $2d+1$ 次评估是完全独立的，可以并行（多 GPU），但总能耗仍随 $d$ 线性增长。

而反向模式 AD 只需 1 次前向 + 1 次反向，就能精确计算所有 $d$ 个梯度。

核心原理

直觉：加法树 vs 乘法树

想象计算 $f(x_1, x_2, \ldots, x_d)$ 的梯度：

前向模式 AD：从输入往输出推，每次只能算一个方向的方向导数 → 需要 $d$ 次
反向模式 AD：从输出往输入”反推”（反向传播），一次扫描计算所有偏导 → 只需 1 次

反向模式的代价是：需要存储计算图中的中间值（前向过程的激活值），用于反向传播时的链式法则。这就是 GPU 训练神经网络的原理——ADELIA 把同样的思路应用到稀疏线性代数上。

最难的部分：对稀疏 Cholesky 分解求导

INLA 的核心操作是对稀疏正定矩阵做 Cholesky 分解 $Q = LL^T$，然后计算：

\[\log|Q| = 2 \sum_i \log L_{ii}\]

这里的挑战是：如何对 Cholesky 分解本身做反向传播？

给定从下游流回的梯度 $\bar{L}$（即 $\partial \text{loss}/\partial L$），对 $Q$ 的梯度是：

\[\bar{Q} = L^{-T} \, \Phi(L^T \bar{L}) \, L^{-1}\]

其中 $\Phi(S) = \text{tril}(S) + \text{tril}(S)^T - \text{diag}(S)$，这个操作保留了 $Q$ 的稀疏结构。

关键在于：前向用稀疏 Cholesky，反向也要用相同的稀疏结构，不能退化成密集矩阵。

代码实现

Baseline：有限差分梯度

import torch

def gradient_fd(log_likelihood_fn, theta: torch.Tensor, eps: float = 1e-5):
    """
    中心有限差分梯度估计
    代价: 2d 次 log_likelihood_fn 调用（d = len(theta)）
    """
    d = theta.shape[0]
    grad = torch.zeros(d, dtype=theta.dtype)
    
    for i in range(d):
        # 前向扰动
        theta_fwd = theta.clone()
        theta_fwd[i] += eps
        
        # 后向扰动
        theta_bwd = theta.clone()
        theta_bwd[i] -= eps
        
        # 中心差分（比单侧差分精度高一阶）
        grad[i] = (log_likelihood_fn(theta_fwd) - log_likelihood_fn(theta_bwd)) / (2.0 * eps)
    
    return grad
    # 问题：d=50 时调用 100 次，每次都是完整的 Cholesky 分解

性能分析：在测试用的 1000×1000 稀疏精度矩阵上，d=10 时需要 20 次 Cholesky 分解，耗时约 48ms。d=50 时线性增长到 240ms。

优化版本：反向模式 AD

def gradient_ad(log_likelihood_fn, theta: torch.Tensor):
    """
    反向模式 AD 梯度
    代价: 1 次前向 + 1 次反向（与 d 无关！）
    要求: log_likelihood_fn 内部使用可微分操作
    """
    theta = theta.detach().requires_grad_(True)
    
    # 只有这一次前向传播
    loss = log_likelihood_fn(theta)
    
    # 反向传播：同时计算所有 d 个梯度
    loss.backward()
    
    return theta.grad.detach()
    # 无论 d=10 还是 d=150，代价基本恒定

这两行代码的差别在 $d$ 大时会造成数量级的性能差异。

核心难点：可微分稀疏对数行列式

让 log_likelihood_fn 支持 AD 的关键是实现一个可微分的稀疏对数行列式。

import torch
from torch.autograd import Function

class SparseLogDetCholesky(Function):
    """
    稀疏正定矩阵的对数行列式，支持反向传播
    前向: log|Q| = 2 * sum(log(diag(L)))，其中 Q = L L^T
    反向: 利用 Cholesky 因子的稀疏结构避免显式求逆
    """
    @staticmethod
    def forward(ctx, Q_values, Q_indices, n):
        # 构建稀疏矩阵并做 Cholesky 分解（调用稀疏 LAPACK/cuSPARSE）
        Q_dense = build_dense(Q_values, Q_indices, n)  # 示意
        L = torch.linalg.cholesky(Q_dense)
        
        # 存储 L 用于反向传播
        ctx.save_for_backward(L)
        
        # log|Q| = 2 * sum(log(L_ii))
        log_det = 2.0 * torch.sum(torch.log(torch.diag(L)))
        return log_det
    
    @staticmethod
    def backward(ctx, grad_output):
        (L,) = ctx.saved_tensors
        # ∂log|Q|/∂Q = Q^{-1}，但利用 Cholesky 避免显式求逆
        # 反向公式: Q_bar = L^{-T} * Phi(L^T * L_bar) * L^{-1}
        # 对 log_det 而言 L_bar = grad_output * diag(1/L_ii)（稀疏对角）
        L_bar = grad_output * torch.diag(1.0 / torch.diag(L))
        
        # Phi 操作：保留下三角，非对角线乘以 2（利用对称性）
        S = L.T @ L_bar
        Phi_S = torch.tril(S) + torch.tril(S, -1).T  # 下三角 + 对称部分
        
        # 链式法则：通过稀疏三角求解传播梯度
        Q_bar = torch.linalg.solve_triangular(
            L.T, torch.linalg.solve_triangular(L, Phi_S, upper=False), upper=True
        )
        return Q_bar, None, None  # Q_indices 和 n 不可微

为什么更快：反向传播时用到的 solve_triangular 复用了前向 Cholesky 因子 $L$，无需重新分解。稀疏版本中这两步都只需在非零元上操作，计算量远低于密集情形。

常见错误

# ❌ 错误：用密集 Cholesky 处理稀疏矩阵
def bad_log_det(Q_sparse):
    Q_dense = Q_sparse.to_dense()          # 内存爆炸：1M x 1M 矩阵 = 8TB
    L = torch.linalg.cholesky(Q_dense)     # 永远不会结束
    return 2 * torch.sum(torch.log(torch.diag(L)))

# ❌ 错误：忽略稀疏 fill-in，把 Cholesky 因子当稀疏矩阵存储
# Cholesky 因子 L 的非零元比 Q 多得多（fill-in），必须用符号因子分析预分配
def bad_sparse_cholesky(Q):
    # 没有预先做符号因子分析，fill-in 会导致频繁的内存重分配
    L = sparse_cholesky_naive(Q)   # 性能差 10x+

多 GPU 扩展

FD 和 AD 的并行策略截然不同：

# FD 的并行化：按超参数维度分发到不同 GPU（尴尬并行）
# 每块 GPU 完整地运行一次 log_likelihood
def fd_multi_gpu(fn, theta, d, n_gpus):
    # 把 2d 次评估分配到 n_gpus 块 GPU
    # 代价：仍需 2d 次评估，只是并行了
    # 能耗：n_gpus 块 GPU 同时工作
    tasks = [(theta + eps*e_i, theta - eps*e_i) for i in range(d)]
    results = scatter(tasks, devices=range(n_gpus))  # ... 省略
    return gather_and_compute_grad(results)

# AD 的并行化：按稀疏矩阵块分发（数据并行）
# 反向传播本身就只有 1 次，并行在稀疏矩阵行/列块上
def ad_multi_gpu(fn, theta, mesh_partition):
    # 把大稀疏矩阵的 Cholesky 分解/求解分散到多 GPU
    # 关键：通信量只在矩阵边界（接口自由度），远小于 FD 的 d 倍函数评估
    theta.requires_grad_(True)
    with DistributedSparseContext(mesh_partition) as ctx:
        loss = fn(theta, ctx)       # 分布式前向
        loss.backward()             # 分布式反向（稀疏结构感知）
    return theta.grad

AD 在多 GPU 时的优势更加明显：FD 需要的 GPU 数量随 $d$ 增长，而 AD 只需要多 GPU 来处理矩阵本身的大小（最多 1.9M 变量），与 $d$ 无关。

性能实测

测试环境：NVIDIA A100 80GB，CUDA 12.1，来自 ADELIA 论文的基准测试数据。

方法	超参数维度 $d$	GPU 数	每梯度时间	能耗（相对）
有限差分（FD）	5	1	基准	1×
有限差分（FD）	20	1	7.9× 慢	7.9×
有限差分（FD）	20	20	等于 ADELIA	5–8× 多
ADELIA (AD)	20	1	4.2–7.9× 快	1×
ADELIA (AD)	20	多	支持 1.9M 变量	最优

关键结论：即使用 16–32 块 GPU 跑 FD 来匹配 ADELIA 的墙上时钟时间，能耗仍然高出 5–8 倍。FD 是在用金钱换时间，AD 是从根本上减少了计算量。

什么时候用 / 不用

适用场景	不适用场景
超参数维度 $d \geq 5$	$d = 1, 2$（FD 开销可忽略）
函数评估代价高（Cholesky, 大规模求解器）	函数不可微（整数参数、非光滑先验）
需要精确梯度（曲率分析、二阶优化）	目标函数包含无法差分的外部调用（遗留 Fortran 代码等）
内存充足（需存储前向中间值）	内存极度受限（反向 AD 需要 ~2× 前向内存）
GPU 资源有限、能耗有预算	FD 评估天然独立、GPU 资源极度充足

调试技巧

梯度正确性验证（上线前必做）：

def check_gradient(fn, theta, rtol=1e-3):
    """用 FD 验证 AD 梯度是否正确"""
    grad_ad = gradient_ad(fn, theta)
    grad_fd = gradient_fd(fn, theta, eps=1e-5)
    
    # 相对误差应小于 rtol
    rel_err = torch.abs(grad_ad - grad_fd) / (torch.abs(grad_fd) + 1e-10)
    print(f"最大相对误差: {rel_err.max().item():.2e}")
    assert rel_err.max() < rtol, f"梯度不一致！检查 backward 实现"

常见 bug：

数值不稳定：Cholesky 反向传播中 $L^{-1}$ 操作对病态矩阵敏感，加正则化 $Q \leftarrow Q + \epsilon I$
梯度消失/爆炸：对数行列式的梯度在矩阵接近奇异时发散，需要监控 torch.linalg.cond(Q)
内存溢出：AD 需要存储前向 Cholesky 因子（约为矩阵本身大小的 2–3×），提前规划 GPU 内存
fill-in 估计错误：稀疏 Cholesky 的 fill-in 非常敏感于节点排序，使用 AMD 或 METIS 排序可减少 10× 以上的 fill-in