Смотрели итоги прошедшего ICLR? Меня заинтересовала довольно провокационная, на первый взгляд, статья от Эплов — ParaRNN. Казалось бы, параллельность РНН — это их главный недостаток, благодаря которому их заменили трансформеры (в большинстве задач).

Так вот, давайте разберемся со всем, на максимально низком уровне, если знаете, что такое RNN и производная — то эта статья для вас.

1. Алгоритм DEER

DEER = Deep Equilibrium Evaluation of Recurrence (Lim et al., 2024). Базовый алгоритм, на котором строится ParaRNN.

1.1. Постановка как задача нахождения корня

Пусть у нас есть обыкновенная RNN с переходной функцией , начальным состоянием и неизвестными состояниями . Введем остаток (residual):

Истинная траектория - это единственное решение уравнения:

Когда говорят «применить RNN к последовательности», имеют в виду стандартную процедуру: взять начальное состояние , применить переходную функцию , получить , потом еще раз применить , получить , и так далее:

Соответственно, получается, что - вектор, у которого все элементы равны 0, опять же потому что при соблюдении рекуррентности и .

1.2. Итерации Ньютона

Соответственно дальше, необходимо найти решение уравнения , или в полном случае — вектор, решающий систему уравнений. Но для начала разберемся со скалярным случаем.

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

Пусть у нас есть гладкая функция и мы хотим найти такое , что . Геометрически — найти точку, где график функции пересекает ось абсцисс.

Идея метода Ньютона строится на простой мысли: в малой окрестности точки гладкая функция почти неотличима от своей касательной. Если мы стоим в текущем приближении (которое, в общем случае, не корень — там ), мы можем сделать вид, что — это ее касательная в этой точке, и для такой линейной функции легко аналитически найти, где она пересекает ось.

Касательная к в точке — это первое слагаемое разложения Тейлора:

Вспомним: Разложение Тейлора — это способ приблизить любую гладкую функцию вблизи точки многочленом:

где каждое следующее слагаемое уточняет приближение, добавляя информацию о все более тонкой особенности формы функции (наклон, кривизна, и т.д.). Логический смысл такой: если функция гладкая, то ее поведение в окрестности точки полностью закодировано в значениях ее производных в этой одной точке — измерив несколько чисел в , мы можем восстановить значения функции рядом. Делитель возникает естественно из требования, чтобы в точке совпадали все производные многочлена и самой функции (он сокращается с факториалом, выскакивающим при -кратном дифференцировании ).

Приравниваем эту линейную аппроксимацию к нулю и находим, где она пересекает ось:

Решаем относительно :

Это и объявляем следующим приближением:

Тут показан графически шаг , и на графике видно, что корень уравнения (нас интересует пересечение функции с осью абсцисс) поменялся с 2 до 1, что показывает улучшение, так как эталонное значение — 0.

Можно переписать это через приращение , что окажется удобнее при обобщении:

То есть «найти такое приращение, чтобы линейная поправка скомпенсировала текущий остаток ».

Многомерный случай: уравнений от переменных

Теперь обобщаем. Вместо одной функции от одной переменной — векторнозначная функция векторного аргумента, и ищем такой вектор , что .

Логика остается такой же, меняются только объекты или их размерности:

Где — якобиан, многомерный аналог обычной производной (об этом подробнее ниже).

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

Линеаризация вокруг точки :

Приравниваем линейную аппроксимацию к нулевому вектору:

И обозначив приращение , получаем линейную систему относительно :

Решив систему и получив , обновляем приближение:

Также можно записать более компактно через обратную матрицу:

— это та же самая формула, просто короче. Запись с — чисто нотационная: на практике обратную матрицу никто никогда не вычисляет, потому что это и дорого, и численно неустойчиво. Вместо этого решают систему напрямую — например, через LU-разложение, или через прямую подстановку, если имеет специальную структуру (что и происходит у нас).

1.3. Применение к нашей задаче с RNN

В случае RNN все ровно по этому шаблону, только размерности конкретные:

• - все скрытые состояния, склеенные в один длинный вектор длины .

• — вектор всех одношаговых остатков, той же длины.

• — якобиан остатка по состоянию.

Применяем тот же ньютоновский шаг:

И тут возникает вопрос: «А разве решить линейную систему размера — это не та же самая последовательная задача? Где здесь параллелизация?»

