有效地对具有不同 sigma（协方差）矩阵的多正态变量集合进行采样

Question

我是 Stan 的新手，希望您能为我指明正确的方向。我会根据自己的情况进行调整，以确保我们在同一页面上...

如果我有一组单变量法线，文档会告诉我：

y ~ normal(mu_vec, sigma);

提供与未向量化版本相同的模型：

for (n in 1:N)
    y[n] ~ normal(mu_vec[n], sigma);

但是矢量化版本（快得多？）。好的，很好，很有道理。

所以第一个问题是：在样本的 mu 和 sigma 都因向量中的位置而异的单变量正常情况下，是否可以利用这种矢量化加速。 IE。如果 mu_vec 和 sigma_vec 都是向量（在前面的例子中 sigma 是一个标量），那么是这样的：

y ~ normal(mu_vec, sigma_vec);

相当于：

for (n in 1:N)
    y[n] ~ normal(mu_vec[n], sigma_vec[n]);

如果有，是否有类似的加速？

好的。那就是热身。真正的问题是如何最好地处理上述等价的多变量。

在我的特定情况下，我有 N 个变量 y 的双变量数据观察值，我将其存储在 N x 2 矩阵中。（对于数量级，N 在我的用例中大约是 1000。）

我相信每个观察的每个分量的平均值是 0，每个分量的标准差是每个观察的 1（我很乐意将它们硬编码为这样的）。但是，我认为相关性 (rho) 因观察而异，作为另一个观察变量 x（存储在 N 元素向量中）的（简单）函数。例如，我们可能会说 rho[n] = 2*inverse_logit(beta * x[n]) - 1 对应 n in 1:N，我们的目标是从我们的数据中了解 beta。 IE。第 n 个观测值的协方差矩阵为：

[1,      rho[n]]
[rho[n], 1     ]

我的问题是，什么是将它们放在 STAN 模型中的最佳方法，这样它才不会那么慢？是否有 multi_normal 分布的矢量化版本，以便我可以将其指定为：

y ~ multi_normal(vector_of_mu_2_tuples, vector_of_sigma_matrices)

或者其他一些类似的表述？或者我需要写：

for (n in 1:N)
    y[n] ~ multi_normal(vector_of_mu_2_tuples[n], vector_of_sigma_matrices[n])

在前面的块中设置了 vector_of_sigma_matrices 和 vector_of_mu_2_tuples 之后？

提前感谢您的指导！

编辑添加代码

使用python，我可以按照我的问题生成如下数据：

import numpy as np
import pandas as pd
import pystan as pys
import scipy as sp
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import seaborn as sns

def gen_normal_data(N, true_beta, true_mu, true_stdevs):
    N = N

    true_beta = true_beta
    true_mu = true_mu
    true_stdevs = true_stdevs

    drivers = np.random.randn(N)
    correls = 2.0 * sp.special.expit(drivers*true_beta)-1.0

    observations = []
    for i in range(N):
        covar = np.array([[true_stdevs[0]**2, true_stdevs[0] * true_stdevs[1] * correls[i]],
                          [true_stdevs[0] * true_stdevs[1] * correls[i], true_stdevs[1]**2]])
        observations.append(sp.stats.multivariate_normal.rvs(true_mu, covar, size=1).tolist())
    observations = np.array(observations)
    
    return {
        'N': N,
        'true_mu': true_mu,
        'true_stdev': true_stdevs,
        'y': observations,
        'd': drivers,
        'correls': correls
    }

然后实际生成数据使用：

normal_data = gen_normal_data(100, 1.5, np.array([1., 5.]), np.array([2., 5.]))

这是数据集的样子（y 的散点图在左窗格中由 correls 着色，在右窗格中由 drivers 着色...所以想法是driver越高，correl越接近1，driver越低，correl越接近-1。所以会期望左窗格中的红点是“左下到右上”，蓝点是“左上到右下”，实际上它们是：

fig, axes = plt.subplots(1, 2, figsize=(15, 6))
x = normal_data['y'][:, 0]
y = normal_data['y'][:, 1]
correls = normal_data['correls']
drivers = normal_data['d']

for ax, colordata, cmap in zip(axes, [correls, drivers], ['coolwarm', 'viridis']):
    color_extreme = max(abs(colordata.max()), abs(colordata.min()))
    sc = ax.scatter(x, y, c=colordata, lw=0, cmap=cmap, vmin=-color_extreme, vmax=color_extreme)
    divider = make_axes_locatable(ax)
    cax = divider.append_axes('right', size='5%', pad=0.05)
    fig.colorbar(sc, cax=cax, orientation='vertical')
