Реформер

Реформер — это специализированная архитектура нейронной сети, предложенная в 2020 году исследователями из компании Google (входит в холдинг Alphabet) для решения проблемы вычислительной сложности, присущей моделям на основе механизма внимания (Transformer). Основная цель разработки — обработка последовательностей данных (текста, изображений, временных рядов) с длинным контекстом, где стандартные трансформеры сталкиваются с квадратичным ростом затрат памяти и времени.

История создания

Архитектура Transformer, представленная в 2017 году, быстро стала стандартом в задачах обработки естественного языка (NLP) и компьютерного зрения. Однако её ключевой компонент — механизм самовнимания (self-attention) — имеет вычислительную сложность O(n²) по длине последовательности n. Это означает, что при увеличении длины входного текста, например, до 100 000 токенов, требования к памяти и времени обучения становятся практически невыполнимыми для обычных графических процессоров (GPU).

В 2020 году группа исследователей из Google Research (Никита Китaев, Łukasz Kaiser, Anselm Levskaya и другие) опубликовала статью «Reformer: The Efficient Transformer». В ней они предложили несколько ключевых модификаций, которые позволили снизить сложность до O(n log n) и значительно уменьшить потребление памяти. Архитектура была реализована в открытой библиотеке машинного обучения Trax.

Ключевые особенности

1. LSH-внимание (Locality-Sensitive Hashing Attention)

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

• Принцип работы: Для каждого токена вычисляется хэш-код на основе его представления. Токены с одинаковым хэшем попадают в одну группу. Внимание вычисляется только внутри группы, а не по всей последовательности.
• Сложность: Снижается с O(n²) до O(n log n), так как количество токенов в каждой корзине в среднем пропорционально log n.
• Ограничение: LSH-внимание является приближённым — оно может не учитывать слабые, но важные связи между токенами из разных корзин. Однако на практике это даёт приемлемую точность для большинства задач.

2. Обратимые слои (Reversible Layers)

В стандартном трансформере каждый слой хранит промежуточные активации (выходы всех нейронов) для обратного распространения ошибки. Это приводит к огромному расходу памяти — объём хранимых данных пропорционален длине последовательности и числу слоёв. Reformer использует обратимые слои, которые позволяют восстановить входные данные из выходных без сохранения промежуточных значений.

• Механизм: Каждый слой разделён на две части: одна часть вычисляет нелинейное преобразование, другая — линейное. Во время обратного прохода вычисления выполняются в обратном порядке, реконструируя входные данные.
• Эффект: Потребление памяти снижается с O(n * L) (где L — число слоёв) до O(n), то есть перестаёт зависеть от глубины сети. Это позволяет строить очень глубокие модели (до десятков слоёв) без переполнения памяти GPU.

3. Кусочно-последовательная обработка (Chunked Feed-Forward)

В стандартном трансформере слои прямой связи (feed-forward) обрабатывают каждый токен независимо, но требуют хранения всех активаций. Reformer разбивает последовательность на небольшие «куски» (chunks) и обрабатывает их по очереди, освобождая память после каждого шага. Это не влияет на точность, так как прямой связи между токенами в этом слое нет.

Архитектура и устройство

Reformer сохраняет общую структуру трансформера: кодировщик (encoder) и декодировщик (decoder), каждый из которых состоит из нескольких блоков. Однако внутренние компоненты модифицированы:

• Входной слой: Токены преобразуются в эмбеддинги (векторные представления) с добавлением позиционных кодировок.
• Блок LSH-внимания: Вместо полного внимания используется LSH-внимание с несколькими раундами хэширования (обычно 4–8) для повышения точности.
• Обратимый слой: Объединяет LSH-внимание и кусочно-последовательную прямую связь в один обратимый блок.
• Выходной слой: Линейное преобразование и softmax для получения вероятностей токенов.

Применение

Reformer ориентирован на задачи, где длина контекста критически важна:

• Обработка длинных текстов: Анализ книг, юридических документов, научных статей (например, обработка полного текста романа «Война и мир»).
• Генерация музыки: Моделирование музыкальных последовательностей с длиной до нескольких тысяч нот.
• Обработка изображений высокого разрешения: Патч-токенизация изображений с последующим вниманием на больших областях.
• Генерация кода: Анализ и генерация программного кода с длинными зависимостями между строками.
• Биоинформатика: Анализ последовательностей ДНК и белков, где длина может достигать миллионов символов.

Преимущества и недостатки

Преимущества

• Эффективность по памяти: Позволяет обрабатывать последовательности в десятки раз длиннее, чем стандартные трансформеры, при тех же вычислительных ресурсах.
• Скорость обучения: LSH-внимание и обратимые слои ускоряют обучение на длинных последовательностях.
• Масштабируемость: Возможность строить очень глубокие модели (до 20–30 слоёв) без переполнения памяти.

Недостатки

• Приближённость внимания: LSH-хэширование может терять редкие, но важные связи между токенами, что снижает качество на некоторых задачах (например, в задачах точного машинного перевода).
• Сложность реализации: Требуется специальная оптимизация для эффективного хэширования на GPU.
• Ограниченная популярность: На практике Reformer уступил место более поздним архитектурам, таким как Longformer, BigBird и Performer, которые предложили альтернативные методы приближённого внимания.

Критика и альтернативы

Несмотря на инновационность, Reformer не стал массовым стандартом. Основная критика связана с тем, что LSH-внимание даёт выигрыш только при очень длинных последовательностях (более 10 000 токенов), а для коротких последовательностей (до 512 токенов) оно менее эффективно, чем полное внимание. Кроме того, были разработаны более простые и быстрые методы:

• Longformer (2020) — использует комбинацию локального и глобального внимания.
• BigBird (2020) — сочетает случайное, локальное и глобальное внимание.
• Performer (2020) — применяет ядерное приближение внимания (FAVOR+).

Тем не менее, Reformer остаётся важной вехой в развитии эффективных трансформеров, продемонстрировавшей практическую возможность обработки сверхдлинных последовательностей.

Интересные факты

• Название «Reformer» является игрой слов: «Reformer» (реформатор) и «Transformer» (трансформер), подчёркивая, что это реформированная версия оригинальной архитектуры.
• В оригинальной статье исследователи показали, что модель может обрабатывать последовательность длиной до 1 миллиона токенов на одном GPU с 16 ГБ памяти.
• Reformer был использован для генерации музыки в проекте Google Magenta, где он создавал композиции длительностью до нескольких минут.

Источники

• Kitaev, N., Kaiser, Ł., Levskaya, A. (2020). «Reformer: The Efficient Transformer». arXiv preprint arXiv:2001.04451.
• Vaswani, A. et al. (2017). «Attention Is All You Need». Advances in Neural Information Processing Systems.
• Beltagy, I., Peters, M. E., Cohan, A. (2020). «Longformer: The Long-Document Transformer». arXiv preprint arXiv:2004.05150.
• Zaheer, M. et al. (2020). «Big Bird: Transformers for Longer Sequences». arXiv preprint arXiv:2007.14062.
• Choromanski, K. et al. (2020). «Rethinking Attention with Performers». arXiv preprint arXiv:2009.14794.