Forgetting Transformer: Softmax Attention with a Forget Gate

🔼 Figure 2 presents a comparison of the Forgetting Transformer (FoX) and the standard Transformer, and also several recurrent models (Mamba-2, HGRN2, DeltaNet), on long-context language modeling. The left panel shows the per-token loss (L(i)) plotted against the token position (i) within a validation sequence. A lower loss indicates better performance. The right panel shows validation perplexity (P(l)) which is the average loss over a given context length (l). Lower perplexity means better performance. The dashed vertical line in both panels denotes the length of the context seen during training. The per-token loss curves are smoothed using a moving average to reduce noise.

read the caption

Figure 2: (left) Per-token loss L⁢(i)𝐿𝑖L(i)italic_L ( italic_i ) at different token position i𝑖iitalic_i. (right) Validation perplexity P⁢(l)𝑃𝑙P(l)italic_P ( italic_l ) over different validation context length l𝑙litalic_l. The vertical dashed line indicates the training context length. The per-token loss is smoothed using a moving average sliding window of 101101101101 tokens.

🔼 This figure visualizes the forget gate weight matrix (F) and the attention score matrix (A) from two different layers of the Forgetting Transformer (FoX) model. The heatmaps show how the forget gate mechanism modulates attention weights. Because the attention score matrix (A) is sparse, only the entries with scores above 0.5 are displayed. This visualization helps to understand how the model’s attention is influenced by the forget gate, allowing it to selectively focus on relevant parts of the input sequence and effectively ‘forget’ less relevant information.

read the caption

Figure 3: Visualization of the forget gate weight matrix 𝑭𝑭{\bm{F}}bold_italic_F and the attention score matrix 𝑨𝑨{\bm{A}}bold_italic_A from two heads in different layers. Since 𝑨𝑨{\bm{A}}bold_italic_A is very sparse, we only show entries with scores larger than 0.50.50.50.5. These results use FoX (Pro). More examples can be found in Appendix F.10.

🔼 This figure displays the results of a needle-in-the-haystack test, which evaluates a language model’s ability to retrieve specific information from a long sequence of text. The test is performed using two different modes: an ’easy mode’ and a ‘standard mode’. In the easy mode, both the question and the answer are explicitly present in the context. In the standard mode, only the answer is present. The model’s success rate in each mode is presented as a heatmap, where different colors represent different accuracy scores. Scores range from 1 to 10 and were determined by GPT-4. The vertical dashed line in each heatmap represents the length of the training context used to train the model. By comparing performance across the two modes and different context lengths, we can gain insight into the model’s long-context retrieval capabilities.

read the caption

Figure 4: Needle-in-the-haystack analysis for different models. We show results for the easy mode on the left and the standard mode on the right. The results are scored on a scale of 1111 to 10101010 by GPT-4o-2024-08-06. The vertical dashed line indicates the training context length.

🔼 This figure displays the results of the ’needle in the haystack’ experiment, a test of long-context retrieval capabilities, using the Forgetting Transformer (FoX) model. It shows the accuracy of the model at retrieving information from a long document. The ’easy mode’ of the experiment includes both the question and answer in the ’needle.’ The figure also presents the per-token loss, which measures how well the model uses context at various distances from the predicted token. Different lines correspond to different numbers of training tokens and learning rates, allowing for the comparison of training data quantity and rate on the performance and loss. The dashed line indicates where the training context ends.

read the caption

Figure 5: FoX (Pro) easy mode needle-in-the-haystack results and per-token loss for different numbers of training tokens and learning rates. The vertical dashed line indicates the training context length. More results can be found in Appendix F.7.

🔼 Figure 6 illustrates the impact of model size, training data size, and training context length on the per-token loss of language models. The plots show the average loss calculated across a 101-token sliding window for each token position. Results are shown only for token positions within the training context length to improve clarity, but Appendix F.6 provides additional results, including an analysis of length extrapolation and results for a 125M parameter model.

read the caption

Figure 6: Per-token loss given different model sizes, numbers of training tokens, and training context lengths. At each token index i𝑖iitalic_i, we report the averaged loss over a window of 101101101101 centered at i𝑖iitalic_i. We only show results within the training context length to reduce visual clutter. See Appendix F.6 for additional results, including length extrapolation and 125M-parameter model results.

