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

Question

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

我遇到过类似的事情jvp(f, input, v=vjp(f, input))，但不太明白，不确定是我想要的。

Answer 1

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

\n

举例来说，假设您有一个R\xe1\xb4\xba \xe2\x86\x92 R\xe1\xb4\xb9如下所示的函数：

\n

import jax.numpy as jnp\nimport numpy as np\n\nnp.random.seed(1701)\nN, M = 10000, 5\nf_mat = np.array(np.random.rand(M, N))\ndef f(x):\n  return jnp.sqrt(f_mat @ x / N)\n

\n

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

\n

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

\n

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

\n

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

\n

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

\n

J1_op = lambda v: jax.jvp(f, (x1,), (v,))[1]\n\nvN = np.random.rand(N)\nnp.allclose(J1 @ vN, J1_op(vN))\n# True\n

\n

类似地，您可以用来计算对向量进行操作jax.vjp的效果，而无需实际计算：J2.TJ2

\n

J2T_op = lambda v: jax.vjp(f, x2)[1](v)[0]\n\nvM = np.random.rand(M)\nnp.allclose(J2.T @ vM, J2T_op(vM))\n# True\n

\n

将它们放在一起并在单位矩阵上进行操作可以得到您想要的完整雅可比矩阵乘积：

\n

def direct(f, x1, x2):\n  J1 = jax.jacfwd(f)(x1)\n  J2 = jax.jacfwd(f)(x2)\n  return J1 @ J2.T\n\ndef indirect(f, x1, x2, M):\n  J1J2T_op = lambda v: jax.jvp(f, (x1,), jax.vjp(f, x2)[1](v))[1]\n  return jax.vmap(J1J2T_op)(jnp.eye(M)).T\n\nnp.allclose(direct(f, x1, x2), indirect(f, x1, x2, M))\n# True\n

\n

除了节省内存之外，这种间接方法也比直接方法快一些，具体取决于所涉及的雅可比矩阵的大小：

\n

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

\n