Крис Ола: backprop как исчисление на графах

Почему это конспект, а не перевод

У блога Криса Олы нет открытой лицензии, поэтому ниже — пересказ своими словами, с нашими пояснениями и кодом. Цитаты короткие и в кавычках. Рисунки графов с подписанными производными смотрите в оригинале; здесь они описаны словами и перерисованы схемами.

Коротко

• обратное распространение ошибкиобратное распространение ошибкиEN backpropagation — это не магия нейросетей, а общий приём быстрого подсчёта производных. В общем виде он называется дифференцированием в обратном режиме (reverse-mode differentiation) и переизобретался десятки раз в разных областях.
• Выражение удобно рисовать как вычислительный графвычислительный графEN computational graphСхема вычисления — узлы-операции и рёбра-значения. По нему идёт прямой проход (значения) и обратный (градиенты).: узлы — промежуточные величины, рёбра — «кто от кого зависит». Производная на ребре — локальная скорость изменения.
• Производная одного узла по другому — это сумма по всем путям между ними произведений производных на рёбрах. Это то же самое цепное правилоцепное правилоEN chain ruleПроизводная композиции — произведение производных звеньев — (f(g(x)))′ = f′(g(x))·g′(x). Математическая основа backprop., только нарисованное.
• Путей бывает экспоненциально много, поэтому их надо «выносить за скобки». Прямой режим считает, как один вход влияет на все узлы; обратный — как все узлы влияют на один выход.
• В обучении сети выход один (функция потерь), а входов-параметров миллионы. Поэтому обратный режим выигрывает в миллион раз. Главный урок Олы: «производные дешевле, чем вы думаете».

Как читать вместе с курсом

Эссе — идеальный мост между уроком 1.2.2 о цепном правиле и модулем 3.2, где вы пишете micrograd: вычислительный граф и обратный проход. Прочитайте его перед уроком 3.2.1: после него Value._backward и топологическая сортировка выглядят очевидными. Математики здесь — только правила суммы и произведения из школьного курса.

1. Зачем вообще говорить о backprop

Ола начинает с масштаба. По его оценке, обратное распространение ускоряет обучение глубокой модели градиентным спуском примерно в десять миллионов раз по сравнению с наивным подсчётом производных. Это разница между неделей обучения и «200 000 лет».

Второй тезис: алгоритм нужен далеко не только в глубоком обучении. Его используют, например, в прогнозе погоды и анализе численной устойчивости, просто под другими именами. Ола ссылается на обзор Гриванка (Griewank, 2010): алгоритм открывали заново как минимум десятки раз. Универсальное имя, не привязанное к приложению, — «дифференцирование в обратном режиме». По сути это приём, чтобы быстро считать производные, и его стоит держать в арсенале любому, кто занимается вычислениями.

2. Вычислительный граф

Возьмём пример Олы: $e = (a + b) \cdot (b + 1)$. Здесь три операции — два сложения и умножение. Дадим имя каждому промежуточному результату:

$$c = a + b, \qquad d = b + 1, \qquad e = c \cdot d.$$

Теперь каждая операция и каждая входная переменная — это узел графа. Если значение одного узла служит входом другого, рисуем стрелку. Ола отмечает, что такие графы постоянно встречаются в информатике (графы зависимостей, графы вызовов) и что именно на этой абстракции был построен фреймворк Theano — предшественник TensorFlow и PyTorch.

Чтобы вычислить выражение, задаём входы и считаем узлы снизу вверх. При $a = 2$, $b = 1$ получаем $c = 3$, $d = 2$, $e = 6$. Это прямой проход (forward pass) из урока 3.1.2.

3. Производные на рёбрах и сумма по путям

Ключ к производным на графе — производные на рёбрах. Если $a$ напрямую влияет на $c$, спрашиваем: если немного изменить $a$, как изменится $c$? Это частная производная $\partial c / \partial a$. Нужны всего два правила:

