Забывающий Трансформер (FoX): Революция в Длинных Контекстах Языкового Моделирования с Эффективным Управлением Памятью

Проблема Управления Памятью в Трансформерах

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

Ограничения Существующих Подходов

Традиционные рекуррентные модели используют forget gates для контроля сохранения памяти, но испытывают трудности с длинными последовательностями из-за фиксированного размера скрытого состояния. Некоторые модификации трансформеров, например ALiBi, добавляют статические позиционные смещения для имитации эффекта недавности, но они не адаптируются к содержимому входных данных. Другие модели, такие как Mamba-2 и GLA, вводят механизмы затвора в линейное внимание, но часто теряют нормализацию и значительно отходят от архитектуры трансформера, что снижает совместимость с оптимизациями.

Архитектура Забывающего Трансформера (FoX)

Исследователи из Mila и Университета Монреаля совместно с MakerMaker AI разработали Забывающий Трансформер (FoX), который внедряет механизм Forgetting Attention, добавляя скалярный forget gate в процесс softmax внимания. Этот затвор динамически корректирует оценки внимания на основе данных, эффективно снижая вес нерелевантных прошлых токенов.

FoX сохраняет совместимость с параллельными вычислениями и эффективным алгоритмом FlashAttention, минимизируя накладные расходы. Были предложены две версии:

• FoX (на базе LLaMA): базовая версия с forget gate.
• FoX (Pro): расширенная версия с нормализацией вывода, выходными затворами и механизмом сдвига токенов, вдохновлённым последними рекуррентными моделями для лучшей чувствительности к контексту.

Технические Детали Forgetting Attention

Значения forget gate вычисляются для каждого временного шага с помощью сигмоидной активации на основе обучаемого линейного преобразования входных данных. Эти скалярные значения влияют на логиты внимания через формулу суммы логарифмов, эффективно модифицируя softmax без необходимости создания больших матриц. Каждый head внимания имеет независимые параметры forget gate.

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

Результаты Экспериментов и Производительность

Тестирование на датасете LongCrawl64 (48 миллиардов токенов) показало, что FoX стабильно превосходит стандартные трансформеры и ведущие рекуррентные модели в задачах языкового моделирования с длинным контекстом. Достигнуты:

• Более резкое снижение потерь на каждом токене.
• Значительно меньшие потери на позиции 64 000 по сравнению с Transformer и LLaMA.
• Лучшие показатели перплексии при увеличении длины контекста, с меньшим падением точности за пределами обучающей длины в 16 384 токена.
• Превосходство над моделями-конкурентами, такими как Mamba-2 и DeltaNet, в способности обобщения.

Обучение проводилось с 760 миллионами параметров и токенизатором TikToken для GPT-2, при этом FoX предпочитал более высокие скорости обучения и меньшие размеры head, что говорит о его устойчивости.

Преимущества и Значение

FoX вводит управляемый данными эффект недавности, улучшая выборочную память трансформеров без увеличения вычислительных затрат и памяти. Модель обобщает статические смещения типа ALiBi через динамические forget gate, повышая адаптивность и точность. Версия Pro особенно эффективна в задачах, требующих чувствительности к контексту.

Это исследование демонстрирует, что динамическое забывание в трансформерах возможно и полезно. Совместимость с FlashAttention позволяет масштабировать модель без потерь эффективности.

Основные Выводы

• Forgetting Attention улучшает softmax внимание с помощью обучаемых forget gate.
• Две архитектурные версии: FoX (LLaMA) и FoX (Pro) с дополнительной нормализацией и затворами.
• FoX превосходит стандартные трансформеры в моделировании длинных контекстов.
• Низкие показатели ошибок и устойчивая перплексия даже на последовательностях свыше 64k токенов.
• Обобщение статических смещений через динамическое управление.
• Эффективность на аппаратном уровне и совместимость с FlashAttention.

Подробнее о работе доступно в оригинальной публикации и репозиториях кода. Следите за обновлениями в Twitter, Telegram и LinkedIn.