Vergeten Transformer: Softmax Attention met een Vergeetpoort

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

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

cs.AI

Samenvatting

Een essentieel onderdeel van moderne recurrente sequentiemodellen is de vergeetpoort. Hoewel Transformers geen expliciete recurrente vorm hebben, laten we zien dat een vergeetpoort op natuurlijke wijze kan worden geïntegreerd in Transformers door de niet-genormaliseerde aandachtsscores op een data-afhankelijke manier te verlagen. We noemen dit aandachtmechanisme de Vergetende Aandacht en het resulterende model de Vergetende Transformer (FoX). We tonen aan dat FoX de Transformer overtreft op het gebied van taalmodellering met lange context, lengte-extrapolatie en downstream taken met korte context, terwijl het op hetzelfde niveau presteert als de Transformer bij downstream taken met lange context. Bovendien is het compatibel met het FlashAttention-algoritme en heeft het geen positionele embeddings nodig. Verschillende analyses, waaronder de naald-in-de-hooiberg-test, laten zien dat FoX ook de superieure lange-contextmogelijkheden van de Transformer behoudt ten opzichte van recurrente sequentiemodellen zoals Mamba-2, HGRN2 en DeltaNet. We introduceren ook een "Pro"-blokontwerp dat enkele veelvoorkomende architectuurcomponenten van recurrente sequentiemodellen integreert en ontdekken dat het de prestaties van zowel FoX als de Transformer aanzienlijk verbetert. Onze code is beschikbaar op 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.