$$\frac{\partial}{\partial a}(a + b) = 1, \qquad \frac{\partial}{\partial u}(u v) = v.$$

Словами: производная суммы по любому слагаемому равна 1, производная произведения по одному множителю равна другому множителю. В нашем графе при $a = 2$, $b = 1$:

flowchart BT
a["a = 2"] -->|"∂c/∂a = 1"| c["c = a + b = 3"]
b["b = 1"] -->|"∂c/∂b = 1"| c
b -->|"∂d/∂b = 1"| d["d = b + 1 = 2"]
c -->|"∂e/∂c = d = 2"| e["e = c · d = 6"]
d -->|"∂e/∂d = c = 3"| e

Граф e = (a + b)(b + 1) с производными на рёбрах (значения при a = 2, b = 1)

А как узнать влияние узлов, которые не соединены напрямую? Ола объясняет через скорости: если $a$ меняется со скоростью 1, то $c$ тоже меняется со скоростью 1, а $c$, меняясь со скоростью 1, меняет $e$ со скоростью 2. Значит, $\partial e / \partial a = 1 \cdot 2 = 2$.

Общее правило: просуммировать по всем путям от одного узла к другому произведения производных на рёбрах пути. От $b$ к $e$ два пути — через $c$ и через $d$:

$$\frac{\partial e}{\partial b} = \underbrace{1 \cdot 2}_{\text{через } c} + \underbrace{1 \cdot 3}_{\text{через } d} = 5.$$

Проверим напрямую: $e = (a + b)(b + 1)$, по правилу произведения $\partial e / \partial b = (b + 1) + (a + b) = 2 + 3 = 5$. Совпало. Ола подчёркивает, что правило «суммы по путям» — это просто другой взгляд на многомерное цепное правило из урока 1.2.2.

4. Беда: путей слишком много

Наивно суммировать пути опасно: их число растёт комбинаторно. Ола рисует граф, где из $X$ в $Y$ ведут три ребра с производными $\alpha, \beta, \gamma$, а из $Y$ в $Z$ — ещё три: $\delta, \epsilon, \zeta$. Путей из $X$ в $Z$ уже $3 \cdot 3 = 9$:

$$\frac{\partial Z}{\partial X} = \alpha\delta + \alpha\epsilon + \alpha\zeta + \beta\delta + \beta\epsilon + \beta\zeta + \gamma\delta + \gamma\epsilon + \gamma\zeta.$$

Добавьте ещё несколько таких «ромбов» подряд — и путей станет $3^k$. Но сумму можно разложить на множители:

$$\frac{\partial Z}{\partial X} = (\alpha + \beta + \gamma)(\delta + \epsilon + \zeta).$$

Вместо девяти умножений и восьми сложений — четыре сложения и одно умножение. Для цепочки из $k$ ромбов вместо $3^k$ слагаемых нужно $3k$ операций. Именно это и делают оба режима дифференцирования: они не перечисляют пути, а сливают их в каждом узле, так что каждое ребро обрабатывается ровно один раз.

Это динамическое программирование

В сноске Ола замечает: если это напоминает динамическое программированиединамическое программированиеEN dynamic programmingРешение задачи через ответы на подзадачи, которые хранятся и переиспользуются. В RL — методы для случая, когда модель среды известна., так и есть. Мы храним в каждом узле готовую частичную сумму по путям и переиспользуем её, вместо того чтобы пересчитывать.

5. Прямой и обратный режимы

Прямой режим (forward-mode) стартует с одного входа и идёт к выходу. В каждом узле он складывает вклады всех входящих рёбер: «насколько этот узел зависит от выбранного входа». Ола замечает, что так вас, по сути, учили считать производные в школе и институте, просто без картинок графов.

Обратный режим (reverse-mode) стартует с выхода и идёт к входам. В каждом узле он складывает вклады всех исходящих рёбер: «насколько выход зависит от этого узла».

