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

问题描述

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

我遇到过类似jvp(f,input,v=vjp(f,input))之类的东西，但是不太了解它，也不知道我想要什么。

解决方法

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

例如，假设您有一个看起来像这样的函数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中。对于您的问题，这些具有不直观的呼叫签名，因为它们主要设计用于正向和反向模式自动区分。但是从广义上讲，您可以将它们视为计算向量J @ v的{​​{1}}和J.T @ v的一种方式，而不必实际地显式计算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

将它们放在一起并在单位矩阵上进行操作，可以为您提供所需要的完整的jacobian矩阵产品：

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,x2,M):
  J1J2T_op = lambda v: jax.jvp(f,jax.vjp(f,x2)[1](v))[1]
  return jax.vmap(J1J2T_op)(jnp.eye(M)).T

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