python - 如何使用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 custom function jacobian gradient

这个问题显示了pytorch中雅可比的计算，但我认为该解决方案不适用于我的问题:Pytorch most efficient Jacobian/Hessian calculation它需要以似乎与分布不兼容的方式堆叠输入。我无法让它发挥作用。

这个要点有一些雅可比行列式的代码片段。原则上，它们类似于上述问题的方法。

Answer 1

PyTorch 1.5.1 引入了 torch.autograd.functional.jacobian功能。这将计算函数 w.r.t. 的雅可比矩阵。输入张量。自 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]])