PaperlayerView exact revision on arXiv

Source-linked AI summary

Forgetting Transformer: Softmax Attention with a Forget Gate

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

arXiv:2503.02130v2cs.LGcs.AIcs.CL

TL;DR

Transformers lack an explicit data-dependent mechanism for forgetting past information, while recurrent models with forget gates can underperform on long-context capabilities. The paper introduces Forgetting Attention and FoX by down-weighting unnormalized softmax attention scores, and reports stronger performance across several settings while retaining long-context retrieval.

  • Problem

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

  • Method

    FoX incorporates a scalar data-dependent forget gate into softmax attention by down-weighting unnormalized attention scores.

  • Results

    FoX outperforms the Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks, while performing on par on long-context downstream tasks.

  • Takeaways & Limitations

    FoX retains Transformer-like long-context retrieval, requires no positional embeddings, and its Pro block design improves both FoX and the Transformer.

  • Takeaways & Limitations

    Experiments are limited to models up to 760M parameters, 48B tokens, a training context length of 16384 tokens, and causal sequence modeling.

Abstract

from arXiv · show

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 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.

1 INTRODUCTION

FoX adds a data-dependent forget gate to softmax attention by down-weighting unnormalized attention scores, targeting the Transformer’s missing mechanism for forgetting past information. It outperforms the Transformer across several tasks while preserving long-context retrieval and compatibility with FlashAttention.

  • Core contribution: FoX applies a data-dependent forget gate to unnormalized softmax attention scores, creating Forgetting Attention.The resulting model is called the Forgetting Transformer (FoX).
  • Empirical findings: FoX outperforms the Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks, while matching it on long-context downstream tasks.
  • Empirical findings: FoX retains strong long-context retrieval, achieving near-perfect needle-in-the-haystack accuracy within the training context length while recurrent models fail.
  • Implementation: FoX requires no positional embeddings and can be implemented with a simple modification to FlashAttention.

2 BACKGROUND: LINEAR ATTENTION WITH A FORGET GATE

The paper motivates forgetting gates by expressing recurrent models in a parallel linear-attention form. A scalar sigmoid gate then down-weights prior attention contributions while retaining softmax-like normalization.

  • Linear attention: Linear attention replaces the exponential dot-product kernel with a feature-based kernel and can be computed either in parallel or recurrently.
  • Forget gate: A scalar forget gate is computed at each timestep as a sigmoid of a learned linear function of the input.
  • Normalization: Unlike many gated linear-attention models, this variant retains a normalization term to preserve similarity with softmax attention.
  • Parallel formulation: The gated linear-attention recurrence has a parallel form with multiplicative factors that down-weight earlier contributions.
  • Connection to softmax attention: Replacing the feature kernel with the exponential dot-product kernel yields softmax attention with a forget gate.

3 FORGETTING TRANSFORMER

Forgetting Attention incorporates scalar, data-dependent forget gates into causal softmax attention through a lower-triangular logit bias. FoX also connects to ALiBi, avoids positional embeddings by default, and supports a Pro architecture and FlashAttention implementation.

  • Forgetting Attention: FoX modifies softmax attention by applying a forget gate to its attention logits before the rowwise softmax.
  • Forgetting Attention: The forget factors form a lower-triangular matrix whose logarithm supplies the causal attention-logit bias.
  • Hardware-aware implementation: The logit-bias formulation can be added to FlashAttention without materializing the L × L bias matrix in high-bandwidth memory.
  • Connection to ALiBi: Forgetting Attention is a data-dependent, learnable version of ALiBi, and experiments verify the superiority of data-dependent forget gates over ALiBi.
  • Architecture design: FoX does not use RoPE or other positional embeddings by default, while FoX (Pro) adds output gates, output normalization, QK-norm, and KV-shift.

4 EMPIRICAL STUDY

