Gated Recurrent Transformer: 3 слоя дают точность 12-слойного GPT-2 Small
Исследователь Amr Hegazy с соавторами опубликовал на Hugging Face препринт о новой архитектуре трансформера, Gated Recurrent Transformer. Проблема, которую она решает: чтобы слои трансформера специализировались на разных задачах (от обработки входа до абстрактного вывода), обычно нужны уникальные веса на каждом слое, а это раздувает память. Простое разделение весов между слоями (depth-sharing) экономит память, но заставляет все слои делать одно и то же и портит качество модели.
Предложенное решение: фиксированные вводный (prelude) и завершающий (coda) блоки обрамляют один общий блок-ядро, который прогоняется R раз подряд (рекуррентно). Чтобы этот единственный блок на каждом проходе вёл себя по-разному, авторы добавили гейт по образцу gated RNN: лёгкая проекция и поэлементный gate, который учитывает скрытое состояние, выход вводного блока и случайный шум, добавляемый заново на каждом шаге. Такой гейт позволяет модели каждый раз по-новому применять один и тот же небольшой набор слоёв к входу, вместо того чтобы держать много разных слоёв ради разнообразия обработки.
Результаты авторы приводят по двум сценариям сравнения. При равном бюджете вычислений (isoFLOPS) 3-слойный Gated Recurrent Transformer сравнялся по точности с базовой моделью GPT-2 Small, у которой 12 слоёв, и обошёл альтернативные методы адаптивной глубины (MoR и heavy-tail depth sampling) во всех девяти проверенных комбинациях масштаба и бюджета; на среднем и крупном масштабе он приближается к качеству обычной (плотной) модели при стандартном бюджете токенов на обучение и превосходит её на среднем масштабе, если этот бюджет удвоить. При равном числе параметров и данных (isoPARAMS) более глубокая рекуррентная версия достигла loss на валидации 2,76 против 2,84 у нерекуррентного аналога.
Авторы также приводят итоговый эффект на крупном масштабе: 63% меньше параметров и 59% меньше пиковой памяти при декодировании ценой роста задержки генерации (при скомпилированной модели) на 10%. Абсолютные размеры моделей для "среднего" и "крупного" масштаба в тексте не указаны, приведены только относительные величины.
Ключевые факты
- Gated Recurrent Transformer оборачивает один общий блок-ядро фиксированными вводным и завершающим блоками и прогоняет его R раз, а не наращивает число уникальных слоёв
- Поэлементный гейт, зависящий от скрытого состояния, выхода вводного блока и добавляемого на каждом шаге шума, заставляет один и тот же блок обрабатывать вход по-разному на каждом проходе
- При равном бюджете вычислений 3-слойная модель сравнялась по точности с 12-слойным GPT-2 Small и обошла методы MoR и heavy-tail depth sampling во всех девяти комбинациях масштаба и бюджета
- При равном числе параметров и данных более глубокая рекуррентная версия дала loss 2,76 против 2,84 у нерекуррентного аналога
- На крупном масштабе метод даёт 63% меньше параметров и 59% меньше пиковой памяти при декодировании ценой роста задержки генерации на 10%
Почему это важно
Стандартный способ дать трансформеру больше выразительности, добавить уникальные слои, но это прямо увеличивает память под веса модели. Gated Recurrent Transformer предлагает получить сопоставимую выразительность за счёт многократного прогона одного общего блока с гейтом, а не за счёт роста числа отдельных слоёв, это отвязывает глубину обработки от количества параметров.
Кому это важно
Разработчикам и исследователям, которые проектируют или дообучают языковые модели и упираются в память под веса и в память при инференсе (decoding), особенно при развёртывании на ограниченном железе, метод напрямую бьёт по этому компромиссу параметры/качество.
Как это применить
Подход меняет саму структуру блока трансформера при обучении с нуля: вместо стека уникальных слоёв, фиксированные prelude/coda и общий блок-ядро, прогоняемый R раз с гейтом на основе скрытого состояния, выхода prelude и случайного шума. Это архитектурное решение, а не надстройка над готовой моделью, поэтому применить его можно на этапе проектирования и обучения новой модели.
Можно ли доверять
Материал, препринт на Hugging Face; в тексте абстракта не указаны ни авторы (кроме данных с самой страницы публикации), ни организация, ни место публикации, ни рецензирование. Заявленные цифры приведены авторами по их собственным экспериментам (isoFLOPS и isoPARAMS сравнения по девяти комбинациям масштаба и бюджета), независимого воспроизведения результатов пока нет.
Риски и подводные камни
Выигрыш в параметрах и памяти при декодировании достаётся ценой роста задержки генерации на 10% на скомпилированной модели. В тексте не приведены абсолютные размеры моделей для "среднего" и "крупного" масштаба, только относительные проценты, что затрудняет независимую оценку. Сравнение дано только с MoR и heavy-tail depth sampling как альтернативными методами адаптивной глубины; код, веса модели или датасет обучения не упомянуты.