Если бы была произвольной плотной матрицей, то да — наивно решение стоило бы , и никакой выгоды бы не было. Но не произвольная. Из-за марковости RNN (каждый шаг видит только предыдущее состояние , а не всю историю) в якобиане подавляющее большинство блоков — нули. Конкретно: в блочной строке номер ненулевые элементы есть только в столбцах и . Получается блочно-бидиагональная структура:

Что такое якобиан вообще

Когда есть обычная функция от одной переменной , ее производная - это одно число, которое говорит «насколько быстро меняется выход при малом изменении входа». Оно играет роль локального коэффициента пропорциональности: если сместить на маленькое , то изменится примерно на .

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

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

Вспомним определение остатка:

Это выражение зависит только от двух переменных: от (через первое слагаемое) и от (через второе). Все остальные в формуле просто не присутствуют. А производная по переменной, которой в формуле нет, равна нулю.

Разберем по случаям, какой блок получается при разных :

Случай 1: . Берем производную по . Только первое слагаемое зависит от , и его производная по самому себе — единичная матрица. Получаем:

Случай 2: . Берем производную по . Первое слагаемое от нее не зависит, второе — это , и его производная — это , вычисленная в точке :

Случай 3: все остальные (то есть и ). Переменная в формуле для просто не встречается. Значит:

Вот и все. Из блоков ненулевых ровно : единичных матриц на главной диагонали и якобианов перехода на поддиагонали. Все остальное — нули. Если расписать всю матрицу:

Где здесь марковость

Ключевая причина, по которой эта структура появилась — марковское свойство RNN. Переходная функция в каждом шаге смотрит только на предыдущее состояние , а не на всю историю . Из-за этого остаток оказывается «локальным» объектом: он зависит только от двух соседних состояний — текущего и предыдущего .

Сколько нам реально нужно памяти

Хотя формально якобиан имеет ячеек, нам нужно хранить только ненулевые блоки. Это:

• единичных матриц — но их даже хранить не нужно, мы знаем, что это и можем подставлять на лету;

• якобианов перехода размера — это чисел, что для , дает около 65 миллионов чисел вместо 65 миллиардов. Уже выполнимо.

Как структура якобиана дает параллелизм

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

Уровень 1: структура позволяет решить систему через прямую подстановку. Возьмем систему и распишем ее построчно. Первая блочная строка матрицы — это , поэтому первое уравнение системы это:

• то есть просто . Получили первый кусок ответа практически даром.

Вторая блочная строка это , поэтому второе уравнение:

Откуда:

И вообще, для произвольного

:

Это и есть линейная рекуррентность, в которую превратилась наша гигантская система . Заметим, что в общем случае решение линейной системы стоит — но мы здесь обошлись без обращения какой-либо матрицы, благодаря тому что блочно-бидиагональна. Система решается простой пробежкой по уравнениям сверху вниз. Это и называется forward substitution.

Если бы дело закончилось здесь, мы бы получили лишь последовательный алгоритм за шагов — каждый зависит от , и пробежать рекурренцию надо строго по порядку. То же самое, что просто прогнать RNN последовательно. Параллелизм рождается на следующем уровне.

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

• Исходная RNN: — функция нелинейна, поэтому распараллелить такую рекуррентность нельзя: приходится считать каждый шаг честно по очереди.

• Рекуррентность для : (где , ) - она линейна. Это значит, что из нее можно получить замкнутую формулу:

Все эти произведения матриц можно вычислить в любом порядке (умножение матриц ассоциативно: ). А значит, можно построить дерево вычислений, в котором мы сначала параллельно считаем все попарные произведения , , , …, затем все четверки , , …, и так далее. За уровней дерева мы получаем все нужные накопленные произведения, и из них собираем все одновременно.

Это и есть параллельный скан (он же parallel prefix sum в обобщенной форме). Аналогия с обычным сложением: если надо сложить миллиард чисел, последовательно — миллиард шагов, а попарным деревом — всего уровней. Тот же прием работает для любой ассоциативной операции, а композиция линейных отображений (умножение их матриц) — ассоциативна.

Итог по сложности: один шаг Ньютона выполняется за параллельной глубины (вместо последовательных шагов), а все применение RNN — за , где iters — число итераций Ньютона.