The empirical study evaluates FoX against Transformers and recurrent sequence models across long-context modeling, retrieval, downstream tasks, and architectural analyses. FoX generally improves long-context and short-context performance while retaining strong retrieval and context utilization.

  • Long-context language modeling: FoX outperforms the standard Transformer within and beyond the training context length on per-token loss and perplexity.This holds with both tested architectures and demonstrates length extrapolation.
  • Long-context language modeling: FoX maintains decreasing per-token loss across the training context, whereas recurrent baselines flatten around 5k tokens and plateau after 10k.The pattern indicates stronger use of distant context by FoX than by the tested recurrent models.
  • Needle in the haystack: FoX achieves near-perfect needle retrieval within the training context, while Mamba-2, DeltaNet, and HGRN2 perform poorly there.FoX and some Transformer variants partially extrapolate beyond the training context, but extrapolation depends on hyperparameters.
  • Downstream tasks: FoX outperforms the Transformer on short-context tasks, with FoX (Pro) performing best among the evaluated models.The evaluation covers language modeling and zero-shot classification or reasoning benchmarks.
  • Analyses: FoX advantages increase with longer training contexts but decrease with larger models and more training tokens.This suggests the usefulness of forgetting depends partly on the relationship between model capacity and context length.
  • Analyses: Data-dependent forget gates perform best among the tested forget-gate designs, and removing FoX components degrades performance.The component analysis also finds that models using neither forget gates nor positional embeddings perform poorly.

5 RELATED WORK

Recent recurrent sequence models often use forget gates, while alternative softmax-attention improvements generally provide only data-independent decay based on relative positions.

  • Recurrent sequence models: Many recurrent sequence models feature forget gates, reflecting renewed interest in recurrent architectures.The cited passage frames forget gates as a recurring design feature in recent recurrent sequence models.
  • Related improvements and alternatives to softmax attention: Several positional embedding methods add distance-based biases to attention logits, implementing data-independent decay.LAS-attention applies multiplicative exponential decay, while RoPE has a similar distance-dependent decay effect.

6 CONCLUSION

The paper proposes FoX, evaluates it against Transformers and recurrent models, and identifies scale, causality, and efficiency as directions for further work.

  • Conclusion: FoX is a Transformer variant with a forget gate that outperforms the Transformer and several recurrent models across long-context language modeling and downstream tasks.The paper also reports that its Pro block design greatly outperforms the basic LLaMA architecture, with or without a forget gate.
  • Limitations and future work: The study is limited to models up to 760M parameters, 48B tokens, 16384-token training contexts, and causal sequence modeling.The authors identify larger-scale evaluation and extension to the non-causal case as important future directions.
  • Limitations and future work: The authors suggest adaptively pruning computation, such as KV-cache eviction, using forget-gate values to reduce training and inference costs.This is presented as potential future work rather than an evaluated result.
  • FoX layer computation: The implementation computes keys, queries, values, forget gates, and output gates before applying normalized attention, gating, and projection.Multi-head FoX computes each head independently and then sums the head outputs.

B.1 MODEL AND TRAINING HYPERPARAMETERS

The appendix specifies model configurations, training conventions, implementation caveats, and the evaluation procedures for long-context modeling and needle retrieval.

  • Model and training hyperparameters: Table 4 lists configuration hyperparameters, including head dimension for FoX and Transformers and layer counts for the model blocks.The head dimension is tuned in {64, 128} for the 760M-parameter FoX and Transformer models.
  • Model and training hyperparameters: The experiments use LLaMA-style initialization, tuned learning rates, bfloat16 mixed-precision training, and no embedding-output parameter sharing.Forget gates are the exception to the otherwise bias-free linear layers.
  • Implementation caveat: A QK-norm implementation accidentally shared scaling parameters across heads, with no observable impact on performance.The intended implementation would use one set of d_head scaling parameters per head.
  • Model parameters and training analysis: Forget-gate bias terms did not matter for performance in preliminary small-scale experiments at statistically significant levels.The authors retained the bias terms because they might be useful in some cases.
  • Needle-in-the-haystack details: The needle-in-the-haystack evaluation uses standard and easy prompts, with results scored by GPT-4o-2024-08-06 on a 1-to-10 scale.The easy mode places both the question and answer inside the context.
  • Long-context analysis: Per-token loss is analyzed by position because its slope indicates how much earlier context contributes to current-token prediction.Randomly rolling LongCrawl64 sequences removes position bias in expectation.

D DATA-INDEPENDENT FORGET GATE INITIALIZATION

The data-independent and fixed forget gates are initialized using decay horizons, while zero initialization performs poorly for those designs.

  • Initialization principle: T(b) is defined so that σ(b)^T(b) = 1/e, making T(b) the number of timesteps needed to reach 1/e decay.This function determines the decay horizon associated with a forget-gate bias.
  • Initialization principle: Biases for fixed and data-independent forget gates are initialized across exponentially spaced horizons from T_min to T_max.For four heads with (T_min, T_max) = (2, 128), the horizons are (2, 8, 32, 128).
  • Relation to ALiBi: A fixed forget gate with (T_min, T_max) is equivalent to ALiBi with minimum slope 1/T_max and maximum slope 1/T_min.The same initialization was tested for data-dependent gates but was not found useful.
  • Initialization comparison: Zero initialization performs extremely poorly for data-independent and fixed forget gates.The data-dependent forget gate instead uses zero initialization in the reported setup.

