Transformador con Mecanismo de Olvido: Atención Softmax con una Compuerta de Olvido

Forgetting Transformer: Softmax Attention with a Forget Gate

March 3, 2025

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

cs.AI

Resumen

Un componente esencial de los modelos recurrentes de secuencias modernos es la puerta de olvido. Aunque los Transformers no tienen una forma recurrente explícita, demostramos que una puerta de olvido puede incorporarse de manera natural en los Transformers al reducir el peso de las puntuaciones de atención no normalizadas de forma dependiente de los datos. Denominamos a este mecanismo de atención "Atención con Olvido" y al modelo resultante "Transformer con Olvido" (FoX). Mostramos que FoX supera al Transformer en modelado de lenguaje de contexto largo, extrapolación de longitud y tareas posteriores de contexto corto, mientras que rinde al mismo nivel que el Transformer en tareas posteriores de contexto largo. Además, es compatible con el algoritmo FlashAttention y no requiere ningún tipo de incrustación posicional. Varios análisis, incluyendo la prueba de la aguja en el pajar, muestran que FoX también conserva las capacidades superiores de contexto largo del Transformer en comparación con modelos recurrentes de secuencias como Mamba-2, HGRN2 y DeltaNet. También presentamos un diseño de bloque "Pro" que incorpora algunos componentes arquitectónicos comunes en modelos recurrentes de secuencias y descubrimos que mejora significativamente el rendimiento tanto de FoX como del Transformer. Nuestro código está disponible en 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.