Transformer con Funzione di Dimenticanza: Softmax Attention con un Cancello di Dimenticanza

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

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

cs.AI

Abstract

Un componente essenziale dei moderni modelli sequenziali ricorrenti è il cancello di dimenticanza (forget gate). Sebbene i Transformer non abbiano una forma ricorrente esplicita, dimostriamo che un cancello di dimenticanza può essere naturalmente incorporato nei Transformer ponderando in modo dipendente dai dati i punteggi di attenzione non normalizzati. Chiamiamo questo meccanismo di attenzione "Forgetting Attention" e il modello risultante "Forgetting Transformer" (FoX). Dimostriamo che FoX supera il Transformer nella modellazione del linguaggio a contesto lungo, nell'estrapolazione della lunghezza e nelle attività downstream a contesto breve, pur mantenendo prestazioni paragonabili al Transformer nelle attività downstream a contesto lungo. Inoltre, è compatibile con l'algoritmo FlashAttention e non richiede alcun embedding posizionale. Diverse analisi, incluso il test "needle-in-the-haystack", mostrano che FoX conserva anche le superiori capacità di gestione del contesto lungo del Transformer rispetto ai modelli sequenziali ricorrenti come Mamba-2, HGRN2 e DeltaNet. Introduciamo inoltre un design a blocco "Pro" che incorpora alcuni componenti architetturali comuni nei modelli sequenziali ricorrenti e scopriamo che migliora significativamente le prestazioni sia di FoX che del Transformer. Il nostro codice è disponibile all'indirizzo 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.