🔼 Figure 7 compares the performance of different forget gate mechanisms in the Forgetting Transformer model. It contrasts data-dependent, data-independent, and fixed forget gates across two model architectures: the LLaMA and Pro architectures. The results show the per-token loss for each configuration, calculated using a 360M parameter model trained on 7.5B tokens of the LongCrawl64 dataset. Smoothing is applied to the loss curves using a moving average over 1001 tokens. The training context length is marked by a dashed vertical line.

read the caption

Figure 7: Data-dependent forget gate (data-dep) vs. data-independent (data-indep) and fixed forget gate. (left and middle-left) Comparison using the LLaMA architecture. (middle-right and right) Comparison using the Pro architecture. We use 360360360360M-parameter models trained on roughly 7.57.57.57.5B tokens on LongCrawl64. All per-token loss curves are smoothed with a moving average sliding window of 1001100110011001 tokens. The vertical dashed line indicates the training context length.

🔼 This table presents the performance of various language models on the LongBench benchmark. Each model has approximately 760 million non-embedding parameters and was trained on roughly 48 billion tokens from the LongCrawl64 dataset. The table shows the results for a range of downstream tasks, allowing for a comparison of the models’ performance across different types of long-context understanding challenges. The best and second-best results for each task are highlighted in bold and underlined, respectively.

read the caption

Table 2: Evalution results on LongBench. All models have roughly 760760760760M non-embedding parameters and are trained on roughly 48484848B tokens on LongCrawl64. Bold and underlined numbers indicate the best and the second-best results, respectively.

🔼 This ablation study investigates the impact of individual components within the Forgetting Transformer (FoX) model. Using 360M parameter models trained on a 7.5B token subset of the LongCrawl64 dataset, the experiment measures the validation perplexity on a 16384-token context length. The bottom half of the table shows incremental additions and removals of components, comparing the results against the baseline FoX (Pro) model. This allows for a precise understanding of each component’s contribution to the model’s overall performance.

read the caption

Table 3: Ablation experiments for FoX. We use 360360360360M-parameter models trained on 7.57.57.57.5B tokens on LongCrawl64. The perplexity is measured over a validation context length of 16384163841638416384 tokens. For the bottom half, all addition (+) or removal (-) of components are relative to FoX (Pro).

🔼 This table details the hyperparameters used for various model configurations in the experiments. It shows the different settings for model size (number of parameters), training data size (number of tokens), the number of layers, and the head dimension (d_head). The head dimension is only relevant to the FoX and Transformer models, and the authors tuned its value within the set {64, 128} for the 760M parameter models. The number of layers (n_layer) indicates the total number of blocks in the model; for the Transformer, each block consists of one attention layer and one multi-layer perceptron (MLP) layer.

read the caption

Table 4: Hyperparameters for different configurations. The head dimension dheadsubscript𝑑headd_{\text{head}}italic_d start_POSTSUBSCRIPT head end_POSTSUBSCRIPT is only applicable to FoX and the Transformer. We tune dheadsubscript𝑑headd_{\text{head}}italic_d start_POSTSUBSCRIPT head end_POSTSUBSCRIPT for the 760M-parameter FoX and Transformer models in {64,128}64128\{64,128\}{ 64 , 128 }. nlayersubscript𝑛layern_{\text{layer}}italic_n start_POSTSUBSCRIPT layer end_POSTSUBSCRIPT counts the number of blocks. For example, for the Transformer each block contains an attention layer and an MLP layer.

🔼 This table shows the peak learning rates used for various language models in the main experiments. All models used in the main experiments had 760 million parameters and were trained on 48 billion tokens. The learning rates were determined through a hyperparameter search. The search involved testing various learning rates, each multiplied by 1, 2, and 5, which are presented in scientific notation using a power of 10 (10ⁱ, 2 × 10ⁱ, and 5 × 10ⁱ, where i is an integer representing a different magnitude of the learning rate).

read the caption

