有没有办法以向量化的方式从 Numpy 的多元正态分布中批量采样？

Question

我目前正在尝试通过 Numpy 运行矢量化批量多元抽样操作。我有k 形状的平均向量[N,] 对应于k 维度的协方差矩阵[N, N]，我试图从多元正态分布中返回k 形状[N,] 的绘制。

我目前有一个执行上述操作的循环，

for batch in range(batch_size):
    c[batch, :] = np.random.multivariate_normal(mean = a[batch, :], cov = b[batch, :, :])

但希望将上述内容整合到矢量化操作中。问题是np.random.multivariate_normal 只能将一维数组作为均值，将二维数组作为协方差。

我可以通过 PyTorch 的多变量普通类进行批量采样，但我正在尝试与一些预先存在的 Numpy 代码集成，并且我更愿意限制发生的转换次数。

谷歌搜索了这个question，这可以通过融化均值来解决，但在我的例子中，我没有使用相同的协方差矩阵并且不能以完全相同的方式处理事情。

非常感谢您的帮助。我认为由于参数限制，我很有可能无法使用 Numpy 分布处理批处理，但我想确保我没有遗漏任何东西。

Answer 1

我在 numpy 中找不到内置函数，但它可以通过执行协方差矩阵 Σ = LLᵀ 的 Cholesky 分解然后利用给定 i.i.d. 的向量 X 的事实来自行实现。标准正态变量，变换 LX + µ 具有协方差 Σ 和均值 µ。

这可以使用例如实现np.linalg.cholesky()（注意这个函数支持批处理模式！），还有np.random.normal()：

# cov:    (*B, D, D)
# mean:   (*B, D)
# result: (*S, *B, D)
L = np.linalg.cholesky(cov)
X = np.random.standard_normal((*S, *B, D, 1))
Y = (L @ X).reshape(*S, *B, D) + mean

这里，为了方便使用，打包成一个函数：

import numpy as np


def sample_batch_mvn(
    mean: np.ndarray,
    cov: np.ndarray,
    size: "tuple | int" = (),
) -> np.ndarray:
    """
    Batch sample multivariate normal distribution.

    Arguments:

        mean: expected values of shape (…M, D)
        cov: covariance matrices of shape (…M, D, D)
        size: additional batch shape (…B)

    Returns: samples from the multivariate normal distributions
             shape: (…B, …M, D)

    It is not required that ``mean`` and ``cov`` have the same shape
    prefix, only that they are broadcastable against each other.
    """
    mean = np.asarray(mean)
    cov = np.asarray(cov)
    size = (size, ) if isinstance(size, int) else tuple(size)
    shape = size + np.broadcast_shapes(mean.shape, cov.shape[:-1])
    X = np.random.standard_normal((*shape, 1))
    L = np.linalg.cholesky(cov)
    return (L @ X).reshape(shape) + mean

现在为了测试这个函数，我们首先需要一批好的协方差矩阵。我们将生成一对来测试一下采样性能：

# Generate N batch of D-dimensional covariance matrices C:
N = 5000
D = 2

L = np.zeros((N, D, D))
L[(..., *np.tril_indices(D))] = \
    np.random.normal(size=(N, D * (D + 1) // 2))
cov = L @ np.swapaxes(L, -1, -2)

这里用于生成协方差矩阵的方法实际上是通过对 Cholesky 因子 L 进行采样来工作的。有了这些因子的先验知识，我们当然不需要在采样函数中计算 Cholesky 分解。但是，为了测试函数的普遍适用性，我们将忘记它们，只传递协方差矩阵 C：

mean = np.zeros(2)
samples = sample_batch_mvn(mean, cov, 1000)

print(samples.shape)   # (1000, 5000, 2)

在我的 PC 上对这 500 万个 2D 向量进行采样大约需要 0.4 秒。

而且，与往常一样，绘图将花费大量精力（这里显示了 5000 个协方差矩阵中前 9 个的一些样本）：

import scipy.stats as stats
import matplotlib.pyplot as plt


fig, axs = plt.subplots(3, 3, figsize=(9, 9))
for ax, i in zip(axs.ravel(), range(5000)):
    cc = cov[i]

    xsamples = samples[:100, i, 0]
    ysamples = samples[:100, i, 1]

    xmin = xsamples.min()
    xmax = xsamples.max()
    ymin = ysamples.min()
    ymax = ysamples.max()
    xpad = (xmax - xmin) * 0.05
    ypad = (ymax - ymin) * 0.05

    xlim = (xmin - xpad, xmax + xpad)
    ylim = (ymin - ypad, ymax + ypad)
    xs = np.linspace(*xlim, num=51)
    ys = np.linspace(*ylim, num=51)
    xy = np.dstack(np.meshgrid(xs, ys))

    pdf = stats.multivariate_normal.pdf(xy, mean, cc)

    ax.contourf(xs, ys, pdf, 33, cmap='YlGnBu')
    ax.plot(xsamples, ysamples, 'r.', alpha=.6,
            markeredgecolor='k', markeredgewidth=0.5)

    ax.set_xlim(*xlim)
    ax.set_ylim(*ylim)

plt.show()

对此的一些启发：

• 关于sampling from a multivariate normal的一些说明
• 皮涅罗和贝茨，1996，Unconstrained Parameterizations for Variance-Covariance Matrices