Короткая формулировка Олы: прямой режим применяет ко всем узлам оператор $\frac{\partial}{\partial X}$ (производная по одному входу $X$), обратный — оператор $\frac{\partial Z}{\partial}$ (производная одного выхода $Z$ по каждому узлу).

flowchart LR
subgraph F["Прямой режим: ∂/∂b для каждого узла"]
  direction LR
  fb["∂b/∂b = 1"] --> fc["∂c/∂b = 1"]
  fb --> fd["∂d/∂b = 1"]
  fc --> fe["∂e/∂b = 2·1 + 3·1 = 5"]
  fd --> fe
end
subgraph R["Обратный режим: ∂e/∂ для каждого узла"]
  direction RL
  re["∂e/∂e = 1"] --> rc["∂e/∂c = 2"]
  re --> rd["∂e/∂d = 3"]
  rc --> ra["∂e/∂a = 2"]
  rc --> rb["∂e/∂b = 2 + 3 = 5"]
  rd --> rb
end

Один проход прямого режима дал производную только по b; один проход обратного — по a и по b сразу

6. В чём выигрыш

На нашем графе прямой режим, запущенный от $b$, дал производные всех узлов по $b$ — и в том числе нужную $\partial e / \partial b$. А обратный режим, запущенный от $e$, дал производные $e$ по всем узлам — и $\partial e / \partial a$, и $\partial e / \partial b$ за один проход. Ола подчёркивает: «по каждому узлу» значит буквально по каждому.

Для этого графа выигрыш всего двукратный. Но представьте функцию с миллионом входов и одним выходом. Прямому режиму понадобится миллион проходов по графу, по одному на вход. Обратному — один. При обучении сети функция потерь — это одно число, а параметров миллионы (сегодня — миллиарды). Поэтому обратный режим, который в нейросетях называют backprop, даёт гигантское ускорение.

Есть и обратная ситуация. Если у функции один вход и много выходов, быстрее прямой режим: один его проход даёт производные всех выходов по этому входу. Ола оставляет это в скобках, но в современных библиотеках оба режима есть: прямой режим называют JVP (произведение якобиана на вектор), обратный — VJP (произведение вектора на якобиан).

Код: оба режима руками и в PyTorch

# Граф e = (a + b) * (b + 1): прямой проход
a, b = 2.0, 1.0
c = a + b        # 3
d = b + 1        # 2
e = c * d        # 6

# Производные на рёбрах (локальные)
dc_da, dc_db = 1.0, 1.0
dd_db = 1.0
de_dc, de_dd = d, c          # для произведения: производная по одному множителю = другой множитель

# Прямой режим (forward-mode): «как вход b влияет на каждый узел»
db_db = 1.0
da_db = 0.0
dc_db_total = dc_da * da_db + dc_db * db_db      # 1
dd_db_total = dd_db * db_db                       # 1
de_db_fwd = de_dc * dc_db_total + de_dd * dd_db_total
print("forward-mode: de/db =", de_db_fwd)        # 5 — но только по b!

# Обратный режим (reverse-mode, backprop): «как каждый узел влияет на выход e»
de_de = 1.0
de_dc_total = de_de * de_dc                       # 2
de_dd_total = de_de * de_dd                       # 3
de_da = de_dc_total * dc_da                       # 2
de_db = de_dc_total * dc_db + de_dd_total * dd_db # 2 + 3 = 5
print("reverse-mode: de/da =", de_da, " de/db =", de_db)   # обе производные за один проход

Вывод:

forward-mode: de/db = 5.0
reverse-mode: de/da = 2.0  de/db = 5.0

То же самое в PyTorch. Обычный backward() — обратный режим, а torch.func.jvp — прямой:

import torch
from torch.func import jvp, grad

def f(x):                       # x = (a, b)
    a, b = x[0], x[1]
    return (a + b) * (b + 1)

