Matrix Calculus, часть 3: прямой и обратный режим автодифференцирования

Это перевод главы 8 конспекта курса MIT «Matrix Calculus (for Machine Learning and Beyond)». Предыдущие части: производная как линейный оператор и градиенты матричных функций; карта курса — в путеводителе. Перевод распространяется под лицензией CC BY-NC-SA 4.0.

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

Читайте после модуля 3.2 «Свой автоград»: в уроках 3.2.1 и 3.2.2 вы своими руками строите обратный режим (как micrograd Карпатого). Эта глава показывает второй способ — прямой режим через дуальные числа, — и объясняет, почему нейросети обучают обратным: у функции потерь один выход и миллиарды входов. Последний раздел (произведения гессиана на вектор) — продвинутый, его можно отложить. Код оригинала — на Julia; мы оставили его с переведёнными комментариями и добавили эквиваленты на Python (блоки «Проверьте в Python» — наше дополнение). Рисунки-графы из оригинала перерисованы схемами.

8. Прямой и обратный режим автоматического дифференцирования

Когда профессор Эдельман впервые услышал об автоматическом дифференцировании (AD), ему было легко представить, что это такое… но представил он неправильно! В голове у него было прямолинейное символьное дифференцирование, применённое к коду: что-то вроде запуска Mathematica или Maple или даже просто автоматического выполнения того, чему он учился на занятиях по анализу. Например, подставлять функции и их области определения из таблицы вроде этой, из первого курса анализа:

Производная	Область определения
$(\sin x)' = \cos x$	$-\infty < x < \infty$
$(\cos x)' = -\sin x$	$-\infty < x < \infty$
$(\tan x)' = \sec^2 x$	$x \ne \frac{\pi}{2} + \pi n,\ n \in \mathbb{Z}$
$(\cot x)' = -\csc^2 x$	$x \ne \pi n,\ n \in \mathbb{Z}$
$(\sec x)' = \tan x \sec x$	$x \ne \frac{\pi}{2} + \pi n,\ n \in \mathbb{Z}$
$(\csc x)' = -\cot x \csc x$	$x \ne \pi n,\ n \in \mathbb{Z}$

А если это не как запуск Mathematica или Maple, то уж точно конечные разности, которые проходят в курсе численных методов (или как мы делали в главе 4).

Оказывается, это определённо не конечные разности: алгоритмы AD в общем случае точные (в точной арифметике, без учёта ошибок округления), а не приближённые. Но и на обычную символьную алгебру это не очень похоже: компьютер на самом деле не строит большое «развёрнутое» символьное выражение, чтобы потом его продифференцировать, как можно было бы делать вручную или в системе компьютерной алгебры. Например, представьте программу, которая вычисляет $\det A$ для матрицы $n \times n$: записать «всё» символьное выражение невозможно, пока программа не запущена и $n$ не известно (например, его вводит пользователь), а наивное символьное выражение в любом случае потребовало бы $n!$ слагаемых. Поэтому системы AD должны справляться с программными конструкциями — циклами, рекурсией, размерами задачи $n$, неизвестными до запуска, — и при этом избегать построения символьных выражений непомерного размера. (Пример, совсем не похожий на формулы, которые дифференцируют на первом курсе, — в разделе 8.1.1.) Разработка систем AD часто оказывается больше делом компиляторов, чем анализа!

8.1. Автоматическое дифференцирование через дуальные числа

(К этой лекции есть Julia-блокнот с результатами разных вычислительных экспериментов; он на странице курса. Выдержки из этих экспериментов приведены ниже.)

Один подход к AD, который можно объяснить относительно просто, — AD в прямом режиме (forward mode): вычисление $f'$ ведётся в паре с вычислением $f$. Каждое промежуточное значение $a$ в программе дополняется ещё одним значением $b$, которое представляет его производную, а правила цепочки протаскивают эти производные через вычисления над значениями программы. Оказывается, это можно рассматривать как замену вещественных чисел (значений $a$) новым видом «дуальных чисел» $D(a, b)$ (значения и производные) с соответствующими правилами арифметики, как объяснено ниже.

8.1.1. Пример: вавилонский квадратный корень

Начнём с простого примера — алгоритма для функции квадратного корня, на котором практический метод автоматического дифференцирования стал для профессора Эдельмана и математическим сюрпризом, и вычислительным чудом. Рассмотрим «вавилонский» алгоритм вычисления $\sqrt{x}$, известный тысячелетиями (а позже оказавшийся частным случаем метода Ньютона для $t^2 - x = 0$): просто повторяйте $t \leftarrow (t + x/t)/2$, пока $t$ не сойдётся к $\sqrt{x}$ с нужной точностью. Каждая итерация — одно сложение и два деления. Для иллюстрации хватит 10 итераций. Вот короткая программа на Julia, реализующая этот алгоритм: начальное приближение 1, затем $N$ шагов (по умолчанию $N = 10$):

julia> function Babylonian(x; N = 10)
           t = (1+x)/2   # один шаг из t=1
           for i = 2:N   # оставшиеся N-1 шагов
               t = (t + x/t) / 2
           end
           return t
       end

Если запустить эту функцию для квадратного корня из $x = 4$, видно, что она сходится очень быстро: всего за $N = 3$ шага она получает правильный ответ (2) почти с тремя знаками после запятой, а задолго до $N = 10$ сходится к 2 в пределах точности компьютерной арифметики (около 16 цифр). На самом деле каждый шаг примерно удваивает число верных цифр:

julia> Babylonian(4, N=1)
2.5

julia> Babylonian(4, N=2)
2.05

julia> Babylonian(4, N=3)
2.000609756097561

julia> Babylonian(4, N=4)
2.0000000929222947

julia> Babylonian(4, N=10)
2.0

Конечно, любой первокурсник знает производную квадратного корня $(\sqrt{x})' = 0{,}5/\sqrt{x}$, которую здесь можно вычислить как 0.5 / Babylonian(x), но мы хотим понять, как получить эту производную автоматически, прямо из самого вавилонского алгоритма. Если мы поймём, как легко и эффективно провести правило цепочки через этот алгоритм, то начнём понимать, как AD дифференцирует и куда более сложные программы, для которых простой формулы производной не известно.

