Transformateur à Oubli : Attention Softmax avec une Porte d'Oubli

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

Auteurs: Zhixuan Lin, Evgenii Nikishin, Xu Owen He, Aaron Courville

cs.AI

Résumé

Un composant essentiel des modèles séquentiels récurrents modernes est la porte d'oubli. Bien que les Transformers ne possèdent pas de forme récurrente explicite, nous montrons qu'une porte d'oubli peut être naturellement intégrée dans les Transformers en pondérant à la baisse les scores d'attention non normalisés de manière dépendante des données. Nous nommons ce mécanisme d'attention l'Attention à l'Oubli et le modèle résultant le Transformer à l'Oubli (FoX). Nous démontrons que FoX surpasse le Transformer dans la modélisation de langage à contexte long, l'extrapolation de longueur et les tâches en aval à contexte court, tout en étant à égalité avec le Transformer sur les tâches en aval à contexte long. De plus, il est compatible avec l'algorithme FlashAttention et ne nécessite aucun encodage positionnel. Plusieurs analyses, dont le test de l'aiguille dans la botte de foin, montrent que FoX conserve également les capacités supérieures du Transformer en contexte long par rapport aux modèles séquentiels récurrents tels que Mamba-2, HGRN2 et DeltaNet. Nous introduisons également une conception de bloc "Pro" qui intègre certains composants architecturaux courants des modèles séquentiels récurrents et constatons qu'elle améliore significativement les performances de FoX et du Transformer. Notre code est disponible à l'adresse 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.