Блог

Предположим, что J это якобиан некоторой функции f относительно некоторых параметров. Существуют ли эффективные способы (в PyTorch или, возможно, Jax) иметь функцию, которая принимает два входных сигнала ( x1 и x2 ) и вычисляет, J(x1)*J(x2).transpose() не создавая экземпляры всех J матриц в памяти?

Я сталкивался с чем-то вроде jvp(f, input, v=vjp(f, input)) , но не совсем понимаю это и не уверен, чего я хочу.

В 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]]
  

Но, как вы заметили, на пути к вычислению этого результата размером 5×5 мы создаем экземпляры двух матриц размером 5×10 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))
# 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
  

Наряду с экономией памяти, этот косвенный метод также работает немного быстрее, чем прямой метод, в зависимости от размеров задействованных якобианов:

 %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