Линейное внимание — один из способов снизить вычислительные затраты при обработке длинных последовательностей. Вместо квадратичного роста сложности по длине контекста оно даёт линейный, что особенно важно для языковых моделей и работы с длинными документами. Однако на практике такие модели часто проигрывают стандартному трансформеру по качеству. Группа студентов мастерской по сервисам и платформам ИИ Инженерно-математической школы НИУ ВШЭ и VK исследовала, можно ли сократить этот разрыв, не возвращаясь к глобальному softmax-вниманию. В статье на Habr они описали гипотезу: отставание связано не только с заменой softmax, но и с архитектурными настройками и схемой обучения, которые обычно переносят из стандартных трансформеров без адаптации.

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

Почему линейное внимание проигрывает и как это исправить

Стандартное softmax-внимание требует хранения матрицы весов размером «длина последовательности × длина последовательности», что при больших контекстах приводит к квадратичному росту памяти и вычислений. Линейное внимание заменяет softmax на простые матричные операции, что позволяет хранить состояние фиксированного размера (аналог KV-cache) и обрабатывать последовательности эффективнее. Но качество таких моделей часто ниже. Исследователи предположили, что причина не только в отсутствии softmax, но и в том, что гиперпараметры и схема обучения, оптимальные для стандартного трансформера, не подходят для линейного механизма.

В частности, они проверили, как на качество влияют:

  • Число голов внимания. В стандартных трансформерах обычно используют 8–16 голов. Для линейного внимания оптимальное число может быть другим.
  • Позиционное кодирование. Абсолютное или относительное кодирование, а также его отсутствие — всё это по-разному влияет на линейные модели.
  • Схема слоёв. Чередование линейных и softmax-слоёв, порядок нормализаций и residual-связей.
  • Предобучение. Использование MLM (masked language modeling) вместо авторегрессионного обучения может дать выигрыш.

Рабочая гипотеза состояла в том, что адаптация этих компонентов под линейное внимание позволит приблизить качество к стандартному трансформеру без возврата к глобальному softmax.

Конкретные шаги для улучшения линейных моделей

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

  • Подберите число голов. Начните с меньшего числа, чем в стандартном трансформере (например, 4–8), и проверьте качество на валидации. Возможно, увеличение числа голов не даст прироста, а только усложнит модель.
  • Экспериментируйте с позиционным кодированием. Попробуйте относительное кодирование (например, RoPE или ALiBi) или его отсутствие. Для линейного внимания абсолютное кодирование может быть менее эффективным.
  • Используйте MLM-предобучение. Авторегрессионное обучение с линейным вниманием может страдать от накопления ошибок. MLM позволяет модели лучше использовать контекст и может сократить разрыв в качестве.
  • Настройте схему слоёв. Попробуйте чередовать линейные и softmax-слои (гибридная архитектура) или изменить порядок нормализации. Это может улучшить стабильность обучения.
  • Проверьте размер состояния. В линейном внимании состояние фиксированного размера — ключевой параметр. Слишком маленькое состояние ограничит выразительность, слишком большое — снизит экономию памяти. Подберите размер под вашу задачу.

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

Типичные ошибки и примеры

При работе с линейным вниманием легко допустить ошибки, которые сводят на нет его преимущества. Вот наиболее распространённые.

  • Слепое копирование гиперпараметров стандартного трансформера. Число голов, размер эмбеддингов, схема обучения — всё это требует адаптации. Например, использование 16 голов для линейного внимания может привести к переобучению и лишним вычислениям.
  • Игнорирование позиционного кодирования. Линейное внимание без позиционного кодирования может плохо справляться с задачами, где важен порядок токенов. Добавление относительного кодирования часто улучшает качество.
  • Использование только авторегрессионного обучения. MLM-предобучение может дать лучший старт, особенно для задач понимания текста.
  • Недостаточный размер состояния. Если состояние слишком мало, модель не сможет запомнить нужную информацию из длинного контекста. Экспериментируйте с размером.
  • Отсутствие гибридных схем. Полностью линейная модель может проигрывать гибриду, где часть слоёв использует softmax. Попробуйте чередование.

Пример из исследования: авторы показали, что изменение числа голов и добавление MLM-предобучения позволило сократить разрыв в качестве на нескольких бенчмарках. Хотя полного паритета достичь не удалось, результаты подтверждают, что архитектурные настройки важны не меньше, чем сам механизм внимания.

Что это значит для практиков

Линейное внимание — не замена стандартному трансформеру «из коробки». Чтобы получить конкурентное качество, нужно потратить время на подбор архитектуры и схемы обучения. Если вы разрабатываете модель для обработки длинных документов или экономии памяти при генерации, линейное внимание может быть оправдано, но только при условии тщательной настройки. Начните с малого: возьмите открытую реализацию линейного внимания, добавьте MLM-предобучение, поэкспериментируйте с числом голов и позиционным кодированием. Сравнивайте результаты с базовым трансформером на валидационной выборке. Не ожидайте мгновенного выигрыша — но если вам критичны затраты на длинные контексты, этот путь может окупиться.