PyTorch/Tensorflow下KDE实现及logprob对x求导方案咨询

关于scikit-learn实现导数计算的说明 #

scikit-learn的KDE基于Numpy实现，本身没有构建计算图，不支持自动微分功能。如果强行基于sklearn的输出结果计算导数，只能手动实现数值微分，存在精度低、高维场景误差大、数值不稳定的问题，不推荐使用。

PyTorch/TensorFlow支持自动微分的KDE方案 #

你可以直接基于深度学习框架的算子手动实现KDE，全链路支持自动微分，也可以使用成熟的第三方封装实现。

1. PyTorch手动实现高斯核KDE（兼容自动微分） #

以下实现和你用到的scikit-learn高斯核KDE逻辑完全对齐：

import torch

def gaussian_kde_logprob(x, x_d, bandwidth=1.0):
    # x: 拟合样本，形状 (样本数, 特征维度)
    # x_d: 查询点，形状 (查询点数量, 特征维度)
    n_samples = x.shape[0]
    n_features = x.shape[1]
    
    # 计算查询点和所有样本的距离平方
    dist_sq = torch.sum((x_d.unsqueeze(1) - x.unsqueeze(0)) ** 2, dim=-1)
    # 高斯核计算
    kernel_vals = torch.exp(-dist_sq / (2 * bandwidth ** 2))
    # 归一化得到概率密度后转对数
    prob = torch.sum(kernel_vals, dim=1) / (n_samples * (bandwidth ** n_features) * (2 * torch.pi) ** (n_features / 2))
    return torch.log(prob)

2. 调用示例 #

# 初始化样本，开启梯度追踪
x = torch.tensor([1,2,3,4,5], dtype=torch.float32).reshape(-1, 1).requires_grad_(True)
x_d = torch.linspace(0, 6, 100).reshape(-1, 1)

# 计算对数概率
logprob = gaussian_kde_logprob(x, x_d, bandwidth=1.0)
# 对任意查询点的logprob求导
logprob[0].backward(retain_graph=True)
# 导数存储在x.grad中
print(x.grad)

其他可选方案 #

• PyTorch生态有封装完善的KDE第三方库，支持多种核函数，原生兼容自动微分
• TensorFlow可以使用tf-probability库中的KernelDensity接口，同样支持自动微分

注意事项 #

• 手动实现的KDE在样本量超过1万时会生成较大的中间距离矩阵，内存不足时可以对查询点分批次计算
• 需要使用非高斯核时，只需替换核函数计算逻辑，只要使用框架原生算子实现，都保留梯度链路
• 不要将sklearn计算得到的logprob直接转成框架张量求导，sklearn的计算过程不在计算图内，无法得到有效梯度

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