Table 5: Peak learning rates for different models for the 760M-parameter/48B-token main experiments. These are tuned using a grid {1×10i,2×10i,5×10i}1superscript10𝑖2superscript10𝑖5superscript10𝑖\{1\times 10^{i},2\times 10^{i},5\times 10^{i}\}{ 1 × 10 start_POSTSUPERSCRIPT italic_i end_POSTSUPERSCRIPT , 2 × 10 start_POSTSUPERSCRIPT italic_i end_POSTSUPERSCRIPT , 5 × 10 start_POSTSUPERSCRIPT italic_i end_POSTSUPERSCRIPT } with different values of i𝑖iitalic_i.

🔼 This table presents a quantitative comparison of different language models in terms of their computational efficiency and training speed. It lists the number of parameters (N), the estimated floating-point operations per token (FLOPs) for a forward pass, the formulas used to calculate these FLOPs, and the training throughput measured in tokens per second on hardware with 4 NVIDIA L40S GPUs. The FLOPs and throughput values provide insights into the computational cost and speed of training each model. The formulas offer further details on the calculation methodology, highlighting the dependencies on the number of parameters (N) and training context length (L). This allows for a comprehensive evaluation of each model’s efficiency and scalability.

read the caption

Table 6: Number of parameters, estimated forward pass FLOPs per token, formulas for FLOPs estimation, and training throughput in tokens per second for different models. Throughput is measured on 4 NVIDIA L40S GPUs. N𝑁Nitalic_N is the number of parameters and L𝐿Litalic_L is the training context length.

🔼 This table presents the results of several language models evaluated on the LM-eval-harness benchmark. The models were trained on the SlimPajama dataset with a training context length of 2048 tokens. All models have approximately 340 million non-embedding parameters and were trained on about 15 billion tokens of SlimPajama data. The evaluation metrics include perplexity (ppl) and accuracy (acc) for various tasks. Length-normalized accuracy (acc-n) is also reported for some tasks. The best and second-best results for each task are highlighted in bold and underlined, respectively. Note that some results (Transformer++, DeltaNet) are taken from a previous publication (Yang et al., 2024) for comparison purposes. The Transformer++ model from Yang et al. (2024) is architecturally identical to the Transformer (LLaMA) model used in this paper.

read the caption

Table 7: Evaluation results on LM-eval-harness for models trained on SlimPajama with a training context length of 2048204820482048 tokens. All models have roughly 340340340340M non-embedding parameters and are trained on roughly 15151515B tokens on SlimPajama. “acc-n” means length-normalized accuracy. Bold and underlined numbers indicate the best and the second best results, respectively. Note the results for Transformer++ and DeltaNet are from Yang et al. (2024). Note that Transformer++ from Yang et al. (2024) and Transformer (LLaMA) in our work have exactly the same architecture.

🔼 This table presents the results of evaluating various language models on the LM-eval-harness benchmark. All models in the comparison have approximately 760 million non-embedding parameters and were trained on roughly 16 billion tokens from the LongCrawl64 dataset. The benchmark assesses performance across several tasks. Length-normalized accuracy (‘acc-n’) is reported for tasks where sequence length affects the score. The best-performing model for each task is indicated in bold and underlined, along with the second-best performing model.

read the caption

Table 8: Evaluation results on LM-eval-harness. All models have roughly 760760760760M non-embedding parameters and are trained on roughly 16161616B tokens on LongCrawl64. “acc-n” means length-normalized accuracy. Bold and underlined numbers indicate the best and the second best results, respectively.

🔼 Table 9 presents the results of several long-context downstream tasks from the LongBench benchmark. Each model in the table has approximately 760M parameters (excluding embeddings) and was trained using roughly 48 billion tokens from the LongCrawl64 dataset. The table compares the performance of the Forgetting Transformer (FoX) with various architectures against several baseline models (including the standard Transformer and recurrent models like Mamba-2, HGRN2, and DeltaNet). The best-performing model for each task is shown in bold and underlined.

read the caption

Table 9: Evalution results on LongBench. All models have roughly 760760760760M non-embedding parameters and are trained on roughly 16161616B tokens on LongCrawl64. Bold and underlined numbers indicate the best and the second best results, respectively.