Summary
The authors introduce the Forgetting Transformer, a transformer variant capable of dynamically focusing on various scopes of history. The mechanism is inspired by the data-dependent forget-gate in recurrent models, which are widely used in modern RNNs. The proposed method demonstrates strong length generalization and long-range capabilities, which the authors evaluated through several benchmarks, including LongCrawl64, Needle, LongBench, and others. Finally, the authors present an efficient implementation that relies on flash attention.
Strengths
1. The proposed transformer variant is simple, elegant, and effective, demonstrating strong performance. If the authors' claims generalize beyond the specific cases presented, the potential impact could be significant.
2. The selected benchmarks and metrics are commendable, incorporating long-range tasks across synthetic and real-world scenarios, such as Needle and LongBench, along with per-token perplexity measurements.
3. The ablation studies in Figure 4 are essential, offering valuable insights into the method’s internal dynamics.
4. Empirical results show notable improvements compared to the baselines.
Weaknesses
I appreciate the ideas and results presented in the paper. However, I have several concerns about the informativeness, consistency, and importance of the method/results:
**W.1. Empirical Evaluation is Insufficient:**
**W.1.1. Current Insufficiency in Empirical Evaluation (Sensitivity to Hyperparameters and Robustness):** The NLP results rely on a **single** instance of the pre-trained forgetting transformer with a specific pre-training setup, which is insufficient to fully support the paper’s claims. Despite impressive performance on length generalization, context tasks, and zero-shot tasks, the robustness and **generality** of these results remain uncertain. I’m not convinced that the results are **not anecdotal**. To strengthen the evaluation, I suggest:
(i) Include **training curves** comparing the variants presented in Figure 1.
(ii) **Additional Datasets:** Conduct pre-training from scratch on more datasets. Even results on a smaller dataset like Wikitext-103 can be valuable. Ideally, it would benefit the community if the authors could compare their method directly with results from original papers. For example, several methods in this domain present results on Wikitext-103 (e.g., HGRN), the Pile (e.g., Pythia, Mamba, RWKV, HGRN, and others), and RedPajama.
(iii) Report results from **multiple seeds** to ensure stability.
(iv) **Robustness Across Models Sizes, Hyperparameters and Settings:** Show that results are not anecdotal by assessing robustness across different hyperparameters and settings, such as varying model sizes (see Mamba and RWKV, for example), other modalities, and different hyperparameters.
W.1.2. Important details about the exact number of parameters, FLOPs, and latency are missing in the tables
___
**W.2. Missing baselines and related works:**
Several transformer variants are similar to the proposed mechanism. The differences between these and the forgetting transformer should be discussed, and some of them should be used as baselines. Examples include:
(i) Mega [1]: This method incorporates Exponential Moving Average before multiplying the Q and K matrices. While this mechanism is not data-dependent, it has several similarities, including manipulation of the attention matrix and adding recency bias.
(ii) COPE [2]: This also manipulates the attention matrix and adds recency bias (in a data-dependent manner). It can be interpreted as a special forget gate.
(iii) Selective Transformer [3]: In a very recent work, a mechanism similar to the forgetting transformer is introduced (see Section 3.3 and Equation 1).
(iv) LAS Attention [4]: This method also adds recency bias (referred to as "local" in the paper) to the transformer by manipulating the attention matrix. It can also be interpreted as a (data-independent) forget gate.
___
I appreciate the ideas, results, and certain aspects of the empirical methodology in the paper. If my concerns are addressed during the discussion period, I will increase my score.
[1] Mega: Moving Average Equipped Gated Attention. Ma et al. ICLR23.
[2] Contextual Position Encoding: Learning to Count What's Important. Golovneva et al.
[3] Selective Attention Improves Transformer. Leviathan et al.
[4] Viewing Transformers Through the Lens of Long Convolutions Layers. Zimerman et al. ICML24.
Questions
1. What is the impact of sharing the new parameters (w_f) across heads?
2. It would be very interesting if standard LLMs could be improved by converting them into a forgetting transformer. For example, you could take Llama 3 7-B, convert it to the forgetting form, fine-tune it, and improve the results. This could be a breakthrough. Have you tried that? What are your thoughts on this?
3. Could you please shed more light on the inner dynamics of the forgetting transformer? For instance, it would be interesting to explore whether different heads focus on different scopes. This could be examined by measuring the F_{ij} values across heads. Additionally, a figure that visualizes the maximal forget values F_{ij} (for j=1, for example) across different positions and inputs could illustrate the forgetting trends (receptive field) learned by the model. While these aspects are partially explored in Figure 2, a more dedicated evaluation considering several inputs, layers, and heads could be very informative.
4. What are the limitations or failure cases of the forgetting transformer? Negative results could be very informative and show the differences between transformer variants.