Forgetting Transformer: Softmax Attention with a Forget Gate
Authors: Zhixuan Lin
Evgenii Nikishin
Xu He
Aaron Courville
Venue: ICLR 2025 Poster (2025)
Keywords: sequence model
long-context sequence modeling
Transformer
softmax attention
linear attention
RNN
language model
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose the Forgetting Transformer, a Transformer variant with a forget gate.

Abstract: 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`](https://github.com/zhixuan-lin/forgetting-transformer).

Published as a conference paper at ICLR 2025
FORGETTING TRANSFORMER: SOFTMAX ATTENTION WITH A FORGET GATE
Zhixuan Lin"
Evgenii Nikishin
Mila & Université de Montréal 2x1 in. csê gma i coo
Mila & Université de Montréal
evgenii ni niki ki sêin 1a quebec
Xu Owen He! MakerMaker AI
Aaron Courville
Mila & Université de Montréal
owen wew.uggaail.comm com
courvila@mila la quebec
ABSTRACT
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 at- tention 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 down- stream tasks. Moreover, it is compatible with the FlashAttention algorithm and does not require any positional embeddings. Several analyses, including the seeell-inntte-haastee
test, show that FoX also retains the Transformer's 's supe- rior long-context capabilities over recurrent sequence models such as Mamba- 2, HGRN2, and DeltaNet. We also introduce a "Pro" block design that incor- porates some common architectural components in recurrent sequence models and find it significantly improves the performance of both FoX and the Trans- former. Our code is available at https: /ttps://github. com/ zhi xmainninni forgetting-trans former
INTRODUCTION
Despite the growing interest in reviving recurrent sequence models (Gu al.. 2021 Peng et al.
2021 Yang et al.) 2023; Gu & Dao 2023 Sun et al. [2023] De et al. [024] Qin et al. [024b Dao &
Gul 2024 Peng et |l.||022 Beck et al. 2024 Zhang al. 2024), these models still underpertorm
the Transformer (Vaswani et al.) 2017) in terms of long-contexi capabilities (Hsieh et al. 2024 Shen
Waleffe al. 2024| Shen et al. 2024 Qin al., (024a), likely due to their relatively small fixed-
sized hidden states (Jelassi et al.) 2024). While the Transformer excels in handling long-conten information, it lacks an explicit mechanism for forgetting past information a data deeendde way. Such mechanism often implemented as some form of the forget gate (Gers al. 2000) is ubiquitous in recurrent sequence models and has proven critical in their success in short-context tasks (Greff al. 2016 Van Der Westhuizen & Lasenby 2018 Peng al.) 2021 Yang al. [2033
Gu & Dao 2023 A natural question to ask is then: can we have forget gate in Transformers?
To address this question, we leverage important fact: many recurrent sequence models with forget gate can be written in parallel linear attention form (Katharopoulos et al.) (020)) analogous softmax attention (Yang et al. 2023 Dao & Gu 2024). In this parallel form, the forget gate mechanism translates into down-weighing the unnormalized attention scores in a data-dependent way. Our key insight that this exact mechanism also applicable to softmax attention. We name this attention mechanism Forgetting Attention and the resulting model the Forgetting Transformer (FoX).
"Coneesoondence to Zhixuan Lin. Work done while Google DeepMind.

---

Published as a conference paper at ICLR 2025
We show that FoX outperforms the Transformer long-conterd language modeling, length trapolation, and short-context downstream tasks, while performing on par with the Transformer on long-context downstream tasks. Notably, it does not require any positional embeddings. It also tains Transformers' long-context retrieval abilities and achieves near-perfect accuracy in the needle- in-the-haystack test (Kamradt (023) within the training context length. In contrast, all the tested recurrent sequence models fail. We also introduce a "Pro" block design that integrates several archi- tectural components commonly used in recurrent sequence models, which significantly improves the performance of FoX and the baseline Transformer. Finally, we show that FoX can be implemented hardware- way with simple modification to the FlashAttention (Dao 2023) algorithm.
BACKGROUND: LINEAR ATTENTION WITH A FORGET GATE
This section introduces the notation used in this work and gives a brief background on linear at- tention. We also introduce gated variant of linear attention and discuss its parallel form, which naturally leads to FoX. Throughout this work, we only consider causal sequence modeling.
2.1 LINEAR ATTENTION
Standard causal softmax attention takes a sequence of input vectors (a:)/=1 and produces sequence of output vectors (o.) where æi, O; RRd, E (1,..., L). Each is computed follows:
q:: Kú, U, = Wq WEX, W. E R
(1)
Σ Kexp (gi, KU
Σラ= exp(g, k2)
O;
(2)
keep (qi, k;)
Σj exp(q k;))
where W, WR. W. E Raddd are projection matrices and kexp (q. k) exp(q' k) is the exponential dot product kernel
Linear attention (Katharopoulos et al. 2020) replaces the exponential dot product kernel
keep (q. k) = exp(q k) with a kernel (q. k) with some feature representation ø : Rd (R+)*
Σ)= kip (g k5 ()(()(())) (k;))) Σj hip qi, kj))
Σj=1 (q))) T((k))
Following Yang et al. (2023). we call this the parallel form of linear attention as it can be computed with matrix multiplications. Alternatvely, linear attention can computed recurrent form:
S = 1 + (k((
()
2, 2,-1 + ((k:))
S.b(g:)
O, 2. ((q:)'
where S, Rdxd' 2, E IRd E (O....,
.......I. L} are computed recurrently, with Sp = 0 and 0.
2.2 LINEAR ATTENTION WITH A FORGET GATE
The recurrent form of linear attention makes natural to introduce a forget gate. Specifically, we can compute scalar forget gate f. aww] / æ, + by) function and wy E RRd by E R are learnable parameters. The recurrent computation then:
each timestep, where is the sigmoid
S. f.S,-1 v,O(())
2: = f + ((k,)
(8)
S,0(qi)
O: 2, o(q))
Note omit the J scaling factor reduce visual clutter. practice we always scale q, k; by Va
2

---

Published as a conference paper at ICLR 2025
Note that this gated variant of linear attention differs from most models in the literature. In particular, most gated variants of linear attention models, such as GLA (Yang et al. 2023) and Mamba-2 (Dao)
& Gu 2024), do not have the normalization term (i.e., there no 2,. and the output is just o, Seo(ge)) We keep the normalization term to maintain similarity with softmax attention.
Crucially, similar to the normalization-fre free version derived in GLA and Mamba-2, can show that this gated variant of linear attention also has a parallel form:
0, Σn=1 Fayo( q:)
(k;)) (
Σj=
(10)
Fu TT(k))
Σ/=1 Fiyhs(q: kj)
where Fij = I ft, with the convention that Fij 1 i j. Our key observation is that
Equation 10 and the softmax attention in Equation 2 are very similar in form. In fact, if we change the kernel K in Equation [0] back to the exponential dot product kernel kexp- obtain softmax attention with a forgel gate. We introduce this formally in the next section.
FORGETTING TRANSFORMER
Our proposed model, the Forgetting Transformer (FoX), features modified softmax attention mechanism with a forget gate. We name this attention mechanism Forgetting Attention. Similar the gated variant of linear attention introduced in the previous section, we first compute scalar forget gate ft = (w I + by) E R for each timestep t. The output of the attention is then
Σj= Fig exp(a k5)0s
Ej exp(g k, + Diy)u,
0,
(11)
Ej=1 Pi explq, k:))
Σj= exp(q, ky + Diy)
where P.j = -II Si and Dij = -log/ Fij = Ei=s+i log ft. This can be written in matrix form:
D log F E ERLXL
(12)
0 soft (QKT D)V E RR
(13)
where F E RLXL is lower triangular matrix whose non-zero entries Fij, i.e, Fy Fij
if i 2 and 0 otherwise. We adopt the convention that log 0 -c. Q.K. V. RLxd are matrices containing qi, v;,
v. o,, {1..... L} as the rows. The softmax operation is applied row- wise. For multi-head attention with H heads, maintain H instances of forget gate parameters !
{w) (h) 1 Jh and (6))) Jhm1 and compute the forget gate values { ( J separately for each head.
Hardware-aware implementation The logit bias form on the rightmost side of Equation can be computed with simple modification to the FlashAttention (Dao) 2023) algorithm. Here we briefly describe the forward pass. The backward pass follows similar idea. First, we compute the cumulative sum G, = Ei= log fi for E fl...., L} and store it in the high-bandwidth memory (HBM) of the GPU. This allows us to compute Dij c, Cj easily later. Whenever compute the attention logit q, k in the GPU's fast shared memory (SRAM) (as in FlashAttention), load e and Cy to SRAM, compute Diy. and add it to the attention logit. The rest of the forward pass remains the same as FlashAttention. This algorithm avoids instantiating the L x L D matrix in the HBM. We provide a detailed algorithm description in Appendix D Moreover, since the forget gates scalars instead of vectors, the additional computation and parameters introduced are negligible.
also
Connection to ALiBi Besides its natural connection to gated linear attention, Forgetting Attention also be seen data-dependent and learnable version of ALiBi (Press et al.) 2021). ALiBi applies a iata--nddepnnee bias byj - (i j)m, to the attention logits, where my, is a fixed slope specific each head h. It easy show that ALiBi equivalent to using a fixed, head-specific, and data-independent forget gate f( = exp(-ma) In Section 4.5
we verify the superiority of
data-dependent forget gates over ALiBi.
Positional embeddings Though we find that using Rotary Position Embeddings (RoPE) (Su et al.) 2024) sometimes slightly improves the performance of FoX, it is not necessary as it is for the stan- dard Transformer (see ablations in Section 4.5). For simplicity, we do not use RoPE or any other positional embeddings for FoX by default.
3

---

Published as a conference paper at ICLR 2025
Linear
SWIGLU MLP
RMSNorm
RM
Forgetting Attention
q1
kt
V
Fox layer
RMSNNrrm
RM5Norm
Linear
ShiftLinear
Linear
Figure 1: Default architecture of FoX. (left) A single FoX block. (right) A single FoX (Pro) layer. All RMSNorms on the right are applied independently each head. is the sigmoid function. @ is element-wise multiplication. Shi ft Linerr implements the computation in Equation 14
Architecture design We test FoX with different architectures. First, replace RoPE in the LLaMA architecture (Touvron et al. 2023) with forget gates and refer to this model as FoX (LLaMA). Second, we test an improved "Pro" architecture with output gates and output normaliza- tion (also used in GLA and Mamba-2). We also use QK-norm (Dehghani et al.) 2023) and apply simplified variant of data-dependent token shift (Peng et al. (2024) to the keys and values (K (KV-shift). Concretely, the keys (k:) are computed follows with additional parameters we Rd:
k; = Wa Rdd, = o x..) E
(14)
k. R.MSNorm(a,)
key k a, key )e)
The values ( are computed in the same way, but without RMSNorm. The overall architecture We refer to the resulting model as FoX (Pro).
is shown in Figure T and detailed in Appendix A
EMPIRICAL STUDY
The advantages of Transformers in long-contest abilities over recurrent sequence models have been verified multiple times (Hsieh et al. 2024 Waleffe et al.) 2024 [hhnn et al. 2024 Qin et al., 2024a) However, forget gates introduce a recency bias. is thus natural ask whether FoX still maintains this advantage. Therefore, our empirical study places special focus on long-context capabilities.
4.1 EXPERIMENTAL SETUP
Dataset and baselines We focus on long-conters language modeling and train all models on
LongCraw164 (Buckman) 2024), long-sequence subset of RedPajama-v2 (Together Computer 2023) pre-tokenized with the Tik tokenizer (OpenAl 2022) for GPT-2 (Radford et al. 2019)) For baselines, we focus on two types of comparisons. First, we compare FoX with the Transformer. also test both the LLaMA and the Pro architecture (referred to as Trans- former (LLaMA) and Transformer (Pro), respectively). Similar to Xiong et al. (2023). we find large RoPE angle 0 for the Transformer for long-context training. Following Xiong jtt al. (2023) we use f 500000. Second, to show the advantage FoX over recurrent sequence models in long-context capabilities,
For the Transformer,
erucial
compare with Mamba-2 (Dao & Gul 2024), HGRN2 (Qin) et al. 2024a), and DeltaNet (Yang et al. 2024). The implementation of all models is based on the Flash Linear Attention repository (Yang & Zhang 2024).
Training setup For our main experiments, train models with 760M (non-embedding) parame- ters on x -tokn (roughly 48B tokens) subset of LongCraw164 with a training context length
Whee output gates are used, reduce the number of parameters in the MI.Ps so the total number of parameters remains the same.
4

---

Published as a conference paper at ICLR 2025
1.65
1.60
Perplexity P())
Loss L(t)
5.0
1.55
4.8
1.50
4.6
4.4
16k
24k
32k
40k
48k
56k
64k
16k
24k
32k
40k
48k
56k
64k
Token index
Validation context length
FoX (Pro)
Transformer (Pro)
Mamba-2
DeltaNet
FoX (LLaMA)
Transformer (LLaMA)
HGRN2
Figure 2: (left) Per-token loss L(t) at different token position i. (right) Validation perplexity P(I) over different validation context length I. The vertical dashed line indicates the training context length. The per-token loss is smoothed using a moving average sliding window of 101 tokens.
of 16384 tokens. For the validation set, we use a 2 x 230-token subset of the LongCraw164 held-out set with sequences of 65536 tokens. We choose a longer validation context length than the training context length to test the length extrapolation abilities
the models. All models are trained with AdamW (Loshchilov) 2017) with (31, (2) (0.9, 0.95). We use linear learning rate warmup from 0 to the learning rate for the first 256 x 22 tokens and then a cosine decay to O. Each training weight decay of 0.1 and gradient clipping of 1.0. We search the learning rate for each model within {1 x 10', x 10 x 10') for different i's until identify locally optimal value. We tune the head dimensions for FoX and the Transformer in {64, 128]. We find that FoX often prefers higher learning rates and more heads/smaller head dimen- sions than the Transformer, and the Pro models often prefer higher learning rates than the LLaMA models. Details of the hyperparameters and experimental setup can be found in Appendix B
batch contains 0.5 x 2 tokens. All models
4.2 LONG-CONTEXT LANGUAGE MODELING
Metrics For main metric, use per-token loss the validation set at different token posi- tions. To be precise, let V be the vocabulary size, y!"' E 13 be the one-hot vector encoding the language modeling target for the i-th token in the j-th validation sequence, and p!") RVV be the corresponding output probabilities of the model, then the per-token loss T((i) at token position i is defined as L(t) 1/4 Σラ= log((p,)) T )T v1 where M is the number of validation sequences.
The per-token loss is particularly meaningful for understanding the long-context capabilities of model. Informally, a monotonically decreasing L(i) with a steep slope indicates the model using the full context well. On the other hand, if L(x) plateaus after some token position k, it indicates the model struggles to use tokens that are tokens away from the current token position for its prediction. This correspondence between the slope of L(i) and the model's context utilization is explained in more detail in Appendix C
Besides per-token loss, also report perplexity over different context lengths. Concretely, perplex- ity P(I) over a context length is defined as P(I) exp() El=1 L(i)). We warn the readers that the slope of P(I) is less meaningful. Since P(2) is the exponential of the cumulative average of L(i), even if L(i) plateaus after some token position k. P(I) may still keep decreasing after giving the wrong impression that the model can make use of the part of the context that is k tokens away.
Results In Figure show the por-token loss L(1) at different token indices i and perplexity P(I) over different validation context lengths As shown in Figure 2 with either architecture, FoX outperforms the standard Transformer both within and beyond (i.e., length extrapolation) the
5

---

Published as a conference paper at ICLR 2025
Layer 8, Head 2,
Layer 8, Head 2,
Layer 16, Head 4,F
Layer 16, Head 4,
0k
0k
Query Position
Query Position
10kk
15k
15k
15k
Sk
10k
15k
0k
15k
Key Position
Key Position
Key Position
Key Position
Score
0.0
0.2
0.6
0.8
1.0
Figure 3: Visualization of the forget gate weight matrix F and the attention score matrix A from two heads different layers. Since A is very sparse, we only show entries with scores larger than 0.1. These results use FoX (Pro). More examples can be found in Appendix F.10
training context length. Similar the Transformer, maintains a monotonically decreasing per- token loss within the training context length, indicating that it utilizes the entire training context for its prediction. contrast, the per-token loss curves all recurrent sequence models start flattening around 5k tokens and plateau after 10k tokens. This indicates that these recurrent models struggle to use the full context effectively for their prediction. In terms of the absolute values of the loss and perplexity, FoX (Pro) also clearly outperforms HGRN2, DeltaNet and Mamba-2.
Visualization of forget gate values and attention map In Figure 3 visualize the forget gate weight matrix and the attention scores A = softmax QK log from two heads different layers. The head the left-hand side exhibits strong decay, and most entries of F are close to zero; accordingly, the attention focuses on local entries. The head on the right-hand side has much weaker decay, and the attention is distributed across the entire context. This shows that FoX can learn to retain information across long contexts when necessary.
4.3 NEEDLE IN THE HAYSTACK
is popular test for the long-context retrieval abil- ities of language models. Besides the standard mode where the "needle" only includes the answer to be retrieved, we also use an "easy mode" (Qin et al. 2024a) where the "needle" placed in the context includes both the question and the answer. This easy mode is particularly suitable for base models that have not been instruction-tuni
The needle-ing the tee--thhhhhsssee haystacyy test (Kamradt 2023
Full details, including the prompts used, are in Appendix B33]
In Figure 4 we show the results for different models. HGRN2 performs even worse than Mamba-2 and we leave it to Appendix F.5 As shown in Figure FoX achieves near-perfect needle retrieval within the truining context length all cases. Transtormer (Pro) and Transformer (LLaMA) also have perfect accuracy within the training context length in the easy mode, though they sometimes fail in the standard mode In contrast, Mamba-2 and DeltaNet (and also HGRN2 in Appendix F.5) perform poorly
within the training context length. FoX (Pro), FoX (LLaMA), and Transformer (Pro) also partially extrapolate beyond the training context length. This is expected given their per-token loss pattern beyond the training context length (see Figure 2). However, find that the extrapolation behaviors of these models could be ayperparameter-ddeendeeee Figure 5 we show that for FoX (Pro), the needle retrieval results and the per-token loss slope beyond the training context length vary depending on the number of training tokens and learning rates. In particular, we find that more training tokens often leads to worse extrapolation, indicating that the models may be gradually "overfitting" to their training context length during training.
For example, in
4.4 DOWNSTREAM TASKS
We use two sets of downstream tasks: set of short-context tasks from L.M-vvalaation-hareee Gao al. 2024) and set of long-context tasks from LongBench (Bai et al.) 2023).
Note these are small models without instraction-tuning. We expect that with more parameters/training tokens instruction-tuning Transformers should also have perfect accuracy within the training context length.

---

Published as a conference paper at ICLR 2025
FoX (Pio), Easy
FoX (LLaMA). Easy
FoX (Pro), Slandard
FoX MA) Standaed
68
04
0.9
0.9
108
106
100
10.3
21-8
16.8
Dopth Porcun
201
Depth Perernt
310
20.3
60.8
406
Depth
00.8
60.0
50.6
00.
500
609
Depth
20.3
00.3
704
20a
201
20.3
6/
90.3
1008
1004
1003
1003
Desemmnt Lengh
Decument Length
Docoment 1 Lengh
Document L Lengh
Transformrr (Pro). Easy
Transoomer (LLaMA), Easy
Transfoemer (Pro), Standard
Transformer (1 MMA Standand
08
04
00
0.
108
106
1va
11.
50.8
20:8
Dapih Poromit
306
201
Dapth Porcent
20.
30a
Depih Peremt
20.
40.8
404
Dapih
20.8
60.8
504
606
200
ma
90.
70:0
200
804
90.8
900
60
00.
1008
1904
1000
100
1200
Document Lengtt
Decument Length
Document Length
Document Length
Mamba 2. Easy
DeltaNet, Easy
Mamba 2, Stundard
NeltaNet, Stundard
68
06
0.
0.9
108
104
10.
10.9
50.8
20.0
304
204
Depth Porcent
20.
Depih Peecent
20.9
20.9
408
Depih
50.8
60.8
Depit
30.6
00.6
50.
609
20.9
708
704
700
50.0
904
800
100
500
90.
1008
1004
1000
1009
Decument Lenglt
Decument Length
Document Lengt
Document Lengh
Score
10
Figure 4: Needle-in-thee beeele-inttee-haystedd analysis for different models. We show results for the easy mode on the left and the standard mode on right. The results are scored on a scale to 10 by GPT- 40-2024-08-06. The vertical dashed line indicates the training context length.
LRx0.00)
1Ro-0.002
LR=00001
L2=0,002
1.9
140t tnkers, RRC (01
Dogth Porcnnt
201
100
tokess. 5,002
Looe M/0
483 tokees, LB 5,002
200
A00
1800
1.5
Dociment Lengh
Decument Length
Document Lengt
Document Lengh
Figure 5: FoX (Pro) easy mode needle-in-the-have -h-thehaaystack results and pok-token loss for different num- bers of training tokens and learning rates. The vertical dashed line indicates the training context length. More results can be found in Appendix F.7
Table 1: Evaluation results on LMM-val-harnedd All models have roughly 760M non-embedding pa- rameters and are trained on roughly 48B tokens on LongCraw164. "acc-n" means length- normml-nd accuracy. Bold and underlined numbers indicate the best and the second best results, respectively.
Model
Wiki. ppl.
LMB.
LMB.
Hella.
Winn.
ARC-c
ARC- acc-n
COPA
SciQA acc1
BoolQ acc]
ppl]
acc?
acc1
ace1
acc1
FoX (Pto)
23.04
14.91
42.75
64.09
38.39
52.33
52.23
26.54
71.00
29.80
85.10
46.57
50.88
Trunsformser (Pro)
24.12
16.16
41.47
64.04
36.60
49.72
51.98
25.26
62.00
29.20
$2.80
60.86
50.39
FoX (LLaMA)
Transformx (LLaMA)
28.14
22.34
38.27
63.22
34,20
49.49
47.98
24.49
66.00
29.40
788 78.90
58.93
49.09
Mumba-2
28.20
21.05
36.50
63.17
35.86
50.59
49.96
25.60
71.00
31.00
80.90
57.49
50.21
HGRN2
DelraNet
29.17
29.14
34.17
62.73
35.28
30.25
47.39
24.32
70.00
29.00
74.30
54.37
47.99

---

Published as a conference paper at ICLR 2025
Table 2: Evalution results LongBench. All models have roughly 760M non-embedding parame- ters and are trained on roughly 4SB tokens on LongCraw164. Bold and underlined numbers indicate the best and the second-best results, respectively.
Single
Leorning
Model
Qpipet
MFQA
JWALDDD
Minippe
Mattrieen
TREC
Thindg4
LCC
Rep.Beuch
FoX (Pro)
18.36
18:88
28-15
15:77
26.99
6.69
12.71
13.51
12.27
A35
31.56
12:34
10.9
9.1
Thansfommer (Ph))
11.42
21.54
22.50
19.58
22.65
6.09
21.92
10.7
8.11
35.0
39,67
30,66
10.79
14.25
(1-aMM)
13.0
+86
61.5
19.11
Trandemm (.LaMA)
18.53
10.9
19.17
25.5
156
15.17
DehaNat
15.09
12.34
350M perameters/...5 taining tokeus
760M parametiers o training tokens
1.85
2.1
Lose M63
Lous L(0
1.80
2.0
1.75
1.70
L.63
512
1024
2048
4096
8192
16384 32768
512
1024 2008 48%6
8192
16384 52768
Tokan index
Token index
Model
Training context length
Model FoX (Pro)
Trining context lenghh
FoX (Pro)
4096
16384
4096
16384
Transformer (Pro)
8192
32768
Tranformer (Pro)
8192
32768
Figure 6: Per-token loss given different model sizes, numbers of training tokens, and training context lengths. At each token index i, we report the averaged loss over a window 101 centered at 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 155M-arrmettr model results.
Short-context tasks We use Wikitext (Merity et al. 2016 LAMBADA (Paperno et al. PiQA (Bisk al. 2020,, HellaSwag (Zellers et al. 2019), WinoGrande (Zellers et al. 2019
2016
ARC-
easy, ARC-challenge (Clark al. 2018), Copa [Roemmele al. 2011). SciQA Auer al. 2023 OpenbookQA (Mihaylov et al. 2018), and BoolQA (Clark et al. 2019). Following Yang et al. 2023)) report perplexity for Wikitext and LAMBADA, length- normalierd accuracy for Hel- laSwag, ARC-challenge, and OpenbookQA, and accuracy for all other tasks (we also report accuracy for LAMBADA). All results are zero-shot. As shown in Table [ FoX outperforms the Transformer with either architecture. FoX (Pro) performs the best among all models.
Long-context tasks We use tasks from LongBench: HotpotQA (Yang el. 2018) 2WikiMul- tihopQA (Ho et 2020), MuSiQue (Trivedi et al. (022). MultiFieldQAAen Narrative (Koiskk ct al. 2018), Qasper (Dasigi et al. 2021) GovReport (Huang et al. 202T). QMSum (Zhong et al, 2021. MultiNews (Fabrr et al. 2019). TriviaQA (Joshi et 2017). SAMSum (Gliwa al. 2019), TREC (Li & Roth 2002). LCC (Guo et al. 2023), and RepoBench-P (Liu et al. 2023). We use the default metrics of LongBench for different tasks, which either F1, Rough-L. accuracy, edit similarity. As shown in Table [ with either architecture, FoX performs on par with the Transformer and better than the recurrent sequence models.
4.5 ANALYSES
We present three sets of analyses. First, we study how the advantages FoX over the Transformer vary with model size and training context length. Second, investigate the contribution of each component in FoX and how RoPE affects performance. Finally, we study the importance of using forget gate that is data-dependent. For these experiments, trained on roughly 16B tokens or 360M-parameter models trained on roughly 7.5B tokens. Experi- mental details be found in Appendix
either 7 models
8

---

Published as a conference paper at ICLR 2025
Table 3: Ablation experiments for FoX. We 360M-parameter models trained 7.5B tokens LongCraw164. The perplexity is measured over a validation context length of 16384 tokens. For the bottom half, all addition (+) or removal (-) of components are relative to FoX (Pro).
Mudel
RoPE
Forget gate
QK-norm
Outgut gate
Output norm
KVshift
Perplexity
Transformer (L.)aMA) we RoPE. Transformer (ULaMA)
29.30
7.49
7.19
FoX (LLaMA)
7.25
7.08
6.88
ToX (Pro)
6.62
QK-norm
6.79
outpat gale
6.86
output norm
6.69
KV-shift
7/
+ RoPE:
6.63
forget gate RoPE Transformer (Pro))
6.82
forget gate
7.40
dets-dap
(Pro), dan-dap
(Pro), dato-dpp
2.1
finad Tise 2016
fxod Tiu 202
dese Lidee,de,,, Tinss
dute-indep, Tass 2048
252
2.1
fxed. Jane
fined Time
2048
266
date date-indeg Yes 256
dat--indap, Tus 2348
Lose A4
dbta-indoen Yaue
das-indaay Yeu
2.0
2.0
1.9
1.8
1.8
1.8
20300
45000 60020
20300
40003
50000
20200
40003
50000
20000 40000
50002
Tokeo todex
Tikkes Indes
Tuken Indin
Taken tndks
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 360M-parameter models trained roughly 7.5B tokens on LongCraw164. All per-token loss curves are smoothed with a moving average sliding window of 1001 tokens. The vertical dashed line indicates the training context length.
Model size and training context length Figure 6 we show the per-token loss for two different model sizes (trained different numbers of tokens) and several training context lengths for FoX (Pro) and Transformer (Pro). As shown in Figure [6] the advantages of FoX over Transformer (1) increase
increase the model size (and training tokens). This indicates that the advantages of having forget gate might depend on the ratio between the model size and the training context length, as larger models can better model long contexts, and thus forgetting may be less important. We also note that long-context training damages short-context performance, which is known effect (Ren al. 2024 Sun al. 2024 likely due to reduced document diversity within training batches.
increase the training context length and (2) decrease
Component analysis We present both (1) an "incremental" style analysis where we incrementally
add/remove components from Transformer (LLaMA) to obtain the complete FoX (Pro) model and (2) a "perturbation" style analysis where we add/remove components from FoX (Pro). The results are shown Table First, mentioned previously, adding RoPE to FoX (LLaMA) and FoX (Pro) results in minor and no improvement, respectively. Second, both types of analyses show that all components in FoX contribute positively. Also note that models that use neither forget gates nor RoPE perform poorly (the first and the last row of the table).
Data-independent and fixed forget gates To show the importance of using a forget gate that is data-dependent, we test a data-independent forget gate p(h)) = o (b())), where the superscript (h)
means for the h-th head. We also test a forget gate that has fixed values (i.e., f((h) a (b())), but we do not update b(h)). As mentioned in Section 3 using a fixed forget gate equivalent to ALiBi. For these inteeinndeenndntt forget gate designs, we find it erucial initialize ((a) properly. In particular, intialize (b(A)) == for the H heads with two hyperparameter T and Tmax such

---

Published as a conference paper at ICLR 2025
that using fixed forget gates with (Tmin: Tmax) is equivalent ALiBi with minimum slope
Timxx
and a maximum slope -ا_- This initialization method is detailed in Appendix D
Tints
In Figure 7 we show the per-token loss of different forget gate designs with the LLaMA and the Pro architecture. For the data-independent and the fixed forget gate designs, we set min 2 and test different values of 7 max- As shown in Figure 7 data-dependent forget gate always has the best performance.
RELATED WORK
Recurrent sequence models
There has been a growing interest in reviving recurrent sequence
models (Katharopoulos et al.
2020 Peng et al. [02 Gu et al., 2021] Orvieto et al.) 2023; Yang
ct al.
2 Ju Dao 2023
Katsch 2023 De et al. 2024|| [Sun al. 20244] Peng et al., 2024 Qin
et al.
2024a Dao & Gu 2024 Beck et al. 2024 Zhang al. [0244] Buckman et al. 2024). Many
recent recurrent sequence models feature some form of the forget gate, which has been shown to be essential in these architectures (Qin et al. 2024b Gu & Dao 2023 Yang ct al.) 2023). Notably, GLA (Yang et al.) 2023) and Mamba-2 (Dao & Gu 2024) show that gated variants of linear attention could be written in form similar to softmax attention, which directly inspired our work. Several works (Ma et al., 2022 2024; Ren et al. (024) combine recurrent layers with quadratic attention. However, unlike our method- which embeds the forget gate into the attention mechanism in these hybrid architectures, the recurrent layers and the quadratic attention layers are independent.
Related improvements and altermatives to softmax attention Several positional embedding methods (Press et al., 2021] Raffel et al., 2020 Chi et al. 2022a (b) add bias terms to the attention logits based on the distances between the keys and queries, which can implement data-ine iata-independent decay. L.A-aattetioon (Zimerman & Wolf) applies multiplicative exponential decay to the attention logits. RoPE (Su et al., (024) also has similar decay effect that becomes stronger with increas- ing relative query/key distances. However, all these methods can only achieve data-independent decay based on the relative distances between the queries and keys. CoPE (Olsson al. 2022) and Selective Attention (Leviathan et al. 2024) modify the current timestep's attention logit based on the sum of transformed logits from some previous timesteps. Our method differs from these in various aspects. For example, in our approach, there is no need to compute sums of transformed logits, which may come with several issues such as potential incompatibility with Attenton Geometric attention (Csordás et al. 2021) and stick-breaking attention (Tan et al.) (024)) use a stick- breaking process compute the attention scores, which has similar data-dependent decay effect to our method. These methods explore in a different direction from ours, as they seek to develop alternatives to softmax attention while our approach is only an improvement on softmax attention.
6 CONCLUSION
We propose the Forgetting Transformer (FoX), a Transformer variant with a forget gate, Our ex- periments show that FoX outperforms the Transformer and several recurrent sequence models on long-context language modeling and various downstream tasks. We also show that our Pro block design greatly outperforms the basic LLaMA architecture, with or without forget gate. We there- fore recommend that future work adopt FoX (Pro) and Transformer (Pro) as baselines in addition to the commonly used L.LMAA architecture.
We discuss several limitations of our work and potential future work. First, due to our limited computing resources, our main experiments only use models up to 760M parameters, 48B tokens, and training context length of 16384 tokens. Thus, an important direction for future work is to test FoX at larger scales. Second, we only consider causal sequence modeling in this work. would be interesting extend Forgetting Attention to the non-causal case. Finally, could potentially prune computation (c.g., KV-cache eviction) adaptively based on the forget gate values, which may greatly reduce training and inference costs.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
ZL thanks Songlin Yang, Jacob Buckman, Zhen Qin, Artem Zholus, Benjamin Thérien, Jonathan Pilault, and Mahan Fathi for their helpful discussion. AC acknowledges funding from Microsoft research. This research was enabled in part by the compute resources, software and technical help provided by Mila (mila. quebec) and the Digital Research Alliance of Canada (alliance. can. ca).
REFERENCES
Sören Auer, Dante AC Barone, Cassiano Bartz, Eduardo G Cortes, Mohamad Yaser Jaradeh, Oliver Karras, Manolis Koubarakis, Dmitry Mouromtsev, Dmitri Pliukhin, Daniil Radyush, et al. The sciqa scientific question answering benchmark for scholarly knowledge. Scientific Reports, 13 (1):7240, 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding, arXiv preprint arXiv:2308. 14508, 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi, Lei Hou, Jie Tang, Yuxiao Dong, and Juanzi Li. Longalign: A recipe for long context alignment of large language models. arXiv preprint arXiv:2401.18058 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04511 2024.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Jacob Buckman. Longcraw164: A Long-Context Natural-Languaae Dataset, 2024. URL https:
//man festai.ts... com arti ceeiicleeeeeee cwwps//l/e/goowlll/ 64
Jacob Buckman, Carles Gelada, and Sean Zhang. Symmetric Power Transformers, 2024.
Ta-Chung Chi, Ting-Han Fan, Peter Ramadge, and Alexander Rudnicky. Kerple: Kernelized rel- ative positional embedding for length extrapolation. Advances in Neural Information Processing Systems, 35:8386-8399 2022a.
Ta-Chung Chi, Ting-Har Han Fan, Alexander Rudnicky, and Peter J Ramadge. Dissecting transformer length extrapolation via the lens of receptive field analysis. arXiv preprint 12X25:22/2.10356, 2022b.
Christopher Clark, Kenton Lee, Ming-Wos Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 19X5::1905.10044 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: /803.05557 2018.
Róbert Csordás, Kazuki Irie, and Jürgen Schmidhuber. The neural data router: Adaptive control flow in transformers improves systematic generalization. arXiv preprint arX1v:2110.07732, 2021.
Tri Dao. Flashattention-2:: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08661, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2 /0600 2024.
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner. A dataset of informattion-sceiing questions and answers anchored in research papers. arXiv preprint 1rXiv:2105.030/1 2021.
11

---

Published as a conference paper at ICLR 2025
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian- Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.1 19427, 2024.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pp. 7480-7512. PMLR, 2023.
Alexander R Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir R Radev. Multi-news; A large- scale multi-document summarization dataset and abstractive hierarchical model. arXiv preprint arXiv:/906.011 01299900011499 2019.
FlagOpen, 2023. http://////////aaaaaaaaaoo/o////
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL htttt::pp//eeeoto org/ records/ 12608602
Felix A Gers, Jürgen Schmidhuber, and Fred Cummins. Learning forget: Continual prediction with Istm. Neural computation, 12(10):2551-247) 2000.
Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A human- annotated dialogue dataset for abstractive summarization. arXiv preprint arXiv: WrXiv:9911.1237,, 2019.
Klaus Greff, Rupesh Srivastava, Jan Koutník, Bas Steumebrink, and Jürgen Schmidhuber. Lstm: A search space odyssey. IEEE transactions neural networks and learning systems, 28(10): 2222-2232, 2016.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00755 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00336. 2021.
Daya Guo, Canwen Xu, Nan Duan, Jian Yin, and Julian McAuley. Longcoder: A long-range pre- trained language model for code completion. In International Conference on Machine Learning, pp. 12098 12107. PMLR, 2023.
Xanh Ho, KhoaKhoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. arXiv preprint 1rXi::0011.0106/ 2020.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:2404.0663 06654, 2024.
Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. arXiv preprint 1rXiv:2104.02112, 2021.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. arXiv preprint arXiv:2402.01032, 2024.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint
2017.
Gregory Kamradt, 2023.
URL bttts::/github... com/ twt.ggaaaaaaaaaaaa
www.eeaaaatttttkcbbbbbbbb...aaaa http:////aaaaaaaaaaaaaaa md
12

---

Published as a conference paper at ICLR 2025
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mnn:: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5166 PMLR, 2020.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint urXiv:2311.01927, 2023.
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge. Transactions the Association for Computational Linguistics, 6:317-328, 2018.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Selective attention improves transformer. arXiv preprint arXiv:2410.02703, 2024.
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Conference on Computational Linguistics, 2002.
Tianyang Liu, Canwen Xu, and Julian McAuley. Repobench: Benchmarking repository-level code auto-completion systems. arXiv preprint arXiv:2306.030// 2023.
I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 2017.
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209. 100655 2022.
Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient Ilm pretraining and inference with unlimited context length. arXiv preprint wxiv:2244.08800 2024.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture 2016.
models. arXiv preprint
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct clectricity? new dataset for open book question answering. arXiv preprint arxiv:1809.02789, 2018.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. http://tttssttaireiiioooo/ooo//t htwpr/dduttttoecccddnnee.....eee
OpenAI, 2021. URL htt.a2202.....ttttss////////
OpenAI, 2022. URL http:p/://ggctb....ooooppaa/////
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando. Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 26670-26698. PMLR, 2023.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint 1:Xivv/606.0600) 2016.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemystaw Kazienko, et al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. arXiv preprint arXiv:2404 05892, 2024.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.021433
2021. 2021.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108, / /2409, 2021.
13

---

Published as a conference paper at ICLR 2025
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrn2: Gated linear mns with state expansion. arXiv preprint arXiv:2404.00900 2024a.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for se- quence modeling. Advances in Neural Information Processing Systems, 36, 2024b.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): 2020.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Sim- ple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522, 2024.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 20/ AAAI spring symposium series, 2011. 1.
Xuyang Shen, Dong Li, Ruitao Leng, Zhen Qin, Weigao Sun, and Yiran Zhong. Scaling laws for linear complexity language models. arXiv preprint arXiv:2406. 16690, 2024.
Daria Soboleva, Faisal Al-Khateeb, Robert
Myers. Jacob R Steeves, Joel Hes- tness, and Nolan Dey. SlimPajama: A 627B token cleaned and dedu-
plicated version slimpajjmme jama-a-aaa----- 627b- cleaned and dedupliicatedd-veerioo---- of reddaa 627B
RedPajama.
https: http:://www.rr http://ww.eeebbbbseetee/bbbb/a
2023. URL http::
Jianlin Su. Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois. Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, al. Learning to (learn test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.04620 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint wxXiv:2307.0862/1 2023.
Shawn Tan, Yikang Shen, Songlin Yang, Aaron Courville, and Rameswar Panda. Stick-breaking attention. arXiv preprint arXiv:2410.179986 2024.
Together Computer. Redpajama: open dataset for training large language models, 2023. URL Lama--ata
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 33111 2023.
Harsh Trivedi, Niranjan Balasubramanian. Tushar Khot, and Ashish Sabharwal. Musique Multihop
questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539-554, 2022.
Jos Van Der Westhuizen and Joan Lasenby. The unreasonable effectiveness of the forget gate. arXiv preprint arXiv:1804. 04844, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Gameet (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https::/pproeeddngs... http://sss..capaaeeeeeeeee////// /iae////////ee//p////01/// fi Sataaa//SSSSSteeeefoooooooaaaaaa
14

---

Published as a conference paper at ICLR 2025
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406. 07887, 2024.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039, 2023.
Songlin Yang and Yu Zhang. Fla: friton-base library for hardware-efficient inplentttnnnnnnnnnnnntnnnnon of linear attention mechanism, January 2024. URL
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficien training. arXiv preprint arXiv:23/2.06635, 2023.
Songlin Yang, Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. arXiv preprint arXiv:2406. 06484, 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: dataset for diverse, explainable multi-hop question answering. arXiv preprint arXiv: 808.00990000,... 2018.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019.
Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, al. Gated slot attention for efficient linear-time sequence modeling. arXiv preprint 2xX1v:2409.07146, 2024.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadal- lah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, et al. Qmsum: new benchmark for query-based multi-domain meeting summarization. arXiv preprint a0Xiv:2104005938, 2021.
Itamar Zimerman and Lior Wolf. Viewing transformers through the lens of long convolutions layers. In Forty-first International Conference on Machine Learning.
15

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
Detailed FoX (Pro) layer computation
17
B Experimental details
17
B.1 Model and training hyperparameters
17
B.2 Model parameters, estimated FLOPs, and throughput
18
B.3 Needle in the haystack details
C
Explanation the relationship between per-token-loss slope and context utilization
19
D
Data-independent forget gate initialization
19
E
Hardware-aware implementation of Forgetting Attention
20
F
Additional results
22
Per-token loss for the ablation studies
22
F2
Transformer (Pro) ablation
22
F3
Short-context training Slim
22
F4
Additional comparison with sliding window attention and Samba
25
F:5
Additional needle-in-the-hhyseee
results
25
F66
Additional results with 125M 1omp-apra22... To-to,n, 360M
360M-param/ param/ 7 7,5B-tokenn and
760M-param/ (6B-token training confiuraatooss-
27
F7
Sensitivity of length extrapolation behaviors to hyperparameters
27
F88
Training curves
32
Stability across random seeds
F.1
Additional visualization of forget gate and attention matrices
33
16

---

Published as a conference paper at ICLR 2025
A DETAILED FOX (PRO) LAYER COMPUTATION
We describe the computation of FoX (Pro) layer (depicted in Figure I rightt in detail. Each FoX layer takes an input sequence (a:)) E IR and produces an output sequence (yt) E IRd. We first describe the computation for the single head case with head dimension dhead. For each time step t, we first compute the key h2 E IR daua query q: R2 value U, E RR forget gate fa E IR and output gate 9: E Rd as follows:
q: = AMSNorrmWWer) ga:z)
(15)
K WETT, at = ouu 21) E
(16)
k; = RMSNorm aa key k + 1o
az
key (k))
(17)
value 5 ( xx) E IR 0 Wizi at
(18)
V: = at value
1 ú (1 af value oJ
(19)
ft = a (w} ti + b,) E R
(20)
gt = o (W)
(21)
This is followed by the computation of the attention output:
Σソ exp(g) k; + Dij)0,
0,
E Bhaa
(22)
Σ exp(q, k, Dsy)
= Ei 1= log fi with Da = 0 for any i. We then apply the output normalization, output gate, and the final output projection:
where Di==
y; W. (RMSNorm(o:) 3 g:) E RRd.
(23)
For the multi-head case, each head maintains an independent copy of the parameters and com- putes its output sequence (y/") (h) ) independently. All RMSNorms are also applied independently to each head y ( from different heads are then summed together to obtain the final output y:- Note that similar to the standard Transformer, even though the computation and parameters of dif- ferent heads are conceptually independent, most computations can be implemented equivalently by properly splittingcooncatenating the intermediate vectors/weight matrices of different heads.
B EXPERIMENTAL DETAILS
B.1 MODEL AND TRAINING HYPERPARAMETERS
Configuration
Tilayess
dmodel
dsead
Peak learning rate
760M params /8B tokens
1536
64 for FoX, 128 for Transformer
See Table
760M params / 16B tokens
24
1536
64 for FoX, 128 for Transformer
1 x 10 3
360M params / 7755 tokens
24
1024
64
2 10
125M params/ 2.7B tokens
12
768
64
2 x 10-3
Table Hyperparameters for different configurations. The head dimension dhead only applicable to FoX and the Transformer. We tune dhead for the 760M-parameter FoX and Transformer models in (64, 128]. Maaye counts the number of blocks. For example, for the Transformer each block contains an attention layer and MLP layer.
We list the hyperparameters for different training configurations used in this work in Table For FoX and Transformer, we follow the HuggingFace LLaMA initialization and initialize all linear layer weights and embedding parameters with (0, 0.02?). Other hyperparameters are either men- tioned in the main text (Section 4.1) follow the default values in the Flash Lienar Attention repos- itory Yang & Zhang 2024
Note that our main experiments use the 660M-paraetter488--okeee
'For the QK-norm implementation, we accidentally shared single of dhess RMSNorm scaling pa- rameters across different heads in each layer for our experiments (normally there should be one set of dhead parameters for each head). We have verified that this has no observable impact on performance. "Commit 1 c55377eeb8800 7bed5ee 34.a344553555 96bd4.
17

---

Published as a conference paper at ICLR 2025
Model
Learning rate
FoX (Pro)
2 x 10 3
Transformer (Pro)
x 10
FoX (LLaMA)
1 10-3
Transformer (LLaMA)
5 X 10
Mamba-2
2 X 10-3
HGRN2
2 x 10-3
DeltaNet
1 X 10-3
Table 5: Peak learning rates for different models for the 160M-parametee/488B-oom main experi- ments. These tuned using grid {1 10 2 10s, 5 x 10') with different values of i.
configuration. The other three configurations for ablation studies and additional analyses. For experiments, we tune the learning rate for each model. We also tune the head dimension for each Transformer and FoX model, along with learning rates. For the other three configurations used for ablation studies and additional analyses, the learning rates are tuned for Transformer (LLaMA) for the 16k training context length setting and transferred to other models and training context lengths. The head dimensions for these three settings are tuned for Transformer (LLaMA) and FoX (LLaMA) for the 16k training context length setting and transferred to Trans- former (Pro) and FoX (Pro) and other training context lengths. The optimal hyperparameters are chosen based on the average training loss over the last 512 x 220 (or 512M) tokens. Note that each token is only visited once by the models during training, so there is no fundamental difference between the training and validation loss.
the B00MM-praatter/48BBBooe
We do not share the parameters between the embedding layer and the output layer. Following the original LLaMA architecture, no bias terms are used in linear layers, except for forget gates Weight decay is not applied to RMSNorm parameters and bias terms in linear layers (again, only used for forget gates) or convolution layers. We use bfloat16 mixed-ppreciste training for all models.
B.2 MODEL PARAMETERS, ESTIMATED FLOPS, AND THROUGHPUT
We report the number of (non-embedding) parameters, estimated FLOPs, and training throughput in Table 6 For the recurrent models, we estimate FLOPs using their recurrent form for simplic- ity. Note that the FlashAttention kernels for FoX (Pro), Transformer (Pro), and FoX (LLaMA) are implemented in Triton by us on top of Flag Attention (FlagOpen) 2023) without significant optimization, while Transformer (I.aaMA) uses the official FlashAttention implementation (Dao) 2023)) in CUDA. Also, we did not focus on optimizing the efficiency of components such as QK- norm and KV-shift. We expect these four models to have similar throughput if they are properly optimized. Since use long training context length compared to model size, recurrent models have a significant advantage in theoretical FLOPs due to their linear complexity. Though an exact FLOPs-matched comparison would be interesting, it will require recalibrating the scaling law for the long-context setting and is beyond the scope of this work.
Model
Params
Forward FLOPP/token
Formula for FLOPs/token
Throughput (tokens/sec)
FoX (Pro)
759M
2.72
10
2N
27k
Transformer (Pro)
757M
2.72
10º
2N
30k
FoX (ILaMA)
757M
2.72 10"
2N
mmagee T.
30k
Transformer (LLaMA)
756M
2.72 10"
2N
38k
Mamba-2
780M
1.65 10"
2N
20 Amoctadua
44k
HGRN2
756M
1.54 10"
2N
5r Ausah
46k
DeltaNet
757M
1.54 x 10P
2N
+ foo
48k
Table 6: Number of parameters, estimated forward pass FLOPs per token, formulas for FLOPs esti- mation, and training throughput tokens per second for different models. Throughput is measured on 4 NVIDIA L40S GPUs. N is the number of parameters and is the training context length.
In preliminary small-scale experiments, we do not find hias terms in forget gates to matter for performance statistically significant way. We keep it might be useful in some cases.
18

---

Published as a conference paper at ICLR 2025
B.3 NEEDLE THE HAYSTACK DETAILS
Our reedle-in-thehaasstad test is based on LongAlign (Bai et al. 2024). which is adapted from the original needle test repositoty (Kamradt, 2023) for HuggingFace models. The prompt for the standard mode has the following structure:
lirrelevant context...
The best thing to in San Francisco is eat sandwich and sit in Dolores Park on a sunny day. [irrelevant context...]
There is important piece of infermation hidden inside the above document. New that you' ve read the document, will quiz you about it. Answer the following question: What is the best thing to da San Francisco? Answer: The best thing to do in San Francisco is
The easy mode is the same, except the needle placed within the context also includes the question:
[irrelevant context...]
What is the best thing do in San Francisco? Answer: The best thing to do San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.
[irrelevant context...]
There an important piece of infermation hidden inside the above document. Now that you' ve read the document, will quiz you about it. Answer the following question: What is the best thing to de in San Francisco? Answer: The best thing to do in San Francisco is
The results scored by GPT-40-2024-08-06 scale from 1 10.
C EXPLANATION ON THE RELATIONSHIP BETWEEN PER-TOKEN-LOSS SLOPE AND CONTEXT UTILIZATION
To understand the relationship between slope of per-token loss and context utilization of the model, we first point out that LongCraw164 applies the preprocessing of randomly "rolling" the sequences to remove any position bias. This means that when given contexts of the same length, the difficulty of predicting tokens at different positions is roughly the same in expectation For example, in expectation, predicting the 100-th token in a sequence given only the previous 90 tokens the context roughly difficult predicting the 90-th token given the full previous 90-token context. Therefore, if L(100) < L(90), indicates that the first 10 tokens in the context contribute the model's predictions for the 100-th token; and larger the difference L(90) L(100) is, the more these distant tokens contribute. On the other hand, if L(100) is roughly the same L(90) (i.e., the graph L(i) plateaus after 100). it means the first 10 tokens do not contribute the model's prediction for the 100-th token, either because they are inherently not useful for this prediction or the model are unable to utilize them.
In summary, the slope of L(i) at token position i reflects how much tokens from roughly steps earlier contribute the model's prediction the current token position.
D DATA-INDEPENDENT FORGET GATE INITIALIZATION
To understand our initialization for the fixed forget gate and the data- data-indepenne inteeenddeeeeettt forget gate, first define function T(b) -og (b) This function is defined such that a(b) T(h) l/c is always true
https: htpes:///huggi angaace.co.
Concretely, this can implemented with np- rol1 with random shift value. Note that without random rolling, given the same number of previous tokens the context, the difficulty of token prediction different positions may still remain relatively uniform.
19

---

Published as a conference paper at ICLR 2025
(i.e., T(b) is the timesteps needed to achieve decay of 1/e). We then initialize 6(4) bid such that T(b(()))
(imit) = exp(log Tiin + (log Taax - log Tmin) ==1) ), where Tuin and maxx are hyperparameters and H is the number of heads. For example, if 4 and (
(2, 128). then we have (T(M))) mit T()))))) (bbm)) ) T())))) = (2, 8, 32, 128). As mentioned in the main text, A fixed forget gate with (Tmin: Tmax) is equivalent to ALiBi with a minimum slope Timx and a maximum slope Tmin We also tested this initialization for the data-dependent forget gate but did not find it useful, so we simply initialize {b(m) h= to zero for the data-dependent forget gate. For the data inta-indeeendent and fixed forget gate, zero iiiiillization performs extremely poorly.
(2)
E HARDWARE-AWARE IMPLEMENTATION OF FORGETTING ATTENTION
Algorithm Forgetting Attention Forward Pass
Require: Matrices K, V Divide Q into T,
RNNxd vector R in HBM, block sizes Be, B,.
N B.
blocks QI...., .Q. of size B, x d each, and divide K, V to
To = B blocks K K.,..., KT. and VI,..., VT., size Be x d each.
2: Divide the output O E RNxd into T, blocks OI...., Or, of size B, x d each, and divide the logsumexp L into T, blocks Li...., Lt, of size B, each.
3: Let c = Devide e into T, blocks ci...... c,
4: Let c Devide ct into T. blocks ct,............ CT.
5:
for 1 Si < T, do
6: Load Q.. cf from HBM to on-chip SRAM.
7: On chip, initialize ( = (0) E B. e(0) = (0) E RD, Rm,, = (-00) E
B.
8:
for 1 j < To do
9:
Load K, V, oj from HBM to on-chip SRAM.
On chip, compute s(D) = Q:KT E B..B. R.
10:
On chip, compute ( = c,17 T - 1(c7)) E B.BB..
11:
On chip, compute s 5(²) + D(1)) B.B.
On chip, compute 5(0) mask(s)) i, E B.B.B.
13:
On chip, compute m. max(m: (1-1) rowmaxx (S ()))) E B. p(5) exp(s())
14:
m()) B. (pointwise). e(()) e 16-1) (m))) o + rowsum P.(0) E R.B
On chip, compute 0.13) diag(em! (1-1) mi)) - j--1 + +,0'
15:
16:
end for
On chip, compute O, diag( ( (T))
17:
18: On chip, compute L, m, (I.) + log (L) (T.))
Write O to HBM the i-th block of O.
20: Write L; to HBM as the i-th block of L.
21:
end for
22: Return the output O and the logsumexp L.
In Algorithm [ and 2 we provide the algorithms for computing the forward pass and backward pass of Forgetting Attention in hardware-aware way. The algorithm is reproduced from the FlashAttention-2 paper (Dao) 2023), with the changes needed to implement Forgetting Attention added and highlighted. In this algorithm, c cumsum(log (log f). The mask operation properly sets some entries of its first operand to -oo to satisfy the causality requirement. Note for the backward pass for ease of presentation combine
assume that pre-computed the cumulative sums
20

---

Published as a conference paper at ICLR 2025
Algorithm 2 Forgetting Attention Backward Pass
Require:
Matrices Q, K O,O, E R in HBM, vector c, dc E TR vector T. E IR- in HBM, block sizes B., Br. Divide Q into T,
N Br
blocks QI...., QT, of size Br x d each, and divide K, V to size B, d each.
T== B blocks Ki,..., KT. and V,,..,
2: Divide O into T, blocks O1...., .r. of size B, x d each, divide dO into T, blocks dO dOr, of size B, x each, and divide L into T, blocks Li...., Lr, of size B, each. 3: Initialize dQ xO) in HBM and divide it into T, blocks .Q....., .dQ,, of size B, d dVT.. each. Divide dK, dV E gxxd in Te blocks dK 1..... dKr. and dV I....,, .......Vy,, T. of size
B. x d each.
4: Let e² c² c. Devide e into T, blocks ci..... dt Devide et into blocks of CT.
Let de de x0) Devide de into T, blocks def, ................................ ddc,, Devide de into T. blocks det...., dei.
5:
6: Compute D = rowsum (d0 o O) E Id (pointwise multiply), write D to HBM and divide into T, blocks , DT, of size B, each.
7: for S j S T. do
8: Load K,, Vj. oj from HBM to on-chip SRAM.
9:
Initialize dK, (0) xd. dV, j = (0) xd on SRAM.
10:
for 1 < i T, do
11:
Load Q O, dO,, 3y dQ L4 D cf from HBM on-chip SRAM.
On chip, compute 5 = K E B..BB.
12:
On chip, compute (1) = c - 1(c)) E ERB.xB.
13:
On chip, compute ( = s(²) + +D(1) E BB.
14:
On chip, compute 5(1) mask(S ;i, E B.B.B.
15:
On chip, compute p!,) exp(S:j - L.) E B.B.B.
16:
On chip, compute dV, < dV, + (P(i)) do, Bexd
18:
On chip, compute d d do,v, E BR.. B.
On chip, compute ds p(J) o (dP()) s - D;) E R B.
19:
20:
Load dQ; from HBM to SRAM, then on chip, update dQ; + dQ, dQ.+ds K, B.xd and write back to HBM.
Load de. from HBM to SRAM, then on chip, update de. - def + +ds (1 1 E R.B. and write back to HBM.
21:
On chip, compute dK, dk; ds (1) T Q, BBxd
22:
On chip, compute de +- det ds()) 1 B.
24:
end for
25:
Write AKj,dV, de to HBM.
26:
end for
27:
Compute de d + d
28: Return dQ, dK, dV, de
the computation of dK, dV, d and the computation of dQ, de in a single algorithm, but in practice these are computed in two different kemels for implementation simplicity.
21

---

Published as a conference paper at ICLR 2025
2.2
2.1
Loss L(1)
2.0
1.9
20000
40000
60000
Token index
Transformer (LLaMA)
FoX (LLaMA) QK-norm output gate
FoX (LLaMA) RoPE
FoX (LLaMA)
QK-nomm + output gate + output
FoX (LLaMA)
FoX (LLaMA) QK-norm + output gate + output norm KV-shiff
FoX (LLaMA) QK-norm
Figure 8: Per-token loss for the incremental-style ablation studies presented in Section 4.5 All
models roughly 360M parameters and trained roughly 7.5B tokens LongCraw164. The vertical line indicates the training context length.
In practice, implement Forgetting Attention based on the Triton (OpenAl 2021) FlashAttention implementation Flag Attention (FlagOpen 2023).
F ADDITIONAL RESULTS
F.1 PER-TOKEN LOSS FOR THE ABLATION STUDIES
In Figure and Figure 9 we show the per-token loss for the ablation studies presented in Table
in Section 4.5 Transformer (LLaMA) without RoPE performs extremely poorly and we show
separately in Figure 10
F.2 TRANSFORMER (PRO) ABLATION
In Figure II parameter/2.7B Bokken setting. We start with Transformer (LLaMA) and add one component at time. Notably, we find that QK-norm seems to helpful for length extrapolation.
we present small-scale ablation study using Transformer (Pro) in the 125M-
F.3 SHORT-CONTEXT TRAINING ON SLIMPAJAMA
To complement our main results in which we perform long-context training on LongCraw164, have also run short-context training on the more commonly used SlimPajama dataset (Soboleval setting in Yang et al.
al. 2023). We follow the (2024). We also use the same hyperparameters and tokenizer as Yang et al. (2024). We train FoX and Transformer with both the LLaMA and the Pro architecture. We also Mamba-2. the strongest recurrent sequence model in our main results.
We show the per-token loss of tested models Figure and downstream task evaluation results in Table We use the same set of tasks as Yang et al. (2024) so our results can be directly compared to those of Yang et al. (2024). As shown in the results, in this short-context training setting FoX (LLaMA) does not have an advantage over the Transformer (LLaMA) except for length extrapo- lation, while FoX (Pro) still outperforms Transformer (Pro) language modeling loss and down- stream tasks. Note that these small-scale experiments without extensive hyperparameter tuning (e.g.. learning rate), so the results might not transfer to larger scales with proper hyperparameter tuning for each model.
22

---

Published as a conference paper at ICLR 2025
2.2
2.1
Loss L(t)
2.0
1.9
1.8
20000
40000
60000
Token index
FoX (Pro)
FoX (Pro)
KV-shift
FoX (Pro) - QK-nomm
FoX (Pro)
RoPE
FoX (Pro) output gate
FoX (Pro)
forget gate + RoPE
FoX (Pro) oulput norm
FoX (Pro) forget gate
Figure 9: Per-token loss for the perturbation-styyl ablation studies presented in Section 4.5 All models are roughly 360M parameters and are trained on roughly 7.5B tokens on LongCraw164. The vertical line indicates the training context length.
E/O)
40000
60000
Token
(LLaMA)
(LLaMA)
Figure 10: Removing RoPE from Transformer (LLaMA) results in poor performance. All models are roughly 360M parameters and are trained on roughly 7.5B tokens on LongCraw164. The vertical line indicates the training context length.
2.5
20000
40000
60000
Tokeu index
Transformer (LLaMA) QK-nom
Transformer (LLaMA) QK-nom ourput gate + output zorm K Vistift
Transformer (LLaMA)) QK-noma output gate
(LLAMMA)
Transformer (I.AMM))
ontput gale outputrocm
Figure 11: Incremental style ablation study for Transformer (Pro). All models roughly 125M parameters and are trained on roughly 2.7B tokens on LongCraw164. The vertical line indicates training context length.
23

---

Published as a conference paper at ICLR 2025
2,00
15
2.55
14
Perpleuity P(C)
2.50
1.3
2445
2.40
11
Token index
Validation context longth
FoX (Pm)
FoX (LLaMA)
Transformer (LLaMA)
Mamba-2
Transfommr (Pro)
Figure 12: Results SlimPajama with training context length of 2048 tokens. All models have roughly 340M non-embedding parameters and are trained on roughly 15B tokens on SlimPajama. The vertical line indicates the training context length.
Table Evaluation results on L.Meeval-harness models trained on SlimPajama with training context length of 2048 tokens. All models have roughly 340M non-embedding parameters and are trained on roughly 15B 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 Transformer++ and DeltaNet are from Yang et al. (2024). Note that Transformer++ from Yang al. (024) and Transformer (LLaMA) in our work have exactly the same architecture.
Model
Wiki. ppl.
LMB.
LMB.
PIQA
Hella. acc-n?
Wino. acc?
ARC-e acc"
ARC-c acc-n"
Avg
Ppl]
acet
acc?
Transformer++ Yang et al. 2024
28.39
42.69
31.00
63,30
34.00
50.40
44,50
24.20
41.23
DeltaNet Yang atal. 2024
28.24
37.37
32.10
64.80
34.30
52.20
80.80
23.20
42.07
FoX (Pro)
25.69
31.98
35.82
65.61
36.39
51.07
45.79
25.09
43.29
Transformer (Pro)
25.92
35.01
65.02
36.09
50.51
46.42
42.74
FoX (LLaMA)
27.86
43.26
32.56
64.80
34.59
50.12
45.12
23.38
41.76
Transformer (I.aMA)
27.98
35.25
32.31
63.71
34.89
48.07
45.33
23.72
41.34
Mamba-2
27.51
41.32
29.83
65.94
35.95
50.20
45.45
23.72
41.85
24

---

Published as a conference paper at ICLR 2025
6.50
1.80
6.25
Perplexity R())
1.75
6.00
5.75
1.70
5.50
1.65
5.25
8k
16k
24k
32k
40k
40k
48k
56k
64k
Token index
Validation context length
FoX (Pro)
FoX (LLaMA)
Samba
Mamba-2
Transformer (Pro)
Transformer (LLaMA)
Transtormer- (LLaMA)
Figure 13: Additional comparison with Samba and Transformer-SWA. (left) Per-token loss I.(i) different token position i. (right) Validation perplexity P(Z) over different validation context length 1. All models have 760M parameters and are trained on roughly 16B tokens. The vertical dashed line indicates the training context length. The per-token loss is typically noisy, so we smooth the curve using a moving average sliding window of 101 tokens. In this plot 1k 1024.
Table 8: Evaluation results on rameters and are trained on roughly 16B tokens on LongCraw164. "acc-n" means length- normalizzed accuracy. Bold and underlined numbers indicate the best and the second best results, respectively.
All models have roughly 760M non con-embedding pa-
Model
Wiki. Ppl]
LMB.
LMB. sco1
PIQA
Hella.
Wino. acet
ARC-e nce1
ARC-e
COPA
OBQJA
SeiQA
BooiQ
ppl
(Pro)
49.72
57.49
Transformer (Pro)
28.17
24.63
36.17
61.53
33.46
50.28
47.81
24.15
67.00
77.90
55.72
28.24
FoX (LLaMA)
31.03
28.41
34.89
61.21
32.27
50.51
46.68
24.06
67.00
29.60
77.30
61.07
48.46
Transformer (LLaMA)
32.33
34.41
32.41
60.94
31.68
49.96
45.62
23.6.
64.00
28.60
74.00
60.06
47.09
Sumba
TransformerS SWA (LLaMA)
33,63
33.04
33.15
60.01
31
51
46.93
23.38
62.00
76.
54.62
Mamba-2
33.26
42.38
27.29
60.83
32.03
50.67
46.21
23.55
64.00
78 40
76. 70
$7.61
46.73
F.4 ADDITIONAL COMPARISON WITH SLIDING WINDOW ATTENTION AND SAMBA
In this section, compare the standard Transformer, FoX, and Mamba-2 with sliding-wing window- iitention-basse Transformer (Transformer-SWe) We also compare with Samba (Ren al. 2024). hybrid architecture combining sliding window attention and Mamba. Both Transtormer-St these experiments, we use
and
Samba use a window size of 2048. configuration in Table Note that as mentioned in Section B all models in this configuration use the same learning rate that is tuned for Transformer (LLaMA), so the results might not be optimal for other models. We show the per-token loss, easy-mode seelle-i--thehhaee context downstream task results, and long-context task results in Figure [3] Figure 14 Table R and Table 9 respectively. Though both Transformer-SWA and Samba perform well short-context tasks, they show an early plateau in the per-token loss, which indicates that they struggle to utilize the long context. Accordingly, they perform poorly in the needle-retrieval task.
760M-parameter// tokee
experiment, short-
F.5 ADDITIONAL NEEDLE-IN
RESULTS
In Figure 15 we show the results of the needle test for HGRN2 in the 200M-paramtter48BBBtao setting.
25

---

Published as a conference paper at ICLR 2025
Table 9: Evalution results on LongBench. All models have roughly 760M non-embedding parame- ters and are trained on roughly 16B tokens on LongCraw164. Bold and underlined numbers indicate the best and the second best results, respectively.
Singl-Daecmmnt Q4
Melli-Decummm Q4
Summatiztiin
Fow-sbot Lecring
Cosk:
Medal
NaralbelS
Quaper
HHapott
Gavilqpet
MallNow
TREC
Sandiam
LEC.
(Ph))
10.45
12.98
20.52
6.77
16.2
548
27.51
10.15
927
63.5
26.07
18.02
5.34
3.4
Tramheese (Poo
8.6
11.9
22.45
9.36
16.11
516
19.8
10.66
12.1
MLT8
26.38
6.17
16.72
(LLMM))
2.45
15.53
17.13
5.25
15.2
375
2L95
10.99
883
200
19.16
10097
538
9.89
Unnateemer (LLAMA)
6/9
11:52
1294
Vannfoemer SWA (LLAMA
Transfirne
Trandformer (LLaMA),
(1..MAA) tnkens
600
Dopph Pereett
2.0
Depeh Percent
860
4.00
468
Document
Decument Lenget
Document Length
Sumba,
Transformer (SWA),
Mamba-2, 168
186
18.
Dapth Panoent
284
38.6
Depth Penoont
28:6
18.
48.
48.0
494
Depii
194
$11
611
Documeni Length
Documeni Lengh
Document Lengh
Figure 14: Easy mode needle-in-ther ineiiinttthhyyee haystack analysis for FoX, the Transformer, Mamba-2, Samba, and the Transformer with sliding window attention. These 760M-parameter models trained on 16B tokens on LongCraw164. The results are scored on scale of 1 (red) to 10 (green) by GPT-4o. The vertical dashed line indicates the training context length.
HGRN2, Easy
HGRN2, Standard
0.0
0.0
10.0
10.0
Depth Deph Pereent
30
30.0
40.0
40.0
Deyth
50.0
60.0
70.0
70.0
90.0
96.0
100.0
100,0
1880
7200
10700
1.000
25000
2000
1880
11000
720e
10:00
22700
2.500
Document Length
Document Length
Figure 15: feedle-in-the-haysee analysis for HGRN2. The results are scored on seale of (red)
to 10 (green) by GPT-4o. The vertical dashed line indicates the training context length.
26

---

Published as a conference paper at ICLR 2025
2.8
2.7
14
2.6
Perplexity P()
Loss L())
13
2.5
12
2.4
11
2.3
2.2
9
8k
16k
24k
32k
40k
48k
56k
64k
0
8k
16k
24k
32k
40k
48k
56k
64k
Token index
Validation context length
FoX (Pro)
FoX (LLaMA)
Mamba-2
DeltaNet
Transformer (Pro)
Transformer (LLaMA)
HGRN2
Figure 16: Results with 125M-parameter models trained on roughly 2.7B tokens. (left) Per-token loss L(t) at different token position i. (right) Validation perplexity P(1) over different validation context length 1. The vertical dashed line indicates the training context length. The per-token loss is typically noisy, so we smooth the curve using a moving average sliding window of 101 tokens. In this plot 1k 1024.
F.6 ADDITIONAL RESULTS WITH 255M-ARAM/2..1 -ATKAM/ 360M-PARAM/755B-TOO AND 760M-PAAAM/66B-OOEK TRAINING CONFIGURATIONS
Besides our main results with 760M parameter model trained on 48B tokens, we also report per-token loss results with 125M- parameter/2.7B-
and 760M- parameter/ 6B tBtt training configurations in this section. The hyperparameters used are given in Appendix B and Table Note that, as mentioned in Appendix the learning rates for these experi- ments are tuned for Transformer (LLaMA) for the 16k training context length setting and transferred to other models and training context lengths, so the reported results may not optimal for some models (e.g., FoX typically prefers higher learning rates than the Transformer).
360Mpparametter7755Bttkkk
Per-token loss for different models in the main experiment In Figure 16 Figure 17 and Fig- ure 18 we show the per-token loss for different models given a training context length 16k tokens for the 1225-parranete/2.77Btrkkk
360M-parameter/7.. Biteee//7.B-ttokeen and 760M-parameters 6B-token training configurations, respectively. These results are consistent with the 760M-parameter/488.
token results in Figure
Per-token loss for different training context lengths In Figure 19 Figure 20 and Figure 21 show the per-token loss for the FoX and Transformer models given different training context lengths for the 125M- 1:5M-parameter222 BB-tok,n, 3.00-parameterr77.. -.5tttkke. and 760M- 1000M---rmmter////// 66-tokenn training configurations, respectively. Consistent with the results in Figure we see that the advan- tages of FoX over the Transformer (1) reduce for larger models and (2) increase for longer training context lengths.
F.7 SENSITIVITY OF LENGTH EXTRAPOLATION BEHAVIORS TO HYPERPARAMETERS
This section presents more results the sensitivity of length extrapolation behaviors to hyperpa- rameters, addition to our results Section 4.3 and Figure 5 Figure 22 and Figure 25 show the mode and sandard-mode needle retrieval results for FoX (Pro) and Transformer (Pro) with dif- ferent numbers of training tokens and learning rates. Figure [4] shows the corresponding per-token loss curves. As shown in these results, length extrapolation is sensitive hyperparameters.
27

---

Published as a conference paper at ICLR 2025
2.2
9.0
8.5
2.1
Perplexity P()
8.0
Loss L(1)
2.0
7.5
7.0
6.5
16k
24k
32k
40k
48k
56k
64k
0
8k
16k
24k
32k
40k
48k
56k
64k
Token index
Validation context length
FoX (Pro)
FoX (LLaMA)
Mamba-2
DeltaNet
Transformer (Pro)
Transformer (LLaMA)
HGRN2
Figure 17: Results with 1o0M-parameter models trained on roughly 7.5B tokens. (left) Per-token loss L(i) at different token position i. (right) Validation perplexity P(T) over different validation context length The vertical dashed line indicates the training context length. The per-token loss is typically noisy, we smooth the curve using moving average sliding window of 101 tokens. In this plot 1024.
1.90
1.85
6.5
1.80
Perplexity P())
1.75
6.0
Loss
1.70
5.5
1.65
1.60
5.0
8k
16k
24k
32k
40k
48k
56k
64k
8k
16k
24k
32k
40k
48k
56k
64k
Token index
Validation context length
FoX (Pro)
FoX (LLaMA)
Mamba-2
DeltaNet
Transformer (Pro)
Transformer (LLaMA)
HGRN2
Figure 18: Results with loss L(t) at different token position i. (right) Validation perplexity P(I) over different validation context length 1. The vertical dashed line indicates the training context length. The per-token loss is typically noisy, so we smooth the curve using a moving average sliding window of 101 tokens. In this plot 1k 1024.
models trained roughly 16B tokens. (left) Per-token
28

---

Published as a conference paper at ICLR 2025
125M parameters/2.11 7B training tokens
125M parameters/2.7 7B training tokens
2.6
2.7
2.6
Li
2.5
2.3
512
1024
2048 4096
8192 16384 32768
512
1024
2048 4096
8192 16384 32768
Token index
Token index
Model
Training longth
Model
Training context length
(LLaMA)
FoX (Pro)
4096
Transfomer (LLaMA)
8192
32768
Transforner (Pro)
8192
32768
Figure 19: Per-token loss given different training context lengths for the 125M-parameter22.75 token setting. (left) Results for the LLaMA models. (right) Results for the Pro models. At each token index i, we report the averaged loss over a window of 101 centered at
350M parameters/7 /.5B training tokens
350M parameters/7.. /7.BB training tokens
2.3
2.2
2.0
2.0
1.9
1.9
512
1024
2048
4096
8192
16384
32768
512
1024
2048 4096
8192
16384 32768
Taken index
Token index i
Model
Training context length
Model
Training context length
FoX (LLaMA)
4096
16384
FoX (Pro)
4096
16384
Transformer (LLaMA)
32768
Transfommer (Pro)
8192
32768
Figure 20: Per-token loss given different training context lengths for the 350M 550M-parameter/77 token setting, (left) Results for the LLaMA models. (right) Results for the Pro models. At each token index i, we report the averaged loss over a window of 101 centered
760M parameters/16B training tokens
760M parameters/ 16B training tokens
2.0
1.85
1.9
L())
1.75
1.8
1.70
1.7
512
16384 32768
Taken index
Token index
Model
Training context length
Model
Training context longth
FoX (LLaMA)
4096
16384
FoX (Pro)
4096
16384
Transforerr (LLaMA)
8192
32768
Transformer (Pro)
8192
32768
Figure 21: Per-token loss given different training context lengths for the
16B token
setting. (left) Results for the LLaMA models. (right) Results for the Pro models. At each token index i, we report the averaged loss over window of 101 centered at
29

---

Published as a conference paper at ICLR 2025
FoX (Pro), 16B tokens, LR-0.001
FoX (Pro), 16B tokens, LR-0,002
10.0
Depth Percent
30.0
40.0
40.0
80.0
90.0
FoX (Pro), 48B tokens, LR-0.001
FoX (Pro), 48B tokens, LR-0,002
Depth Percent
20.0
50.0
70.0
100.0
1000
41800
7200
10000
1340
19600
22700
25000
28600
1380
1eog
4180
7200
1000
22700
21800
30900
3200
Document Length
Document Length
Transformer (Pro),
tokens, LR=0.001
Transformer (Pro).
tokens. LR=0.002
0.0
Depth Percent
20.0
50.0
70.0
80.0
100.0
Transformer (Pro), 48B tokens. LR-0.001
Transformer (Pro), 48B tokens. LR-0.002
10.0
Depth Percent
38.0
48.0
60.0
90.0
1880
Fested
22700
2900
1000
2200
13400
16600
199603
237e00
25t00
2000l
3200
Document Length
Document Length
Figure 22: FoX (Pro) and Transformer (Pro) easy mode needle-in-the the-thystach results for different numbers of training tokens and learning rates. The vertical dashed line indicates the training context length.
30

---

Published as a conference paper at ICLR 2025
FoX (Pro), 16B tokens, LR-0.001
FoX (Pro), 16B tokens, LR-0,002
10.0
Depth Percent
30.0
40.0
60.0
80.0
90.0
FoX (Pro), 48B tokens, LR-0.001
FoX (Pro), 48B tokens, LR-0,002
Depth Percent
20.0
50.0
70.0
100.0
1000
41800
7200
10000
1340
19600
22700
25000
28600
1380
1eop
4180
7200
1000
22700
21000
30900
3200
Document Length
Document Length
Transformer (Pro), L6B tokens, LR=0.001
Transformer (Pro).
tokens. LR=0.002
0.0
Depth Percent
20.0
50.0
70.0
80.0
100.0
Transformer (Pro), 48B tokens. LR-0.001
Transformer (Pro), 48B tokens. LR-0.002
10.0
Depth Percent
38.0
48.0
60.0
90.0
1880
Fested
22700
2900
1000
2200
13400
16900
19660
237e00
25t00
2000l
3200
Document Length
Document Length
Figure 23: FoX (Pro) and Transformer (Pro) standard mode needle-in insllle-ii-the--eee haystack results for dif- ferent numbers of training tokens and learning rates. The vertical dashed line indicates the training context length.
31

---

Published as a conference paper at ICLR 2025
1.9
FoX (Pro), tokens,
Transformer (Pro). tokens, 0R-0.001
FoX (Pro), 16B tokens, LR-0.002
Transformer (Pro), 16B tokens, LR-0.002
FoX (Pro), 48B tokens, LR-0,001
Trinsformer (Pro), 48B tokens, LR-0.001
FoX (Pro), 48B tokens, LR=0.002
Transformer (Pro), 48B tokens, LR-0.002
Loss M(8)
Loss I(i)
1.7
1.7
1.6
1.6
8k
16k
24k
32k
40k
48%
56k
64k
24k
32k
40k
48k
56k
Token index
Taken index
Figure 24: FoX (Pro) and Transformer (Pro) per-token loss for different numbers of training tokens and learning rates. The vertical dashed line indicates the training context length.
2.6
2.4
2.2
Loss
2.0
1.8
1.6
3
4
5
Training tokens
1e10
FoX (Pro)
Mamba-2
FoX (LLaMA)
HGRN2
Transformer (Pro)
DeltaNet
Transformer (LLaMA)
Figure 25:
Training curves of different models presented Figure These eurves show the training loss averaged every 512 x 220 tokens. All models have 760M parameters.
F.8 TRAINING CURVES
we show the training curves for all models presented in Figure 2 Note that different models use different peak learning rates, so their learning curves have different shapes.
In Figure 25
F.9 STABILITY ACROSS RANDOM SEEDS
Although it is computationally impractical for multiple seeds for all results, have run three seeds for our 360M- -arrmmerreeter FoX (LLaMA) model to show that the variance across seeds small. As shown in Figure 26, the variance across seeds is small.
32

---

Published as a conference paper at ICLR 2025
2.2
2.1
Loss L(
2.0
1.9
8k
16k
24k
32k
40k
48k
56k
64k
Token index
FoX (LLaMA, seed 0)
FoX (LLaMA, seed 2)
FoX (LLaMA, seed I)
Figure 26: Result stability across seeds with 360M-parameter FoX (LLaMA). All models are trained on roughly 7.5B tokens. The vertical dashed line indicates the training context length. The per-token loss is typically noisy, so we smooth the curve using a moving average sliding window of 101 tokens. In this plot 1k 1024.
F.10 ADDITIONAL VISUALIZATION OF FORGET GATE AND ATTENTION SCORE MATRICES
In Figure 27 and Figure we show the forget gate matrices and the attention score matrices A from 16 heads distributed in 4 layers. Note that since these matrices are large (16384 16384), if only near-diagonal entries of F are non-zero the visualization will look almost all black.
33

---

Published as a conference paper at ICLR 2025
Layer Head 1
Layer I, Head 2, P
Layar Head 3,
Layer 1, Head 4,
0k
6k
0k
Ok
Queny Postions
Query Postions
10k
Quary
b0k
10k
Quany 30k
15k
Key
Key
Layer 8, Tlead I.
Layer 8, Hllnnd 2.
Layer 8, Head P,
Layer R, Head 4,
0k
0k
0k
0k
Quee Postions
Quey Poslians
10k
20k
10k
20k
15k
15k
Key
Key
Layer 16, Ilead F.
Layer 16, Ilead F
Layer 16, Ilead 3
Layer 16, Hlead f.
0k
6k
0k
0k
Quey Puslius
Quecy Postiuns
10k
10k
10k
70k
15k
15k
Key
Key
Key Positions
Key
Layer 24, Head
Laver 24, Head 2,
Laver 24, Head 3,
Laver 24, Head 4,
0k
6k
0k
6k
Quaty Postions
Quecy Bustiona
10k
10k
10k
10k
15k
15k
Key Posiltins
Key Positions
Key Positions
Key Positions
Score
0.0
0.2
0.4
0.6
0.8
1.0
Figure 27: Visualization of the forget gate weight matrix F from heads 4 different layers. These results use FoX (Pro).
34

---

Published as a conference paper at ICLR 2025
Layer Head J.A
Layer I, Hend 2,
Layer Head 3,A
Layer J. Hend 4,A
6k
Ok
0k
Ok
Quey Posttam
Quany Poolions
5k
10k
10k
10k
10k
15k
19k
15k
15k
10kk
10k
Key Postnoms
Key Posttions
Key Postiions
Key Postions
Layer Ilead 1.4
Layer 8, Illed 2.A
Layer Illddd 3,4
Layer R, 1led 4.4
6%
0k
0k
0k
Quay Postuons
5k
5k
Quay Postunn
5k
10k
10k
10k
15k
19k
15k
15k
15k
0k
Sk
L0k
15k
15k
5k
L0k
15k
Key Positoms
Key Positions
Key Postiions
Key Postions
Layer Head 1..4
Layer 16, Head 2.A
Laver Head 3.A
Laver 16, Hoad
6k
6k
0k
Quary Potliiss
5k
5k
Quacy Posiious
5k
5k.
10k
Quary
30k
Quary
15k
15k
15k
19k
0k
10k
15k
0k
5k
10k
15k
0k
10k
15k
5k
10k
15k
Key Posiiom
Key Posilions
Key Postioon
Kay Postions
Head
Layer Head
Head
Layer
Ok
Query Poslbons
5k
5k
Query Posilbons
5k.
5k.
10k
10k
Query
15k
15k
15k
15k
0k
5k
100L 15k
0k
5k
10k
15k
06
5k
10k
15k
08
5k
10k Postioons
1.kk
Key Posiiimm
Kry Posdiims
Key Postiions
Kry
Score
Figure 28: Visualization of the attention score matrix A from 16 heads in different layers. These results use FoX (Pro). Since A is very sparse, we only show entries larger than 0.1.
35


Reviews:
[{"review_id": "nm7CsUBsAs", "reviewer": "Reviewer_Ldrv", "summary": "The paper presents an addition of a forget gate to standard causal attention. In particular, the paper connects the addition of a simple multiplicative forget gate to the recurrent formulation of attention with the parallel formulation of attention and shows how it results in a simple addition of a data-dependent additive bias to softmax attention. Importantly the bias depends on all previous inputs and not just the current one. A thorough experimental evaluation shows that the proposed forgetting attention performs as well as other methods on short context tasks with possibly an improved performance on long context tasks where it performs the best in 9 out of 14 benchmarks. There are also ablations comparing the proposed forget gate to learnable or fixed decay rates where it is shown that the proposed method performs better than either option.", "questions": "My main question is also slightly mentioned at the end of the weaknesses section, the transformer seems to be significantly underperforming all methods in practically every benchmark. Is this an artifact of the long context training set? Could it be undertrained or the hyperparameters not tuned sufficiently?\n\nSimilarly, how would a transformer without any positional encoding perform on these benchmarks and especially on the needle-in-the-haystack test?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The writing of the background leading to the formulation of the forgetting transformer is very good. It is a pleasure to read and it flows naturally.\n- The idea is simple both to implement and understand.\n- There is thorough experimental evaluation and ablation studies. I especially appreciate the comparison with data independent and fixed positional biases.", "weaknesses": "The paper doesn't have a lot of weaknesses. The main question is how convincing is the experimental evaluation which even though it is thorough in the sense that it evaluates most if not all possible questions it is\n1. on smallish models (as mentioned by the authors themselves as a limitation)\n2. lacking a baseline without RoPE or any other positional encoding\n\nThe first one is less significant given the computational requirements but the second may be important especially since the evaluation is focused heavily on really long sequences for which RoPE may not necessarily be the best choice as well as sequence length generalization for which any positional encoding will hurt performance.\n\nThe results on table 1, for the transformer's Wikipedia perplexity for instance, are quite low which calls slightly into question the validity of the conclusions. Could it be the long context training set?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pO1usgOf9A", "reviewer": "Reviewer_bWGZ", "summary": "This paper proposed to incorporate a forgetting gate, which is working as a decaying mechanism, into standard attention mechanism. The forgetting gate was directly applied to the unnormalized attention scores to down-weight long-term dependencies. \n\nThe authors pre-trained a 760M model on 16B tokens. Compared with Transformer with full attention and some recurrent architectures, such as Mamba2, HGRN2 and DeltaNet, the proposed Forgetting Transformer architecture obtained better performance on long-context retrieval task such as needle-in-haystack, and other widely used long- and short-context benchmarks.", "questions": "Why does the impact QK-norm on Forgetting Transformer vary on different tasks?\n\nWhat are the initialization methods for the parameter $w_f$ and $b_f$?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written, and the proposed architecture is well-motivated and easy to implement.", "weaknesses": "However, there are weaknesses that the authors need to address:\n\n1. The scale of training dataset is relatively small (16B tokens).\n\n2. The forgetting gate works as a decaying mechanism to down-weight historical contextual information. Hence, there are some important baselines missed in the experimental comparisons. First, one straight-forward baseline is Transformer with sliding window attention. The second is chunk-wise attention with moving average mechanism (such as Megalodon).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4msAN15uLB", "reviewer": "Reviewer_vgjA", "summary": "This paper proposes a modified causal attention mechanism for transformers which incorporates a _forget gate_, inspired by forget gates in recurrent neural networks. The forget gate can take the place of traditional positional encodings, so the overall architecture remains quite simple. The authors show that the forget gate significantly improves length generalization when compared to a standard transformer architecture with rotary positional encodings, and it is competitive with non-transformer architectures such as Mamba 2.", "questions": "The exposition was clear and I have no questions at this time.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a simple, easily implementable idea that seems to improve upon prior state of the art in autoregressive language modeling. The results on length generalization are particularly compelling when compared to RoPE.", "weaknesses": "The authors use QK-norm in combination with their proposed method for many experiments, finding that it significantly improves performance. The paper would be improved if there were some investigation of why QK-norm seems to be important for optimal results, especially since QK-norm was introduced in the context of Vision Transformers rather than language models. It would also be nice to see results with and without QK-norm on the standard transformer baseline, so we can see if it is actually more important to use QK-norm with Forgetting Transformers than it is with normal transformers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2yqSVXxAGV", "reviewer": "Reviewer_VGjT", "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.", "questions": "1.\tWhat is the impact of sharing the new parameters (w_f) across heads?\n2.\tIt 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?\n3. 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.\n4. What are the limitations or failure cases of the forgetting transformer? Negative results could be very informative and show the differences between transformer variants.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1.\tThe 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. \n2.\tThe 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.\n3.\tThe ablation studies in Figure 4 are essential, offering valuable insights into the method’s internal dynamics.\n4.\tEmpirical 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:\n\n**W.1. Empirical Evaluation is Insufficient:** \n\n**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:\n\n   (i) Include **training curves** comparing the variants presented in Figure 1. \n\n   (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. \n\n   (iii) Report results from **multiple seeds** to ensure stability.\n\n   (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.\n\n  W.1.2. Important details about the exact number of parameters, FLOPs, and latency are missing in the tables\n\n___\n\n**W.2. Missing baselines and related works:**\n 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:\n\n   (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.\n\n   (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. \n\n   (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).\n\n   (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.\n\n___\n\nI 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.\n\n[1] Mega: Moving Average Equipped Gated Attention. Ma et al. ICLR23.\n\n[2] Contextual Position Encoding: Learning to Count What's Important. Golovneva et al.\n\n[3] Selective Attention Improves Transformer. Leviathan et al.\n\n[4]  Viewing Transformers Through the Lens of Long Convolutions Layers. Zimerman et al. ICML24.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)