8.1.2. Простой AD в прямом режиме

Основная идея проведения правила цепочки через программу очень проста: замените каждое число двумя числами — одно хранит значение, другое следит за производной этого значения. Значения вычисляются как и раньше, а производные — по правилу цепочки для элементарных операций вроде $+$ и $/$.

В Julia эту идею можно реализовать, определив новый тип чисел — назовём его D, — который хранит значение val и производную deriv.

julia> struct D <: Number
           val::Float64
           deriv::Float64
       end

(Подробное объяснение синтаксиса Julia можно найти в других местах, но, надеемся, основная идея понятна, даже если неясен каждый знак препинания.) Величина x = D(a,b) этого нового типа имеет две компоненты: x.val = a и x.deriv = b, которыми мы будем представлять значения и производные соответственно. Код Babylonian использует только две арифметические операции, $+$ и $/$, так что нам достаточно перегрузить встроенные (из Base) определения этих операций в Julia, добавив новые правила для типа D:

julia> Base.:+(x::D, y::D) = D(x.val+y.val, x.deriv+y.deriv)
       Base.:/(x::D, y::D) = D(x.val/y.val, (y.val*x.deriv - x.val*y.deriv)/y.val^2)

Присмотревшись, вы увидите, что значения просто складываются и делятся обычным образом, а производные вычисляются по правилу суммы (складываем производные входов) и правилу частного соответственно. Нужен ещё один технический приём: определить правила преобразования (conversion) и продвижения типов (promotion), которые говорят Julia, как сочетать значения D с обычными вещественными числами в выражениях вроде $x + 1$ или $x/2$:

julia> Base.convert(::Type{D}, r::Real) = D(r,0)
       Base.promote_rule(::Type{D}, ::Type{<:Real}) = D

Это просто говорит, что обычное вещественное число $r$ сочетается со значением D, сначала превращаясь в D(r,0): значение равно $r$, а производная — 0 (производная любой константы равна нулю).

С этими определениями можно подставить значение D в нашу неизменённую функцию Babylonian, и она «волшебным образом» вычислит производную квадратного корня. Попробуем для $x = 49 = 7^2$:

julia> x = 49
49

julia> Babylonian(D(x,1))
D(7.0, 0.07142857142857142)

Видно, что функция правильно вернула значение 7.0 и производную 0.07142857142857142, что действительно совпадает с квадратным корнем $\sqrt{49}$ и его производной $0{,}5/\sqrt{49}$:

julia> (√x, 0.5/√x)
(7.0, 0.07142857142857142)

Почему мы подали на вход D(x,1)? Откуда взялась 1? Это просто то, что производная входа $x$ по самому себе равна $(x)' = 1$, — это начальная точка для правила цепочки.

На практике всё это (и многое другое) уже реализовано в пакете Julia ForwardDiff.jl (и во многих похожих пакетах на разных языках). Пакет прячет детали реализации и явно предоставляет функцию для вычисления производной. Например:

julia> using ForwardDiff

julia> ForwardDiff.derivative(Babylonian, 49)
0.07142857142857142

По сути, однако, это то же самое, что наша маленькая реализация D, только сделанная более общо и изощрённо (правила цепочки для большего числа операций, поддержка большего числа числовых типов, частные производные по нескольким переменным и т. д.): как и мы, ForwardDiff дополняет каждое значение вторым числом, следящим за производной, и протаскивает обе величины через вычисление.

Ту же идею можно было реализовать и специально для вавилонского алгоритма, написав новую функцию dBabylonian, которая по ходу вычисления следит и за переменной $t$, и за её производной $t' = dt/dx$:

julia> function dBabylonian(x; N = 10)
           t = (1+x)/2
           t′ = 1/2
           for i = 1:N
               t = (t+x/t)/2
               t′= (t′+(t-x*t′)/t^2)/2
           end
           return t′
       end

julia> dBabylonian(49)
0.07142857142857142

Это ровно те же вычисления, что и вызов Babylonian(D(x,1)) или ForwardDiff.derivative(Babylonian, 49), но требующие гораздо больше человеческого труда: пришлось бы делать это для каждой программы, которую мы пишем, вместо того чтобы один раз реализовать новый тип чисел.

