Forgetting Transformer: Softmax Attention with a Forget Gate
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.
67cc697fa029f09af72cca01
Zhixuan Lin
extracted_confirmed
2025-03-08T16:00:21.933Z
6694cc1009326cb83f2d11bb
/avatars/1ddaaed70a16ac475a9404848aef5d48.svg
Zhixuan Lin
user
zhixuan-lin
67cc697fa029f09af72cca02
Evgenii Nikishin
admin_assigned
2025-03-10T09:54:09.945Z
64234eadd654afd6931a288b
https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/UTU-XcO_ssKpIYr5MBujK.jpeg
Evgenii Nikishin
user
nikishin
67cc697fa029f09af72cca03
Xu Owen He
extracted_pending
2025-03-08T16:00:01.392Z
66906c4e37eadb9c577984d3
/avatars/b81765472942fdf94c0ee885ca62df2d.svg
Owen He
user
littleowen
67cc697fa029f09af72cca04
Aaron Courville
Forgetting Transformer: Softmax Attention with a Forget Gate
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.
A Forgetting Transformer (FoX) improves performance on long-context language modeling and short-context tasks by incorporating a forget gate mechanism into attention scores, while maintaining compatibility with FlashAttention and superior long-context handling.
forget gate Transformers Forgetting Attention Forgetting Transformer (FoX) long-context language modeling length extrapolation short-context downstream tasks FlashAttention positional embeddings Pro block design Mamba-2 HGRN2 DeltaNet