Вопрос или проблема
Предположим, у меня есть $f:\mathbb{R}^{d_i}\to\mathbb{R}^{d_o}$. Пусть $X \in \mathbb{R}^{n \times d_i}$, и я применяю $f$ к каждой строке $X$, получая $Y = f(X) \in \mathbb{R}^{n \times d_o}$. Я хотел бы вычислить тензор $Z$, который определяется как
$$Z_{i,j,k} = \frac{\partial Y_{i,j}}{\partial X_{i,k}}$$
с использованием pytorch. Как мне следует определить v, чтобы это достичь?
def get_row_wise_jacobian(X, f):
n = X.shape[0]
d_i = X.shape[1]
x = X.clone().detach().requires_grad_(True)
y = f(x)
d_o = y.shape[1]
v = ?
y.backward(v)
Z = x.grad
return Z
Если f(x) – это глубокая сеть, вы можете использовать реализацию дифференциальной сети из этого репозитория. Эта дифференциальная сеть добавляет некоторые вычисления к графу прямой передачи для вычисления аналитического якобиана по отношению к входным данным сети. В отличие от реализации autograd, вычислительные затраты минимальны, и якобиан вычисляется с использованием прямого дифференцирования. Кроме того, примеры из Репозитория показывают, что вам следует использовать активирующую функцию softplus для получения более гладких якобианов.
Ответ или решение
Вопрос о вычислении якобиана рядов с использованием библиотеки PyTorch является весьма актуальным в области машинного обучения и численного анализа. Рассмотрим ситуацию, когда у нас имеется функция ( f: \mathbb{R}^{d_i} \to \mathbb{R}^{d_o} ), которая применяется ко всем строкам матрицы ( X \in \mathbb{R}^{n \times d_i} ) с целью получить новую матрицу ( Y = f(X) \in \mathbb{R}^{n \times d_o} ). Мы хотим вычислить тензор ( Z ), заданный формулой
[
Z{i,j,k} = \frac{\partial Y{i,j}}{\partial X_{i,k}}
]
Для достижения этой цели в PyTorch требуется изменить код функции, чтобы правильно определить градиенты. Рассмотрим, как можем это сделать.
Определение функции
Начнём с определения функции get_row_wise_jacobian, которая будет вычислять якобиан. В этой функции нужно правильно настроить переменную v, которая будет передаваться в метод .backward().
import torch
def get_row_wise_jacobian(X, f):
n = X.shape[0] # количество строк в X
d_i = X.shape[1] # размерность входного пространства
x = X.clone().detach().requires_grad_(True) # создание тензора x с градиентами
y = f(x) # вычисление Y
d_o = y.shape[1] # размерность выходного пространства
# Создаем тензор v с той же формой, что и y, и инициализируем единичные значения
v = torch.zeros_like(y) # начальная форма v
v[:, :] = 1 # устанавливаем значение 1 для всех элементов
# Запуск обратного распространения градиентов
y.backward(v) # вычисление градиентов
Z = x.grad # получаем градиенты, которые являются якобианом
return Z # возвращаем якобианОбъяснение кода
-
Инициализация: Мы сначала клонируем входной тензор
Xи устанавливаем для него требование градиента, что позволяет вычислить градиенты после вызова функцииf. -
Вычисление выходного тензора: Затем мы вычисляем
Y(выход) через функцию ( f ). -
Инициализация ( v ): Здесь мы создаем тензор ( v ) такой же формы, как и ( Y ), и инициализируем его значениями единиц. Это проходит по всем измерениям выходного тензора, тем самым указывая, что мы хотим получить градиенты по всем выходным компонентам.
-
Обратный проход: Выполняем обратное распространение градиентов, позволяющее получить частные производные каждого элемента
Yотносительно каждого элементаX. -
Возвращение якобиана: Наконец, получаем градиенты и возвращаем их в качестве конечного результата. Этот тензор ( Z ) будет иметь размерность ( n \times d_o \times d_i ).
Заключение
Композиция функции во сочетании с функцией актуализации, такой как softplus, может привести к более гладким якобианам, что может быть полезно в дальнейших вычислениях. Если требуется минимизировать вычислительную нагрузку, стоит рассмотреть варианты специализированных библиотек, таких как дифференциальная сеть, упомянутая в вашем вопросе. Это поможет в более эффективном вычислении якобианов для глубоких сетей.
Данный подход обеспечивает высокую гибкость и масштабируемость, что делает его идеальным для применения в различных задачах машинного обучения и численных вычислений.