动手实践：可扩展变分推断

计算精确后验 $p(\mathbf{z}|\mathbf{x})$ 通常是难以处理的。尽管马尔可夫链蒙特卡洛 (MCMC) 提供了一种替代方法，但对于大型数据集来说可能很慢。变分推断将推断问题重新构造为一个优化问题，旨在找到一个近似分布 $q(\mathbf{z})$，使其到真实后验的KL散度最小化，这等同于最大化证据下界 (ELBO)。

$$\mathcal{L}(q) = \mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})] - \mathbb{E}_{q(\mathbf{z})} [\log q(\mathbf{z})]$$

对于大型数据集，在每个优化步骤中计算关于整个数据集 $\mathbf{x}$ 的期望 $\mathbb{E}_{q(\mathbf{z})} [\log p(\mathbf{x}, \mathbf{z})]$ 计算成本过高。随机变分推断 (SVI) 通过使用小批量数据来估计 ELBO 的梯度来处理这个问题，从而使我们能够使用 Adam 等随机优化技术。黑箱变分推断 (BBVI) 通过提供一种通用的梯度计算方法，进一步扩展了这种能力，即使 ELBO 梯度无法解析计算，它也能依赖于得分函数估计器。

在这个实践练习中，我们将使用 Pyro 概率编程库为贝叶斯逻辑回归模型实现 SVI。这将展示如何有效处理更大的数据集。

目标

在一个模拟数据集上，使用随机变分推断 (SVI) 实现并训练贝叶斯逻辑回归模型。我们将在训练期间监测 ELBO，并检查模型参数 (parameter)的学习到的近似后验分布。

设置

首先，确保你已安装所需的库。你主要需要 torch 和 pyro-ppl。

# 标准库
import torch
import torch.nn as nn
from torch.distributions import constraints
import numpy as np
import matplotlib.pyplot as plt

# Pyro 库
import pyro
import pyro.distributions as dist
from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam

# 确保结果可复现
pyro.set_rng_seed(101)

生成模拟数据

让我们创建一个简单的二分类数据集。我们将模拟数据，其中分类线性依赖于两个特征，并经过 sigmoid 函数处理。

# 生成合成数据
N = 1000  # 数据点数量
D = 2     # 特征数量
true_weights = torch.tensor([1.5, -2.0])
true_bias = torch.tensor([0.5])

# 生成特征（正态分布）
X_data = torch.randn(N, D)

# 计算线性组合和概率
linear_combination = X_data @ true_weights + true_bias
prob = torch.sigmoid(linear_combination)

# 根据概率生成二元标签
y_data = torch.bernoulli(prob).float() # 确保标签是浮点数，以满足潜在的损失计算需求

print(f"Generated X_data shape: {X_data.shape}")
print(f"Generated y_data shape: {y_data.shape}")

定义贝叶斯逻辑回归模型

在 Pyro 中，我们将概率模型定义为一个以数据为输入的 Python 函数。在此函数内部，我们使用 pyro.sample 语句来声明随机变量（带有先验的参数 (parameter)和观测似然）。

def logistic_regression_model(features, labels):
    # 权重和偏置的先验
    # 假设标准正态先验（均值=0，标准差=1）
    # 权重的形状应与特征数量 D 匹配
    weights = pyro.sample("weights", dist.Normal(torch.zeros(D), torch.ones(D)))
    # 偏置是一个单一的标量值
    bias = pyro.sample("bias", dist.Normal(0., 1.))

    # 定义线性组合（logits）
    # 使用 pyro.plate 表示数据点之间的条件独立性
    with pyro.plate("data", size=features.shape[0]):
        logits = torch.sigmoid(features @ weights + bias)
        # 使用伯努利分布定义似然
        # obs=labels 将分布与观测数据关联起来
        pyro.sample("obs", dist.Bernoulli(logits), obs=labels)

这个 logistic_regression_model 函数定义了生成过程：从先验中抽取权重 (weight)和偏置 (bias)，使用逻辑函数计算每个数据点的概率，然后根据这些概率从伯努利分布中抽取观测标签。

定义引导函数（变分分布）

引导函数 $q(\mathbf{z})$ 指定了我们将用于近似后验的分布族。对于平均场变分推断 (VI)，我们假设潜在变量（在本例中是权重 (weight)和偏置 (bias)）在近似后验中是独立的。我们通常使用具有可学习参数 (parameter)的分布（例如，均值和标准差是待优化的参数的正态分布）。

def mean_field_guide(features, labels):
    # 定义权重的变分参数
    # 'weight_loc' 和 'weight_scale' 将在优化过程中学习
    weight_loc = pyro.param("weight_loc", torch.randn(D))
    # 尺度参数必须为正，因此我们使用 softplus 变换和约束
    weight_scale = pyro.param("weight_scale", torch.randn(D))
    weight_scale_constrained = torch.nn.functional.softplus(weight_scale)

    # 定义偏置的变分参数
    bias_loc = pyro.param("bias_loc", torch.randn(1))
    bias_scale = pyro.param("bias_scale", torch.randn(1))
    bias_scale_constrained = torch.nn.functional.softplus(bias_scale)

    # 从引导分布中抽样潜在变量
    # 这些名称必须与模型函数中使用的名称匹配
    pyro.sample("weights", dist.Normal(weight_loc, weight_scale_constrained))
    pyro.sample("bias", dist.Normal(bias_loc, bias_scale_constrained))

在这里，pyro.param 声明可学习参数。weight_loc、weight_scale、bias_loc 和 bias_scale 分别是我们用于权重和偏置的近似正态分布的参数。我们使用 softplus 函数来确保尺度参数保持正值。

