Transformador de Esquecimento: Atenção Softmax com um Portão de Esquecimento

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

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

cs.AI

Resumo

Um componente essencial dos modelos de sequência recorrentes modernos é o portão de esquecimento. Embora os Transformers não tenham uma forma recorrente explícita, mostramos que um portão de esquecimento pode ser naturalmente incorporado aos Transformers ao reduzir o peso das pontuações de atenção não normalizadas de maneira dependente dos dados. Nomeamos esse mecanismo de atenção como Atenção de Esquecimento e o modelo resultante como Transformer de Esquecimento (FoX). Demonstramos que o FoX supera o Transformer em modelagem de linguagem de contexto longo, extrapolação de comprimento e tarefas subsequentes de contexto curto, enquanto apresenta desempenho semelhante ao Transformer em tarefas subsequentes de contexto longo. Além disso, ele é compatível com o algoritmo FlashAttention e não requer nenhum embedding posicional. Várias análises, incluindo o teste de "agulha no palheiro", mostram que o FoX também mantém as capacidades superiores de contexto longo do Transformer em comparação com modelos de sequência recorrentes como Mamba-2, HGRN2 e DeltaNet. Também introduzimos um design de bloco "Pro" que incorpora alguns componentes arquitetônicos comuns em modelos de sequência recorrentes e descobrimos que ele melhora significativamente o desempenho tanto do FoX quanto do Transformer. Nosso código está disponível em 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.