如何使用pytorch获取jacobian的多元正态分布的对数概率

Question

我从多元正态分布中提取了样本，并希望获得其对数概率相对于均值的梯度。由于样本很多，因此需要一个雅可比行列式：

import torch

mu = torch.ones((2,), requires_grad=True)
sigma = torch.eye(2)
dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, sigma)

num_samples=10
samples = dist.sample((num_samples,))
logprobs = dist.log_prob(samples)

现在，我想获得logprobs相对于中每个条目的派生形式mu。

一个简单的解决方案是python循环：

grads = []
for logprob in logprobs:
    grad = torch.autograd.grad(logprob, mu, retain_graph=True)
    grads.append(grad)

如果叠放毕业生，则结果是所需的雅可比矩阵。是否还对此提供了内置的矢量化支持？

相关问题/互联网资源：

这是一个巨大的话题，有很多相关文章。不过，我认为这个特定问题（关于分布）尚未得到回答：

• 这个问题与我的问题基本相同（但没有示例代码和解决方案尝试），可悲的是它没有得到答案：Pytorch自定义函数jacobian渐变

• 这个问题显示了在pytorch中进行雅可比计算，但我认为该解决方案不适用于我的问题：pytorch最有效的Jacobian / Hessian计算它需要以似乎与发行版不兼容的方式堆叠输入。我无法使其工作。

• 该要点有一些针对Jacobian的代码片段。原则上，它们类似于上述问题的方法。

Answer 1

PyTorch 1.5.1 引入了torch.autograd.function.jacobian函数。这将计算输入张量函数的雅可比行列式。由于jacobian需要 python 函数作为第一个参数，因此使用它需要一些代码重组。

import torch

torch.manual_seed(0)    # for repeatable results

mu = torch.ones((2,), requires_grad=True)
sigma = torch.eye(2)
num_samples = 10

def f(mu):
    dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, sigma)
    samples = dist.sample((num_samples,))
    logprobs = dist.log_prob(samples)
    return logprobs

grads = torch.autograd.functional.jacobian(f, mu)

print(grads)

tensor([[-1.1258, -1.1524],
        [-0.2506, -0.4339],
        [ 0.5988, -1.5551],
        [-0.3414,  1.8530],
        [ 0.4681, -0.1577],
        [ 1.4437,  0.2660],
        [ 1.3894,  1.5863],
        [ 0.9463, -0.8437],
        [ 0.9318,  1.2590],
        [ 2.0050,  0.0537]])