В совместном исследовании Варшавского университета, Google Research и OpenAI представлена новая концепция оптимизации нейросетей, описанная в статье «Sparse is Enough in Scaling Transformers». Известный ИИ-исследователь и блогер Янник Кильхер (Yannic Kilcher) подробно разбирает, как архитектурные модификации Scaling Transformer и Terraformer позволяют радикально сократить вычислительные затраты. Главная идея работы заключается в том, что селективная активация параметров позволяет достичь многократного ускорения инференса без потери качества языкового моделирования.
🚀 Эпоха разреженности: как ускорить инференс в десятки раз 0:00
Традиционные трансформеры требуют колоссальных вычислительных ресурсов при генерации текста. Как отмечает Янник Кильхер, авторы статьи предлагают набор базовых блоков для внедрения разреженности (sparsity), трансформируя стандартную архитектуру сначала в Scaling Transformer, а затем, с добавлением дополнительных оптимизаций, в Terraformer.
Особое преимущество данные модели демонстрируют при так называемом однопоточном (unbatched) декодировании — обработке единичных запросов в режиме реального времени, когда на GPU нельзя запустить параллельный пакет вычислений. Эффект масштабирования выглядит следующим образом:
- Для модели с 800 млн параметров базовый трансформер тратит на декодирование 0,16 секунды, а Scaling Transformer ускоряет процесс в 2,6 раза.
- Для модели с 17 млрд параметров базовый вариант требует 3,6 секунды, тогда как Scaling Transformer справляется за 0,18 секунды, демонстрируя 20-кратное ускорение.
- В отдельных конфигурациях, по словам Янника Кильхера, прирост скорости генерации токенов может достигать внушительных 37 раз.
Самым удивительным фактом исследователь называет то, что разреженные слои демонстрируют ту же перплексию (качество работы со словесной последовательностью), что и стандартные плотные аналоги. По мнению Кильхера, это указывает на избыточность стандартных архитектур для многих прикладных задач.
🧩 Шаг 1: Разреженный полносвязный слой (Sparse Feed-Forward) 6:49
В классическом трансформере слой Feed-Forward (FFN) принимает последовательность эмбеддингов токенов и обрабатывает каждый независимо через две последовательные матрицы с функцией активации ReLU между ними. Поскольку ReLU естественным образом зануляет часть выходов, исследователи решили превентивно исключить лишние вычисления из графа нейросети.
Пошаговый алгоритм настройки разреженного FFN-слоя включает следующие параметры и операции:
- Разделение (компартментализация): Вся исходная весовая матрица $d_{model} \times d_{ff}$ разбивается на отдельные независимые блоки (компартменты).
- Активация столбцов: Внутри каждого компартмента в один момент времени разрешено активировать строго один столбец параметров, который и загружается в память. Соответственно, связанная строка во второй матрице тоже становится активной, а остальные игнорируются.
- Вычисление маски (Контроллер): Для определения того, какие именно столбцы необходимо активировать, применяется вспомогательная низкоранговая (low-rank) сеть-контроллер. Она временно сжимает размерность вектора, а затем разворачивает её до размерности FFN-слоя, выдавая вероятности для каждого элемента.
- Обучение через Gumbel-Softmax: Во время обучения для сквозного прохождения градиента используется трюк со straight-through Gumbel-Softmax распределением. При инференсе сеть жестко выбирает один максимальный элемент (Argmax).
Янник Кильхер подчеркивает важный вычислительный компромисс: работа низкорангового контроллера создает дополнительную нагрузку. Метод эффективен только тогда, когда затраты на вычисление маски меньше, чем выигрыш от линейного снижения объема умножения матриц за счет разреженности.
🔍 Шаг 2: Модуль мультипликативного внимания и свертки (MultConv) 19:22
Оптимизация FFN-слоя переносит вычислительное «бутылочное горлышко» на механизм внимания. Прямой перенос трюка с ReLU здесь невозможен, так как матрицы проекций запросов (Queries), ключей (Keys) и значений (Values) не содержат обилия нулевых элементов.
Для решения этой проблемы авторы внедряют комбинированный модуль MultConv, полностью заменяющий плотные слои проекций:
- Мультипликативный слой (Mult): Вместо трех раздельных тяжелых матричных умножений создается один общий мультипликативный слой для Q, K и V. Вектор токена подвергается поэлементному умножению с бродкастингом (трансляцией размерностей) матрицы весов, а затем проецируется в низкоранговое пространство с параметрами $s \times m = d_{model}$, где $s$ — число голов внимания, а $m$ — уменьшенная размерность. Авторы математически доказывают теорему, что такая структура способна выразить любую произвольную перестановку признаков.
- Сверточный слой (Conv): На выходе мультипликативного слоя формируется тензор размерности $[s \times l \times m]$, где $l$ — длина последовательности. К нему применяется стандартная свертка, которая обрабатывается алгоритмами как двумерное изображение. Свертка позволяет обмениваться информацией между соседними головами внимания (по оси $s$) и соседними токенами (по оси $l$).
По мнению Кильхера, это не только ускоряет вычисления, но и привносит в модель встроенное понимание позиционности токенов и взаимосвязи между головами внимания еще до самого расчета механизма Self-Attention.
⚡ Шаг 3: Переход к архитектуре Terraformer и оптимизация длинных последовательностей 44:02
Во второй части статьи исследователи объединяют наработки Scaling Transformer с концепциями из модели Reformer, создавая архитектуру Terraformer, ориентированную на сверхдлинные контексты.
Интеграция включает в себя три ключевых шага:
- Отказ от раздельного внимания: Разработчики полностью убирают классический блок Encoder-Decoder attention. Вместо этого они просто конкатенируют (склеивают) представления энкодера перед токенами декодера, упрощая граф вычислений.
- Хеширование LSH (Locality Sensitive Hashing): Для избавления от квадратичной сложности вычисления парных скалярных произведений используется алгоритм LSH-внимания. Токены распределяются по хеш-корзинам, и расчеты производятся только в случае совпадения хешей (коллизий), отсекая заведомо нерелевантные связи.
- Обратимые блоки (Reversibility): Информационный поток разделяется на две параллельные ветви. Изменения, вносимые функциями слоев, накладываются на альтернативную ветвь, сохраняя вторую нетронутой. Это делает блоки обратимыми: во время обратного прохода (backpropagation) сеть способна аналитически восстановить промежуточные активации, что исключает необходимость хранить их в памяти GPU.
В качестве финального штриха авторы добавляют рекуррентные связи (Simple Recurrent Units, SRU) внутрь FFN-блоков. Янник Кильхер называет это решение «крайне странным выбором». Он отмечает, что на стандартных языковых задачах вроде датасета C4 прирост от рекуррентности минимален. Однако она резко улучшает экстраполяционные способности модели на синтетических тестах, позволяя сети, обученной на последовательностях из 128 цифр, успешно обрабатывать контекст длиной в 256 цифр.