E HARDWARE-AWARE IMPLEMENTATION OF FORGETTING ATTENTION

Forgetting Attention is implemented with FlashAttention-style tiled computation, incorporating cumulative forget-gate terms without materializing the full L × L matrix. The implementation keeps the added computation and parameters negligible while supporting forward and backward passes.

  • The forward pass processes Q, K, and V in blocks, computes outputs and logsumexp values on chip, and writes them back to HBM.The block structure avoids constructing the full attention-related matrix in high-bandwidth memory.
  • The hardware-aware algorithms adapt FlashAttention-2 by adding forget-gate terms to attention logits during on-chip block computation.The cumulative sums c = cumsum(log f) are precomputed, and the causal mask remains enforced.
  • The backward pass computes gradients for Q, K, V, and cumulative forget-gate terms through tiled SRAM kernels.The paper notes that these computations may be split into two kernels for implementation simplicity.
  • Forgetting Attention is implemented using Triton FlashAttention, and its scalar forget gates add negligible computation and parameters.The algorithm avoids instantiating the L × L D matrix in HBM.

F.1 PER-TOKEN LOSS FOR THE ABLATION STUDIES

The ablations show that removing RoPE from the Transformer causes very poor performance, while QK-norm appears helpful for length extrapolation in the Pro architecture.

  • QK-norm seems helpful for length extrapolation in Transformer (Pro).The result comes from an incremental ablation that adds components one at a time in the 125M-parameter/2.7B-token setting.

F.3 SHORT-CONTEXT TRAINING ON SLIMPAJAMA

Short-context SlimPajama experiments compare FoX and Transformer variants with recurrent and attention-based baselines. FoX (Pro) outperforms Transformer (Pro), whereas FoX (LLaMA) shows no general advantage except for length extrapolation.

  • The experiments evaluate per-token loss and downstream tasks while comparing FoX, Transformer, Mamba-2, and additional attention or hybrid baselines.The figure and table captions specify the short-context training setup and evaluation formats.
  • FoX (Pro) outperforms Transformer (Pro) in language-modeling loss and downstream tasks under short-context training.The models use roughly 340M non-embedding parameters, roughly 15B SlimPajama tokens, and a 2048-token training context.
  • FoX (LLaMA) has no advantage over Transformer (LLaMA) in this setting except for length extrapolation.The comparison follows the same task set and hyperparameters as the referenced prior work.
  • The long-context comparisons include Transformer-SWA and Samba, both using a window size of 2048.These experiments use the 760M-parameter/16B-token configuration.
  • HGRN2 is evaluated with a needle-in-the-haystack test in the 760M-parameter/48B-token setting.The results are presented separately from the short-context SlimPajama experiments.

F.6 ADDITIONAL RESULTS WITH 125M-PARAM/2.7B-TOKEN, 360M-PARAM/7.5B-TOKEN, AND 760M-PARAM/16B-TOKEN TRAINING CONFIGURATIONS

Additional experiments across model sizes and training budgets confirm the main per-token-loss trends, while probing context-length effects, hyperparameter sensitivity, training stability, and internal gate patterns.

  • Additional per-token-loss results: The reported per-token-loss results are consistent across 125M/2.7B-token, 360M/7.5B-token, and 760M/16B-token configurations and with the 760M/48B-token results.These evaluations use a 16k-token training context.
  • Training context and scale: FoX’s advantage over the Transformer decreases for larger models but increases with longer training context lengths.The comparison covers the 125M/2.7B-token, 360M/7.5B-token, and 760M/16B-token settings.
  • Hyperparameter sensitivity: Length extrapolation is sensitive to the number of training tokens and the learning rate.This sensitivity is measured with easy-mode and standard-mode needle retrieval together with per-token-loss curves.
  • Training dynamics: Training curves differ in shape because the models use different peak learning rates.The curves cover the models presented in the main comparison.
  • Seed stability: Variance across three seeds is small for the 360M-parameter FoX (LLaMA) model.All models in this stability check are trained on roughly 7.5B tokens.
  • Matrix visualizations: The visualizations compare forget-gate matrices F and attention-score matrices A across 16 heads in four layers.Because the matrices are 16384 × 16384, near-diagonal nonzero entries in F can make the visualization appear almost black.
Loading 2503.02130v2…