Улучшить производительность автограда Якобиана

Question

Мне интересно, как следующий код может быть быстрее. На данный момент это кажется неоправданно медленным, и я подозреваю, что я неправильно использую API Autograd. Результат, который я ожидаю, - это каждый элемент timeline оценивается по якобиану f, который я получаю, но это занимает много времени:

import numpy as np
from autograd import jacobian


def f(params):
    mu_, log_sigma_ = params
    Z = timeline * mu_ / log_sigma_
    return Z


timeline = np.linspace(1, 100, 40000)

gradient_at_mle = jacobian(f)(np.array([1.0, 1.0]))

Я ожидал бы следующее:

1. jacobian(f) возвращает функцию, которая представляет вектор градиента по параметрам.
2. jacobian(f)(np.array([1.0, 1.0])) Якобиан оценивается в точке (1, 1). Для меня это должно быть похоже на векторизованную функцию numpy, поэтому она должна выполняться очень быстро, даже для массивов длиной 40 КБ. Однако это не то, что происходит.

Даже что-то вроде следующего имеет такую ​​же низкую производительность:

import numpy as np
from autograd import jacobian


def f(params, t):
    mu_, log_sigma_ = params
    Z = t * mu_ / log_sigma_
    return Z


timeline = np.linspace(1, 100, 40000)

gradient_at_mle = jacobian(f)(np.array([1.0, 1.0]), timeline)

Answer 1

С https://github.com/HIPS/autograd/issues/439 я понял, что есть недокументированная функция autograd.make_jvp который вычисляет якобиан в режиме ускоренной перемотки вперед.

Ссылка гласит:

Дана функция f, векторы x и v в области f, make_jvp(f)(x)(v) вычисляет как f(x), так и якобиан из f, вычисленный в точке x, умноженный справа на вектор v.

Чтобы получить полный якобиан f, вам просто нужно написать цикл для оценки make_jvp(f)(x)(v) для каждого v в стандартной основе домена f. Наш оператор Якобиана обратного режима работает аналогично.

Из вашего примера:

import autograd.numpy as np
from autograd import make_jvp

def f(params):
    mu_, log_sigma_ = params
    Z = timeline * mu_ / log_sigma_
    return Z

timeline = np.linspace(1, 100, 40000)

gradient_at_mle = make_jvp(f)(np.array([1.0, 1.0]))

# loop through each basis
# [1, 0] evaluates (f(0), first column of jacobian)
# [0, 1] evaluates (f(0), second column of jacobian)
for basis in (np.array([1, 0]), np.array([0, 1])):
    val_of_f, col_of_jacobian = gradient_at_mle(basis)
    print(col_of_jacobian)

Выход:

[  1.           1.00247506   1.00495012 ...  99.99504988  99.99752494
 100.        ]
[  -1.           -1.00247506   -1.00495012 ...  -99.99504988  -99.99752494
 -100.        ]

Это работает в ~ 0,005 секунд на Google Collab.

Редактировать:

Функции как cdf не определены для регулярного jvp пока, но вы можете использовать другую недокументированную функцию make_jvp_reversemode где это определено. Использование аналогично, за исключением того, что выводом является только столбец, а не значение функции:

import autograd.numpy as np
from autograd.scipy.stats.norm import cdf
from autograd.differential_operators import make_jvp_reversemode


def f(params):
    mu_, log_sigma_ = params
    Z = timeline * cdf(mu_ / log_sigma_)
    return Z

timeline = np.linspace(1, 100, 40000)

gradient_at_mle = make_jvp_reversemode(f)(np.array([1.0, 1.0]))

# loop through each basis
# [1, 0] evaluates first column of jacobian
# [0, 1] evaluates second column of jacobian
for basis in (np.array([1, 0]), np.array([0, 1])):
    col_of_jacobian = gradient_at_mle(basis)
    print(col_of_jacobian)

Выход:

[0.05399097 0.0541246  0.05425823 ... 5.39882939 5.39896302 5.39909665]
[-0.05399097 -0.0541246  -0.05425823 ... -5.39882939 -5.39896302 -5.39909665]

Обратите внимание, что make_jvp_reversemode будет немного быстрее чем make_jvp постоянным фактором из-за использования кэширования.