在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)
但是,正如您所指出的那样,在计算这个5x5结果的过程中,我们实例化了两个5x10,000矩阵。我们该如何解决这个问题呢?
答案在于jax.jvp
和jax.vjp
。它们的调用签名对于您的问题来说有些不太直观,因为它们主要设计用于正向模式和反向模式自动微分。但总体而言,您可以将它们视为一种计算J @ v
和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))
同样地,您可以使用
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))
将它们结合起来并对单位矩阵进行操作,即可得到您所需的完整雅可比矩阵乘积。
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))
除了节省内存外,这种间接方法也比直接方法快得多,具体取决于所涉及的雅可比矩阵的大小:
%time direct(f, x1, x2)
%time indirect(f, x1, x2, M)