Примечание переводчика: в строке для t′ стоит производная частного $x/t$ по $x$: $(1 \cdot t - x \cdot t')/t^2$. Заметьте также, что цикл здесь делает $N$ шагов, а в Babylonian — $N - 1$; на ответ при $N = 10$ это не влияет, алгоритм давно сошёлся.

Проверьте в Python (дополнение переводчика)

Тот же приём на Python: класс D с перегруженной арифметикой. Функцию babylonian мы не меняем.

class D:
    """Дуальное число a + bϵ: значение и производная."""
    def __init__(self, val, deriv=0.0):
        self.val, self.deriv = float(val), float(deriv)

    @staticmethod
    def lift(x):                      # обычное число r превращаем в D(r, 0)
        return x if isinstance(x, D) else D(x, 0.0)

    def __add__(self, other):         # правило суммы
        o = D.lift(other)
        return D(self.val + o.val, self.deriv + o.deriv)
    __radd__ = __add__

    def __sub__(self, other):
        o = D.lift(other)
        return D(self.val - o.val, self.deriv - o.deriv)

    def __rsub__(self, other):
        return D.lift(other) - self

    def __mul__(self, other):         # правило произведения
        o = D.lift(other)
        return D(self.val * o.val, self.deriv * o.val + self.val * o.deriv)
    __rmul__ = __mul__

    def __truediv__(self, other):     # правило частного
        o = D.lift(other)
        return D(self.val / o.val, (o.val * self.deriv - self.val * o.deriv) / o.val**2)

    def __rtruediv__(self, other):
        return D.lift(other) / self

    def __repr__(self):
        return f"{self.val} + {self.deriv}ϵ"


def babylonian(x, N=10):
    t = (1 + x) / 2              # один шаг из t = 1
    for i in range(2, N + 1):    # оставшиеся N-1 шагов
        t = (t + x / t) / 2
    return t


print(babylonian(4, N=3))        # 2.000609756097561
print(babylonian(D(49, 1)))      # 7.0 + 0.07142857142857142ϵ
print(0.5 / 49**0.5)             # 0.07142857142857142 — производная √x при x = 49
eps = D(0, 1)
print(eps * eps)                 # 0.0 + 0.0ϵ — то самое ϵ² = 0

8.1.3. Дуальные числа

Есть приятный алгебраический способ думать о нашем новом типе чисел $D(a, b)$ вместо точки зрения «значение и производная». Помните, как комплексное число $a + bi$ строится из двух вещественных $(a, b)$: вводится особая новая величина $i$ (мнимая единица), удовлетворяющая $i^2 = -1$, и все остальные правила комплексной арифметики следуют отсюда? Аналогично можно считать $D(a, b)$ числом $a + b\epsilon$, где $\epsilon$ — новая «бесконечно малая единица», удовлетворяющая $\epsilon^2 = 0$. Такая точка зрения называется дуальными числами.

Из элементарного правила $\epsilon^2 = 0$ сразу следуют остальные алгебраические правила дуальных чисел:

$$\begin{aligned} (a + b\epsilon) \pm (c + d\epsilon) &= (a \pm c) + (b \pm d)\epsilon, \\ (a + b\epsilon) \cdot (c + d\epsilon) &= (ac) + (bc + ad)\epsilon, \\ \frac{a + b\epsilon}{c + d\epsilon} &= \frac{a + b\epsilon}{c + d\epsilon} \cdot \frac{c - d\epsilon}{c - d\epsilon} = \frac{(a + b\epsilon)(c - d\epsilon)}{c^2} = \frac{a}{c} + \frac{bc - ad}{c^2}\epsilon. \end{aligned}$$

Коэффициенты при $\epsilon$ в этих правилах соответствуют правилам суммы/разности, произведения и частного из дифференциального исчисления!

На самом деле это в точности правила, которые мы реализовали выше для типа D. Не хватало только правил вычитания и умножения, которые теперь можно добавить:

julia> Base.:-(x::D, y::D) = D(x.val - y.val, x.deriv - y.deriv)
       Base.:*(x::D, y::D) = D(x.val*y.val, x.deriv*y.val + x.val*y.deriv)

Ещё приятно добавить правило «красивой печати», чтобы Julia показывала дуальные числа как a + bϵ, а не как D(a,b):

julia> Base.show(io::IO, x::D) = print(io, x.val, " + ", x.deriv, "ϵ")

Как только в Julia реализовано правило умножения дуальных чисел, $\epsilon^2 = 0$ следует из частного случая $a = c = 0$ и $b = d = 1$:

julia> ϵ = D(0,1)
0.0 + 1.0ϵ

julia> ϵ * ϵ
0.0 + 0.0ϵ

julia> ϵ^2
0.0 + 0.0ϵ

(Мы не определяли правило для степеней $D(a, b)^n$ — как же было вычислено ϵ^2? Ответ: по умолчанию Julia вычисляет $x^n$ повторным умножением, так что хватило определить правило для $*$.) Теперь производную вавилонского алгоритма в $x = 49$ можно вычислить, как выше:

julia> Babylonian(x + ϵ)
7.0 + 0.07142857142857142ϵ

где «бесконечно малая часть» — это производная $0{,}5/\sqrt{49} = 0{,}0714\cdots$

Приятно, что точка зрения дуальных чисел прямо соответствует нашему понятию производной как линеаризации:

$$f(x + \epsilon) = f(x) + f'(x)\epsilon + \text{(члены высших порядков)},$$

а правило дуальных чисел $\epsilon^2 = 0$ соответствует отбрасыванию членов высших порядков.

8.2. Наивное символьное дифференцирование

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

Представьте, что мы подаём в код Babylonian «символьную переменную» $x$, выполняем алгоритм и записываем большое алгебраическое выражение для результата. После одного шага, например, получим $(x + 1)/2$. После двух шагов — $((x + 1)/2 + 2x/(x + 1))/2$, что упрощается до отношения двух многочленов («рациональной функции»):

$$\frac{x^2 + 6x + 1}{4(x + 1)}.$$

Продолжать вручную довольно утомительно, но, к счастью, это может сделать компьютер (как показано в Julia-блокноте). Три вавилонские итерации дают

$$\frac{x^4 + 28x^3 + 70x^2 + 28x + 1}{8\left(x^3 + 7x^2 + 7x + 1\right)},$$

четыре итерации —

$$\frac{x^8 + 120x^7 + 1820x^6 + 8008x^5 + 12870x^4 + 8008x^3 + 1820x^2 + 120x + 1}{16\left(x^7 + 35x^6 + 273x^5 + 715x^4 + 715x^3 + 273x^2 + 35x + 1\right)},$$

а пять итераций — огромное выражение:

Пять итераций (прокрутите формулу вбок)

$$\scriptsize \frac{x^{16} + 496x^{15} + 35960x^{14} + 906192x^{13} + 10518300x^{12} + 64512240x^{11} + 225792840x^{10} + 471435600x^{9} + 601080390x^{8} + 471435600x^{7} + 225792840x^{6} + 64512240x^{5} + 10518300x^{4} + 906192x^{3} + 35960x^{2} + 496x + 1}{32\left(x^{15} + 155x^{14} + 6293x^{13} + 105183x^{12} + 876525x^{11} + 4032015x^{10} + 10855425x^{9} + 17678835x^{8} + 17678835x^{7} + 10855425x^{6} + 4032015x^{5} + 876525x^{4} + 105183x^{3} + 6293x^{2} + 155x + 1\right)}$$

Заметьте, как быстро они растут — степень многочленов удваивается на каждой итерации! Если теперь взять символьные производные этих функций по обычным правилам анализа и упростить (с помощью компьютера), то производная одной итерации равна $\frac{1}{2}$, двух итераций —

$$\frac{x^2 + 2x + 5}{4\left(x^2 + 2x + 1\right)},$$

трёх итераций —

$$\frac{x^6 + 14x^5 + 147x^4 + 340x^3 + 375x^2 + 126x + 21}{8\left(x^6 + 14x^5 + 63x^4 + 100x^3 + 63x^2 + 14x + 1\right)},$$

четырёх итераций —

Производная четырёх итераций

$$\scriptsize \frac{x^{14} + 70x^{13} + 3199x^{12} + 52364x^{11} + 438945x^{10} + 2014506x^{9} + 5430215x^{8} + 8836200x^{7} + 8842635x^{6} + 5425210x^{5} + 2017509x^{4} + 437580x^{3} + 52819x^{2} + 3094x + 85}{16\left(x^{14} + 70x^{13} + 1771x^{12} + 20540x^{11} + 126009x^{10} + 440986x^{9} + 920795x^{8} + 1173960x^{7} + 920795x^{6} + 440986x^{5} + 126009x^{4} + 20540x^{3} + 1771x^{2} + 70x + 1\right)}$$

а производная пяти итераций — чудовище, которое можно прочитать, только сильно увеличив масштаб:

Производная пяти итераций

$$\scriptsize \frac{x^{30} + 310x^{29} + 59799x^{28} + 4851004x^{27} + 215176549x^{26} + 5809257090x^{25} + 102632077611x^{24} + 1246240871640x^{23} + 10776333438765x^{22} + 68124037776390x^{21} + 321156247784955x^{20} + 1146261110726340x^{19} + 3133113888931089x^{18} + 6614351291211874x^{17} + 10850143060249839x^{16} + 13883516068991952x^{15} + 13883516369532147x^{14} + 10850142795067314x^{13} + 6614351497464949x^{12} + 3133113747810564x^{11} + 1146261195398655x^{10} + 321156203432790x^{9} + 68124057936465x^{8} + 10776325550040x^{7} + 1246243501215x^{6} + 102631341330x^{5} + 5809427001x^{4} + 215145084x^{3} + 4855499x^{2} + 59334x + 341}{32\left(x^{30} + 310x^{29} + 36611x^{28} + 2161196x^{27} + 73961629x^{26} + 1603620018x^{25} + 23367042639x^{24} + 238538538360x^{23} + 1758637118685x^{22} + 9579944198310x^{21} + 39232152623175x^{20} + 122387258419860x^{19} + 293729420641881x^{18} + 546274556891506x^{17} + 791156255418003x^{16} + 894836006026128x^{15} + 791156255418003x^{14} + 546274556891506x^{13} + 293729420641881x^{12} + 122387258419860x^{11} + 39232152623175x^{10} + 9579944198310x^{9} + 1758637118685x^{8} + 238538538360x^{7} + 23367042639x^{6} + 1603620018x^{5} + 73961629x^{4} + 2161196x^{3} + 36611x^{2} + 310x + 1\right)}$$

Это ужасный способ вычислять производные! (Впрочем, существуют более изощрённые подходы к эффективному символьному дифференцированию, например алгоритм «$D^*$», которые избегают явных гигантских формул, используя повторяющиеся подвыражения.)

Чтобы было ясно: подход дуальных чисел (без учёта ошибок округления) вычисляет ответ в точности так, как если бы подставил некоторое конкретное $x$ в эти безумные выражения, но слова «как если бы» здесь очень важны. Как видите, мы эти выражения не строим и тем более не вычисляем. Мы просто получаем результаты, равные тем значениям, которые получили бы, если бы их построили.

8.3. Автоматическое дифференцирование через вычислительные графы

Перейдём к автоматическому дифференцированию через вычислительные графы. В этом разделе мы рассматриваем следующий простой мотивирующий пример.

Пример. Определим функции

$$\begin{cases} a(x, y) = \sin x, \\ b(x, y) = \frac{1}{y} \cdot a(x, y), \\ z(x, y) = b(x, y) + x. \end{cases}$$

Найдите $\frac{\partial z}{\partial x}$ и $\frac{\partial z}{\partial y}$.

Решить задачу можно несколькими способами. Во-первых, конечно, можно вычислить символьно, заметив, что

$$z(x, y) = b(x, y) + x = \frac{1}{y} a(x, y) + x = \frac{\sin x}{y} + x,$$

откуда

$$\frac{\partial z}{\partial x} = \frac{\cos x}{y} + 1 \quad \text{и} \quad \frac{\partial z}{\partial y} = -\frac{\sin x}{y^2}.$$

Однако можно воспользоваться и вычислительным графом (рисунок ниже), в котором ребро из узла $A$ в узел $B$ подписано производной $\frac{\partial B}{\partial A}$.

flowchart LR
X(("x")) -->|"cos x"| A(("a = sin x"))
A -->|"1/y"| B(("b = a/y"))
B -->|"1"| Z(("z = b + x"))
X -->|"1"| Z
Y(("y")) -->|"−a/y²"| B

Вычислительный граф для примера: выход z(x, y) вычисляется из двух входов x и y через промежуточные a и b. Узлы подписаны значениями, рёбра — производными значения по предыдущему значению.

Как найти производные по этому ориентированному ациклическому графу (DAG)? Один взгляд («прямой») — идти по путям от входов к выходам и (слева) умножать по ходу, складывая несколько путей. Например, проделав это для путей из $x$ в $z(x, y)$, получим

$$\frac{\partial z}{\partial x} = 1 \cdot \frac{1}{y} \cdot \cos x + 1 = \frac{\cos x}{y} + 1.$$

Аналогично для путей из $y$ в $z(x, y)$:

$$\frac{\partial z}{\partial y} = 1 \cdot \frac{-a(x, y)}{y^2} = \frac{-\sin x}{y^2},$$

и если на рёбрах числовые производные, этот алгоритм работает. Можно пойти и «обратным» путём — идти по путям назад (умножая справа налево) — и получить тот же результат. Заметьте, что ничего волшебного в том, что здесь всё скалярное, нет: можно представить, что это функции того вида, что мы встречаем в этом курсе, и проделать те же вычисления! Принципиально здесь важна только ассоциативность. Однако для векторных функций порядок, в котором перемножаются веса рёбер, жизненно важен (поскольку векторные и матричные функции в общем случае не коммутируют).

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

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

flowchart LR
B1(("B₁")) --> A(("A"))
B2(("B₂")) --> A
B3(("B₃")) --> A
A -->|"∂f(A)/∂A"| F(("f(A)"))

Узел A получает значения от узлов B₁, B₂, B₃ и передаёт f(A) дальше.

Пусть мы знаем произведение по пути $P$ по всем рёбрам до ребра из $B_2$ в $A$ включительно. Каким будет новое произведение по пути, когда мы сдвинемся вправо от $A$? Это $f'(A) \cdot P$! Значит, нужна структура данных, которая отображает

$$(\text{значение}, \text{произведение по пути}) \mapsto (f(\text{значение}), f' \cdot \text{произведение по пути}).$$

В некотором смысле это ещё один взгляд на дуальные числа: принимаем произведения по путям и выдаём значения. В любом случае мы перегружаем программу, которая легко вычисляет $f(\text{значение})$, и добавляем к ней $f' \cdot (\text{произведение по пути})$.

Можно спросить, с чего программа начинается: мы описали, как она работает «в середине», а каким должно быть начальное значение? Единственное, при котором метод работает, — $(x, 1)$. Затем на каждом шаге выполняется указанное выше отображение

$$(\text{значение}, \text{произведение по пути}) \mapsto (f(\text{значение}), f' \cdot \text{произведение по пути}),$$

и в конце мы получаем производные.

Как теперь объединять стрелки? Иными словами, пусть в двух узлах слева у нас значения $(a, p)$ и $(b, q)$, как на схеме ниже.

flowchart LR
P["(a, p)"] -->|"∂z/∂a"| Z(("z = f(a, b)"))
Q["(b, q)"] -->|"∂z/∂b"| Z

Два входящих ребра: значения a, b и накопленные произведения по путям p, q.

Здесь мы думаем об $a$, $b$ не как о числах, а как о переменных. Каким должно быть новое выходное значение? Мы хотим сложить два произведения по путям и получить

$$\left(f(a, b),\ \frac{\partial z}{\partial a} p + \frac{\partial z}{\partial b} q\right).$$

То есть наша перегруженная структура данных на самом деле переводит пары $(a, p)$ и $(b, q)$ в пару $\left(f(a, b),\ \frac{\partial z}{\partial a} p + \frac{\partial z}{\partial b} q\right)$. Разумеется, это обобщается на случай, когда слева в графе много узлов.

Если придумать такую структуру для всех простых вычислений (сложения/вычитания, умножения и деления) и если программе больше ничего не нужно, то всё готово! Вот как структура определяется для сложения/вычитания, умножения и деления.

• Сложение/вычитание (рисунок «граф сложения/вычитания»): входы $(a_1, p = 1)$ и $(a_2, q = \pm 1)$ дают узел $\left(z = a_1 \pm a_2,\ \frac{\partial z}{\partial a_1} \cdot 1 + \frac{\partial z}{\partial a_2} \cdot (\pm 1)\right)$.
• Умножение (рисунок «граф умножения»): входы $(a_1, p = a_2)$ и $(a_2, q = a_1)$ дают узел $\left(z = a_1 a_2,\ \frac{\partial z}{\partial a_1} \cdot a_2 + \frac{\partial z}{\partial a_2} \cdot a_1\right)$.
• Деление (рисунок «граф деления»): входы $(a_1, p = a_2/a_2^2)$ и $(a_2, q = -a_1/a_2^2)$ дают узел $\left(z = a_1/a_2,\ \frac{\partial z}{\partial a_1} \cdot \frac{1}{a_2} - \frac{\partial z}{\partial a_2} \cdot \frac{a_1}{a_2^2}\right)$.

Примечание переводчика: обозначения на этих рисунках оригинала довольно вольные. Смысл такой: вес ребра от $a_1$ к $z$ — это $\partial z / \partial a_1$ (для умножения $a_2$, для деления $1/a_2$), от $a_2$ к $z$ — $\partial z / \partial a_2$ (для умножения $a_1$, для деления $-a_1/a_2^2$); эти веса умножаются на накопленные во входах произведения по путям и складываются. Запись $p = a_2/a_2^2$ — это просто $1/a_2$.

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

8.3.1. Обратный режим автоматического дифференцирования на графах

В обратном режиме стрелки идут в другую сторону, и в этом разделе мы разберёмся, как это устроено. В прямом режиме всё сводилось к вопросу «от чего мы зависим?», то есть к вычислению производной в правой части схемы выше по функциям в узлах слева. В обратном режиме вопрос скорее такой: «на что мы влияем?» или «на что мы повлияем позже?»

Двигаясь «назад», нужно знать, на какие узлы влияет данный узел. Например, для узла $A$ мы хотим знать узлы $B_i$, на которые $A$ влияет, то есть которые от него зависят. Теперь схема выглядит так:

flowchart LR
X["(x, ∂z/∂x)"] -.-> A["(a, ∂z/∂a)"]
A --> B1["(b₁, ∂z/∂b₁)"]
A --> B2["(b₂, ∂z/∂b₂)"]
A --> B3["(b₃, ∂z/∂b₃)"]
B1 -.-> Z["(z, 1)"]
B2 -.-> Z
B3 -.-> Z

Обратный режим: узел a влияет на b₁, b₂, b₃, а те — (через другие узлы, пунктир) на выход z. Вычисление начинается с (z, 1) справа.

Итак, в конце у нас есть финальный узел $(z, 1)$ (крайний справа), с которого всё начинается. На этот раз все умножения идут справа налево, поскольку мы в обратном режиме. Наша цель — вычислить узел $(x, \partial z / \partial x)$. Значит, если мы знаем, как заполнить член $\frac{\partial z}{\partial a}$, то сможем двигаться по этим вычислительным графам справа налево (то есть в обратном режиме). Формула для $\frac{\partial z}{\partial a}$ такова:

$$\frac{\partial z}{\partial a} = \sum_{i=1}^{s} \frac{\partial b_i}{\partial a} \frac{\partial z}{\partial b_i},$$

где $b_i$ — узлы, на которые влияет узел $A$. Это снова просто правило цепочки из анализа, но можно смотреть на это и как на суммы произведений всех весов в той части графа, на которую влияет $A$.

flowchart LR
X(("x")) -->|"a"| P(("p"))
Y(("y")) -->|"b"| P
P -->|"c"| Q(("q"))
Q -->|"d"| Z(("z"))

Граф «источники — сток»: два входа x, y сходятся в p, дальше одна цепочка p → q → z.

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

Если $x$, $y$ — наши источники, а $z$ — сток, мы хотим вычислить сумму произведений весов по путям от источников к стокам. В прямом режиме пришлось бы вычислить пути $dca$ и $dcb$, а это четыре умножения (и потом их сложить). В обратном режиме достаточно вычислить $a\,\underline{cd}$ и $b\,\underline{cd}$ и сложить; заметьте, что обратный режим (поскольку $cd$ вычисляется только один раз) требует всего 3 умножения. В общем случае так можно эффективнее решать определённые типы задач, например задачу «источники — сток».

Связь с micrograd (комментарий переводчика)

Формула $\frac{\partial z}{\partial a} = \sum_i \frac{\partial b_i}{\partial a} \frac{\partial z}{\partial b_i}$ — это ровно строчка a.grad += local_derivative * b.grad в методе backward из урока 3.2.2. Знак += — это и есть сумма по узлам $b_i$, на которые влияет $a$; а обход в обратном топологическом порядке гарантирует, что все $\partial z / \partial b_i$ уже посчитаны, когда очередь доходит до $a$.

8.4. Прямой и обратный режимы: сравнение

Кратко подытожим относительные плюсы и минусы двух подходов к вычислению производных (вручную или с помощью программ AD). С математической точки зрения два подхода — зеркальные отражения друг друга, но с вычислительной они сильно различаются, потому что программы обычно выполняются «вперёд» во времени, от входов к выходам.

Пусть мы дифференцируем функцию $f: \mathbb{R}^n \mapsto \mathbb{R}^m$, которая отображает $n$ скалярных входов ($n$-мерный вход) в $m$ скалярных выходов ($m$-мерный выход). Первое ключевое различие прямого и обратного режимов — как вычислительная стоимость растёт с числом (размерностью) входов и выходов:

• Стоимость прямого дифференцирования (от входов к выходам) пропорциональна $n$ — числу входов. Это идеально для функций с $n \ll m$ (мало входов, много выходов).
• Стоимость обратного дифференцирования (от выходов к входам) пропорциональна $m$ — числу выходов. Это идеально для функций с $m \ll n$ (мало выходов, много входов).

До этой главы мы впервые увидели такую зависимость в разделе о стоимости умножения матриц (часть 1, раздел 2.5.1), а затем в методе сопряжённых уравнений (глава 6); в одной из следующих лекций увидим её ещё раз (чувствительность решений ОДУ, глава 9). Случай немногих выходов крайне распространён в крупномасштабной оптимизации (будь то машинное обучение, инженерное проектирование или другие приложения): параметров оптимизации много ($n \gg 1$), а выход один ($m = 1$) — целевая функция (функция потерь), — или несколько выходов: цель и ограничения. Поэтому обратный режим (backpropagation) — основной подход в крупномасштабной оптимизации и таких приложениях, как обучение нейросетей.

Есть, однако, и другие практические соображения:

• Прямое дифференцирование идёт в том же порядке, что и вычисление самой функции, — от входов к выходам. Похоже, поэтому AD в прямом режиме проще реализовать (например, как наша реализация в разделе 8.1) и сделать эффективным.
• Обратное дифференцирование идёт в противоположном обычному вычислению направлении. Из-за этого AD в обратном режиме гораздо сложнее реализовать, и оно добавляет к вычислению функции много накладных расходов на память. Сначала вы вычисляете функцию от входов к выходам, но вы (или система AD) ведёте запись («ленту», tape) всех промежуточных шагов вычисления; затем вы выполняете вычисление в обратном порядке («прокручиваете ленту назад»), обратно распространяя производные.

Из-за этих практических преимуществ, даже в случае многих ($n > 1$) входов и одного ($m = 1$) выхода, практики говорят нам, что прямой режим оказывается эффективнее, пока $n$ не станет достаточно большим (возможно, даже пока $n > 100$, в зависимости от дифференцируемой функции и реализации AD). (Вас может заинтересовать и пост в блоге Криса Ракаукаса Engineering Trade-offs in AD, в основном о реализациях обратного режима.)

Если $n = m$ и ни у одного подхода нет преимущества в масштабировании, обычно предпочитают меньшие накладные расходы и простоту прямого дифференцирования. Такой случай возникает при вычислении явных матриц Якоби для нелинейного поиска корней (глава 6) или матриц вторых производных — гессианов (глава 12), для которых часто используют прямой режим… или даже сочетание прямого и обратного, как обсуждается ниже.

Конечно, прямой и обратный — не единственные варианты. Правило цепочки ассоциативно, так что возможных порядков много (например, начать с обоих концов и встретиться посередине или наоборот). Трудная¹ задача, которая часто требует гибридных схем, — вычислить якобианы (или гессианы) за минимальное число операций, используя любую специфичную для задачи структуру (например, разреженность: многие элементы могут быть нулями). Это и другие темы AD обсуждаются гораздо подробнее, чем в этом конспекте, в книге Evaluating Derivatives (2-е изд.) Гриванка и Вальтер (2008).

¹ На самом деле необычайно трудная: «NP-полная» (Naumann, 2006).

8.4.1. Прямой поверх обратного: вторые производные

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

Вычисление гессиана. Рассмотрим, например, функцию $f(x): \mathbb{R}^n \to \mathbb{R}$, отображающую $n$ входов $x$ в один скаляр. Первую производную $f'(x) = (\nabla f)^T$ лучше всего вычислять обратным режимом, если $n \gg 1$ (много входов). Но теперь рассмотрим вторую производную — производную функции $g(x) = \nabla f$, которая отображает $n$ входов $x$ в $n$ выходов $\nabla f$. Понятно, что $g'(x)$ — матрица Якоби $n \times n$, которая называется гессианом $f$ (в главе 12 он обсуждается гораздо шире). Поскольку у $g(x)$ одинаковое число входов и выходов, ни у прямого, ни у обратного режима нет внутреннего преимущества в масштабировании, так что для $g'$ обычно выбирают прямой режим благодаря его практической простоте, а $\nabla f$ по-прежнему вычисляют обратным. То есть мы вычисляем $\nabla f$ обратным режимом, а затем вычисляем $g' = (\nabla f)'$, применяя прямое дифференцирование к алгоритму вычисления $\nabla f$. Это называется алгоритмом «прямой поверх обратного» (forward-over-reverse).

Ещё более явное применение схемы «прямой поверх обратного» — произведения гессиана на вектор (Hessian–vector products). Во многих приложениях оказывается, что нужен только произведение $(\nabla f)' v$ гессиана $(\nabla f)'$ на произвольный вектор $v$. В этом случае можно вообще не вычислять (и не хранить) матрицу гессиана явно, а стоимость будет пропорциональна лишь стоимости одного вычисления функции $f(x)$. Приём такой: вспомним (из раздела 2.2.1), что для любой функции $g$ линейная операция $g'(x)[v]$ — это производная по направлению, равносильная производной одной переменной $\frac{\partial}{\partial\alpha} g(x + \alpha v)$ при $\alpha = 0$. Здесь мы просто применяем это правило к функции $g(x) = \nabla f$ и получаем формулу произведения гессиана на вектор:

$$(\nabla f)' v = \left. \frac{\partial}{\partial\alpha} \left( \left. \nabla f \right|_{x + \alpha v} \right) \right|_{\alpha = 0}.$$

Вычислительно внутреннее вычисление градиента $\nabla f$ в произвольной точке $x + \alpha v$ эффективно делается обратным алгоритмом (adjoint, backpropagation). А внешнюю производную по одному входу $\alpha$ лучше всего брать прямым режимом.² Поскольку гессиан симметричен (это в большой общности обсуждается в главе 12), тот же алгоритм работает и для произведений вектора на гессиан $v^T (\nabla f)' = [(\nabla f)' v]^T$ — этим мы воспользуемся в следующем примере.

² В The Autodiff Cookbook — части документации JAX — этот алгоритм обсуждается в разделе о произведениях гессиана на вектор. Там отмечено, что можно поменять местами производные $\partial/\partial\alpha$ и $\nabla_x$ и применить режим «обратный поверх прямого», но на практике это, видимо, менее эффективно: «поскольку у прямого режима меньше накладных расходов, чем у обратного, а внешнему оператору дифференцирования приходится дифференцировать большее вычисление, чем внутреннему, держать прямой режим снаружи — лучше всего». Там же предложена ещё одна альтернатива: по тождеству $(\nabla f)' v = \nabla(v^T \nabla f)$ можно применить режим «обратный поверх обратного» и взять градиент $v^T \nabla f$, но накладные расходы тут ещё больше.

Скалярные функции от градиентов. Есть ещё одна частая ситуация, в которой прямое и обратное дифференцирование сочетают, но которая выглядит тоньше: дифференцирование скалярной функции от градиента другой скалярной функции. Рассмотрим пример.

Пример. Пусть $f(x): \mathbb{R}^n \mapsto \mathbb{R}$ — скалярная функция $n \gg 1$ входов с градиентом $\left. \nabla f \right|_x = f'(x)^T$, и пусть $g(z): \mathbb{R}^n \mapsto \mathbb{R}$ — ещё одна такая функция с градиентом $\left. \nabla g \right|_z = g'(z)^T$. Рассмотрим скалярную функцию $h(x) = g(\left. \nabla f \right|_x): \mathbb{R}^n \mapsto \mathbb{R}$ и найдём $\left. \nabla h \right|_x = h'(x)^T$.

Обозначим $z = \left. \nabla f \right|_x$. По правилу цепочки $h'(x) = g'(z)(\nabla f)'(x)$, но мы хотим избежать явного вычисления большой матрицы гессиана $(\nabla f)'$ размера $n \times n$. Вместо этого, как обсуждалось выше, воспользуемся тем, что такое произведение вектора на гессиан равносильно (по симметрии гессиана) транспонированному произведению гессиана $(\nabla f)'$ на вектор $\nabla g = g'(z)^T$, а оно равносильно производной по направлению:

$$\left. \nabla h \right|_x = h'(x)^T = \left. \frac{\partial}{\partial\alpha} \left( \left. \nabla f \right|_{x + \alpha \left. \nabla g \right|_z} \right) \right|_{\alpha = 0},$$

где дифференцирование идёт по одному скаляру $\alpha \in \mathbb{R}$. Значит, как и для любого произведения гессиана на вектор, $h$ и $\nabla h$ можно вычислить так:

1. Вычислить $h(x)$: вычислить $z = \left. \nabla f \right|_x$ обратным режимом и подставить в $g(z)$.
2. Вычислить $\nabla h$:
   1. вычислить $\left. \nabla g \right|_z$ обратным режимом;
   2. реализовать $\left. \nabla f \right|_{x + \alpha \left. \nabla g \right|_z}$ обратным режимом, а затем продифференцировать по $\alpha$ прямым режимом при $\alpha = 0$.

Это алгоритм «прямой поверх обратного»: прямой режим эффективно используется для производной по единственному входу $\alpha \in \mathbb{R}$ в сочетании с обратным режимом для дифференцирования по $x, z \in \mathbb{R}^n$.

Ниже — пример кода на Julia, реализующий этот процесс «прямой поверх обратного» как раз для такой функции $h(x) = g(\nabla f)$. Здесь прямое дифференцирование по $\alpha$ выполняет пакет ForwardDiff.jl из раздела 8.1, а обратное дифференцирование по $x$ или $z$ — пакет Zygote.jl. Сначала импортируем пакеты и определим простые функции-примеры $f(x) = 1/\Vert x \Vert$ и $g(z) = (\sum_k z_k)^3$, а также вычисление $h$ через Zygote:

julia> using ForwardDiff, Zygote, LinearAlgebra
julia> f(x) = 1/norm(x)
julia> g(z) = sum(z)^3
julia> h(x) = g(Zygote.gradient(f, x)[1])

Теперь вычислим $\nabla h$ схемой «прямой поверх обратного»:

julia> function ∇h(x)
           ∇f(y) = Zygote.gradient(f, y)[1]
           ∇g = Zygote.gradient(g, ∇f(x))[1]
           return ForwardDiff.derivative(α -> ∇f(x + α*∇g), 0)
       end

Теперь подставим случайные числа и сравним с проверкой конечными разностями:

julia> x = randn(5); δx = randn(5) * 1e-8;

julia> h(x)
-0.005284687528953334

julia> ∇h(x)
5-element Vector{Float64}:
 -0.006779692698531759
  0.007176439898271982
 -0.006610264199241697
 -0.0012162087082746558
  0.007663756720005014

julia> ∇h(x)' * δx    # производная по направлению
-3.0273434457397667e-10

julia> h(x+δx) - h(x)  # проверка конечной разностью
-3.0273433933303284e-10

Проверка конечной разностью совпадает примерно до 7 значащих цифр — на большее рассчитывать нельзя: код «прямой поверх обратного» работает!

Проверьте в Python (дополнение переводчика)

То же на PyTorch: torch.func.grad — обратный режим, torch.func.jvp — прямой (произведение якобиана на вектор, то есть производная по направлению). Числа будут другими — у Julia и PyTorch разные генераторы случайных чисел.

import torch
from torch.func import grad, jvp

torch.manual_seed(0)
torch.set_default_dtype(torch.float64)        # двойная точность, как в Julia

f = lambda x: 1 / torch.linalg.norm(x)        # f(x) = 1/‖x‖
g = lambda z: z.sum() ** 3                    # g(z) = (Σ z_k)³
h = lambda x: g(grad(f)(x))                   # h(x) = g(∇f(x))

def grad_h(x):
    grad_f = grad(f)                          # ∇f — обратный режим
    v = grad(g)(grad_f(x))                    # ∇g в точке z = ∇f(x) — тоже обратный режим
    _, hvp = jvp(grad_f, (x,), (v,))          # d/dα ∇f(x + α v) при α = 0 — прямой режим
    return hvp

x = torch.randn(5)
dx = torch.randn(5) * 1e-8
print((grad_h(x) @ dx).item())                # производная по направлению
print((h(x + dx) - h(x)).item())              # проверка: совпадает примерно до 7 значащих цифр

Задача. Распространённый вариант этой процедуры, который часто встречается в машинном обучении, — функция $f(x, p) \in \mathbb{R}$, отображающая входные «данные» $x \in \mathbb{R}^n$ и «параметры» $p \in \mathbb{R}^N$ в скаляр. Обозначим через $\nabla_x f$ и $\nabla_p f$ градиенты по $x$ и по $p$.

Пусть, как и раньше, есть функция $g(z): \mathbb{R}^n \mapsto \mathbb{R}$, и определим $h(x, p) = g(\left. \nabla_x f \right|_{x,p})$. Мы хотим вычислить $\nabla_p h = (\partial h / \partial p)^T$, что потребует «смешанных» производных $f$ и по $x$, и по $p$.

Покажите, что $\nabla_p h$ можно вычислить как

$$\left. \nabla_p h \right|_{x,p} = \left. \frac{\partial}{\partial\alpha} \left( \left. \nabla_p f \right|_{x + \alpha \left. \nabla g \right|_z,\, p} \right) \right|_{\alpha = 0},$$

где $z = \left. \nabla_x f \right|_{x,p}$. (Принципиально, что так мы никогда не вычисляем матрицу смешанных производных $f$ размера $n \times N$.)

Придумайте простые функции-примеры $f$ и $g$, реализуйте эту формулу схемой «прямой поверх обратного» на Julia по образцу выше (прямой режим для $\partial/\partial\alpha$ и обратный — для $\nabla$) и проверьте результат приближением конечными разностями.

Комментарий переводчика: такие «градиенты от градиентов» встречаются, например, в штрафе на норму градиента (gradient penalty) при обучении GAN и в метаобучении (MAML), где шаг обучения сам дифференцируется по параметрам.

Словарик

EN	RU
automatic differentiation (AD)	автоматическое дифференцирование
forward mode / reverse mode	прямой режим / обратный режим
dual number, infinitesimal unit $\epsilon$	дуальное число, бесконечно малая единица $\epsilon$
operator overloading	перегрузка операторов
conversion and promotion rules	правила преобразования и продвижения типов
symbolic differentiation	символьное дифференцирование
finite differences	конечные разности
rational function	рациональная функция
computational graph, DAG	вычислительный граф, ориентированный ациклический граф
path product	произведение по пути
source / sink	источник / сток
tape	лента (запись промежуточных шагов)
storage overhead	накладные расходы на память
Hessian, Hessian–vector product	гессиан, произведение гессиана на вектор
forward-over-reverse	«прямой поверх обратного»
Jacobian–vector product (jvp), vector–Jacobian product (vjp)	произведение якобиана на вектор, вектора на якобиан

Назад к карте курса: путеводитель по Matrix Calculus.