x = torch.tensor([2.0, 1.0])
# Прямой режим: производная по направлению v = (0, 1), то есть по b
_, df_db = jvp(f, (x,), (torch.tensor([0.0, 1.0]),))
print("jvp по b:", df_db.item())          # 5.0
# Обратный режим: весь градиент сразу
print("grad:", grad(f)(x))                # tensor([2., 5.])

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

7. «Это же просто цепное правило?»

Ола признаётся, что, впервые поняв backprop, подумал: «Да это же просто цепное правило! Почему до него так долго доходили?» И он такой не один. Действительно, если спросить «есть ли умный способ считать производные в нейросети прямого распространения», ответ найти несложно.

Но, по мнению Олы, во времена изобретения всё было сложнее, чем кажется. Тогда не занимались именно такими сетями, какие изучаем мы, и вовсе не было очевидно, что учить их нужно через производные. Это становится очевидным только когда знаешь, что производные считаются быстро. Получалась циклическая зависимость: чтобы захотеть быстрый алгоритм, надо верить в обучение градиентом, а чтобы верить в обучение градиентом, надо иметь быстрый алгоритм.

Хуже того, любое звено этого круга легко было отбросить с ходу: «застрянем в локальных минимумах», «производных слишком много, это дорого». Мы не начинаем сразу перечислять такие возражения только потому, что уже знаем: подход работает. Это преимущество взгляда назад: когда вопрос правильно поставлен, главная работа уже сделана.

8. Чему учит эссе

Ола подводит три итога.

1. Производные дешевле, чем вы думаете. Настолько дешевле, что людям пришлось много раз открывать это заново. Градиент по всем параметрам стоит примерно как несколько прямых проходов — в курсе вы это проверите на micrograd.
2. Backprop — линза для понимания обучения. Глядя на то, как производные текут по модели, можно понять, почему некоторые модели плохо оптимизируются. Классический пример — затухание градиентов в рекуррентных сетях (урок 3.4.3 и конспект эссе Олы про LSTM).
3. Общий алгоритмический урок. Backprop и прямой режим держатся на паре приёмов — линеаризации (локальные производные на рёбрах) и динамическом программировании (слияние путей). Поняв их, можно эффективно считать и другие выражения с производными; Ола обещал вернуться к этому в отдельном посте.

Своё изложение Ола называет очень абстрактным и советует для конкретики главу 2 книги Майкла Нильсена «Neural Networks and Deep Learning», где backprop разобран именно на нейросетях.

Что сделать руками

1. Нарисуйте на бумаге граф для $f = (x \cdot y + z) \cdot x$ при $x = 3$, $y = -1$, $z = 2$. Подпишите производные на рёбрах и найдите $\partial f / \partial x$ суммой по путям. (Ответ: путей два, $f = x^2 y + x z$, $\partial f / \partial x = 2xy + z = -4$.) Проверьте через torch.
2. Постройте цепочку из пяти «ромбов» как в разделе 4 со случайными производными на рёбрах. Посчитайте $\partial Z / \partial X$ перебором всех $3^5 = 243$ путей и через произведение сумм. Числа должны совпасть.
3. Возьмите torch.nn.Linear(1000, 1) и сравните время: один backward() против 1000 вызовов jvp по базисным векторам. Во сколько раз отличается?
4. После этого конспекта переходите к уроку 3.2.1: класс Value — это узел графа, _backward — локальные производные на рёбрах, а топологическая сортировка — порядок, в котором обратный режим обходит граф.

Словарик

English	Русский
backpropagation	обратное распространение ошибки
reverse-mode differentiation	дифференцирование в обратном режиме
forward-mode differentiation	дифференцирование в прямом режиме
computational graph	вычислительный граф
node / edge	узел / ребро
partial derivative	частная производная
sum rule / product rule	правило суммы / правило произведения
multivariate chain rule	многомерное цепное правило
sum over paths	сумма по путям
factoring paths	вынесение путей за скобки
dynamic programming	динамическое программирование
linearization	линеаризация
JVP / VJP	произведение якобиана на вектор / вектора на якобиан
vanishing gradients	затухание градиентов