fig.tight_layout()

使用蛮力方法，我可以建立一个如下所示的 STAN 模型：

model_naked = pys.StanModel(
    model_name='naked',
    model_code="""
data {
    int<lower=0> N;
    vector[2] true_mu;
    vector[2] true_stdev;
    real d[N];
    vector[2] y[N];
}

parameters {
    real beta;
}

transformed parameters {
}

model {
    real rho[N];
    matrix[2, 2] cov[N];
    
    for (n in 1:N) {
        rho[n] = 2.0*inv_logit(beta * d[n]) - 1.0;
 
        cov[n, 1, 1] = true_stdev[1]^2;
        cov[n, 1, 2] = true_stdev[1] * true_stdev[2] * rho[n];
        cov[n, 2, 1] = true_stdev[1] * true_stdev[2] * rho[n];
        cov[n, 2, 2] = true_stdev[2]^2;
    }
    
    beta ~ normal(0, 10000);
    for (n in 1:N) {
        y[n] ~ multi_normal(true_mu, cov[n]);
    }
}
"""
)

这很合身：

fit_naked = model_naked.sampling(data=normal_data, iter=1000, chains=2)f = fit_naked.plot();
f.tight_layout()

但我希望有人能为我指明“边缘化”方法的正确方向，在这种方法中，我们将双变量法线分解为一对可以使用相关性混合的独立法线。我需要这个的原因是在我的实际用例中，的两个维度都是肥尾的。我很乐意将其建模为 student-t 分布，但问题是 STAN 只允许指定一个 nu （不是每个维度一个），所以我想我需要找到一种方法将 multi_student_t 分解成一对独立的 student_t，这样我就可以为每个维度分别设置自由度。

Answer 1

单变量正态分布确实接受其任何或所有参数的向量，并且比循环遍历 N 观察以使用标量参数调用它 N 次要快。

然而，加速只会是线性的，因为计算都是一样的，但如果你只调用一次，它只需要分配一次内存。总体时间受您必须执行的函数评估次数的影响更大，每次 MCMC 迭代最多 2^10 - 1（默认情况下），但您是否达到最大树深度取决于后验分布的几何形状您正在尝试从中抽样，而这又取决于一切，包括您所依据的数据。

双变量正态分布可以 written 作为第一个变量的边际单变量正态分布和给定第一个变量的第二个变量的条件单变量正态分布的乘积。在 Stan 代码中，我们可以利用逐元素乘法和除法来编写其对数密度，如

target += normal_lpdf(first_variable | first_means, first_sigmas);
target += normal_lpdf(second_variable | second_means + 
  rhos .* first_sigmas ./ second_sigmas .* (first_variable - first_means),
  second_sigmas .* sqrt(1 - square(rhos)));

不幸的是，Stan 中更通用的多元正态分布没有输入协方差矩阵数组的实现。

Answer 2

这并不能完全回答您的问题，但您可以通过删除一堆冗余计算并稍微转换比例以使用 tanh 而不是比例逆 logit 来提高程序的效率。我会摆脱缩放比例，只使用较小的 beta，但我保留了它，这样它应该会得到相同的结果。

data {
  int<lower=0> N;
  vector[2] mu;
  vector[2] sigma;
  vector[N] d;
  vector[2] y[N];
}
parameters {
  real beta;
}
transformed data {
  real var1 = square(sigma[1]);
  real var2 = square(sigma[2]);
  real covar12 = sigma[1] * sigma[2];
  vector[N] d_div_2 = d * 0.5;
}
model {
  // note: tanh(u) = 2 * inv_logit(u / 2) - 1
  vector[N] rho = tanh(beta * d_div_2);
  matrix[2, 2] Sigma;
  Sigma[1, 1] = var1;
  Sigma[2, 2] = var2;
  // only reassign what's necessary with minimal recomputation
  for (n in 1:N) {
    Sigma[1, 2] = rho[n] * covar12;
    Sigma[2, 1] = Sigma[1, 2];
    y[n] ~ multi_normal(true_mu, Sigma);
  }
  // weakly informative priors fit more easily
  beta ~ normal(0, 8);
}

您还可以通过计算 Cholesky 分解作为 rho 和其他固定值的函数来分解并使用它——它在多变量法线中节省了求解器步骤。

您的另一个选择是直接写出 multi-student-t 而不是使用我们的内置实现。内置可能不会快很多，因为整个操作在很大程度上由矩阵求解主导。

有效地对具有不同 sigma（协方差）矩阵的多正态变量集合进行采样

Efficiently sample a collection of multi-normal variables with varying sigma (covariance) matrix

stan

rstan

pystan