基于GPyTorch多元正态分布计算效应量类似物及分区问题

核心理论基础（针对多元正态条件分布） #

对于p维多元正态分布 $p(\boldsymbol{\beta}) = \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma})$，将变量划分为单个特征$\beta_j$和其余特征$\boldsymbol{\beta}_{-j}$后：

• 条件期望：$E(\boldsymbol{\beta}{-j} | \beta_j) = \boldsymbol{\mu}{-j} + \boldsymbol{\Sigma}{j,-j}^T \Sigma{jj}^{-1} (\beta_j - \mu_j)$
• 条件方差：$V(\boldsymbol{\beta}{-j} | \beta_j) = \boldsymbol{\Sigma}{-j,-j} - \boldsymbol{\Sigma}{j,-j}^T \Sigma{jj}^{-1} \boldsymbol{\Sigma}_{j,-j}$
• 论文补充材料中的$\alpha_j$（对应$\beta_j$的KLD）可通过下三角矩阵的行列式高效计算，无需直接求逆协方差矩阵。

代码实现（基于PyTorch/GPyTorch） #

假设你已通过SAASBO模型获取了后验MVN的属性：

model.eval()
posterior = model.posterior(X_test)
mvn = posterior.mvn
mu = mvn.mean.squeeze()  # 假设为[p]维度的均值向量
Sigma = mvn.covariance_matrix.squeeze()  # [p,p]协方差矩阵
Lambda = mvn.precision_matrix.squeeze()  # [p,p]精度矩阵
L = mvn.scale_tril.squeeze()  # [p,p]下三角缩放矩阵（Sigma = L @ L.T）

1. 矩阵分区工具函数 #

用于提取协方差/精度矩阵中与$\beta_j$相关的子矩阵：

import torch

def partition_matrix(mat, j):
    """
    划分p维方阵为与beta_j相关的子矩阵
    参数:
        mat: [p,p] 协方差/精度矩阵
        j: 目标特征索引（0-based）
    返回:
        mat_jj: [1,1] 单个元素的子矩阵
        mat_j_minusj: [1, p-1] beta_j与其余特征的交叉项
        mat_minusj_j: [p-1, 1] 转置交叉项
        mat_minusj_minusj: [p-1, p-1] 其余特征的子矩阵
    """
    p = mat.shape[0]
    mat_jj = mat[j, j].unsqueeze(0).unsqueeze(0)
    # 提取beta_j对应的行/列，排除自身
    idx_minusj = torch.cat([torch.arange(j), torch.arange(j+1, p)], dim=0)
    mat_j_minusj = mat[j, idx_minusj].unsqueeze(0)
    mat_minusj_j = mat_j_minusj.T
    mat_minusj_minusj = mat[idx_minusj][:, idx_minusj]
    return mat_jj, mat_j_minusj, mat_minusj_j, mat_minusj_minusj

2. 条件分布计算（协方差矩阵版） #

def conditional_mvn_cov(Sigma, mu, j, beta_j):
    """
    基于协方差矩阵计算条件分布N(beta_{-j} | beta_j)的期望和方差
    参数:
        Sigma: [p,p] 协方差矩阵
        mu: [p] 均值向量
        j: 目标特征索引
        beta_j: 标量，beta_j的取值
    返回:
        cond_mu: [p-1] 条件期望
        cond_cov: [p-1, p-1] 条件方差
    """
    sigma_jj, sigma_j_minusj, sigma_minusj_j, sigma_minusj_minusj = partition_matrix(Sigma, j)
    mu_minusj = mu[torch.cat([torch.arange(j), torch.arange(j+1, p)], dim=0)]
    
    # 计算条件期望
    sigma_jj_inv = 1.0 / sigma_jj
    cond_mu = mu_minusj + (sigma_minusj_j @ sigma_jj_inv @ (beta_j - mu[j]).unsqueeze(0).T).squeeze(-1)
    # 计算条件方差
    cond_cov = sigma_minusj_minusj - sigma_minusj_j @ sigma_jj_inv @ sigma_j_minusj
    return cond_mu, cond_cov

3. 条件分布计算（精度矩阵版，更高效） #

利用精度矩阵的稀疏性（SAASBO模型特性），避免协方差求逆：

def conditional_mvn_precision(Lambda, mu, j, beta_j):
    """
    基于精度矩阵计算条件分布N(beta_{-j} | beta_j)的期望和方差
    参数:
        Lambda: [p,p] 精度矩阵（Sigma^{-1}）
        mu: [p] 均值向量
        j: 目标特征索引
        beta_j: 标量，beta_j的取值
    返回:
        cond_mu: [p-1] 条件期望
        cond_cov: [p-1, p-1] 条件方差
    """
    lambda_jj, lambda_j_minusj, lambda_minusj_j, lambda_minusj_minusj = partition_matrix(Lambda, j)
    mu_minusj = mu[torch.cat([torch.arange(j), torch.arange(j+1, p)], dim=0)]
    
    # 条件方差为Lambda_{-j,-j}的逆，用Cholesky分解求逆更稳定
    L_chol = torch.linalg.cholesky(lambda_minusj_minusj)
    cond_cov = torch.cholesky_inverse(L_chol)
    # 计算条件期望
    cond_mu = mu_minusj - (cond_cov @ lambda_minusj_j @ (beta_j - mu[j]).unsqueeze(0).T).squeeze(-1)
    return cond_mu, cond_cov

4. 计算$\alpha_j$（对应$\beta_j$的KLD） #

利用下三角矩阵L的行列式特性（行列式为对角线元素乘积），高效计算：

def compute_alpha_j(L, j):
    """
    计算论文补充材料中的alpha_j，即beta_j对应的KLD值
    参数:
        L: [p,p] 下三角缩放矩阵（Sigma = L @ L.T）
        j: 目标特征索引
    返回:
        alpha_j: 标量，KLD值
    """
    p = L.shape[0]
    # 计算det(Sigma) = (prod(L.diag()))^2
    det_Sigma = (torch.prod(L.diag())) ** 2
    # 计算det(Sigma_{-j,-j})
    idx_minusj = torch.cat([torch.arange(j), torch.arange(j+1, p)], dim=0)
    L_minusj = L[idx_minusj][:, idx_minusj]
    det_Sigma_minusj = (torch.prod(L_minusj.diag())) ** 2
    # Sigma_jj = L[j,j]^2
    sigma_jj = L[j, j] ** 2
    # 计算alpha_j
    alpha_j = 0.5 * torch.log(det_Sigma / (det_Sigma_minusj * sigma_jj))
    return alpha_j.item()

批量计算所有特征的KLD #

# 计算p=10个特征的KLD值
kld_values = [compute_alpha_j(L, j) for j in range(10)]
print("各特征的KLD值（alpha_j）:", kld_values)

内容的提问来源于stack exchange，提问作者iammteah