Трансформер с механизмом забывания: Softmax-внимание с гейтом забывания

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

Авторы: Zhixuan Lin, Evgenii Nikishin, Xu Owen He, Aaron Courville

cs.AI

Аннотация

Важным компонентом современных рекуррентных моделей последовательностей является забывающий вентиль. Хотя Transformers не имеют явной рекуррентной формы, мы показываем, что забывающий вентиль может быть естественным образом интегрирован в Transformers путем понижения весов ненормализованных оценок внимания в зависимости от данных. Мы называем этот механизм внимания "Забывающее внимание" (Forgetting Attention), а результирующую модель — "Забывающий Transformer" (Forgetting Transformer, FoX). Мы демонстрируем, что FoX превосходит Transformer в задачах моделирования языка с длинным контекстом, экстраполяции длины и задачах с коротким контекстом, при этом показывая сопоставимые результаты с Transformer в задачах с длинным контекстом. Более того, он совместим с алгоритмом FlashAttention и не требует использования позиционных эмбеддингов. Несколько анализов, включая тест "иголка в стоге сена", показывают, что FoX также сохраняет превосходные способности Transformer в работе с длинным контекстом по сравнению с рекуррентными моделями последовательностей, такими как Mamba-2, HGRN2 и DeltaNet. Мы также представляем дизайн "Pro" блока, который включает некоторые общие архитектурные компоненты рекуррентных моделей последовательностей, и обнаруживаем, что он значительно улучшает производительность как FoX, так и Transformer. Наш код доступен по адресу https://github.com/zhixuan-lin/forgetting-transformer.

English

An essential component of modern recurrent sequence models is the forget gate. While Transformers do not have an explicit recurrent form, we show that a forget gate can be naturally incorporated into Transformers by down-weighting the unnormalized attention scores in a data-dependent way. We name this attention mechanism the Forgetting Attention and the resulting model the Forgetting Transformer (FoX). We show that FoX outperforms the Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks, while performing on par with the Transformer on long-context downstream tasks. Moreover, it is compatible with the FlashAttention algorithm and does not require any positional embeddings. Several analyses, including the needle-in-the-haystack test, show that FoX also retains the Transformer's superior long-context capabilities over recurrent sequence models such as Mamba-2, HGRN2, and DeltaNet. We also introduce a "Pro" block design that incorporates some common architectural components in recurrent sequence models and find it significantly improves the performance of both FoX and the Transformer. Our code is available at https://github.com/zhixuan-lin/forgetting-transformer.