# Как уместить 64 000 токенов в одну видеокарту: разбор нейросети Reformer

Источник: https://www.youtube.com/watch?v=i4H0kjxrias
Канал: Yannic Kilcher
Опубликовано: 22.01.2020

---

Архитектура Transformer произвела настоящую революцию в сфере глубокого обучения, однако ее практическое применение жестко ограничено колоссальными требованиями к вычислительным ресурсам. В своем видеообзоре исследователь искусственного интеллекта Янник Кильчер подробно разбирает научную работу «Reformer: The Efficient Transformer», авторы которой предложили изящные математические методы для радикального снижения потребления памяти. Этот подход позволяет нейросетям обрабатывать гигантские текстовые последовательности и изображения на стандартном коммерческом оборудовании.

## 🧠 Проблема квадратичной сложности в стандартных трансформерах
[[JUMP:0:00]]

В классических моделях Transformer информация передается последовательно от слоя к слою. На каждом этапе нейросеть вычисляет механизм под названием «внимание» (attention), распределяя ресурсы по трем ключевым векторам: запросам (queries), ключам (keys) и значениям (values). Запросы формируют требования текущего узла к предыдущему слою, ключи представляют доступную информацию, а значения определяют непосредственно передаваемый контент.

Главная вычислительная проблема кроется во взаимодействии запросов и ключей. Чтобы направить информацию в следующий слой, алгоритм обязан сопоставить каждый отдельный запрос со всеми имеющимися ключами с помощью операции скалярного произведения. Если длина последовательности составляет d элементов, количество необходимых вычислительных операций вырастает до величины порядка d². Столь стремительный квадратичный рост создаёт критическую нагрузку на память, превращая обработку длинных текстов в неподъемную задачу для большинства вычислительных систем.

## 🪣 Хеширование LSH: как упорядочить хаос
[[JUMP:06:02]]

Для решения пространственной проблемы авторы Reformer предложили использовать технику бакетизации — распределения данных по виртуальным «корзинам». Идея заключается в том, чтобы группировать похожие векторы вместе. Если распределить векторы запросов и ключей по одинаковым бакетам, то сравнивать их придется только внутри этих локальных групп, полностью игнорируя остальную часть огромного пространства. 

В основе этого математического трюка лежит алгоритм локально-чувствительного хеширования (Locality-Sensitive Hashing, или LSH). Его фундаментальное правило формулируется следующим образом:

* Если расстояние между двумя векторами мало, они должны с высокой вероятностью оказаться в одном бакете.
* Если расстояние между векторами велико, вероятность их попадания в общую корзину должна стремиться к минимуму.

Янник Кильчер приводит простой пример работы LSH на основе расстояния Жаккара через субсемплирование битов. Если у нас есть векторы в виде битовых строк, мы можем выбрать первые два бита в качестве идентификатора бакета. В архитектуре Reformer для вычисления косинусного расстояния применяется более продвинутый метод — случайные проекции гиперплоскостей. Через начало координат случайным образом проводятся гиперплоскости, разделяющие пространство. Каждому вектору присваивается бинарный код (плюс или минус) в зависимости от того, по какую сторону от конкретной плоскости он находится. 

Поскольку функция Softmax в механизме внимания математически доминируется исключительно самыми большими элементами (то есть наиболее близкими векторами), всеми остальными далекими значениями можно пренебречь. Благодаря LSH отпадает необходимость просчитывать полную матрицу взаимодействий, что колоссально экономит ресурсы.

## ✂️ Разделение на блоки (Chunking) для стабилизации памяти
[[JUMP:17:09]]

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

Чтобы стабилизировать нагрузку, в Reformer ввели механизм разделения на блоки фиксированного размера (chunking). Алгоритм искусственно нарезает отсортированную последовательность на равные части. Если какой-то бакет из-за разделения оказывается разорван между блоками, модель задействует специальное правило: элементы могут обращаться ко входам как внутри своего блока, так и внутри соседнего. Это жесткое ограничение позволяет зафиксировать объем памяти на уровне квадрата от размера блока, а не всей длины последовательности.

## 🔄 Обратимые слои: избавляемся от хранения активаций
[[JUMP:20:40]]

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

Reformer решает эту задачу с помощью внедрения обратимых слоев (reversible layers). Если математическая функция слоя обратима, то промежуточные активации нет смысла хранить — их можно в любой момент вычислить в обратном порядке на этапе обратного прохода. Архитектура заимствует подход из реверсивных остаточных сетей (RevNet), где данные разделяются на два параллельных потока, обновляемых поочередно:

1.  Первый слой берет два потока, обновляет второй с помощью функции слоя и передает данные дальше.
2.  Следующий слой оставляет полученную вторую часть неизменной, но аналогичным образом обновляет первый поток.

Такая симметричная шахматная структура позволяет легко запустить вычисления вспять. Янник Кильчер указывает, что концепция обратимости не нова и давно применяется, например, в нормализующих потоках (normalizing flows). При этом ведущий замечает, что авторы Reformer почему-то проигнорировали и не упомянули в библиографии обширную литературу по потоковым моделям, ограничившись лишь ссылкой на базовую статью о RevNet.

## 🚀 Результаты: терабайты контекста на одной видеокарте
[[JUMP:26:28]]

Сочетание LSH-внимания и обратимых слоев полностью убирает квадратичную зависимость из формулы требований к памяти, превращая Reformer в невероятно масштабируемый инструмент. Для сравнения, популярная модель BERT жестко ограничена лимитом в 512 входных токенов, а продвинутая XLNet способна незначительно раздвигать эти рамки лишь за счет циклического переноса скрытых состояний.

В то же время Reformer без труда оперирует контекстом в 64 000 токенов, требуя для такой масштабной операции всего 16 гигабайт видеопамяти — объем, доступный на любой флагманской потребительской видеокарте. Столь мощный прорыв открывает невиданные ранее возможности: например, обработку огромных изображений из датасета ImageNet буквально попиксельно, что раньше казалось абсолютно немыслимым из-за ограничений классического механизма Transformer.