使用 SVI 设置推断

现在我们配置 SVI 算法。我们需要一个优化器（Adam 常用）和损失函数 (loss function) (Trace_ELBO，它计算负 ELBO 的蒙特卡洛估计)。

# 设置优化器
optimizer = Adam({"lr": 0.01}) # 学习率

# 设置推断算法
svi = SVI(model=logistic_regression_model,
          guide=mean_field_guide,
          optim=optimizer,
          loss=Trace_ELBO())

训练循环

我们运行 SVI 优化循环。在每一步中，svi.step 使用小批量数据计算 ELBO 梯度的估计值，并通过优化器更新引导函数中定义的参数 (parameter)（如 weight_loc、weight_scale 等）。

# 训练配置
num_iterations = 2000
batch_size = 100
n_data = X_data.shape[0]
elbo_history = []

# 在开始训练前清除 Pyro 的参数存储
pyro.clear_param_store()

print("SVI 训练开始...")
for j in range(num_iterations):
    # 创建小批量索引
    indices = torch.randperm(n_data)[:batch_size]
    mini_batch_features = X_data[indices]
    mini_batch_labels = y_data[indices]

    # 计算损失并执行梯度步
    loss = svi.step(mini_batch_features, mini_batch_labels)
    elbo_history.append(-loss) # 存储 ELBO（负损失）

    if j % 200 == 0:
        print(f"[Iteration {j+1}/{num_iterations}] ELBO: {-loss:.2f}")

print("SVI 训练完成。")

结果与评估

训练结束后，Pyro 参数 (parameter)存储会包含引导函数的参数优化值。我们可以访问它们以了解近似后验。

# 获取引导函数的学习参数
learned_params = {}
for name, param_val in pyro.get_param_store().items():
    learned_params[name] = param_val.detach().numpy()

print("\n学习到的变分参数：")
print(f"权重均值 (loc)：{learned_params['weight_loc']}")
print(f"权重标准差 (scale)：{np.log(1 + np.exp(learned_params['weight_scale']))}") # 应用 softplus 逆变换
print(f"偏置均值 (loc)：{learned_params['bias_loc']}")
print(f"偏置标准差 (scale)：{np.log(1 + np.exp(learned_params['bias_scale']))}")

# 与真实值比较
print(f"\n真实权重：{true_weights.numpy()}")
print(f"真实偏置：{true_bias.numpy()}")

你应该会看到，学习到的均值（weight_loc、bias_loc）与用于生成数据的真实权重 (weight)和偏置 (bias)相当接近。学习到的标准差（经过变换的 weight_scale，经过变换的 bias_scale）为我们提供了这些参数不确定性的度量。

可视化 ELBO 收敛情况

让我们绘制训练期间记录的 ELBO 值。我们期望 ELBO 增加并趋于平稳，这表明优化收敛。

# 绘制 ELBO 历史
plt.figure(figsize=(10, 4))
plt.plot(elbo_history)
plt.title("SVI 训练期间的 ELBO 收敛")
plt.xlabel("迭代次数")
plt.ylabel("ELBO")
plt.grid(True)
plt.show()

证据下界 (ELBO) 在训练期间通常会增加，这表明变分分布正在更好地近似真实后验。随着优化收敛，曲线通常会趋于平坦。

可视化后验近似

我们还可以将权重的近似后验分布与它们的真实值进行可视化比较。

# 可视化权重的近似后验
from scipy.stats import norm

# 创建子图
fig, axs = plt.subplots(1, D, figsize=(12, 5), sharey=True)
fig.suptitle('权重的近似后验分布')

for i in range(D):
    # 获取当前权重的均值和标准差
    loc = learned_params['weight_loc'][i]
    scale = np.log(1 + np.exp(learned_params['weight_scale'][i])) # 应用 softplus

    # 生成均值附近的 x 值
    x = np.linspace(loc - 3*scale, loc + 3*scale, 200)
    # 计算概率密度函数 (PDF)
    y = norm.pdf(x, loc, scale)

    # 绘制 PDF
    axs[i].plot(x, y, label=f'近似后验 q(w_{i})', color='#1c7ed6')
    # 绘制真实权重值作为垂直线
    axs[i].axvline(true_weights[i].numpy(), color='#f03e3e', linestyle='--', label=f'真实 w_{i}')
    axs[i].set_title(f'权重 {i+1}')
    axs[i].set_xlabel('值')
    axs[i].legend()
    axs[i].grid(True, linestyle=':', alpha=0.6)

axs[0].set_ylabel('概率密度')
plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # 调整布局以避免标题重叠
plt.show()

通过 SVI 获得的模型权重的近似后验分布。图表显示了以真实参数值（红色虚线）为中心学习到的正态分布（蓝色曲线），其宽度表示了变分近似捕捉到的不确定性。

讨论

这个动手示例展示了 SVI 的核心工作流程：定义模型和引导函数，设置 SVI 优化器和损失，以及使用小批量数据运行训练循环。SVI 使我们能够将贝叶斯推断应用于比精确方法或标准 MCMC 可处理的更大的数据集，这得益于其随机优化方法。

我们在这里使用了简单的平均场引导函数。对于更复杂的后验几何结构，尝试更高级的变分族（在第 3.7 节中简要介绍）可能会产生更好的近似效果，但这通常会增加优化过程中的计算复杂性。SVI、BBVI 和 MCMC 方法之间的选择通常取决于具体模型、数据集大小、所需精度和可用计算资源，这是我们在第 3.8 节中考察过的一个权衡问题。