Research map / Architectures and efficiency / Attention and transformers

Forgetting Transformer: Softmax Attention with a Forget Gate

ICLR 2025 · Architectures and efficiency › Attention and transformers · approach: Linear attention approximations · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Forgetting Attention, a softmax attention variant that adds a data-dependent forget-gate decay to attention logits, yielding the Forgetting Transformer (FoX) that outperforms Transformers and recurrent models on long-context language modeling and short-context tasks.

Problem

Transformers lack an explicit data-dependent mechanism for forgetting past information, while recurrent sequence models with forget gates underperform Transformers on long-context capabilities.

Approach

Derive the parallel form of gated linear attention and apply the same exact forget-gate down-weighting D_ij = sum log f_l to softmax attention logits, making softmax attention forget-gated, compatible with FlashAttention and requiring no positional embeddings; add a "Pro" block with output gates, output norm, QK-norm, and KV-shift.

Key results

• FoX (Pro) achieves an average of 50.88 on LM-eval-harness versus 50.39 for Transformer (Pro) and 49.09 for Transformer (LLaMA) at 760M params / 48B tokens.
• FoX achieves near-perfect needle retrieval within the 16k training context, while Mamba-2 and DeltaNet perform poorly even within the training context.
• On ablations, the forget gate is the most important component: removing it from FoX (Pro) raises perplexity from 6.62 to 7.40 (360M, 7.5B tokens).

Measured effects reported in the paper

Surprising findings

• Adding RoPE to FoX (Pro) gives no improvement (6.63 vs 6.62), while RoPE helps the Transformer.
• More training tokens often leads to worse length extrapolation in FoX (Pro) needle retrieval.

Limitations

• Experiments limited to models up to 760M parameters, 48B tokens, and 16k training context.
• Only causal sequence modeling is considered.
• Extrapolation behavior is hyperparameter-dependent.

Experimental setup

Datasets: LongCrawl64 RedPajama-v2 SlimPajama LM-eval-harness tasks LongBench LongAlign needle prompts

Benchmarks: LongCrawl64 validation per-token loss needle-in-the-haystack LM-eval-harness LongBench

Metrics: per-token loss perplexity accuracy acc-n F1 Rouge-L edit similarity needle score 1-10

Baselines: Transformer (LLaMA) Transformer (Pro) Mamba-2 HGRN2 DeltaNet Transformer++

Builds on: GLA Mamba-2 FlashAttention-2 ALiBi LLaMA

Closest papers by meaning

• Memorization Capacity of Multi-Head Attention in Transformers (ICLR 2024 · similarity 0.88)
• Theory, Analysis, and Best Practices for Sigmoid Self-Attention (ICLR 2025 · similarity 0.88)
• Beyond Attention: Breaking the Limits of Transformer Context Length with Recurrent Memory (AAAI 2024 · similarity 0.88)
• Gated Slot Attention for Efficient Linear-Time Sequence Modeling (NeurIPS 2024 · similarity 0.88)
• Gated Linear Attention Transformers with Hardware-Efficient Training (ICML 2024 · similarity 0.88)
• Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.88)