计算 Jacobian x Jacobian.T 的有效方法答案

【问题标题】：Efficient way to compute Jacobian x Jacobian.T计算 Jacobian x Jacobian.T 的有效方法
【发布时间】：2020-12-12 23:33:37
【问题描述】：

假设J 是某个函数f 关于某些参数的雅可比行列式。是否有有效的方法（在 PyTorch 或 Jax 中）拥有一个函数，该函数接受两个输入（x1 和 x2）并计算 J(x1)*J(x2).transpose() 而不在内存中实例化整个 J 矩阵?

我遇到过jvp(f, input, v=vjp(f, input)) 之类的东西，但不太明白，也不确定是不是我想要的。

【问题讨论】：

标签： python pytorch backpropagation autograd jax

【解决方案1】：

在 JAX 中，您可以使用 jax.jacfwd 或 jax.jacrev 计算完整的雅可比矩阵，也可以使用 jax.jvp 和 jax.vjp 计算雅可比运算符及其转置。

例如，假设您有一个函数Rᴺ → Rᴹ，看起来像这样：

import jax.numpy as jnp
import numpy as np

np.random.seed(1701)
N, M = 10000, 5
f_mat = np.array(np.random.rand(M, N))
def f(x):
  return jnp.sqrt(f_mat @ x / N)

给定两个向量 x1 和 x2，您可以使用 jax.jacfwd 计算每个向量的雅可比矩阵

import jax
x1 = np.array(np.random.rand(N))
x2 = np.array(np.random.rand(N))
J1 = jax.jacfwd(f)(x1)
J2 = jax.jacfwd(f)(x2)
print(J1 @ J2.T)
# [[3.3123782e-05 2.5001222e-05 2.4946943e-05 2.5180108e-05 2.4940484e-05]
#  [2.5084497e-05 3.3233835e-05 2.4956826e-05 2.5108084e-05 2.5048916e-05]
#  [2.4969209e-05 2.4896170e-05 3.3232871e-05 2.5006309e-05 2.4947023e-05]
#  [2.5102483e-05 2.4947576e-05 2.4906987e-05 3.3327218e-05 2.4958186e-05]
#  [2.4981882e-05 2.5007204e-05 2.4966144e-05 2.5076926e-05 3.3595043e-05]]

但是，正如您所注意到的，在计算这个 5x5 结果的过程中，我们实例化了两个 5x10,000 矩阵。我们该如何解决这个问题？

答案在jax.jvp 和jax.vjp。就您的问题而言，这些调用签名有些不直观，因为它们主要设计用于正向模式和反向模式自动区分。但从广义上讲，您可以将它们视为一种为向量 v 计算 J @ v 和 J.T @ v 的方法，而无需显式实际计算 J。

例如，您可以使用jax.jvp 计算J1 对向量进行操作的效果，而无需实际计算J1：

J1_op = lambda v: jax.jvp(f, (x1,), (v,))[1]

vN = np.random.rand(N)
np.allclose(J1 @ vN, J1_op(vN))
# True

同样，您可以使用jax.vjp 计算J2.T 对向量进行操作的效果，而无需实际计算J2：

J2T_op = lambda v: jax.vjp(f, x2)[1](v)[0]

vM = np.random.rand(M)
np.allclose(J2.T @ vM, J2T_op(vM))
# True

将这些放在一起并在单位矩阵上进行运算，即可得到您所追求的完整雅可比矩阵乘积：

def direct(f, x1, x2):
  J1 = jax.jacfwd(f)(x1)
  J2 = jax.jacfwd(f)(x2)
  return J1 @ J2.T

def indirect(f, x1, x2, M):
  J1J2T_op = lambda v: jax.jvp(f, (x1,), jax.vjp(f, x2)[1](v))[1]
  return jax.vmap(J1J2T_op)(jnp.eye(M)).T

np.allclose(direct(f, x1, x2), indirect(f, x1, x2, M))
# True

除了节省内存外，这种间接方法也比直接方法快一点，具体取决于所涉及的 jacobian 的大小：

%time direct(f, x1, x2)
# CPU times: user 1.43 s, sys: 14.9 ms, total: 1.44 s
# Wall time: 886 ms
%time indirect(f, x1, x2, M)
# CPU times: user 311 ms, sys: 0 ns, total: 311 ms
# Wall time: 158 ms

【讨论】：