Правило цепочки и якобиан

Якобиан

Для отображения $\mathbf{f} : \R^n \to \R^m$ линейная аппроксимация — это матрица, а не вектор. Она называется якобианом:

$J_{ij} = \frac{\partial f_i}{\partial x_j}, \qquad J \in \R^{m \times n}$

Строк столько, сколько выходов, столбцов — сколько входов. Проверка формы та же, что для матриц в блоке 1: якобиан — это матрица линейного отображения, которое приближает $\mathbf{f}$ около точки.

Частные случаи, которые полезно узнавать:

• $f : \R^n \to \R$ — якобиан это строка, транспонированный градиент;
• $\mathbf{f} : \R \to \R^m$ — якобиан это столбец;
• линейное отображение $\mathbf{f}(\mathbf{x}) = A\mathbf{x}$ — якобиан равен $A$ в любой точке. Линейная аппроксимация линейной функции есть она сама.

Правило цепочки в матричном виде

$$J_{f \circ g}(\mathbf{x}) = \htmlData{k=outer}{J_f\bigl(g(\mathbf{x})\bigr)}\,\htmlData{k=inner}{J_g(\mathbf{x})}$$

Это буквально то же самое, что в одномерном случае, — только вместо умножения чисел умножение матриц. Внешний якобиан обязательно вычисляется в точке $g(\mathbf{x})$, а не в $\mathbf{x}$: это самая частая ошибка в выкладках.

Порядок множителей проверяется размерностями. Если $g : \R^n \to \R^k$ и $f : \R^k \to \R^m$, то $J_g$ имеет форму $k \times n$, $J_f$ — форму $m \times k$, и только произведение $J_f J_g$ вообще определено. Как и в блоке 1, порядок записи обратен порядку применения.

Сумма по путям

Если переменная влияет на результат несколькими способами, вклады складываются:

$\frac{\partial L}{\partial x} = \sum_{\text{пути от } x \text{ к } L} \ \prod_{\text{звенья пути}} \text{локальная производная}$

Умножение вдоль пути, сложение по путям. Это полная форма правила цепочки, и именно её забывают, когда переменная используется дважды.

Простейший пример: $L = x \cdot x$. Есть два пути от $x$ к произведению — через первый аргумент и через второй. По каждому локальная производная равна $x$, сумма даёт $2x$. Тот же ответ, что и по правилу для $x^2$, но полученный механически, без знания формулы.

Обратный проход по шагам

Ниже — граф вычисления $L = \tanh(wx + b)^2$. Нажимайте «шаг назад»: сначала выход получает затравку $\partial L/\partial L = 1$, потом каждый узел передаёт свою производную входам, умножая на локальную.

x знач 2 град —

w знач -1.5 град —

b знач 1 град —

wx × знач -3 град —

n + знач -2 град —

t tanh знач -0.96 град —

L x² знач 0.93 град —

шаг назад до конца сбросить

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

Почему это дёшево

Ключевое наблюдение, из которого выросло всё глубокое обучение. Чтобы получить производные по всем входам, достаточно одного обратного прохода — его стоимость примерно равна стоимости прямого.

Сравните с численным дифференцированием: там на каждый параметр нужно как минимум два вычисления функции. Для сети с $10^9$ параметрами это $2 \cdot 10^9$ прямых проходов против одного обратного.

Причина такой асимметрии — в том, с какой стороны перемножать цепочку якобианов. Разберём это в уроке про VJP: то, что backprop идёт справа налево, а не слева направо, и есть источник всей экономии.

Диамант: где ошибаются вручную

Рассмотрим $L = e^x + x^2$. Граф ветвится: $x$ идёт и в экспоненту, и в квадрат, а потом ветви сходятся в сумме.

x знач 0.8 град —

eˣ exp знач 2.23 град —

x² x² знач 0.64 град —

L + знач 2.87 град —

шаг назад до конца сбросить

Пройдите обратный проход и посмотрите на градиент $x$ после каждого шага: он накапливается, а не перезаписывается. В коде это +=, и забытый плюс — классическая ошибка при написании автограда вручную.

Теорема о неявной функции — только формулировка

Иногда связь задана уравнением $F(x, y) = 0$, а не формулой $y = f(x)$. Теорема утверждает: если $\partial F/\partial y \ne 0$, то локально $y$ выражается через $x$ как дифференцируемая функция, и

$\frac{dy}{dx} = -\frac{\partial F/\partial x}{\partial F/\partial y}$

Доказательство пропускаем. Знать формулировку стоит, потому что она объясняет, как дифференцируют через решение уравнения, не решая его: так устроены implicit layers, deep equilibrium models и дифференцирование через оптимизационную задачу.