如何在Python中计算两组高维分子特征的多元JS/KL散度？

我有两个二维数组，尺寸分别为N×3968和M×3968，对应两组分子的N个和M个特征向量，每个向量包含3968个浮点特征。需要衡量这两组非标准多元分布（不服从多元正态、多项式等分布）之间的散度，目前遇到以下问题：

• 现有工具（dit包、scipy）的JS散度/距离函数仅支持一维分布
• 用scipy按axis=0计算得到3968个一维距离结果后，存在疑问：
  1. 取均值作为多元散度估计是否合理？已知均值无法捕捉特征间依赖关系，那取乘积呢？
  2. 是否可通过降维技术消除特征间的非线性依赖（类似PCA处理线性依赖），再按低维空间各维度的方差解释占比「pv」加权平均对应维度的JS距离，以此计算多元散度？

核心结论与解决方案 #

1. 一维JS距离的均值/乘积均不可行 #

• 均值：完全忽略特征间的依赖关系，相当于假设所有特征独立，这在分子特征场景下几乎不成立（分子特征往往存在复杂关联，比如官能团相关特征），会严重偏离真实散度。
• 乘积：无理论依据，会放大极端值影响，无法反映依赖关系本质，结果不具备统计意义。

2. 降维后加权的思路可行，但需选对降维方法 #

PCA仅能处理线性依赖，而分子特征间多为非线性关联，需选择适合非线性降维的方法：

• 核PCA：通过核函数将数据映射到高维空间后做PCA，可捕捉非线性依赖
• UMAP：相比t-SNE保留更多全局结构，可计算各维度方差占比用于加权
• 自编码器（AutoEncoder）：通过神经网络学习非线性降维，可自定义损失保留关键特征

具体步骤：

1. 合并两组数据进行非线性降维，得到低维嵌入（比如降到50-100维，根据方差解释率调整）
2. 计算低维空间每个维度的方差解释占比「pv」（核PCA可直接输出，UMAP/自编码器可通过维度方差占比计算）
3. 对低维每个维度分别计算两组数据的JS距离，按「pv」加权求和得到多元散度估计值

3. 直接计算多元JS/KL散度的Python实现 #

若不想降维，可通过**核密度估计（KDE）**先估计多元概率密度，再计算散度：

import numpy as np
from scipy.stats import gaussian_kde
from scipy.special import rel_entr

def kl_divergence(p, q):
    # 计算KL散度 D(p||q)
    return np.sum(rel_entr(p, q))

def js_divergence(p, q):
    # 计算JS散度
    m = 0.5 * (p + q)
    return 0.5 * kl_divergence(p, m) + 0.5 * kl_divergence(q, m)

# 示例数据：X为N×3968，Y为M×3968
X = np.random.randn(100, 3968)
Y = np.random.randn(150, 3968)

# 生成KDE模型
kde_X = gaussian_kde(X.T)
kde_Y = gaussian_kde(Y.T)

# 用合并后的样本点计算密度（加小值避免0）
combined = np.vstack([X, Y])
p = kde_X.evaluate(combined.T) + 1e-10
q = kde_Y.evaluate(combined.T) + 1e-10

js_div = js_divergence(p, q)
print(f"多元JS散度: {js_div}")

注意：3968维数据的KDE会面临维数灾难，计算慢且估计误差大，降维后计算更实用；可通过bw_method参数调整KDE带宽优化效果。

4. 替代方案：最大均值差异（MMD） #

对于非标准多元分布，MMD是更常用的分布差异衡量方法，无需估计概率密度，直接通过核函数计算样本均值差异：

from sklearn.metrics.pairwise import rbf_kernel

def mmd(X, Y, sigma=1.0):
    Kxx = rbf_kernel(X, X, gamma=1/(2*sigma**2))
    Kyy = rbf_kernel(Y, Y, gamma=1/(2*sigma**2))
    Kxy = rbf_kernel(X, Y, gamma=1/(2*sigma**2))
    return np.mean(Kxx) + np.mean(Kyy) - 2*np.mean(Kxy)

mmd_score = mmd(X, Y)
print(f"MMD值: {mmd_score}")

MMD在高维数据上表现更稳定，分子特征场景中常用RBF核或线性核，可根据数据调整核函数和参数。

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