Backprop и vector-Jacobian product

Backprop — это правило цепочки и ничего больше

Сеть — это композиция: $L = f_n \circ f_{n-1} \circ \cdots \circ f_1(\mathbf{x})$. По правилу цепочки якобиан композиции есть произведение якобианов:

$J = J_n J_{n-1} \cdots J_1$

Всё. Никакой отдельной «теории обратного распространения» не существует — есть правило цепочки и вопрос, в каком порядке перемножать эту цепочку. Ответ на него и есть весь backprop.

Порядок умножения решает всё

Матричное произведение ассоциативно, поэтому считать можно с любого конца. Но стоимость разная, и разница катастрофическая.

Пусть сеть отображает $n$ параметров в один скаляр лосса, а промежуточные слои имеют ширину около $k$.

Слева направо (от входа) — это forward mode. Первое произведение имеет форму $k \times n$, и все промежуточные результаты — матрицы такого размера. Чтобы получить все производные, нужно $n$ проходов: по одному на каждый входной параметр.

Справа налево (от лосса) — это reverse mode, он же backprop. Первый множитель это $1 \times k$ — строка, а не матрица. Умножение строки на матрицу даёт строку, и так до конца: все промежуточные результаты остаются векторами.

Отсюда правило: если выходов мало, а входов много — считайте с конца. В обучении выход один (скалярный лосс), входов миллиарды, поэтому reverse mode выигрывает в $n$ раз. Если бы функция потерь была векторной той же размерности, что параметры, преимущества не было бы.

VJP: то, что действительно вычисляется

Ключевой момент: якобиан никогда не собирается. Вместо него вычисляется произведение вектора на якобиан:

$$\htmlData{k=out}{\frac{\partial L}{\partial \mathbf{x}}} = \htmlData{k=v}{\frac{\partial L}{\partial \mathbf{y}}}\,\htmlData{k=J}{J}$$

Каждая операция обязана уметь ровно одно: получив входящий градиент, вернуть исходящий. Как именно — её дело, и почти всегда это делается без построения якобиана.

Примеры, где выигрыш очевиден:

• линейный слой $\mathbf{y} = W\mathbf{x}$: якобиан по $\mathbf{x}$ равен $W$, и VJP это $W^\top \mathbf{v}$ — одно матрично-векторное умножение;
• поэлементная функция $y_i = g(x_i)$: якобиан диагонален, и VJP это поэлементное умножение. Собрать диагональную матрицу $k \times k$ было бы чистым расточительством;
• softmax: якобиан $K \times K$, но после кросс-энтропии он сократился в $\mathbf{p} - \mathbf{y}$ — предыдущий урок целиком про этот случай.

Стоимость и память

Обратный проход стоит примерно столько же, сколько прямой — константа порядка 2–3. Но платить приходится памятью: локальные производные зависят от значений с прямого прохода, а значит их надо сохранить.

Отсюда две вещи, знакомые всякому, кто обучал большие модели:

• память растёт с глубиной, а не только с шириной. Активации всех слоёв живут до конца обратного прохода;
• gradient checkpointing — сознательный обмен: часть активаций не хранят, а пересчитывают. Плата — примерно один дополнительный прямой проход, выигрыш — память как корень из глубины.

Проход по шагам

Граф двухслойной сети с одним скрытым нейроном. Обратите внимание на порядок обхода: градиент по $w_1$ становится известен последним, потому что до него нужно пройти всю цепочку.

x знач 1.5 град —

w₁ знач 0.8 град —

w₂ знач -1.2 град —

y знач 1 град —

w₁x × знач 1.2 град —

h tanh знач 0.83 град —

o × знач -1 град —

o − y − знач -2 град —

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

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

Forward mode тоже нужен

Reverse mode выигрывает при «много входов, один выход». В обратной ситуации — один вход, много выходов — выигрывает forward mode, и он вычисляет Jacobian-vector product, $J\mathbf{v}$, а не $\mathbf{v}^\top J$.

Где это встречается на практике: произведение гессиана на вектор. $Hv$ считается как forward-режим, применённый к обратному проходу, — двумя проходами, без сборки гессиана. Так работают методы, которым нужна кривизна, но не вся матрица.

Что стоит унести

Три утверждения, которые вместе объясняют, почему обучение больших моделей вообще возможно:

1. Backprop — правило цепочки, посчитанное справа налево.
2. Якобианы не собираются: каждая операция реализует VJP.
3. Обратный проход стоит как прямой по времени и как глубина по памяти.

Ни одно из них не про нейросети. Всё это свойства композиции функций, и в этом смысле блок 2 закончен — осталось собрать своими руками.