ArXiv: 2503.02130
🎯 Pitch
Transformers lack an explicit forgetting mechanism, yet a simple, data-dependent gate on attention scores boosts long-context language modeling and achieves near-perfect needle-in-a-haystack retrieval where Mamba-2, HGRN2, and DeltaNet all fail.
1. Executive Summary
This paper proposes the Forgetting Transformer (FoX), a Transformer variant that incorporates a forget gate into softmax attention by down-weighting unnormalized attention scores in a data-dependent way—operationalized as a cumulative product of learned scalar gates that bias the attention logits toward more recent tokens. Evaluated on long-context language modeling using the LongCrawl64 dataset with models up to 760M parameters, FoX achieves superior performance over the standard Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks while retaining near-perfect needle-in-the-haystack retrieval accuracy within the training context length of 16,384 tokens, establishing that a data-dependent forget gate improves Transformer performance and long-context retention over both standard Transformers and recurrent sequence models such as Mamba-2, HGRN2, and DeltaNet. The paper also introduces a "Pro" block design—incorporating output gates, output normalization, QK-norm, and data-dependent key-value token shifting—that significantly improves both FoX and the baseline Transformer, with the advantage of FoX over the Transformer growing as training context length increases and shrinking as model size increases, indicating that forget gates are most beneficial when the ratio of model capacity to context length is low.
2. Context and Motivation
The Core Problem: Transformers Lack Explicit Forgetting Mechanisms
The fundamental problem this paper addresses is deceptively simple: standard Transformers have no built-in mechanism for selectively forgetting or down-weighting past information based on the content of the current sequence. This is not merely an architectural curiosity—it has direct consequences for how Transformers process long sequences and how they compare to alternative architectures.
To understand why this matters, we need to look at two parallel threads in sequence modeling research that have been unfolding over the past several years.
On one side, recurrent sequence models—including LSTMs, the more recent Mamba family (Gu & Dao, 2023; Dao & Gu, 2024), gated linear attention models like GLA (Yang et al., 2023), and others (Qin et al., 2024b; Peng et al., 2024; Beck et al., 2024)—have undergone a renaissance. A defining feature of virtually all successful modern recurrent models is the forget gate (Gers et al., 2000), a learned, data-dependent mechanism that decides how much of the previous hidden state to retain versus discard at each timestep. This mechanism has been shown to be essential for their performance, including in large-scale language modeling settings (Greff et al., 2016; Van Der Westhuizen & Lasenby, 2018; Gu & Dao, 2023; Yang et al., 2023). The forget gate gives recurrent models a principled way to implement recency bias—the intuition that, all else being equal, more recent information is generally more relevant for predicting the next token—while still allowing the model to retain information from arbitrarily far back in the sequence when the content warrants it.
On the other side, Transformers (Vaswani et al., 2017) have no such mechanism. Standard softmax attention computes the relevance of every previous token to the current query purely based on the dot-product similarity between query and key vectors, without any built-in notion of temporal distance or selective decay. The only way Transformers can implement something resembling forgetting is through the content of the keys and queries themselves—if a key vector happens to be orthogonal to a query vector, its influence is naturally suppressed. But this is entirely content-driven, not temporally structured, and it provides no explicit inductive bias toward recency.
This gap might seem academic, but it has practical consequences. Consider a long document where the topic shifts gradually. A Transformer must learn entirely from data that tokens from the earlier topic become less relevant to predictions in the later topic, encoding this knowledge implicitly in its attention patterns. A model with a forget gate, by contrast, has a structural prior that naturally down-weights distant tokens unless there is a strong content-based reason to attend to them. This prior can make learning more sample-efficient, particularly when the model capacity is limited relative to the context length.
Why This Matters: The Stubborn Superiority of Transformers on Long-Context Tasks
The practical urgency of this problem comes from a tension in the current literature. Despite the architectural elegance and theoretical efficiency of modern recurrent sequence models (which scale linearly with sequence length rather than quadratically), Transformers consistently outperform them on tasks requiring genuine long-context utilization (Hsieh et al., 2024; Waleffe et al., 2024; Shen et al., 2024; Qin et al., 2024a). This has been attributed to the limited capacity of the fixed-size hidden states in recurrent models (Jelassi et al., 2024)—essentially, recurrent models compress the entire history into a state vector of finite dimension, whereas Transformers can attend directly to every token in the context, providing lossless access to distant information.
So we have an uncomfortable tradeoff: recurrent models have a structural mechanism (the forget gate) that provides useful inductive bias and enables linear-complexity processing, but they underperform on long-context tasks because their compressed state loses information. Transformers preserve access to the full context but lack the inductive bias of temporal decay and selective forgetting, potentially making them less sample-efficient and more prone to being distracted by irrelevant distant context.
The natural question—and the one this paper directly asks—is: can we have both? Specifically, can we incorporate a data-dependent forget gate into the Transformer's attention mechanism, gaining the inductive bias benefits without sacrificing the direct-access property that makes Transformers excel at long-context tasks?
The Key Insight: Forget Gates in Linear Attention Have a Parallel Form
The conceptual bridge that makes this possible comes from recent work on gated linear attention. Models like GLA (Yang et al., 2023) and Mamba-2 (Dao & Gu, 2024) showed that recurrent models with forget gates can be rewritten in a parallel form that closely resembles softmax attention. This is not just a mathematical curiosity—it reveals a structural isomorphism between the two families of models.
Consider a gated linear attention model with a scalar forget gate at each timestep (the paper develops this in Section 2.2). In its recurrent form, this model updates a state matrix and normalization vector as:
where is a feature map for the linear attention kernel. The output is .
When unrolled into parallel form, this becomes:
where is the cumulative product of forget gates between position and position . In words: the forget gate mechanism translates into a multiplicative decay factor on the unnormalized attention scores, with the decay being stronger when more forget gates (each ) are multiplied together—that is, when the key is further in the past.
The Critical Leap: This Multiplicative Decay Works for Softmax Attention Too
The paper's central insight is that this decay factor is completely independent of the choice of kernel . You can take the exact same factor—computed as a cumulative product of data-dependent scalar forget gates—and apply it to the exponential dot-product kernel that defines standard softmax attention. The result is Forgetting Attention:
This is mathematically elegant. The forget gate contributes a learned, data-dependent additive bias to the attention logits, where is computed as from the input at position . The sigmoid constrains to , meaning , so the bias is always non-positive—it can only decrease the attention logit for a given key, never increase it. And because is a sum over intermediate positions, the decay is automatically monotonically decreasing with distance: keys that are further in the past have more negative log-forget-gate terms summed into their bias, giving them lower effective attention scores on average.
Where Prior Approaches Fall Short
The paper identifies several categories of prior work that attempt to address related problems but fall short in specific ways.
Positional embedding methods that add distance-based biases. ALiBi (Press et al., 2021) adds a data-independent, linearly decreasing bias to the attention logits, where is a fixed head-specific slope. This implements a fixed exponential decay: since the bias decreases linearly with distance, the attention scores (after softmax) decay exponentially with distance. T5's relative position bias (Raffel et al., 2020) and KERPLE (Chi et al., 2022a) follow a similar philosophy of adding learned or fixed biases based on relative distance. RoPE (Su et al., 2024), while not an additive bias, also induces a distance-dependent decay through its rotation mechanism.
The fundamental limitation of all these approaches is that they are data-independent: the decay applied to a key depends only on its distance from the query, not on the content of the tokens in between. This means the model cannot learn to, say, sharply down-weight information when a section break or topic shift occurs, or maintain strong attention to an important fact mentioned much earlier if it remains relevant. The decay pattern is structurally baked in, limiting the model's flexibility.
LAS-attention (Zimerman & Wolf) applies multiplicative exponential decay to the attention logits, which is a step toward FoX's approach, but again the decay is data-independent—it does not condition the decay rate on the actual content of the sequence.
Methods that condition attention on aggregated past information. Selective Attention (Leviathan et al., 2024) and CoPE (Olsson et al., 2022) modify the current query's attention logits based on aggregating information from previous timesteps. While these introduce data-dependence, they do so through a different mechanism that requires computing sums of transformed logits from prior positions, which may complicate hardware-efficient implementations and does not directly implement a multiplicative decay structure.
Alternative attention formulations with data-dependent decay. Stick-breaking attention (Tan et al., 2024) and geometric attention (Csordás et al., 2021) use stick-breaking processes to compute attention scores, which naturally induce a data-dependent decay effect similar to FoX. However, these works frame themselves as alternatives to softmax attention—they replace the softmax operator entirely rather than augmenting it. FoX takes a more conservative approach: it keeps all the desirable properties of softmax attention (normalization, differentiability, compatibility with existing infrastructure) and simply adds a bias term. This makes it easier to integrate into existing Transformer codebases and training pipelines.
Hybrid architectures that combine recurrence with attention. Models like Mega (Ma et al., 2022), Megalodon (Ma et al., 2024), and Samba (Ren et al., 2024) combine recurrent layers (with forget gates) and quadratic attention layers in the same architecture. However, in these hybrid designs, the recurrent layers and the attention layers are independent modules—the forget gate operates in the recurrent layers, while the attention layers remain standard softmax attention. FoX is fundamentally different: it embeds the forget gate directly into the attention mechanism itself, so every attention head has its own learned forgetting behavior. This means FoX can be used as a drop-in replacement for standard attention in any Transformer architecture without adding separate recurrent modules.
The gap in empirical understanding. Beyond specific architectural limitations, the paper identifies a broader empirical gap: while forget gates are ubiquitous in recurrent models, there has been no systematic study of whether and how such a mechanism could benefit Transformers specifically. The recurrent model literature has established that forget gates are critical (Qin et al., 2024b; Gu & Dao, 2023; Yang et al., 2023), but this evidence comes from models whose fundamental computation is different from softmax attention. It is not obvious a priori that adding a forget gate to a Transformer would help—Transformers already have direct access to all tokens and can in principle learn to ignore distant information through the query-key dot products. A forget gate might be redundant, or worse, it might impose too strong a recency bias that prevents the model from attending to important distant tokens.
How This Paper Positions Itself
The paper frames its contribution not as proposing an entirely new model class, but as a targeted architectural improvement to the Transformer: take the forget gate mechanism that has proven essential in recurrent models, translate it into the attention bias form that is compatible with softmax attention, and demonstrate that it provides consistent improvements across language modeling, downstream tasks, and long-context retrieval.
This positioning has several strategic implications for how the paper is structured:
First, the paper emphasizes compatibility. Forgetting Attention is implemented as a simple modification to FlashAttention (Dao, 2023)—computing a cumulative sum of log-forget-gate values once and adding the appropriate bias during the attention logit computation in SRAM. The additional parameters (a weight vector and bias per head) and computation (a sigmoid followed by a cumulative sum per head) are negligible relative to the overall model. This is important because it means FoX is not just theoretically interesting but practically deployable at scale without requiring specialized kernels or significantly increased training cost.
Second, the paper positions FoX as an alternative to positional embeddings. A striking finding is that FoX does not require any positional embeddings—neither learned absolute positions nor RoPE. The forget gate mechanism itself provides sufficient positional information through the temporal decay structure it imposes on the attention scores. This is validated empirically: adding RoPE to FoX provides only minor or no improvement (Table 3). This positions FoX as a unified solution that handles both position encoding and adaptive context utilization simultaneously.
Third, the paper makes a direct connection to ALiBi as a special case. As noted in Section 3, if the forget gate is fixed and data-independent with , then Forgetting Attention reduces exactly to ALiBi with slope . This is a powerful conceptual unification: it shows that ALiBi is the simplest possible instance of the broader Forgetting Attention framework, and that the key missing ingredient is data-dependence. The ablation in Section 4.5 (Figure 7) directly tests this, showing that data-dependent forget gates consistently outperform both fixed and data-independent variants across architectures.
Fourth, the paper introduces a "Pro" block design as a complementary contribution. Recognizing that the Transformer architecture itself has received less optimization attention compared to recurrent models (which have incorporated advances like output gates, normalization schemes, and gating mechanisms), the paper proposes a set of architectural improvements—output gates, output normalization, QK-norm (Dehghani et al., 2023), and a simplified data-dependent token shift for keys and values (inspired by Peng et al., 2024). Importantly, this Pro block design benefits both FoX and the standard Transformer, but the combination of FoX + Pro yields the strongest results overall. This is methodologically important because it prevents conflating the benefits of the forget gate with the benefits of better architecture design.
Fifth, the paper explicitly targets the long-context regime. The training context length of 16,384 tokens and validation length of 65,536 tokens (used for length extrapolation testing) are modest by modern standards but large enough to reveal differences in how models utilize long contexts. The choice of the per-token loss at different positions as the primary metric—rather than just aggregate perplexity—reflects the paper's focus on understanding context utilization patterns rather than just reporting summary statistics. As Appendix C explains, a monotonically decreasing per-token loss with a steep slope indicates the model is effectively using tokens from across the full context for its predictions, while a plateau indicates the model's effective context window is shorter than the nominal context length.
The Practical and Theoretical Stakes
The paper addresses a problem with both practical and theoretical dimensions. Practically, as language models are increasingly deployed in settings requiring long-context understanding—document summarization, multi-hop question answering, code repository understanding—the ability to efficiently and effectively utilize long contexts becomes critical. A model that can learn to selectively forget irrelevant past information while retaining relevant context could be more robust to distractors, more sample-efficient during training, and potentially more interpretable (since the forget gate values provide an explicit signal of what the model considers worth retaining).
Theoretically, the paper contributes to a broader unification of recurrent and attention-based sequence models. By showing that the forget gate mechanism can be naturally expressed in the parallel form of both linear and softmax attention, it reveals a deeper structural connection between these model classes. This connection has been explored from the linear attention side (GLA, Mamba-2 showing that gated linear attention looks like softmax attention), but FoX approaches it from the opposite direction: starting with softmax attention and showing that it can be augmented with the key mechanism that makes recurrent models successful. Together, these lines of work suggest that the distinction between "recurrent" and "attention-based" models may be less fundamental than previously thought—both can be expressed as forms of token-to-token interaction modulated by cumulative gating operations, differing primarily in the choice of similarity kernel and the dimensionality of the state.
The paper thus sits at the intersection of two research trajectories: the effort to improve Transformers by incorporating insights from recurrent architectures, and the effort to improve recurrent architectures by incorporating the direct-access properties of Transformers. FoX represents a step toward a unification where the best properties of both families can coexist in a single attention mechanism.
3. Technical Approach
3.1 Reader Orientation
This paper develops the Forgetting Transformer (FoX), a modified Transformer architecture where every attention head learns a scalar "forget gate" at each position that systematically down-weights how much attention can be paid to tokens further in the past relative to the current position. The system solves the problem that standard Transformers lack any structural mechanism for recency bias or selective forgetting—they must learn entirely from data which past tokens are irrelevant—by introducing a cumulative multiplicative decay on attention scores that is computed on-the-fly from the sequence content itself, giving the model an inductive bias toward attending to recent information while preserving the ability to retain access to distant tokens when the content warrants it.
3.2 Big-Picture Architecture (Diagram in Words)
The FoX system consists of four major components integrated into a standard Transformer backbone:
-
Forget Gate Computation — At each position in the sequence, a small neural network (one weight vector and one bias per attention head) produces a scalar value
$f_t \in (0, 1)$by passing the current token's representation through a sigmoid nonlinearity. This scalar encodes "how much to retain from the past at this timestep." -
Cumulative Decay Matrix Construction — The forget gate values are accumulated multiplicatively across positions to form a lower-triangular matrix
$F$where$F_{ij} = \prod_{l=j+1}^i f_l$represents the total decay applied to the attention from position$i$(query) to position$j$(key). Taking the logarithm converts this into an additive bias matrix$D_{ij} = \sum_{l=j+1}^i \log f_l$. -
Forgetting Attention Module — Standard softmax attention is augmented by adding the log-decay bias
$D_{ij}$into the pre-softmax attention logits. Specifically,$O = \text{softmax}(QK^\top + D)V$, so each query-key dot product receives a position-dependent, content-driven penalty that increases with the number of intermediate forget gates (i.e., with distance). This module is implemented as a modification to FlashAttention, avoiding materialization of the$L \times L$bias matrix. -
Architectural Variants (LLaMA and Pro) — The forgetting attention mechanism can be dropped into a standard LLaMA-style Transformer (replacing RoPE positional embeddings) or into an enhanced "Pro" block that adds output gates, output normalization, QK-norm, and data-dependent token shifting for keys and values. The Pro block represents a complementary set of architectural improvements drawn from the recurrent sequence modeling literature.
Information flows as follows: a token sequence enters the FoX layer → each token produces query, key, value vectors via linear projections → a forget gate is computed from each token's input via a sigmoid-gated linear layer → the forget gates are accumulated into a cumulative log-decay vector → during attention computation, each query-key dot product is biased by the total log-decay between those positions → the biased attention scores are softmax-normalized and used to weight the values → the result passes through optional output gates and normalization before proceeding to the next layer.
3.3 Roadmap for the Deep Dive
- First, the exact mathematical form of Forgetting Attention (Equation 11 and variants), building from the gated linear attention derivation to show why the logit-bias formulation is the natural way to embed a forget gate into softmax attention.
- Second, the forget gate computation itself—how
$f_t$is produced from the input, parameterization choices, and multi-head treatment, since this is the only new learned component. - Third, the hardware-aware implementation via a modified FlashAttention algorithm, because the paper explicitly claims compatibility with efficient attention computation, and understanding this is necessary to assess practical deployability.
- Fourth, the connection to ALiBi as a special case, showing that fixed data-independent forget gates recover existing positional encoding methods precisely.
- Fifth, the architectural variants (LLaMA-style vs. Pro block), including the detailed computation of the Pro block components (output gates, QK-norm, KV-shift) and their roles.
- Sixth, the relationship to positional embeddings, explaining why FoX does not require RoPE or learned positions.
3.4 Detailed, Sentence-Based Technical Breakdown
This is primarily an architectural innovation paper whose core idea is that the forget gate mechanism from gated linear attention models can be transcribed into a logit-bias form that applies directly to standard softmax attention, yielding a Transformer variant that learns data-dependent temporal decay while retaining full compatibility with efficient attention implementations and the direct-access property that makes Transformers excel at long-context tasks.
The Mathematical Form of Forgetting Attention
The paper derives Forgetting Attention from the parallel form of gated linear attention. In a gated linear attention model with feature map $\phi: \mathbb{R}^d \to (\mathbb{R}^+)^{d'}$, the output at position $i$ can be written as:
where $F_{ij} = \prod_{l=j+1}^i f_l$ is the cumulative product of forget gate values $f_l = \sigma(w_f^\top x_l + b_f) \in (0, 1)$ for positions between $j$ and $i$ (with $F_{ii} = 1$, i.e., no decay between a position and itself), $q_i, k_j, v_j \in \mathbb{R}^d$ are projected query, key, and value vectors, and $\phi(q_i)^\top \phi(k_j)$ is the linear attention kernel between query $i$ and key $j$.
What it computes: a weighted average of value vectors where the weight for key $j$ in computing the output for query $i$ is $F_{ij}$ times the query-key similarity $\phi(q_i)^\top \phi(k_j)$, normalized by the sum of similarly weighted similarities. The factor $F_{ij}$ is the product of all forget gate values between positions $j$ and $i$; since each $f_l \leq 1$, this product decreases (or stays constant) as $j$ moves further into the past relative to $i$, implementing a cumulative decay.
Why this form: the product structure means that each intermediate token between $j$ and $i$ has an independent, multiplicative effect on how much information from $j$ survives to influence $i$. This is fundamentally different from a single distance-dependent decay (as in ALiBi) because the decay can vary based on what those intermediate tokens actually are—a section break or topic shift can produce forget gates near zero, sharply truncating the influence of everything before it, while a sequence of highly related tokens can produce forget gates near one, preserving information from far in the past.
The crucial transition from linear to softmax attention is then to simply replace the kernel $\phi(q_i)^\top \phi(k_j)$ with the exponential dot-product kernel $\exp(q_i^\top k_j)$:
Taking the logarithm of $F_{ij}$ to move the multiplicative factor into an additive term inside the exponential yields the logit-bias form:
where $D_{ij} = \log F_{ij} = \sum_{l=j+1}^i \log f_l$ with $D_{ii} = 0$ (achieved by the convention $\log 1 = 0$).
What it computes: standard softmax attention where the pre-softmax logit for key $j$ and query $i$ is the dot-product similarity $q_i^\top k_j$ plus a bias $D_{ij}$ that is the cumulative sum of log-forget-gate values between those positions. Since each $f_l \in (0, 1)$, each $\log f_l \leq 0$, so $D_{ij} \leq 0$—the bias can only reduce the effective attention score, never increase it, and the reduction is larger (more negative) when more intermediate forget gates have small values.
Why this form: the additive logit bias is exactly the computation that FlashAttention and other efficient attention implementations already support (they compute $\text{softmax}(QK^\top + \text{mask})$ where the mask handles causality). Adding $D_{ij}$ requires only that the bias matrix be available in SRAM during the attention computation—a requirement that can be satisfied by precomputing a cumulative sum vector. This is the key insight that makes Forgetting Attention practically implementable without instantiating the $L \times L$ bias matrix.
In matrix form, the full attention computation is:
where $Q, K, V \in \mathbb{R}^{L \times d}$ are matrices with query, key, and value vectors as rows, $D \in \mathbb{R}^{L \times L}$ is a strictly lower-triangular matrix (plus zeros on the diagonal) whose entry $(i, j)$ for $i > j$ is $D_{ij} = \sum_{l=j+1}^i \log f_l$ (with entries for $i \leq j$ set to $-\infty$ via the causal mask), and the softmax is applied row-wise.
A critical detail on the convention: The paper adopts $\log 0 = -\infty$. This is necessary because if any forget gate $f_l$ equals exactly 0 (which can happen with sigmoid saturation), then $F_{ij} = 0$ for all $j < l \leq i$, meaning the model completely cuts off information from before position $l$. In the logit-bias formulation, $\log 0 = -\infty$, which when exponentiated inside the softmax yields exactly zero attention weight to those positions—the correct behavior.
Why not put the forget gate elsewhere? A natural alternative would be to decay the value vectors directly ($F_{ij} v_j$) or to apply the forget gate to the attention weights post-softmax. The logit-bias approach is preferred because it integrates cleanly with the softmax normalization—the decay competes with the content-based similarities in the same exponent space, so a key with very high dot-product similarity can overcome a moderate decay bias, while a key with low similarity will be further suppressed by the decay. This mirrors how gated linear attention works: the forget gate modulates the effective attention score before normalization, not the normalized weights directly.
Forget Gate Computation
The forget gate at each position is computed from the input at that position via a simple learned mapping. For a single attention head, the computation is:
where $x_t \in \mathbb{R}^d$ is the input vector at position $t$ (the layer's input, before any projection), $w_f \in \mathbb{R}^d$ is a learned weight vector, $b_f \in \mathbb{R}$ is a learned scalar bias, and $\sigma$ is the sigmoid function $\sigma(z) = 1 / (1 + \exp(-z))$.
What it computes: a scalar in $(0, 1)$ that represents the proportion of past information retained at this timestep. A value of $f_t \approx 1$ means "keep almost everything from before this position"; a value of $f_t \approx 0$ means "nearly wipe out all information from before this position."
Why a scalar per head rather than a vector? In many recurrent models (such as LSTMs and the original GLA), the forget gate is a vector of the same dimensionality as the hidden state, allowing per-dimension forgetting. The paper chooses a scalar for two practical reasons. First, the computation and memory cost is dramatically lower—a scalar gate adds only $d + 1$ parameters per head (the weight vector $w_f$ and bias $b_f$), compared to vector gates that would require a full $d \times d$ weight matrix. Second, and more importantly for efficiency, the scalar gate makes the cumulative decay $F_{ij}$ a scalar that can be applied uniformly to all dimensions of the dot product $q_i^\top k_j$, which is essential for the FlashAttention implementation where the bias $D_{ij}$ must be a single scalar added to each logit.
Multi-head treatment: For a model with $H$ attention heads, the paper maintains separate forget gate parameters $\{w_f^{(h)}\}_{h=1}^H$ and $\{b_f^{(h)}\}_{h=1}^H$ for each head $h$. This means different heads can learn different forgetting behaviors—some heads might implement strong local attention (forget gates consistently near 0, creating short decay horizons), while others might maintain near-perfect memory (forget gates near 1, allowing attention across the full context). Figure 3 in the paper visualizes exactly this diversity: some heads show near-diagonal attention patterns (strong local focus), while others show diffuse attention across thousands of tokens (weak decay).
Initialization: For the data-dependent forget gate, the biases $\{b^{(h)}\}_{h=1}^H$ are initialized to zero. The paper notes that for data-independent forget gates, a specialized initialization based on geometric spacing of decay timescales is critical (detailed in Section 3.4 - Connection to ALiBi), but for the full data-dependent variant, zero initialization works adequately because the model can learn appropriate biases during training.
A subtle implementation note: The paper's default architecture does not use bias terms in linear layers except for forget gates (following the LLaMA convention). However, preliminary small-scale experiments indicated that adding a bias term to the forget gate computation was not statistically significant for performance; it is retained mainly as a convenience.
Hardware-Aware Implementation (Modified FlashAttention)
A central practical claim of the paper is that Forgetting Attention is compatible with FlashAttention, the standard algorithm for computing exact softmax attention in GPU SRAM without materializing the full attention matrix. This section explains exactly how the modification works.
The challenge: Standard FlashAttention computes $\text{softmax}(QK^\top)V$ in blocks loaded into SRAM, using online softmax rescaling to avoid writing the intermediate $L \times L$ attention matrix to HBM (high-bandwidth memory, i.e., GPU global memory). Forgetting Attention requires adding a bias $D_{ij}$ to each logit $q_i^\top k_j$. Naively, $D_{ij}$ is an $L \times L$ matrix, which would defeat the purpose of FlashAttention if it needed to be fully materialized.
The solution: The paper observes that $D_{ij}$ can be decomposed as $D_{ij} = c_i - c_j$ where $c_i = \sum_{l=1}^i \log f_l$ is a cumulative sum vector of length $L$. This works because:
with the convention that $\sum_{l=1}^0 \log f_l = 0$ and $D_{ii} = c_i - c_i = 0$.
What this decomposition achieves: Instead of storing an $L \times L$ bias matrix, the system only needs to store a single vector $c \in \mathbb{R}^L$ in HBM. During the attention computation, whenever a block of logits $S_{ij}^{(j)} = Q_i K_j^\top$ is computed in SRAM, the algorithm also loads the relevant segments of $c$ (specifically, $c_i$ values for the query block and $c_j$ values for the key block), computes $D_i^{(j)} = c_i 1^\top - 1 (c_j)^\top$ (a rank-1 update in SRAM), and adds it to the logit block. The rest of the FlashAttention algorithm—online softmax rescaling, block-wise output accumulation, and log-sum-exp tracking—proceeds identically to the standard version.
The forward pass (Algorithm 1 in Appendix E): The algorithm follows the FlashAttention-2 pattern (Dao, 2023) with the following modifications:
-
Precomputation: Before the attention computation begins, the cumulative sum vector
$c = \text{cumsum}(\log f)$is computed and stored in HBM. This is an$O(L)$operation that produces a single vector of length$L$. -
Block loading: When blocks
$Q_i$(query block$i$, size$B_r \times d$) and$K_j$(key block$j$, size$B_c \times d$) are loaded from HBM to SRAM, the corresponding segments$c^{\text{q}}_i$(size$B_r$) and$c^{\text{k}}_j$(size$B_c$) of the cumulative sum vector are also loaded. -
Logit computation with bias: The raw logits
$S_i^{(j)} = Q_i K_j^\top \in \mathbb{R}^{B_r \times B_c}$are computed in SRAM. Then the bias$D_i^{(j)} = c^{\text{q}}_i 1^\top - 1 (c^{\text{k}}_j)^\top \in \mathbb{R}^{B_r \times B_c}$is computed on-chip (this is a rank-1 outer product that is very fast in SRAM) and added element-wise to$S_i^{(j)}$. The causal mask is then applied. -
Online softmax: The biased logits undergo the standard FlashAttention online softmax procedure: running maximum
$m_i^{(j)}$is updated, exponentiated logits$\tilde{P}_i^{(j)} = \exp(S_i^{(j)} - m_i^{(j)})$are computed, the normalization factor$\ell_i^{(j)}$is updated with a running sum incorporating the rescaling factor$\exp(m_i^{(j-1)} - m_i^{(j)})$, and the output block$O_i^{(j)}$is accumulated with the same rescaling. The key point is that all of these operations treat the biased logits exactly as they would treat standard logits—the bias is just an extra additive term that is folded in before the softmax nonlinearity. -
Final normalization: After all key blocks are processed,
$O_i = \text{diag}(\ell_i^{(T_c)})^{-1} O_i^{(T_c)}$produces the final output for query block$i$.
The backward pass (Algorithm 2 in Appendix E): The backward pass follows the standard FlashAttention backward algorithm (recomputing attention logits and softmax weights from saved statistics) with the same $D_{ij} = c_i - c_j$ bias added at the recomputation step. Additionally, the backward pass must compute gradients with respect to the cumulative sum vector $c$ (since the forget gate parameters $w_f$ and $b_f$ need gradients). Because $D_{ij} = c_i - c_j$, the gradient with respect to $c_i$ collects contributions from all positions $j$ that query $i$ attends to (through $+dS_{ij}$), and the gradient with respect to $c_j$ collects contributions from all positions $i$ that attend to key $j$ (through $-dS_{ij}$). These gradients are accumulated in SRAM during the two-pass backward computation and written to HBM as two vectors $dc^{\text{q}}$ and $dc^{\text{k}}$, which are then summed to get the full gradient $dc = dc^{\text{q}} + dc^{\text{k}}$. From $dc$, gradients with respect to the individual $\log f_t$ are computed via reverse cumulative sum, and from there to $w_f$ and $b_f$ via the sigmoid's derivative and the input $x_t$.
Why this matters for practicality: The additional computation for Forgetting Attention consists of:
- Precomputation of
$c$: one cumulative sum per head per layer,$O(L)$with negligible constant - Loading
$c$segments into SRAM: requires 2 additional vectors of size$B_r$and$B_c$per block—trivial compared to the$Q, K, V$matrices of size$B_r \times d$or$B_c \times d$ - Computing the rank-1 bias: one outer product and one addition in SRAM per block—negligible compared to the matrix multiplication
$Q_i K_j^\top$ - Backward gradient computation for
$dc$: accumulates contributions during the standard recomputation passes, adding minor arithmetic
The paper reports that their Triton implementation achieves throughput of approximately 27k tokens/sec for FoX (Pro) compared to 30k for Transformer (Pro) and 38k for Transformer (LLaMA) on 4 NVIDIA L40S GPUs (Table 6). The authors note that this implementation is "without significant optimization" and that the Transformer (LLaMA) uses the highly optimized CUDA FlashAttention from Dao (2023), so the throughput gap may narrow substantially with engineering effort. The point is that the overhead is modest and not fundamental.
Implementation detail on the mask operation: In Algorithms 1 and 2, the mask(S_i^{(j)}, i, j) function sets entries to $-\infty$ where the query position is less than the key position (to enforce causality). This is the standard causal mask, applied after the bias addition, and is unchanged from FlashAttention.
Connection to ALiBi
The paper demonstrates that ALiBi (Press et al., 2021) is a special case of Forgetting Attention where the forget gates are data-independent and fixed (not learned). This connection provides both conceptual clarity and a concrete ablation baseline.
ALiBi recap: ALiBi adds a bias $b_{ij} = -(i - j) m_h$ to the attention logits, where $m_h$ is a head-specific slope that geometrically increases across heads. In the original paper, slopes are set such that the ratio between consecutive heads is $2^{8/H}$, giving a range of decay rates.
FoX-to-ALiBi reduction: If the forget gate at every position in head $h$ is a constant $f^{(h)} = \exp(-m_h)$ (independent of the input $x_t$), then:
which is exactly the ALiBi bias. The key properties: $f^{(h)} \in (0, 1)$ since $m_h > 0$, the forget gate is identical across all positions, and the resulting decay is a pure exponential function of distance.
Initialization for data-independent and fixed forget gates: The paper introduces a careful initialization scheme for experiments comparing data-dependent and data-independent forget gates (Section 4.5). Rather than directly setting $b^{(h)}$ values, they introduce a reparameterization using a function $T(b) = \frac{1}{-\log \sigma(b)}$. This function has the property that $\sigma(b)^{T(b)} = 1/e$ always—$T(b)$ is the number of timesteps needed for a fixed forget gate $\sigma(b)$ to achieve $1/e$ decay.
For $H$ heads, the biases are initialized as $b^{(h)}_{\text{init}} = \sigma^{-1}\left(\exp(-1 / T^{(h)})\right)$ where $T^{(h)}$ is geometrically spaced:
For example, with $(T_{\min}, T_{\max}) = (2, 128)$ and $H = 4$, we get $T^{(1)}, T^{(2)}, T^{(3)}, T^{(4)} = 2, 8, 32, 128$, meaning the heads are initialized with characteristic decay timescales ranging from 2 tokens to 128 tokens.
Why this initialization matters: The paper notes that for data-independent and fixed forget gates, "zero initialization performs extremely poorly" (Appendix D). This is because a forget gate of $\sigma(0) = 0.5$ produces extremely rapid decay—after just 2 timesteps, the attention weight for a past token is multiplied by $0.5^2 = 0.25$, effectively preventing long-range attention. By spacing the decay timescales geometrically from very short ($T_{\min} = 2$, corresponding to $f \approx 0.606$) to very long ($T_{\max}$, corresponding to $f \approx 0.992$ for $T_{\max} = 128$), the initialization ensures that some heads can attend across the full context length while others implement local focus.
The ALiBi equivalence with fixed forget gates and $(T_{\min}, T_{\max})$: A fixed forget gate with $(T_{\min}, T_{\max})$ is equivalent to ALiBi with a minimum slope $1/T_{\max}$ and a maximum slope $1/T_{\min}$. This equivalence is used in Section 4.5 (Figure 7) to directly compare data-dependent forget gates against ALiBi under identical architectural conditions.
Why data-dependence matters (conceptual): The reduction to ALiBi reveals exactly what Forgetting Attention adds. ALiBi applies the same decay at all positions—a position 100 tokens away from the query always gets the same attenuation regardless of what happened in those intervening 100 tokens. Forgetting Attention allows the decay between two positions to depend on the content of the tokens between them. If the intervening tokens are highly relevant (part of the same paragraph, same topic), the forget gates may be near 1, preserving information. If a clear boundary occurs (a section break, a topic shift), the forget gate at that boundary can be near 0, sharply truncating the influence of everything before it. This is a qualitative change in capability, not just a quantitative improvement in decay rate selection.
The Pro Block Architecture
The paper introduces a "Pro" block design (Section 3, Figure 1, and Appendix A) that incorporates several architectural components commonly found in modern recurrent sequence models. The Pro block represents a set of improvements orthogonal to the forget gate itself—they benefit both FoX and the standard Transformer, though the paper's experiments show they amplify the advantages of the forget gate.
Motivation: The paper argues that recurrent sequence models have benefited from architectural innovations beyond just the forget gate—output gates, normalization schemes, and gating mechanisms that improve training dynamics. These innovations have not been systematically applied to Transformer attention blocks. The Pro block is an attempt to do so, providing a stronger baseline for both FoX and the Transformer.
Full Pro block computation (single head): Starting from the input sequence $(x_i)_{i=1}^L$ at a given layer, a Pro block computes the output $(y_i)_{i=1}^L$ as follows (from Appendix A):
Step 1 - Query computation with QK-norm:
where $W_q \in \mathbb{R}^{d_{\text{head}} \times d}$ is the query projection matrix and RMSNorm is Root Mean Square Layer Normalization applied to the query vector. QK-norm (Dehghani et al., 2023) applies normalization to queries and keys before computing dot products, which has been shown to stabilize training and improve length extrapolation.
Step 2 - Key computation with data-dependent token shift:
where $W_k \in \mathbb{R}^{d_{\text{head}} \times d}$ is the key projection matrix, $w_k \in \mathbb{R}^d$ is an additional learned vector for the shift gate, and $\tilde{k}_{t-1}$ is the key candidate from the previous timestep (not the final key). The shift gate $\alpha^{\text{key}}_t \in (0,1)$ interpolates between the previous key candidate and the current key candidate. This is a simplified variant of data-dependent token shift from Peng et al. (2024). The intuition is similar to a forget gate but applied in key space: the model can choose to carry forward information from the previous position's key representation rather than computing an entirely new key, which may help with smooth key representations across related tokens.
Step 3 - Value computation with data-dependent token shift (no normalization):
where $W_v \in \mathbb{R}^{d_{\text{head}} \times d}$ is the value projection matrix and $w_v \in \mathbb{R}^d$ is the value shift gate vector. Note the absence of RMSNorm here—values are not normalized before attention aggregation. The shift mechanism is identical in form to the key shift but uses independent parameters, allowing different shift behaviors for keys versus values.
Step 4 - Forget gate computation:
This is identical to the LLaMA version of FoX—the forget gate computation is unchanged between architectures.
Step 5 - Output gate computation:
where $W_g \in \mathbb{R}^{d_{\text{head}} \times d}$ is a learned matrix. This is an output gate in the style of GLA and Mamba-2: it produces a per-dimension gating vector that modulates the attention output before it is combined with the MLP output. This is fundamentally different from the forget gate: the forget gate controls how much past information enters the attention computation, while the output gate controls how much of the attention output flows to the next layer.
Step 6 - Forgetting Attention:
with $D_{ij} = \sum_{l=j+1}^i \log f_l$. This is the core Forgetting Attention mechanism, unchanged from the LLaMA version.
Step 7 - Output normalization and gating:
where $W_o \in \mathbb{R}^{d \times d_{\text{head}}}$ projects from head dimension back to model dimension, $\text{RMSNorm}$ normalizes the attention output (output normalization, also used in GLA and Mamba-2), and $\odot$ is element-wise multiplication with the output gate $g_i$.
Step 8 - Multi-head aggregation: For multi-head attention, each head $h$ computes its output $y_i^{(h)}$ independently with separate parameters for all projections and gates. The final output is the sum across heads: $y_i = \sum_{h=1}^H y_i^{(h)}$. All RMSNorm operations are applied independently per head (the scaling parameters are per-head and per-RMSNorm instance).
Parameter accounting: When output gates are used, the paper reduces the number of parameters in the MLP layers so that the total parameter count remains approximately the same across architectures. This is detailed in Table 6: FoX (Pro) has 759M non-embedding parameters versus 757M for FoX (LLaMA)—a negligible difference. The output gate adds $H \times d \times d_{\text{head}}$ parameters (for $W_g$), while the KV-shift adds $2d$ parameters per head (for $w_k$ and $w_v$), so the total head-level parameter increase is offset by MLP reduction.
A note on QK-norm implementation: The paper acknowledges (Appendix A, footnote 4) that they accidentally shared a single set of $d_{\text{head}}$ RMSNorm scaling parameters across all heads within a layer for their experiments, rather than having separate parameters per head. They verified that this "has no observable impact on performance."
Why the Pro block works: The paper does not provide a theoretical justification, but the components have known benefits from prior work:
- QK-norm stabilizes training by bounding the norm of query and key vectors, preventing the attention logits from growing too large in magnitude
- Output normalization normalizes the attention output, which may help with gradient flow in deep networks
- Output gates provide a learned mechanism for the model to modulate how much each head contributes to the final representation
- KV-shift provides a form of local smoothness in key and value representations, which may be particularly helpful when the model needs to track gradual topic evolution across a long context
The ablation study in Table 3 (Section 4.5) tests each component's contribution for FoX and Transformer separately, finding that all components contribute positively for FoX (perplexity decreases as each is added), while for the Transformer the pattern is similar but slightly different—for example, adding output gates to Transformer (LLaMA) + QK-norm actually slightly worsens perplexity in the incremental ablation, suggesting that the output gate interacts differently with forgetting attention than with standard attention.
Positional Embeddings and Their Relationship to Forgetting Attention
One of the paper's striking claims is that FoX does not require any positional embeddings—neither learned absolute positions nor Rotary Position Embeddings (RoPE). This section explains why that is and what the ablation results show.
The core argument: The forget gate mechanism inherently provides positional information through the temporal decay structure. Specifically, the bias $D_{ij}$ depends on the number of tokens between $i$ and $j$ and on the forget gate values at those intermediate positions. Even in the data-dependent case, the expected value of $D_{ij}$ is monotonically decreasing with $|i - j|$ (since each $\log f_l \leq 0$), providing a strong inductive bias that keys further in the past receive lower attention scores. This bias is sufficient for the model to distinguish positions without any explicit position encoding.
Empirical evidence from Table 3:
- FoX (LLaMA) without RoPE: perplexity 7.25 at 16384-token validation
- FoX (LLaMA) with RoPE: perplexity 7.25 (no improvement)
- FoX (Pro) without RoPE: perplexity 6.62
- FoX (Pro) with RoPE: perplexity 6.63 (negligible improvement of 0.01)
For the Transformer, by contrast:
- Transformer (LLaMA) without RoPE: perplexity 29.30 (catastrophic)
- Transformer (LLaMA) with RoPE: perplexity 7.49
This is dramatic: the Transformer cannot function without positional embeddings, while FoX performs essentially identically with or without RoPE. The forget gate mechanism is providing all the positional information the model needs.
Why does the Transformer fail without positional embeddings? Without any position encoding, the softmax attention $\text{softmax}(QK^\top)$ is a permutation-equivariant operation on the set of keys—shuffling the order of tokens in the context produces exactly the same attention weights for each key (though the keys themselves are different since they come from different positions). This means the model has no way to distinguish "the word 'bank' that appeared 5 tokens ago" from "the word 'bank' that appeared 500 tokens ago" based on attention patterns alone—it must rely entirely on the content of the key vectors. For language modeling, position information is critical (word order matters), so the model fails catastrophically without it.
Why does FoX succeed without positional embeddings? The forget gate bias $D_{ij}$ breaks permutation equivariance because it depends on the order of tokens through the cumulative product structure. Even if two keys $k_5$ and $k_{500}$ have identical vectors (which they wouldn't, but hypothetically), the bias applied to $k_5$ when queried from position $i$ would be $\sum_{l=6}^i \log f_l$, while the bias applied to $k_{500}$ would be $\sum_{l=501}^i \log f_l$—a much more negative number (in expectation) because it sums over more intermediate forget gates. This means the model naturally pays more attention to tokens that are closer in position, which is a fundamental positional signal.
Why might RoPE sometimes help slightly? RoPE encodes relative position through rotation of query and key vectors: $q^\top R(i) R(j)^\top k = q^\top R(i-j) k$, where $R(\theta)$ is a rotation matrix. This provides fine-grained relative position information that the forget gate's scalar decay does not capture—specifically, RoPE can modulate how different dimensions of the query-key dot product vary with relative distance, potentially allowing more nuanced position-dependent attention patterns. The fact that RoPE helps FoX (LLaMA) slightly (7.19 → 7.08 when adding QK-norm → 6.88 when adding output gate + output norm, though some of this is from the other components) but not FoX (Pro) at all suggests that the Pro block's additional gating and normalization mechanisms already provide sufficient representational flexibility to capture whatever positional nuance RoPE would add.
The paper's default: "For simplicity, we do not use RoPE or any other positional embeddings for FoX by default" (Section 3). This is a significant simplification—it removes a commonly tuned hyperparameter (RoPE base frequency) and reduces model complexity.
Summary of Design Choices and Their Justifications
-
Scalar per-head forget gate (over vector gates): dramatically reduces parameter count (
$d+1$vs.$d^2$per head) and enables the$D_{ij} = c_i - c_j$decomposition that makes FlashAttention compatibility possible. The tradeoff is less expressive per-dimension decay control, but apparently this is not needed for language modeling performance. -
Logit-bias formulation (
$D_{ij}$added to pre-softmax logits) over post-softmax decay: keeps the forget gate's effect within the softmax normalization, so very high content-based similarities can overcome moderate decay, and the overall attention distribution remains a proper probability distribution. Post-softmax decay would produce non-normalized outputs. -
Sigmoid nonlinearity for forget gates: constrains values to
$(0, 1)$, guaranteeing that$\log f_t \leq 0$and thus$D_{ij} \leq 0$—the forget gate can only reduce attention, never amplify it. This is the correct semantics for a forgetting mechanism. -
Cumulative sum decomposition (
$D_{ij} = c_i - c_j$): the critical implementation insight that makes the method practical. It avoids materializing an$L \times L$matrix and fits cleanly into the block-based FlashAttention algorithm, requiring only$O(L)$additional HBM storage. -
Zero initialization for data-dependent forget gate biases: works well empirically for the learned case, unlike the data-independent case where careful geometric initialization is essential (because fixed gates cannot adapt during training).
-
Pro block components are independent of the forget gate: output gates, QK-norm, output normalization, and KV-shift are architectural improvements that any attention mechanism can benefit from. The paper shows they amplify FoX's advantages but also improve the Transformer baseline, which is methodologically clean—it prevents attributing general architectural gains to the forget gate specifically.
-
No positional embeddings by default: simplifies the architecture and removes a hyperparameter while achieving equal or better performance, validating that the temporal structure imposed by the forget gate provides sufficient positional signal.
4. Key Insights and Innovations
Innovation 1: Forget Gates Are Kernel-Agnostic — The Isomorphism Between Gated Linear Attention and Softmax Attention
The paper's most fundamental conceptual move is the recognition that the forget gate mechanism from gated linear attention models does not depend on the specific similarity kernel being used. The cumulative decay factor $F_{ij} = \prod_{l=j+1}^i f_l$ appears as a multiplicative weight on the kernel evaluation in the parallel form of gated linear attention (Equation 10)—and that multiplicative weight can be applied to any kernel, including the exponential dot-product kernel $\exp(q^\top k)$ that defines standard softmax attention.
This is a genuine insight rather than an obvious observation because it inverts the direction of prior work's reasoning. GLA (Yang et al., 2023) and Mamba-2 (Dao & Gu, 2024) showed that gated linear attention could be written in a parallel form that resembles softmax attention, but they used this observation to motivate efficient implementations of recurrent models—essentially saying "look, recurrent models can be computed like Transformers." FoX makes the opposite move: it says "if the forget gate works in the parallel form of linear attention, and the parallel form looks similar to softmax attention, then softmax attention can adopt the forget gate directly without giving up its kernel." The isomorphism between the two model classes is symmetric, and the paper exploits it in the previously unexplored direction.
The significance of this insight extends beyond the specific architectural proposal. It reveals that the distinction between "recurrent models with forget gates" and "Transformers with softmax attention" is less fundamental than the literature often implies. Both can be expressed as token-to-token interaction matrices modulated by cumulative gating operations; they differ primarily in the choice of similarity kernel (linear vs. exponential dot product) and the dimensionality of the state representation (finite in recurrent models, growing with sequence length in Transformers). By showing that the cumulative gating mechanism is kernel-agnostic, the paper opens the door to systematically studying which kernel properties matter for which tasks, with the forget gate mechanism held constant across comparisons. This is a reframing of the model design space, not just a new model.
The evidence for this insight's validity is primarily structural—it follows from the algebraic equivalence in Equations 10 and 11—but the empirical results in Figure 2 and Table 1 confirm that this is not just a formal curiosity: the forget gate applied to the softmax kernel yields consistent improvements over both the standard softmax kernel (Transformer) and linear-kernel models (Mamba-2, HGRN2, DeltaNet). The mechanism works regardless of the kernel, establishing the kernel-agnostic nature of forget gates as an empirical fact rather than merely a theoretical possibility.
Innovation 2: Data-Dependent Decay as a New Axis of Attention Modulation — Beyond ALiBi and Positional Embeddings
Section 3 established the mechanics of how FoX reduces to ALiBi when forget gates are fixed and data-independent. But the conceptual innovation is deeper: the paper identifies data-dependent temporal decay as a distinct and previously underexplored axis of attention modulation, orthogonal to both content-based attention (the $QK^\top$ term) and position-based attention (the ALiBi/RoPE-style bias term).
Prior work on improving softmax attention has largely focused on two axes. The first axis is better positional information: ALiBi (Press et al., 2021) adds fixed distance-dependent biases; RoPE (Su et al., 2024) encodes relative position via rotation; T5 (Raffel et al., 2020) uses learned relative position biases; KERPLE (Chi et al., 2022a) learns kernelized position biases. All of these are data-independent in the sense that the positional contribution to the attention logit for a given key-query distance is the same regardless of what tokens appear in between. The second axis is better content-based interaction: various attention variants modify the $QK^\top$ computation itself or the subsequent normalization.
FoX introduces a third axis: content-conditioned temporal decay. The forget gate $f_t$ is a function of the token $x_t$, so the decay applied to information passing through position $t$ depends on what that position actually contains. This means the model can learn to associate certain types of content (section breaks, topic shifts, dialogue turns) with sharp forgetting, and other types of content (continuation signals, anaphoric references) with information preservation. This is qualitatively different from saying "tokens 100 positions away are decayed by a factor that depends on the distance 100"—it says "tokens 100 positions away are decayed by a factor that depends on what happened in those 100 positions."
The evidence for why this matters comes from Figure 7 and Table 3. The data-dependent forget gate consistently outperforms both fixed and data-independent variants across both the LLaMA and Pro architectures. For FoX (LLaMA) with data-dependent gates, the per-token loss at 64K validation length (extrapolating 4× beyond training) continues decreasing, while for data-independent gates, the loss plateaus earlier. This is not just a performance improvement—it demonstrates that data-dependence enables qualitatively different behavior (sustained context utilization beyond the training length) that fixed-decay approaches cannot replicate regardless of hyperparameter tuning.
The connection to ALiBi is particularly revealing. ALiBi can be seen as the optimal fixed-weight solution when you have no information about content—you set a decay rate per head and apply it uniformly. FoX shows that allowing the decay rate to vary based on content is strictly more powerful, and that this additional power translates to meaningful improvements in language modeling and downstream task performance. This establishes a hierarchy: data-dependent decay > data-independent learned decay > fixed decay > no decay.
Innovation 3: Length Extrapolation as an Emergent Property of Learned Forgetting — The "Overfitting to Context Length" Finding
One of the paper's most intriguing empirical findings is not that FoX extrapolates well (it does, as shown in Figure 2), but that extrapolation behavior is sensitive to hyperparameters and training duration in a systematic way that suggests models gradually "overfit" to their training context length. This finding appears in Section 4.3 and Appendix F.7, and it represents a diagnostic contribution to understanding length generalization in Transformers.
The specific pattern, shown in Figure 5 and Figures 22-24, is that FoX (Pro) models trained on 16B tokens with a 16,384-token context length extrapolate better to 32,768-token needle evaluation than models trained on 48B tokens with the same context length. With more training tokens, the per-token loss beyond the training context length becomes flatter—the model learns to rely more heavily on tokens within the training window and stops benefiting from information that is further away than what it saw during training. The needle retrieval results mirror this: 16B-token models partially succeed at 32K document length, while 48B-token models largely fail beyond 16K.
This is a genuinely surprising result that complicates the standard narrative of "more training is always better." It suggests that length generalization is not just a property of the architecture but a dynamic that evolves during training, with models potentially learning to tune their effective context window to match the training distribution. The forget gate values themselves likely adapt during training: early in training, the model may maintain diverse forget gate behaviors (some heads with weak decay for long-range attention), but as training proceeds and the model discovers that tokens beyond 16K are never useful during training, it may shift toward uniformly stronger decay, sacrificing extrapolation ability for better in-distribution performance.
This finding has implications beyond FoX. It suggests that the common practice of evaluating length extrapolation at a single checkpoint may be misleading—the checkpoint that performs best in-distribution may not be the one that extrapolates best. It also suggests that techniques for improving length extrapolation (such as continuing training with longer sequences or periodically evaluating on longer contexts) may need to be integrated into the training process rather than treated as architectural properties to be tuned once.
The paper does not fully investigate this phenomenon—it is presented as an observation rather than a controlled study—but it opens a diagnostic window into how models learn to use context that is valuable independent of FoX's architectural merits. The fact that this pattern appears for Transformer (Pro) as well (Figure 23) reinforces that it is a general phenomenon rather than a FoX-specific quirk.
Innovation 4: The Pro Block as a Systematic Porting of Recurrent Model Innovations to Transformers
The Pro block design is more than an architectural improvement—it represents a methodological contribution: the systematic identification and integration of components from the recurrent sequence modeling literature that had not been applied to standard Transformer attention blocks. This is distinctive because it reframes the relationship between recurrent and attention-based architectures as one of mutual learning rather than competition.
The specific components the Pro block adopts are revealing. Output gates (from LSTM, GLA, Mamba-2), output normalization (from GLA, Mamba-2), QK-norm (from vision Transformers, Dehghani et al., 2023), and data-dependent token shift (from Peng et al., 2024) were all developed in contexts where Transformers were not the primary architecture being studied. The Pro block demonstrates that these innovations are architecture-agnostic—they improve softmax attention as much as they improve linear attention or state-space models. This is not obvious a priori: output gates and normalization might have interacted with the specific properties of recurrent state updates (bounded state size, linear dynamics) in ways that don't transfer to the unbounded, non-linear computation of softmax attention. The empirical results in Figure 2, Tables 1-3, and Figures 19-21 show they do transfer, and strongly.
The incremental ablation in Table 3 is particularly informative for understanding the contribution of each component. Moving from FoX (LLaMA) + QK-norm + output gate + output norm (perplexity 6.80) to the full FoX (Pro) with KV-shift (perplexity 6.62) shows that the token shift mechanism is not just a minor detail—it contributes meaningfully. The ablation of FoX (Pro) - KV-shift (perplexity 6.80) versus FoX (Pro) (6.62) confirms this. The key insight is that data-dependent token shift in key and value space provides a complementary form of temporal smoothing that is different from the forget gate's temporal decay—the forget gate controls whether past information influences the present, while KV-shift controls how the model represents keys and values in a way that enables smooth transitions between adjacent positions.
The practical significance of the Pro block is that it establishes stronger baselines for future work. The paper explicitly recommends that "future work adopt FoX (Pro) and Transformer (Pro) as baselines in addition to the commonly used LLaMA architecture" (Section 6). This is important because comparing against a weak baseline (standard LLaMA) can make modest architectural improvements appear more significant than they are. By providing the Pro block as a shared foundation, the paper enables more rigorous comparisons between attention variants.
Innovation 5: Per-Token Loss as a Diagnostic Tool for Context Utilization
While not a model-level innovation, the paper's use of per-token loss at different positions as the primary metric for evaluating long-context models is a methodological contribution that addresses a common pitfall in the literature. Appendix C provides a formal justification: because LongCrawl64 applies random rolling to remove position bias, the difficulty of predicting tokens at different positions is uniform in expectation, so differences in per-token loss directly reflect the model's ability to utilize the available context.
This is significant because the standard metric—perplexity over the full context—can be misleading. As the paper warns (Section 4.2): "even if $L(i)$ plateaus after some token position $k$, $P(l)$ may still keep decreasing after $k$, giving the wrong impression that the model can make use of the part of the context that is $k$ tokens away." This is a specific, quantifiable pitfall: perplexity $P(l) = \exp(\frac{1}{l} \sum_{i=1}^l L(i))$ is a cumulative average, so improvements in early-position loss can mask a plateau in later-position loss. A model that achieves very low loss on the first 1000 tokens but completely ignores the next 15,000 tokens might still have a decreasing perplexity curve, creating a false impression of effective long-context utilization.
The per-token loss curves in Figure 2 provide a striking illustration of why this matters. Mamba-2 and DeltaNet show per-token loss curves that flatten starting around 5K tokens and plateau after ~10K tokens—yet their perplexity curves (right panel) continue to decrease gradually. The perplexity curves alone might suggest these models are making some use of the full context, when in fact they are effectively operating with a ~10K-token context window. The per-token loss reveals the true behavior unambiguously.
This diagnostic approach is not novel in isolation (per-token loss has been plotted before), but the paper's systematic use of it—as the primary metric for all model comparisons, across all training configurations and architectures—elevates it to a recommended practice. The needle-in-the-haystack test provides complementary evidence (Figure 4 shows recurrent models fail at retrieval while FoX succeeds), but the per-token loss provides a more continuous and fine-grained signal about context utilization that does not require designing specific retrieval prompts. This is a practical contribution to evaluation methodology that the community would benefit from adopting more widely.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. All main experiments use LongCrawl64 (Buckman, 2024), a long-sequence subset of RedPajama-v2 (Together Computer, 2023) pre-tokenized with the TikToken tokenizer (OpenAI, 2022) for GPT-2 (Radford et al., 2019). The training split consists of a 45 × 2³⁰-token subset (approximately 48B tokens), each sequence containing 16,384 tokens. The validation set uses a 2 × 2³⁰-token subset of the LongCrawl64 held-out data with sequences of 65,536 tokens—4× longer than the training context length—specifically chosen to test length extrapolation. The preprocessing includes random rolling of sequences to remove positional bias, ensuring that token prediction difficulty at different positions is uniform in expectation (Appendix C). For short-context experiments, the paper also trains on SlimPajama (Soboleva et al., 2023) using a 340M-parameter/15B-token/2048-context-length configuration matching Yang et al. (2024).
-
Base model(s). The primary models have 760M non-embedding parameters, using 24 layers with a model dimension
d_model = 1536. The head dimension is tuned per-model: FoX variants used_head = 64, while Transformer variants used_head = 128. The paper also studies three smaller configurations for ablation and scaling analysis: 125M parameters (12 layers,d_model = 768), 360M parameters (24 layers,d_model = 1024), and a 760M-parameter variant trained on only 16B tokens. The choice of model scale reflects the authors' available computational resources, which they explicitly acknowledge as a limitation (Section 6: "due to our limited computing resources, our main experiments only use models up to 760M parameters, 48B tokens, and a training context length of 16384 tokens"). The base Transformer model uses the LLaMA architecture (Touvron et al., 2023) with RoPE at a large base angle of θ = 500,000, following Xiong et al. (2023)'s finding that this is crucial for long-context training. -
Metrics. The primary metric is per-token loss
L(i)at each token positioni, defined as the cross-entropy loss averaged across all validation sequences for tokens at that specific position:L(i) = (1/M) ∑ⱼ -log[(pᵢ⁽ʲ⁾)ᵀ yᵢ⁽ʲ⁾], whereMis the number of validation sequences. The paper argues (and Appendix C justifies) that the slope ofL(i)is the most meaningful diagnostic for context utilization: a monotonically decreasing curve with steep slope indicates the model is using information from tokens across the full context, while a plateau indicates the effective context window is shorter than the nominal length. PerplexityP(l) = exp((1/l) ∑ᵢ₌₁ˡ L(i))over different validation context lengthslis also reported, but the paper explicitly warns that perplexity "may still keep decreasing after [a plateau in per-token loss], giving the wrong impression that the model can make use of the part of the context that is k tokens away" (Section 4.2). For the needle-in-the-haystack test, retrieval quality is scored by GPT-4o-2024-08-06 on a 1–10 scale. For downstream tasks, the paper uses standard metrics per benchmark: perplexity for WikiText and LAMBADA, length-normalized accuracy for HellaSwag, ARC-challenge, and OpenbookQA, accuracy for all other LM-eval-harness tasks, and F1/Rouge-L/accuracy/edit similarity as appropriate for LongBench tasks. -
Baselines. The paper compares against three categories of models:
- Transformer variants: Transformer (LLaMA) — the standard LLaMA architecture with RoPE at θ = 500,000 — and Transformer (Pro) — the LLaMA architecture augmented with the Pro block components (output gates, output normalization, QK-norm, KV-shift) but no forget gate.
- Recurrent sequence models: Mamba-2 (Dao & Gu, 2024), HGRN2 (Qin et al., 2024a), and DeltaNet (Yang et al., 2024). All implementations are based on the Flash Linear Attention repository (Yang & Zhang, 2024).
- Additional comparisons (Appendix F.4): Samba (Ren et al., 2024), a hybrid architecture combining sliding window attention and Mamba, and a Transformer with sliding window attention (Transformer-SWA) using a window size of 2048. These are tested at the 760M-parameter/16B-token scale. The paper also compares against prior published results: Transformer++ and DeltaNet from Yang et al. (2024) for the SlimPajama experiments (Table 7).
-
Generation budget / compute accounting. The paper measures compute in terms of estimated FLOPs per token and training throughput (Table 6). For FoX and Transformer models, the forward FLOPs per token are approximately 2.72 × 10⁹, computed as 2N + 2 n_layers d_model L where N is the number of non-embedding parameters and L is the training context length (16,384). For recurrent models, FLOPs are substantially lower (1.54–1.65 × 10⁹ per token) due to their linear complexity in sequence length; the paper notes that "an exact FLOPs-matched comparison would be interesting, [but] will require recalibrating the scaling law for the long-context setting and is beyond the scope of this work" (Appendix B.2). Training throughput is measured on 4 NVIDIA L40S GPUs: Transformer (LLaMA) achieves 38k tokens/sec using the official CUDA FlashAttention implementation, while FoX variants achieve 27–30k tokens/sec using a Triton implementation that the authors describe as "without significant optimization." All models are trained with a single epoch over the training data (each token is visited once), so there is no fundamental distinction between training and validation loss for convergence assessment.
-
Cross-validation / statistical protocol. The paper does not use cross-validation for strategy selection (unlike the compute-optimal paper). Instead, for the main 760M-parameter/48B-token experiments, learning rates and head dimensions are tuned independently for each model using a grid search: learning rates within {1 × 10ⁱ, 2 × 10ⁱ, 5 × 10ⁱ} and head dimensions within {64, 128}. The optimal configuration is selected based on the average training loss over the last 512 × 2²⁰ (512M) tokens. For the smaller-scale ablation configurations (125M, 360M, 760M/16B), learning rates are tuned for Transformer (LLaMA) at the 16K training context length and transferred to other models and context lengths—a cost-saving approximation that the paper acknowledges may not be optimal for all models. Stability across random seeds is assessed for one configuration (360M-parameter FoX (LLaMA)) with three seeds, showing "variance across seeds is small" (Appendix F.9, Figure 26).
Main Quantitative Results
Long-Context Language Modeling (Section 4.2, Figures 2, 16–18, 19–21)
The headline result is that FoX outperforms the standard Transformer on long-context language modeling within the training context length, extrapolates better beyond it, and substantially outperforms all tested recurrent sequence models on both metrics.
Per-token loss within training context length (Figure 2, left panel). At the 760M-parameter/48B-token scale with a 16,384-token training context, FoX (Pro) achieves the lowest per-token loss across all positions, with a monotonically decreasing curve that continues to descend through the full 16K context. The per-token loss at position 16,384 is approximately 1.52 for FoX (Pro) versus 1.54 for Transformer (Pro), 1.55 for FoX (LLaMA), and 1.57 for Transformer (LLaMA). The slope of the FoX curves is steeper for the Pro architecture than the LLaMA architecture, indicating the Pro block amplifies context utilization regardless of the attention mechanism. All recurrent models show fundamentally different behavior: Mamba-2, HGRN2, and DeltaNet all exhibit per-token loss curves that flatten starting at roughly 5,000 tokens and plateau after approximately 10,000 tokens. At position 16,384, Mamba-2 achieves approximately 1.60, HGRN2 approximately 1.61, and DeltaNet approximately 1.64—all substantially higher than FoX (Pro)'s ~1.52. This plateau pattern is the key evidence that recurrent models "struggle to use the full context effectively for their prediction" (Section 4.2).
Length extrapolation beyond training context (Figure 2, left panel, 16K–64K range). Beyond the training context length (vertical dashed line at 16,384), FoX (Pro) and FoX (LLaMA) continue to show decreasing per-token loss, reaching approximately 1.48 and 1.50 respectively at position 64,000. Transformer (Pro) also extrapolates (reaching ~1.50 at 64K), while Transformer (LLaMA) degrades more noticeably (reaching ~1.55 at 64K). By contrast, the recurrent models' loss curves remain essentially flat from 16K to 64K, confirming that they gain no benefit from tokens beyond their effective context window. The extrapolation advantage is quantified more precisely in Figures 19–21 across three model scales and four training context lengths (4K, 8K, 16K, 32K), showing that FoX's extrapolation advantage over Transformer grows as the training context length increases—a finding the paper summarizes as "the advantages of FoX over Transformer (1) increase as we increase the training context length and (2) decrease as we increase the model size" (Section 4.5).
Perplexity over validation context length (Figure 2, right panel). At validation context length 65,536, FoX (Pro) achieves perplexity of approximately 4.45 versus 4.55 for Transformer (Pro), 4.65 for FoX (LLaMA), and 4.70 for Transformer (LLaMA). Mamba-2 achieves approximately 5.00, HGRN2 approximately 5.10, and DeltaNet approximately 5.30. However, the paper's warning about perplexity as a metric is borne out: even though Mamba-2 and DeltaNet show flat per-token loss beyond ~10K tokens, their perplexity curves continue to decrease gradually (from ~5.50 at 16K to ~5.00 at 64K for Mamba-2), creating a misleading impression of continued context utilization. The per-token loss curves provide the unambiguous diagnostic.
Scaling with training context length (Figure 6). For 760M-parameter models trained on 16B tokens, FoX (Pro) maintains a steeper per-token loss slope than Transformer (Pro) at all training context lengths (4K, 8K, 16K, 32K). The advantage is most pronounced at 32K training context length, where FoX (Pro) reaches approximately 1.70 at position 32K versus approximately 1.74 for Transformer (Pro). At shorter training contexts (4K, 8K), the gap narrows. This pattern is consistent across the 360M-parameter/7.5B-token configuration (Figure 20) and the 125M-parameter/2.7B-token configuration (Figure 19), though with wider confidence intervals at the smaller scale. The interaction between model size and training context length—FoX's advantage shrinking with larger models—is visible in the 760M/48B results (Figure 2) versus 360M/7.5B results (Figure 17): the gap in per-token loss at 16K is approximately 0.02–0.03 for the larger model versus 0.05–0.08 for the smaller model.
Short-context baseline on SlimPajama (Appendix F.3, Figure 12, Table 7). When trained on SlimPajama with a 2,048-token context (340M parameters, 15B tokens), FoX (LLaMA) does not show an advantage over Transformer (LLaMA) within the training context—both achieve per-token loss of approximately 2.43 at position 2,048—but FoX (LLaMA) extrapolates noticeably better beyond 2,048 tokens (reaching ~2.40 at 8K versus ~2.43 for Transformer). FoX (Pro) outperforms Transformer (Pro) both within and beyond the training context length (per-token loss of ~2.42 at 2K versus ~2.44). On downstream tasks under this configuration (Table 7), FoX (Pro) achieves the highest average across WikiText, LAMBADA, PIQA, HellaSwag, WinoGrande, ARC-easy, and ARC-challenge (43.29 average versus 42.74 for Transformer (Pro) and 41.23 for Transformer++ from Yang et al., 2024). The paper notes that "these are small-scale experiments without extensive hyperparameter tuning (e.g., learning rate)" and that "results might not transfer to larger scales" (Appendix F.3).
Needle-in-the-Haystack (Section 4.3, Figures 4, 5, 14, 15, 22–24)
The headline result: FoX achieves near-perfect needle retrieval accuracy within the training context length in both easy and standard modes, while all tested recurrent sequence models (Mamba-2, DeltaNet, HGRN2) fail even within the training context length.
Easy mode results (Figure 4, left panel). With a training context length of 16,384 tokens, FoX (Pro) scores 10/10 (green) on GPT-4o evaluation for all document lengths up to 16,384 and all depth percentages (where the needle is placed within the document). FoX (LLaMA) shows the same near-perfect pattern within 16K, with only minor degradation beyond the training context (some 8–9/10 scores at 19,600–25,800 document lengths). Transformer (Pro) also achieves perfect scores within 16K but shows more degradation beyond (scores of 7–9 at longer documents). Transformer (LLaMA) is slightly worse, with scores of 8–10 at the longest documents. Mamba-2 shows a clear failure pattern: within the training context, it achieves 10/10 only for short documents (under ~7,200 tokens) and when the needle is at depth 0–30%, with scores dropping to 1–3 for longer documents or deeper needle placement. DeltaNet performs even worse, with most of the heatmap showing scores of 1–2 for documents longer than ~4,000 tokens. HGRN2 (Figure 15, Appendix F.5) performs the worst, with near-universal failure (scores of 1–2) across all document lengths and depths.
Standard mode results (Figure 4, right panel). The standard mode—where the needle contains only the answer, not the question—is substantially harder. FoX (Pro) maintains near-perfect scores (9–10) within 16K for all depths and document lengths. FoX (LLaMA) shows slightly more degradation (scores of 8–10 at longer documents). Transformer (Pro) shows noticeable failures: at document length 16,384, several depth percentages score 6–8, and beyond 16K, scores drop to 4–6 for many configurations. Transformer (LLaMA) is worse still, with scores of 2–4 at the longest documents. Mamba-2 and DeltaNet both fail comprehensively, with scores of 1–2 across almost all configurations.
Length extrapolation sensitivity (Figure 5, Figures 22–24). The paper reports a striking finding about length extrapolation: for FoX (Pro), models trained on 16B tokens extrapolate better to 32K documents than models trained on 48B tokens. With 16B tokens and learning rate 0.001, the easy-mode needle heatmap shows scores of 9–10 at 32K document length for most depths; with 16B tokens and learning rate 0.002, scores drop to 7–9 at 32K. With 48B tokens and learning rate 0.001, scores at 32K are mostly 7–8; with 48B tokens and learning rate 0.002, scores drop to 4–8. The corresponding per-token loss curves (Figure 24) show that the 48B-token models have flatter slopes beyond the training context length than the 16B-token models—indicating they have learned to rely less on tokens beyond the training window. The paper interprets this as evidence that "the models may be gradually 'overfitting' to their training context length during training" (Section 4.3). Similar sensitivity to hyperparameters is observed for Transformer (Pro) (Figures 22–23), suggesting this is a general phenomenon rather than FoX-specific.
Sliding window and hybrid model results (Appendix F.4, Figure 14). Transformer-SWA (sliding window attention with window size 2048) and Samba (a hybrid of sliding window attention and Mamba) both perform poorly in the needle test at the 760M/16B scale. Both show heatmaps with scores of 1–3 across most configurations beyond ~7,200 document length, consistent with their per-token loss curves in Figure 13 that plateau early. This confirms that sliding window attention, despite being computationally efficient, fundamentally cannot access information beyond its window size.
Downstream Tasks (Section 4.4, Tables 1, 2, 8, 9)
Short-context tasks (Table 1). On the LM-evaluation-harness suite, evaluated zero-shot at the 760M/48B scale, FoX (Pro) achieves the highest average accuracy of 50.88. Transformer (Pro) is second at 50.39, FoX (LLaMA) third at 49.82, and Transformer (LLaMA) fourth at 49.09. Mamba-2 (50.21) ranks between Transformer (Pro) and FoX (LLaMA)—notable as the only recurrent model competitive on short-context tasks. HGRN2 (49.45) and DeltaNet (47.99) trail further behind. The pattern is consistent across individual tasks: FoX (Pro) achieves the best or second-best score on 9 of the 12 tasks. The largest margins are on LAMBADA accuracy (FoX (Pro) at 42.75 versus DeltaNet at 34.27), COPA (FoX (Pro) at 71.00 versus Transformer (Pro) at 62.00), and SciQA (FoX (Pro) at 85.10 versus HGRN2 at 75.60). The Pro architecture consistently improves performance for both FoX and Transformer across all tasks—the average gain from LLaMA to Pro is approximately 1–1.5 percentage points.
Long-context tasks (Table 2). On LongBench, evaluated at the same scale, the picture is more nuanced. FoX (Pro) achieves the highest average score across the 14 tasks, though the paper characterizes performance as "on par" with the Transformer (the exact averages are not reported in a single number, but FoX (Pro) wins on 6 tasks, Transformer (Pro) on 5, and they tie or are close on the remaining 3). FoX (Pro) shows particular strength on summarization tasks: GovReport (22.71 vs. 19.88 for Transformer (Pro)), QMSum (13.51 vs. 10.70), and MultiNews (12.27 vs. 8.11). On code tasks, Transformer (Pro) has an advantage: LCC (10.79 for Transformer (Pro) vs. 10.90 for FoX (Pro)) and RepoBench-P (14.25 vs. 9.10). The recurrent models perform substantially worse on most long-context tasks, particularly on multi-document QA and summarization: Mamba-2 achieves 9.31 on GovReport versus 22.71 for FoX (Pro), and DeltaNet achieves 8.19 versus 22.71. This is consistent with the per-token loss and needle results showing recurrent models cannot effectively utilize the full context.
Additional comparison with Samba and Transformer-SWA (Tables 8 and 9). At the 760M/16B scale, Samba performs competitively on short-context tasks (average 48.12, between FoX (Pro) at 48.33 and Transformer (Pro) at 48.24) but poorly on long-context tasks, particularly on summarization (GovReport: 9.42 for Samba vs. 27.51 for FoX (Pro)). Transformer-SWA (LLaMA) is the weakest on both short-context (46.72 average) and long-context tasks (GovReport: 7.47), confirming that sliding window attention's limited context window directly harms performance on tasks requiring long-range understanding.
Ablation Studies and Robustness Checks
Model size and training context length interaction (Section 4.5, Figure 6, Figures 19–21): The advantage of FoX over Transformer systematically shrinks as model size increases and grows as training context length increases. At the 125M-parameter/2.7B-token scale with 16K training context (Figure 19), FoX (Pro) achieves per-token loss approximately 0.05–0.08 lower than Transformer (Pro) at positions beyond 8K. At the 360M/7.5B scale (Figure 20), the gap narrows to approximately 0.03–0.05. At the 760M/16B scale (Figure 21), the gap is approximately 0.02–0.03. For training context length variation, at 4K context length (Figure 6, top-left), FoX (Pro) and Transformer (Pro) are nearly indistinguishable; at 32K context length, the gap is approximately 0.04 at the 760M/16B scale and approximately 0.08 at the 360M/7.5B scale. The paper interprets this as evidence that "the advantages of having a 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" (Section 4.5). A negative side-effect noted: long-context training damages short-context performance ("likely due to reduced document diversity within training batches"), visible in Table 8 where the 760M/16B models achieve lower average scores than the 760M/48B models on LM-eval-harness (e.g., FoX (Pro) 48.33 vs. 50.88).
Pro block component analysis (Table 3, Figures 8, 9): Using 360M-parameter models trained on 7.5B tokens, the paper presents both incremental and perturbation-style ablations.
Incremental starting from Transformer (LLaMA) (Table 3, top half, Figure 8): Transformer (LLaMA) without RoPE achieves perplexity 29.30 at 16,384-token validation—catastrophically bad, confirming positional information is essential. Adding RoPE drops perplexity to 7.49. Adding forget gates (+ FoX (LLaMA)) improves to 7.25. Adding QK-norm further improves to 7.19. Adding output gates and output normalization drops to 7.08 and 6.88 respectively. Adding KV-shift to reach the full FoX (Pro) achieves 6.80. The final addition of RoPE to FoX (Pro) yields 6.82—essentially no improvement over 6.80 without RoPE, confirming the forget gate makes positional embeddings redundant.
Perturbation from FoX (Pro) (Table 3, bottom half, Figure 9): Starting from FoX (Pro) at perplexity 6.62, removing QK-norm degrades to 6.79, removing output gates degrades to 6.86, removing output norm degrades to 6.69, and removing KV-shift degrades to 6.80. Adding RoPE yields 6.63 (negligible gain). Removing the forget gate while keeping RoPE (i.e., Transformer (Pro)) degrades to 6.82. Removing both forget gate and RoPE collapses to 7.40. All components contribute positively for FoX, with the forget gate (+ RoPE removal) being the largest single contributor (7.40 → 6.82, a 0.58 perplexity improvement) and the remaining Pro components contributing smaller but consistent gains (0.20 total from 6.82 to 6.62).
The per-token loss curves (Figures 8, 9) reveal where these improvements manifest: the Pro components primarily improve loss at longer token positions (beyond ~8K), while the forget gate improves loss across all positions.
Data-dependent vs. data-independent vs. fixed forget gates (Figure 7): Using 360M-parameter models on 7.5B tokens, the data-dependent forget gate consistently achieves the lowest per-token loss across all positions for both LLaMA and Pro architectures. For FoX (LLaMA) with data-dependent gates, per-token loss reaches approximately 1.86 at 64K validation length; with fixed gates at T_max = 256, loss is approximately 1.95; at T_max = 2048, approximately 1.92; at T_max = 16,384, approximately 1.88. The data-independent (learned but input-independent) variant performs identically to the fixed variant for each T_max setting. For FoX (Pro), data-dependent gates achieve approximately 1.82 at 64K; fixed/data-independent gates at T_max = 16,384 achieve approximately 1.83—the gap is smaller than for the LLaMA architecture, suggesting the Pro block's additional gating mechanisms partially compensate for the lack of data-dependence. The critical pattern is in the slope beyond the training context length (vertical dashed line at 16K): data-dependent gates maintain a steeper decreasing slope, while fixed/data-independent gates show flatter curves, indicating poorer length extrapolation. The paper also notes that "for the data-independent and the fixed forget gate designs, we set T_min = 2 and test different values of T_max" (Section 4.5), and that "zero initialization performs extremely poorly" for these variants, necessitating the geometric initialization scheme described in Appendix D.
Forget gate + RoPE interaction (Table 3, Figure 8): For FoX (LLaMA), adding RoPE provides a minor improvement (7.25 → 7.19 with QK-norm, and further improvements with output gate + output norm to 6.88). For FoX (Pro), adding RoPE yields 6.63 versus 6.62 without RoPE—within measurement noise. This ablation establishes that RoPE provides diminishing returns as more Pro components are added, and is completely unnecessary for the full Pro architecture. The per-token loss curves (Figure 8) show that RoPE's benefit for FoX (LLaMA) is concentrated at shorter token positions (<8K tokens), with negligible effect on long-range extrapolation.
Transformer (Pro) ablation (Figure 11): At the 125M-parameter/2.7B-token scale, incrementally adding Pro components to Transformer (LLaMA) reveals that QK-norm is particularly helpful for length extrapolation. Transformer (LLaMA) achieves per-token loss of approximately 2.38 at 8K (beyond the 16K training context); adding QK-norm reduces this to ~2.34; adding output gate + output norm reduces to ~2.33; adding KV-shift reduces to ~2.32. The cumulative improvement from Pro components is approximately 0.03–0.06 in per-token loss, smaller than the benefit of the forget gate in FoX but still meaningful.
Sliding window attention vs. full attention (Figures 13, 14, Tables 8, 9): At the 760M/16B scale, Transformer-SWA (window size 2048) achieves per-token loss that plateaus at approximately 1.72 after ~2,048 tokens, confirming it cannot utilize tokens beyond its window. Samba (hybrid sliding window + Mamba) plateaus at approximately 1.70 after ~4,096 tokens—better than pure sliding window but far worse than full attention. On needle retrieval (Figure 14), both models fail beyond their effective window, with Samba showing slightly better performance (scores of 4–6 at 16K for some depths) than Transformer-SWA (scores of 1–2). On downstream tasks (Table 8), Samba achieves competitive short-context performance (48.12 average) but poor long-context results (Table 9: GovReport 9.42 vs. FoX (Pro) 27.51). This ablation establishes that the full quadratic attention with forget gates is genuinely more effective than architectural hybrids for long-context tasks at this scale.
Short-context training on SlimPajama (Appendix F.3, Figure 12, Table 7): In the 340M/15B/2K-context setting, FoX (LLaMA) does not outperform Transformer (LLaMA) within the training context length (per-token loss ~2.43 at position 2,048 for both), but extrapolates better beyond it (reaching ~2.40 versus ~2.43 at 8K). FoX (Pro) outperforms Transformer (Pro) both within and beyond the training context (per-token loss ~2.42 versus ~2.44 at 2K, and ~2.39 versus ~2.42 at 8K). The downstream results (Table 7) show FoX (Pro) with the highest average (43.29) among tested models, though the paper acknowledges limited hyperparameter tuning. This ablation demonstrates that the forget gate's benefits are most pronounced in long-context regimes; at very short contexts (2K tokens), the advantage is marginal or absent.
Training curves and convergence (Figure 25): All models show stable training loss curves that decrease smoothly over the full 48B tokens. FoX variants and Transformer variants converge to similar final loss values, though FoX (Pro) achieves the lowest training loss (approximately 1.50 averaged over the last 512M tokens). The recurrent models converge more slowly and to higher final loss values: Mamba-2 at approximately 1.58, HGRN2 at approximately 1.62, DeltaNet at approximately 1.67. Different learning rates (tuned per model) result in different curve shapes, with FoX models using higher learning rates (2 × 10⁻³ for FoX (Pro), 1 × 10⁻³ for FoX (LLaMA)) than Transformers (1 × 10⁻³ for Transformer (Pro), 5 × 10⁻⁴ for Transformer (LLaMA)).
Stability across random seeds (Figure 26): For the 360M-parameter FoX (LLaMA) configuration trained on 7.5B tokens with three different random seeds, the per-token loss curves are nearly identical (within approximately 0.01 of each other at all positions). This supports the reliability of the reported results, though the paper notes it is "computationally impractical for us to run multiple seeds for all our results" (Appendix F.9).
Visualization of forget gate behavior (Figures 3, 27, 28): The forget gate weight matrix F and attention score matrix A from FoX (Pro) heads show diverse learned behaviors. Some heads (e.g., Layer 8, Head 2 in Figure 3) exhibit strong decay with most F entries near zero off-diagonal and attention focused on very local entries (scores > 0.1 only within ~100 positions of the diagonal). Other heads (e.g., Layer 16, Head 4) show much weaker decay with F entries remaining near 1 for thousands of positions and attention distributed across the entire 16K context. Appendix F.10 provides visualizations for 16 heads across 4 layers, revealing a consistent pattern: early layers (Layer 1) tend to have more local attention, middle layers (Layers 8, 16) show a mix of local and global attention, and later layers (Layer 24) show predominantly global attention with very weak decay. This hierarchical organization emerges purely from training—the model learns to allocate different forgetting behaviors to different heads at different depths, with lower layers focusing on local context and higher layers integrating information across the full sequence.
Critical Assessment
The experimental evidence provides strong support for the paper's core claim that adding a data-dependent forget gate to softmax attention improves language modeling performance, length extrapolation, and long-context retrieval over standard Transformers. The evidence is consistent across multiple model scales (125M, 360M, 760M parameters), training data volumes (2.7B, 7.5B, 16B, 48B tokens), context lengths (4K, 8K, 16K, 32K), and evaluation protocols (per-token loss, needle-in-the-haystack, downstream tasks). The effect is not marginal: FoX (Pro) achieves per-token loss at position 16K of ~1.52 versus ~1.54 for Transformer (Pro) at the 760M/48B scale (Figure 2), and the gap is substantially larger at smaller model scales or longer context lengths (Figures 19–21).
However, the paper's claims require careful qualification regarding scale, generalization, and the specific conditions under which the advantages manifest.
Claim: FoX outperforms the Transformer on long-context language modeling. This is supported within the tested regime (up to 760M parameters, 48B tokens, 16K context), but the paper's own analysis reveals that the advantage shrinks systematically with model size. At 125M parameters, the gap is ~0.05–0.08 in per-token loss; at 760M, it is ~0.02 (Figures 19–21). The paper does not test at scales where this trend would plausibly reverse or vanish (e.g., 7B parameters or larger), leaving open the question of whether the forget gate provides meaningful benefit at the scales of frontier models. The observation that "larger models can better model long contexts, and thus forgetting may be less important" (Section 4.5) is speculative—it could equally be that the current training setups and hyperparameters were not fully optimized for FoX at larger scales, or that the advantage manifests differently (e.g., in faster convergence rather than final performance). An experiment that would have clarified this: a scaling law study tracking the FoX-vs-Transformer gap across an order of magnitude more model sizes.
Claim: FoX retains the Transformer's superior long-context capabilities over recurrent models. This is strongly supported. The per-token loss curves (Figure 2), needle-in-the-haystack results (Figure 4), and LongBench performance (Table 2) all show FoX matching or exceeding Transformer baselines while recurrent models consistently fail at tasks requiring genuine long-context utilization. The needle test is particularly decisive: Mamba-2, HGRN2, and DeltaNet all show retrieval scores of 1–2 for documents longer than ~7K tokens, while FoX achieves 9–10 throughout the 16K training context. However, the recurrent model comparisons are not FLOPs-matched—the recurrent models have fundamentally lower per-token FLOPs (1.54–1.65 × 10⁹ vs. 2.72 × 10⁹ for FoX, per Table 6) due to linear complexity. The paper acknowledges this limitation explicitly ("an exact FLOPs-matched comparison would be interesting, [but] will require recalibrating the scaling law for the long-context setting and is beyond the scope of this work"). This means the comparison is not entirely fair: in a FLOPs-matched setting, recurrent models could be scaled to larger parameter counts or trained on more tokens, potentially closing part of the gap. The paper does not provide evidence that FoX would outperform recurrent models under equal computational budget.
Claim: FoX does not require any positional embeddings. This is supported by the ablation in Table 3: FoX (Pro) with RoPE achieves 6.63 perplexity versus 6.62 without RoPE—effectively identical performance. This is a genuinely strong result. However, the claim is validated only on LongCrawl64 with long-context training. The SlimPajama results (Appendix F.3, 2K context) do not include a RoPE ablation for FoX, so it is unclear whether positional embeddings become necessary at very short context lengths or on different data distributions. Additionally, the paper does not test whether FoX without positional embeddings can learn tasks that require precise fine-grained position discrimination (e.g., token-level tasks like character-level language modeling or code completion where exact distance matters), as opposed to the coarse positional signal provided by cumulative decay.
Claim: The Pro block significantly improves both FoX and the Transformer. Supported robustly by Table 3 and Figures 8–9. However, the paper does not provide a FLOPs comparison for the Pro block's additional computation. The Pro block adds output gates (a matrix multiplication per head), output normalization (RMSNorm per head), QK-norm (RMSNorm per query and key), and KV-shift (two additional linear layers per head). While the paper parameter-matches by reducing MLP size, the additional forward-pass operations may affect throughput. The reported throughput in Table 6 shows Transformer (Pro) at 30k tokens/sec versus 38k for Transformer (LLaMA)—a 21% reduction—but it is unclear how much of this is due to the Pro components versus the unoptimized Triton implementation. A more detailed FLOP-speed analysis would help practitioners decide whether the Pro block's accuracy gains are worth the computational cost.
The length extrapolation "overfitting" finding is preliminary but important. Figure 5 and Appendix F.7 show that more training tokens can actually worsen length extrapolation, with 16B-token models extrapolating better than 48B-token models. This is a genuinely interesting observation, but it is presented as an empirical finding without controlled experiments to isolate the mechanism. Critical missing experiments include: training with explicit length regularization (e.g., periodically evaluating on longer sequences and using that signal for early stopping), testing whether the effect is due to optimizer state rather than model parameters (by comparing fresh optimizer at 48B vs. continued training from 16B), and measuring whether the forget gate values themselves systematically change during training (e.g., do the mean and variance of forget gate values shift toward stronger decay as training progresses?). Without these, the claim of "overfitting to context length" remains plausible but unsubstantiated.
The short-context downstream task results are promising but not decisive. Table 1 shows FoX (Pro) with the highest average accuracy (50.88) on LM-eval-harness at the 760M/48B scale, but the gaps between models are small relative to the variance across tasks. For example, on BoolQA, Transformer (Pro) achieves 60.86 versus FoX (Pro) at 46.57—a 14-point gap in the opposite direction. The paper does not report statistical significance for any downstream task comparison, and with 12 tasks, some of the apparent advantages could be noise. The SlimPajama results (Table 7) provide a partial replication with a different training distribution, and the pattern largely holds (FoX (Pro) > Transformer (Pro) > Transformer (LLaMA)), but with similarly small margins (43.29 vs. 42.74 vs. 41.34 average).
Missing experiments that would strengthen the case:
-
Scaling law study. The interaction between model size and FoX's advantage (Figure 6) is the most important open question. A systematic study tracking the FoX-vs-Transformer gap at 5+ model sizes (from ~100M to ~3B parameters) would clarify whether the advantage asymptotically vanishes, persists, or takes a different form (e.g., faster convergence rather than better final loss). The current three data points (125M, 360M, 760M) are suggestive but insufficient to establish a trend.
-
FLOPs-matched comparison with recurrent models. The table stakes for comparing architectures with fundamentally different complexity characteristics. At equal FLOPs, could a larger Mamba-2 or DeltaNet match FoX's long-context performance? This is acknowledged as future work but is critical for the paper's narrative about Transformer superiority.
-
Training data ablation. All experiments use LongCrawl64, which is derived from RedPajama-v2. It is unclear whether the forget gate's benefits depend on properties of this specific data distribution (e.g., average document length, topic coherence, repetition patterns). Training on a different long-context corpus (e.g., books, code, scientific papers) would test data dependence.
-
Forget gate analysis over training. A mechanistic study tracking how forget gate values, forget gate diversity across heads, and the effective decay timescales evolve during training would provide insight into how FoX learns to use context. Do forget gate values change when the model transitions from short-context to long-context utilization? Do they become more or less diverse? The current analysis (Figures 3, 27, 28) is purely end-of-training.
-
Comparison against other positional embedding schemes. The paper compares against ALiBi (via the fixed forget gate equivalence) and RoPE, but not against other methods that provide data-independent decay (e.g., KERPLE, Chi et al., 2022a) or learned relative position biases (e.g., T5, Raffel et al., 2020). Showing that data-dependent decay outperforms a broader set of fixed-decay methods would strengthen the claim that data-dependence specifically (not just having decay) is the critical factor.
-
Instruction-tuned or larger-scale needle test. The needle-in-the-haystack results for the standard mode show Transformer (Pro) failing in some configurations within the training context length (Figure 4). The paper notes that "these are small models without instruction-tuning. We expect that with more parameters/training tokens or instruction-tuning Transformers should also have perfect accuracy within the training context length" (Section 4.3, footnote). This is speculation—testing on a larger or instruction-tuned model would confirm whether the Transformer's standard-mode failures are a scale issue or a fundamental limitation that FoX addresses.
-
The cost of the forget gate. While the paper argues the additional parameters and computation are negligible, a direct measurement of training time overhead at larger scales or with optimized CUDA kernels would make this claim more concrete. The current Triton implementation is "without significant optimization" and shows a 21% throughput reduction for FoX (Pro) versus Transformer (LLaMA), but this confounds kernel maturity with algorithmic overhead.
Summary of the evidence-to-claims mapping:
The paper's central innovation—that a data-dependent cumulative decay bias improves softmax attention—is well-supported by consistent empirical patterns across configurations. The practical value proposition—that FoX offers a simple drop-in improvement over standard Transformers—is supported for the tested regime (hundreds of millions of parameters, tens of billions of tokens, ~16K context) but is not yet established for the scales at which most production language models operate. The most robust finding is the qualitative behavior difference: FoX with data-dependent forget gates maintains a steeper and more sustained per-token loss slope than either standard Transformers or fixed-decay variants, indicating fundamentally better context utilization that manifests across scales, architectures, and evaluation protocols. The most important open question is whether this advantage persists, diminishes, or transforms at the scales where language models are typically deployed.
6. Limitations and Trade-offs
The "Overfitting to Context Length" Phenomenon Constrains Reliable Length Extrapolation
The assumption or constraint. Section 4.3 and Appendix F.7 document a finding that appears to complicate one of the paper's key claimed advantages: models trained longer on more data actually extrapolate worse beyond the training context length than models trained on less data. Specifically, FoX (Pro) models trained on 16B tokens with a 16,384-token context length achieve better needle retrieval accuracy at 32K document length than identically-architected models trained on 48B tokens. The paper states that "more training tokens often leads to worse extrapolation, indicating that the models may be gradually 'overfitting' to their training context length during training" (Section 4.3).
The consequence. Length extrapolation—the ability to process sequences longer than those seen during training—is one of the paper's marquee claims for FoX (Section 4.2, Figure 2). If this ability is not a stable architectural property but rather a transient phenomenon that peaks early in training and then degrades, then FoX's extrapolation advantage over Transformers may not be reliable at scale. A practitioner training a production model would face an uncomfortable tradeoff: stop training early to preserve extrapolation ability (sacrificing in-distribution performance) or train to convergence and accept that the model will not extrapolate meaningfully. Worse, the optimal stopping point for extrapolation is not predictable without running expensive needle evaluations during training, and it may depend on hyperparameters like learning rate (Figures 5, 22–24 show the extrapolation pattern shifting with both token count and learning rate). The paper does not offer a principled criterion for when to stop training to maximize extrapolation, nor does it demonstrate that the extrapolation advantage survives at scales beyond 760M parameters or 48B tokens.
What evidence exists in the paper. Figure 5 and Figures 22–24 (Appendix F.7) are the primary evidence. For FoX (Pro) in easy-mode needle evaluation at 32K document length (16K beyond the training context), the 16B-token model with learning rate 0.001 achieves scores of 9–10 across most depths; the 48B-token model with the same learning rate achieves scores of 7–8; and the 48B-token model with learning rate 0.002 achieves scores of 4–8. The corresponding per-token loss curves (Figure 24, left panel) show that the 48B-token models have flatter slopes beyond 16K than the 16B-token models—the stronger in-distribution models are worse at using information beyond their training horizon. The paper also notes that the same sensitivity pattern appears for Transformer (Pro) (Figures 22–23), suggesting this is a general property of training dynamics rather than a FoX-specific issue. The paper does not measure this effect at scales beyond 760M/48B, nor does it track the evolution of forget gate values or effective context windows during training to understand why the overfitting occurs.
Mitigation status. The paper acknowledges this finding but offers no mitigation beyond describing it as an observation. Section 4.3 flags it as a finding that extrapolation behaviors "could be hyperparameter-dependent," and Section 6 (conclusion) does not list it as a limitation, focusing instead on the need to test at larger scales. The phenomenon remains unexplained and unaddressed—it is not incorporated into any training recipe or architectural modification. A practitioner reading this paper would be aware that FoX's extrapolation ability is fragile but would receive no guidance on how to preserve it.
No Demonstration at Frontier Model Scales
The assumption or constraint. All experiments use models up to 760M non-embedding parameters trained on up to 48B tokens with a training context length of 16,384 tokens (Section 4.1, Table 4). The paper explicitly acknowledges: "due to our limited computing resources, our main experiments only use models up to 760M parameters, 48B tokens, and a training context length of 16384 tokens. Thus, an important direction for future work is to test FoX at larger scales" (Section 6). The scale gap between these experiments and frontier language models (billions to hundreds of billions of parameters, trillions of training tokens, context lengths of 128K+) is between 2 and 4 orders of magnitude on key axes.
The consequence. The paper's own analysis reveals that the advantage of FoX over the standard Transformer systematically shrinks as model size increases. Section 4.5 and Figures 19–21 show this trend across three model sizes: at 125M parameters, the per-token loss advantage of FoX (Pro) over Transformer (Pro) at long positions is approximately 0.05–0.08; at 360M, it narrows to approximately 0.03–0.05; at 760M, it is approximately 0.02–0.03. The paper interprets this as suggesting "the advantages of having a 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" (Section 4.5). If this trend continues—and the paper provides no evidence that it plateaus or reverses—then at scales of 7B parameters or larger, the forget gate's benefit could become negligible or even negative (imposing a recency bias that a sufficiently capable model would be better off learning purely from data). Three data points (125M, 360M, 760M) are insufficient to extrapolate a scaling trend, and the paper does not provide a scaling law analysis that would allow practitioners to predict whether FoX is worth adopting at their target scale. Additionally, the training context length of 16K tokens is small by contemporary standards (128K–1M tokens). The paper's finding that FoX's advantage grows with training context length (Section 4.5, Figure 6) suggests a more optimistic scaling direction, but the interaction between model size, context length, and the forget gate benefit is not resolved.
What evidence exists in the paper. Figures 19, 20, 21, and the main Figure 2 provide the per-token loss scaling data. Table 3 provides the perplexity scaling within a single model size. The paper does not experiment at any scale beyond 760M/48B/16K. The trend of shrinking advantage is visible but not quantified as a scaling law.
Mitigation status. The authors are transparent about this limitation in Section 6 and flag larger-scale testing as "an important direction for future work." However, the paper's claims are stated without scale qualification in the abstract and introduction ("FoX outperforms the Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks"), which could mislead readers who do not carefully examine the model size and training token counts. The paper does not provide any theoretical argument for why the forget gate advantage would persist at scale (e.g., that the forgetting mechanism addresses a fundamental inductive bias gap that no amount of scale can compensate for), making the scaling behavior purely an empirical question.
FLOPs-Matched Comparison Against Recurrent Models Is Absent
The assumption or constraint. The paper compares FoX against recurrent sequence models (Mamba-2, HGRN2, DeltaNet) at equal parameter counts (~760M non-embedding parameters) and equal training data (48B tokens), but not at equal computational cost. As Table 6 documents, the recurrent models have substantially lower forward FLOPs per token due to linear complexity: Mamba-2 at 1.65 × 10⁹ FLOPs/token versus FoX (Pro) at 2.72 × 10⁹ FLOPs/token—a 1.65× difference. The paper states: "an exact FLOPs-matched comparison would be interesting, [but] will require recalibrating the scaling law for the long-context setting and is beyond the scope of this work" (Appendix B.2).
The consequence. The paper's central comparative claim—that FoX "retains the Transformer's superior long-context capabilities over recurrent sequence models such as Mamba-2, HGRN2, and DeltaNet" (Abstract)—is not tested under equal computational budget. In a FLOPs-matched setting, the recurrent models could be scaled to approximately 1.65× more parameters or trained on 1.65× more tokens for the same total compute. Given that the recurrent models already show competitive performance on short-context tasks at equal parameters (Table 1: Mamba-2 achieves 50.21 average vs. 50.88 for FoX (Pro)) and that their primary weakness is long-context utilization (which is the dimension where additional parameters or training tokens might provide the most benefit), it is plausible that a FLOPs-matched recurrent model would close a substantial portion of the performance gap. The paper also does not account for the recurrent models' inference efficiency advantage: at very long sequences (where the quadratic cost of attention dominates), recurrent models' linear complexity becomes an exponentially larger advantage in wall-clock time and total FLOPs. A fair deployment comparison would need to consider total cost of ownership (training + inference), not just parameter-matched training comparisons.
What evidence exists in the paper. Table 6 provides the FLOPs per token and throughput numbers that quantify the computational asymmetry. The per-token loss curves (Figure 2) and needle results (Figure 4) show recurrent models failing at long-context tasks, but these are parameter-matched, not FLOPs-matched. The paper does not provide any FLOPs-controlled comparison or scaling law analysis that would allow readers to estimate what performance recurrent models would achieve at equal compute.
Mitigation status. The paper acknowledges this limitation explicitly in Appendix B.2 but does not attempt even a partial mitigation (e.g., comparing against a scaled-up recurrent model at a single FLOPs-matched point, or extrapolating from scaling trends in the literature). The framing throughout the main text (e.g., "FoX also retains the Transformer's superior long-context capabilities over recurrent sequence models" in Section 1) presents the comparison as resolved when it is not. A skeptical reader could argue that the paper's conclusions about Transformer superiority over recurrent models are confounded by the unequal computational budget, and that the question of which architecture class is more compute-efficient for long-context tasks remains open.
Single Dataset and Model Family — No Evidence of Cross-Domain or Cross-Architecture Generalization
The assumption or constraint. All long-context experiments are conducted on a single dataset (LongCrawl64, a subset of RedPajama-v2) using a single model family (PaLM 2 is not used; the base architecture is the LLaMA-style Transformer, initialized with the HuggingFace LLaMA scheme). The short-context SlimPajama experiments (Appendix F.3) provide a partial second data point but are at a smaller scale (340M parameters) and use a fundamentally different context length (2K tokens). The paper does not test on other long-context corpora (e.g., books, code, scientific papers, multilingual text) or with other base Transformer architectures (e.g., models using parallel attention + MLP rather than sequential, models with different normalization schemes, models pre-trained from different initializations).
The consequence. The forget gate mechanism introduces a specific inductive bias: information from the past is multiplicatively decayed based on the content of intervening tokens. Whether this bias is universally beneficial for language modeling, or whether its benefits are specific to the statistical properties of web-scraped text corpora like RedPajama-v2, is unknown. For example, code repositories have strong long-range dependencies (a function definition may be referenced hundreds of lines after it is defined, with intervening code that is structurally unrelated but syntactically coherent). Scientific papers have explicit section structure where section boundaries typically signal topic shifts—exactly the kind of content that FoX's forget gates might learn to recognize. Conversely, narrative text (novels) has gradual thematic evolution where sharp forgetting at section breaks might be harmful. Without testing across domains, a practitioner cannot know whether FoX's forget gate will adapt appropriately to their target distribution or whether it will impose an overly aggressive recency bias in domains where distant context is consistently relevant. Similarly, the Pro block components (output gates, QK-norm, KV-shift) were validated only on the same corpus and model family. Their benefits may interact with specific properties of the LLaMA architecture (pre-norm, SwiGLU MLP, sequential attention-then-MLP block structure) and may not transfer to other design choices (post-norm, parallel blocks, different activation functions).
What evidence exists in the paper. The main results (Figures 2–4, Tables 1–2) are all on LongCrawl64-trained models. The SlimPajama results (Figure 12, Table 7) provide a limited cross-domain check at 340M/15B/2K scale but are acknowledged as "small-scale experiments without extensive hyperparameter tuning" where "results might not transfer to larger scales" (Appendix F.3). The paper does not test on code, books, scientific text, or multilingual corpora. No experiments replace the LLaMA backbone with an alternative Transformer architecture.
Mitigation status. The paper does not claim cross-domain generalization and does not flag single-corpus evaluation as a limitation in Section 6. The SlimPajama experiments provide suggestive evidence that the benefits persist in a different web-text distribution at short context lengths, but this is insufficient to establish robustness. A practitioner deploying FoX on a domain substantially different from RedPajama-v2 would be operating without evidence.
The Pro Block's Overhead Is Not Isolated from Implementation Immaturity
The assumption or constraint. The paper reports training throughput in Table 6: FoX (Pro) achieves 27k tokens/sec, Transformer (Pro) achieves 30k tokens/sec, and Transformer (LLaMA) achieves 38k tokens/sec on 4 NVIDIA L40S GPUs. The authors note that "the FlashAttention kernels for FoX (Pro), Transformer (Pro), and FoX (LLaMA) are implemented in Triton by us on top of Flag Attention without significant optimization, while Transformer (LLaMA) uses the official FlashAttention implementation 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" (Appendix B.2).
The consequence. The paper's practical value proposition—that FoX is a drop-in modification to standard Transformers with negligible overhead—rests on the claim that the throughput gap is an artifact of kernel maturity rather than a fundamental algorithmic cost. However, this claim is unverified. The current measurements show a 29% throughput reduction from Transformer (LLaMA) to FoX (Pro) (38k → 27k tokens/sec), and a 21% reduction from Transformer (LLaMA) to Transformer (Pro) (38k → 30k tokens/sec). Even if the FlashAttention kernel is fully optimized (closing the gap between FoX (Pro) and Transformer (Pro) implementations), the Pro block components themselves—output gates (an additional matrix multiplication per head), QK-norm (additional RMSNorm operations on queries and keys), KV-shift (additional linear layers and interpolation), and output normalization (additional RMSNorm on attention outputs)—represent real computational work beyond what a standard LLaMA block performs. The paper reduces MLP parameters to keep total parameter count equal (Table 6 shows FoX (Pro) at 759M vs. Transformer (LLaMA) at 756M), but this does not equalize FLOPs or latency: the Pro block's operations are in the attention path (which has quadratic complexity in sequence length) while the reduced MLP parameters save computation in the feedforward path (which is linear in sequence length). At long context lengths, the attention path dominates, so the Pro block's additional attention-path operations may cost more than the MLP reduction saves. The paper provides no FLOP-breakdown by component and no latency measurement at different context lengths that would allow a practitioner to estimate the real-world cost of adopting the Pro block for their specific sequence length.
What evidence exists in the paper. Table 6 provides throughput numbers but without component-level breakdowns or scaling with context length. The paper states the expectation of throughput parity but provides no evidence—even a FLOP-count analysis of the additional operations—to support it.
Mitigation status. The authors transparently describe the implementation limitations and express optimism about optimization potential. However, the paper does not attempt even an engineering-order-of-magnitude estimate: counting the additional FLOPs from the Pro components as a fraction of total attention FLOPs, for example. The recommendation that "future work adopt FoX (Pro) and Transformer (Pro) as baselines" (Section 6) implicitly assumes the Pro block's cost is acceptable, but a practitioner considering this recommendation has no data on what that cost actually is at their target scale and sequence length.
The Forget Gate Initialization and Learning Dynamics Are Underexplored — Robustness to Hyperparameters Is Unclear
The assumption or constraint. The forget gate mechanism introduces new hyperparameters (the initialization of biases b_f, the learning rate for forget gate parameters, potential interaction with the learning rate of other parameters) and new training dynamics (the effective context window is learned rather than architecturally fixed). The paper's experiments tune learning rates and head dimensions independently per model and per architecture (Table 5), but they do not systematically explore the sensitivity of FoX's performance to forget-gate-specific hyperparameters. Crucially, the paper reports that for data-independent and fixed forget gates, "zero initialization performs extremely poorly" (Appendix D), requiring a carefully designed geometric initialization scheme across heads. For data-dependent forget gates (the main proposal), zero initialization is used and "works well," but the paper does not test alternative initializations, does not explore whether the forget gate learning rate should differ from the main model learning rate, and does not measure the variance in final performance across different forget gate initializations.
The consequence. In practice, the success of FoX may depend on hyperparameter choices and initialization schemes that are not obvious a priori. The fact that fixed forget gates completely fail with zero initialization but work with geometric initialization (Appendix D) suggests that the forget gate biases b_f are in a sensitive part of the parameter space: if initialized too close to zero, the model starts with an effective context length of ~2 tokens (since σ(0) = 0.5 and 0.5^2 = 0.25) and may never recover the ability to attend to longer contexts. The data-dependent variant apparently escapes this trap by learning appropriate biases during training, but the paper does not demonstrate that this learning is robust—for instance, does the model reliably discover long-range attention if the training data requires it, or can it get stuck in a local optimum of strong local attention? The interaction between learning rate and the forget gate's behavior is visible in Figures 5 and 24: FoX (Pro) with learning rate 0.001 at 48B tokens extrapolates better than with learning rate 0.002, suggesting that the forget gate's learned decay rates are sensitive to optimization hyperparameters. A practitioner training FoX at a new scale or on new data would need to tune these hyperparameters, but the paper provides no guidance on how to do so efficiently (e.g., which metrics to monitor, what failure modes to look for, what reasonable ranges are for forget-gate-specific hyperparameters).
What evidence exists in the paper. Appendix D documents the sensitivity of data-independent forget gates to initialization. Figures 5 and 24 show sensitivity of length extrapolation to learning rate and training tokens. The paper does not present a systematic study of forget gate initialization schemes for the data-dependent variant, forget gate learning rate sensitivity, or the variance of final performance across multiple initializations (beyond the three-seed stability check in Figure 26, which uses the same initialization and hyperparameters).
Mitigation status. The paper provides the geometric initialization scheme for data-independent forget gates (Appendix D) as a solution to the zero-initialization failure mode, but does not extend this analysis to the data-dependent case or provide diagnostic tools for practitioners. The observation that FoX "often prefers higher learning rates and more heads/smaller head dimensions than the Transformer" (Section 4.1) is a useful hint but is not quantified (how much higher? is this consistent across scales?) and is based on the tuned hyperparameters in Table 5 without a controlled sensitivity analysis.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper makes a reframing contribution rather than a paradigm shift: it shows that the forget gate—a mechanism previously considered the exclusive domain of recurrent architectures—is kernel-agnostic and can be cleanly ported to softmax attention without sacrificing any of the Transformer's structural advantages. The magnitude of this reframing is best understood along three axes.
First, it dissolves a false dichotomy. The literature has largely treated "recurrent models with forget gates" and "Transformers with softmax attention" as competing design philosophies, with a growing body of work arguing for hybrid architectures that bolt recurrent components alongside attention layers (Ma et al., 2022; 2024; Ren et al., 2024). FoX demonstrates that the forget gate does not require recurrence—it can live entirely within the attention mechanism itself, expressed as an additive logit bias D_{ij} = \sum_{l=j+1}^i \log f_l derived from a cumulative product of scalar gates. This insight, combined with the parallel work showing that gated linear attention looks like softmax attention (Yang et al., 2023; Dao & Gu, 2024), completes a two-way isomorphism: recurrent models can be written to resemble Transformers, and Transformers can be augmented with the signature mechanism of recurrent models. The practical upshot is that future research should not ask "should we use attention or recurrence?" but rather "which kernel and which gating mechanism best serve this task?"—with the understanding that these are independent design axes.
Second, it establishes data-dependent temporal decay as a distinct axis of attention modulation, orthogonal to content-based similarity (QK^T) and position-based bias (RoPE, ALiBi). Prior work on improving attention largely operated within one of these two axes: either designing better positional encodings (Press et al., 2021; Su et al., 2024; Chi et al., 2022a) or modifying the similarity computation (various kernel attention variants). FoX introduces a third axis—content-conditioned decay—that interpolates between them: it is positional in that it naturally encodes distance through cumulative summation, but it is also content-driven in that the decay rate at each position depends on what token appears there. Figure 7 provides the critical evidence for why this axis matters: data-dependent forget gates consistently outperform both fixed-decay (pure positional) and data-independent learned decay variants, and the advantage is most pronounced in length extrapolation beyond the training context. This opens a new dimension for attention design that was previously invisible in the standard positional-vs-content dichotomy.
Third, it provides a diagnostic tool and a warning about length extrapolation. The finding that models can "overfit" to their training context length—with more training tokens actually worsening extrapolation (Figure 5, Figures 22–24)—is more significant than the paper treats it. This is not a FoX-specific quirk; it appears for Transformer (Pro) as well. It implies that length extrapolation is not a stable architectural property that can be assessed at a single checkpoint, but rather a dynamic that evolves during training and may peak before convergence. The per-token loss versus perplexity distinction (Section 4.2, Appendix C) is central here: aggregate metrics like perplexity mask the plateau in context utilization that the per-token loss reveals. This finding should change how the community evaluates long-context models: rather than reporting perplexity at the final checkpoint, papers should report per-token loss curves at multiple checkpoints during training to characterize whether length generalization is improving or degrading. Research directions that become more attractive include dynamic training schedules that periodically evaluate (and optimize for) length extrapolation, auxiliary losses that explicitly penalize context-length overfitting, and architectures that decouple the effective context window from the training context length more robustly than current approaches. Directions that become less attractive include the naive assumption that more training data at a fixed context length will naturally improve length generalization—the evidence here suggests the opposite may be true.
The reconciliation of prior contradictions. The paper does not resolve a major contradiction in the literature in the way that, say, the Chinchilla scaling laws reconciled conflicting results about model size versus data scaling. Rather, it resolves a subtler tension: the ubiquity of forget gates in recurrent models versus their complete absence in Transformers, despite both model classes performing similar sequence modeling tasks. The resolution is that previous work was looking at the forget gate in the wrong place—as part of recurrent state updates rather than as a general mechanism for modulating token-to-token interactions. Once the forget gate is reformulated as a cumulative decay on attention scores, it becomes clear that there was no fundamental incompatibility; the mechanism was always applicable to softmax attention, just in a parallel rather than recurrent form.
Research directions that become higher priority:
- Attention kernel design space exploration. If forget gates are kernel-agnostic, then systematic comparisons of different kernels (exponential dot product, linear, other learned kernels) with the forget gate held constant become the obvious next step. FoX opens the door to studying what the softmax kernel provides over linear kernels when both have access to data-dependent decay.
- Verifier and gating mechanism co-design. The parallel to the compute-optimal test-time scaling work is instructive: just as that work found verifier quality to be the primary bottleneck for search-based inference scaling, future work on attention gating should recognize that the forget gate's learned decay rates interact with what the model chooses to encode in key and value vectors. Improving the forget gate (e.g., making it multi-dimensional, using more expressive functions of the input) may yield diminishing returns unless the key/value representations are also optimized for the decay structure they will undergo.
- Training dynamics of learned context windows. The overfitting-to-context-length finding elevates a mechanistic study of how forget gate values evolve during training from a curiosity to a priority. If models learn to tune their effective context window to match the training distribution, this has implications for curriculum learning (start with short contexts, gradually increase), data mixing (include some longer-sequence data even when training at shorter contexts to prevent overfitting), and early stopping criteria.
Research directions that become lower priority:
- Pure positional encoding research as a solution to long-context modeling. If a learned scalar forget gate can provide sufficient positional signal without any explicit positional embeddings (Table 3: FoX (Pro) with RoPE at 6.63 perplexity versus 6.62 without), then the search for ever-more-sophisticated positional encoding schemes may be solving a problem that a simpler, content-aware mechanism already handles. The marginal value of a new positional encoding over the forget gate baseline is likely to be small and decreasing.
- Hybrid recurrent-attention architectures that keep the two mechanisms separate. FoX shows that the forget gate can be embedded directly into attention, making the case for separate recurrent layers alongside attention layers weaker. Why add a recurrent module when you can fold its key mechanism into the attention computation itself, with negligible parameter overhead (a weight vector and bias per head) and no additional architectural complexity? Future hybrid designs should justify why their recurrence is doing something FoX-style gating within attention cannot, rather than assuming recurrence is necessary for temporal processing.
Follow-Up Research This Work Enables
Scaling law for the forget gate benefit: does FoX's advantage persist, vanish, or transform at 1B–7B+ parameters? The paper's own data (Figures 19–21) shows the FoX-vs-Transformer gap shrinking from ~0.05–0.08 per-token loss at 125M parameters to ~0.02 at 760M parameters. This trend has three data points, which is insufficient to fit a functional form. A strong follow-up trains FoX (Pro) and Transformer (Pro) at five or more model sizes spanning 100M to 3B parameters (or larger, if resources permit), holding training data, context length (16K or 32K), and hyperparameter tuning budget constant. The measured quantity should be the per-token loss gap at the final training context position as a function of parameter count, fit to a power law or exponential decay. If the gap extrapolates to zero at ~3–7B parameters, then FoX is a small-to-medium-scale technique and the community should stop exploring it for frontier models. If the gap asymptotes to a non-zero value, then the forget gate provides an irreducible inductive bias benefit that scale cannot compensate for—a much stronger result that would justify investment in optimized kernels and large-scale training. The paper's finding that the advantage grows with context length (Figure 6) complicates the picture: the scaling law should be measured at multiple context lengths (e.g., 4K, 16K, 64K) to disentangle whether the model-size trend and the context-length trend interact multiplicatively.
Controlled study of context-length overfitting: what mechanism causes longer training to degrade extrapolation? The observation that 48B-token models extrapolate worse than 16B-token models (Figure 5) is striking but unexplained. A controlled experiment would track the following quantities at regular intervals during training for both FoX (Pro) and Transformer (Pro): (a) per-token loss at positions both within and beyond the training context length; (b) the distribution of forget gate values across heads and layers (mean, variance, effective decay timescale T(b) = 1 / (-log σ(b)) for each head); (c) the effective context window, defined as the distance at which the average attention weight (or cumulative forget gate product) drops below some threshold (e.g., 1/e); and (d) needle-in-the-haystack retrieval accuracy at 1.5× and 2× the training context length. The hypothesis to test is that forget gate values systematically shift toward stronger decay (lower f_t, more negative log f_t) as training progresses, effectively shrinking the model's learned context window to match the training distribution. If this hypothesis holds, it predicts that (1) the variance of forget gate values across heads decreases over training (heads converge to similar decay rates), (2) the mean forget gate value decreases, and (3) models trained with higher learning rates show this shift earlier and more dramatically. A follow-up intervention would test whether an auxiliary loss penalizing the difference between the model's effective context window and the training context length (encouraging the model to maintain long-range attention even when not strictly necessary for in-distribution loss) prevents the overfitting and preserves extrapolation. The beyond-training-context needle accuracy would be the primary metric.
Multi-dimensional or vector-valued forget gates: does per-dimension decay improve over scalar decay? The paper uses scalar forget gates (f_t \in \mathbb{R}) per head, citing parameter efficiency and FlashAttention compatibility (the D_{ij} = c_i - c_j decomposition requires scalar gates). But many recurrent models use vector forget gates (LSTMs, the original GLA formulation) to allow per-dimension forgetting—different features in the key/value space might require different decay rates. A natural extension is to test whether this expressivity matters. The challenge is making vector gates compatible with efficient attention: if f_t \in \mathbb{R}^{d_head}, then F_{ij} \in \mathbb{R}^{d_head} and the bias D_{ij} is a vector added element-wise to q_i^T k_j—which is a scalar. One approach is to tie groups of d_head / g dimensions to share a scalar gate (reducing to g scalar gates per head, with g = 1 recovering the paper's design and g = d_head being full vector gates). A study sweeping g from 1 to d_head at a moderate scale (e.g., 360M parameters on 7.5B tokens) would measure whether per-dimension decay provides benefits beyond scalar decay, and whether the benefit saturates at small g. The primary metric is per-token loss slope within and beyond the training context, since vector gates might enable more nuanced long-range retention beyond what scalar gates can express. The cost is implementation complexity: vector gates require storing g cumulative sum vectors per head rather than 1, and the backward pass must compute per-dimension gradients.
Cross-domain stress test: does the forget gate's recency bias help or hurt when long-range dependencies are the rule? All of the paper's experiments use web-scraped text (RedPajama-v2 / LongCrawl64), where the relevance of distant context is highly variable and often low. In domains where long-range dependencies are consistently important—code repositories (a function call may reference a definition thousands of lines earlier, and those references must be preserved precisely), scientific papers (where the methods section is essential context for interpreting results tables hundreds of tokens later), legal documents (where definitions from the preamble constrain interpretation of clauses throughout)—an overly aggressive recency bias could be harmful. A cross-domain study would train FoX (Pro) and Transformer (Pro) at equal parameter count and data volume on three domains: web text (LongCrawl64, as a baseline), code (e.g., The Stack or a repository-level code dataset), and scientific text (e.g., S2ORC or arXiv papers). The primary metrics are per-token loss within the training context (does FoX underperform the Transformer in domains where distant context is always relevant?) and a retrieval task specific to each domain (e.g., for code, predicting the body of a function given its signature and a call site thousands of tokens later; for science, answering a question about a result given the methods section as context). This study would establish whether the forget gate's inductive bias is universally beneficial or whether it must be tuned per-domain (or even disabled) to avoid imposing unhelpful decay where long-range dependencies dominate. If FoX underperforms the Transformer in code or science, the next step is testing whether initializing forget gate biases to favor very long decay timescales (e.g., T(b_init) = 10,000 for all heads) recovers Transformer-level performance while preserving FoX's advantages in mixed-dependency regimes.
Forget gate analysis as an interpretability tool: does the learned decay structure reveal document structure? The visualization in Figure 3 shows that some heads learn strong local attention (near-diagonal attention patterns) while others maintain global attention (diffuse attention across the full context). This suggests that forget gate values might serve as an unsupervised signal for document structure—section breaks, topic shifts, or dialogue turns should produce forget gates near zero (sharp forgetting), while within-section tokens should produce forget gates near one (information preservation). A concrete study would take a trained FoX (Pro) model, run it on long documents with known structure (e.g., Wikipedia articles with section boundaries, transcripts with speaker turns, or code files with function boundaries), and measure whether the forget gate values at structure boundaries are systematically lower than at within-structure positions. A simple metric: the mean forget gate value in a window of ±5 tokens around section boundaries versus at random within-section positions, tested for statistical significance across many documents. If the forget gates reliably signal structure, this has practical applications: (a) text segmentation without supervised labels (use forget gate minima as boundary candidates); (b) KV-cache eviction strategies that prune tokens before sharp forget gates, since information before a strong reset is unlikely to be needed later; and (c) training data filtering (documents where forget gates show no structure might be poorly organized and less useful for training). This direction is enabled by FoX because standard Transformers do not have an explicit forgetting signal—the attention weights confound content similarity and temporal relevance, making it harder to isolate structure from content.
Hardware-aware co-design: can the cumulative sum precomputation be fused into the FlashAttention kernel for zero-overhead forgetting? The paper's current implementation precomputes the cumulative sum vector c in a separate kernel and loads segments of it alongside Q, K, V blocks into SRAM. A more aggressive optimization would fuse the cumulative sum computation into the FlashAttention kernel itself, computing c_i values on-the-fly from log f_t values during the forward pass. This is challenging because FlashAttention processes blocks out of order (the query blocks are tiled over key blocks), but the cumulative sum at a given position depends on all previous positions. A practical intermediate goal: implement the forward pass where each thread block, when processing query block i, computes the partial cumulative sum for just the positions in that block by loading the preceding block's final cumulative sum value. This would eliminate the separate c precomputation kernel and the HBM traffic for storing/loading the cumulative sum vector (reducing from O(L) additional HBM reads to O(Tr) where Tr = L / B_r is the number of query blocks). The metric is wall-clock time overhead relative to standard FlashAttention, measured at multiple sequence lengths (4K, 16K, 64K) to characterize how the overhead scales. If the overhead can be reduced to <2% across sequence lengths, the practical barrier to adopting FoX in production training pipelines largely disappears. This is an engineering contribution rather than a research contribution, but it is directly enabled by the paper's D_{ij} = c_i - c_j formulation and would significantly impact adoption.
Practical Applications and Downstream Use Cases
Long-document processing pipelines where context utilization is critical and compute is constrained. The paper's most directly actionable result for practitioners is that FoX (Pro) achieves equivalent or better long-context performance to Transformer (Pro) while reducing the effective context-length overfitting that degrades Transformers with more training. For applications like document summarization (where the model must extract key information from 10K–50K token documents), multi-hop question answering over large corpora, or repository-level code understanding, deploying FoX (Pro) instead of a standard Transformer could yield accuracy improvements on long-context tasks: Table 2 shows FoX (Pro) achieving 22.71 on GovReport vs. 19.88 for Transformer (Pro) at 760M parameters, a ~14% relative improvement, and 13.51 vs. 10.70 on QMSum, a ~26% relative improvement. These gains are achieved with the same parameter count and comparable training FLOPs (Table 6: 2.72 × 10^9 FLOPs/token for both), meaning the deployment cost is unchanged—the same hardware can serve a more accurate model. The caveat is that these numbers are at 760M parameters; practitioners should verify the scaling trend for their target model size before committing, given the evidence that the FoX advantage shrinks with scale.
Training data generation and self-improvement pipelines where diverse context windows matter. For organizations generating synthetic training data using language models (e.g., distilling long-context reasoning, generating multi-turn dialogue, or creating instruction-following examples with complex context dependencies), the observation that FoX's length extrapolation ability varies with training duration (Figure 5) has a practical implication: early checkpoints may generalize better to longer contexts than final checkpoints. If the generation task involves contexts longer than the model's training length (e.g., generating summaries of 32K-token documents using a model trained at 16K context), using an intermediate checkpoint rather than the final one could improve output quality. FoX amplifies this because its forget gate mechanism provides an additional knob: practitioners could also directly manipulate the forget gate biases at inference time (e.g., adding a positive offset to b_f to increase all f_t values, effectively lengthening the model's context window) without retraining. While the paper does not test this intervention, the architecture makes it straightforward—unlike a standard Transformer where context window behavior is distributed across all attention weights. A principled generation pipeline would: (1) estimate the effective context window of each checkpoint using per-token loss on long sequences; (2) select the checkpoint that maximizes accuracy at the target generation length; and (3) optionally adjust forget gate biases to further tune the decay rate.
Edge deployment where positional embedding simplicity reduces implementation burden. FoX's elimination of positional embeddings (Table 3: no performance difference with or without RoPE for FoX (Pro)) simplifies the inference stack for on-device or embedded deployment. Rotary Position Embeddings require computing sin/cos values and applying rotations to query and key vectors at every layer, every token—this is a small but non-trivial fraction of inference FLOPs and adds implementation complexity (particularly in quantized or custom-hardware settings where trigonometric functions may not be optimized). By removing RoPE entirely, FoX (Pro) reduces the number of operations in the attention path and eliminates a source of numerical precision issues (RoPE rotations can accumulate errors in low-precision formats). The benefit is small in absolute terms but relevant in environments where every millisecond and every milliwatt matters: on-device language models for keyboards, voice assistants, or real-time translation. The paper demonstrates that this simplification costs nothing in accuracy at the tested scale, though practitioners should verify on their specific hardware and model size.
KV-cache eviction strategies guided by forget gate values. A speculative but promising application that the paper explicitly mentions (Section 6: "we could potentially prune computation (e.g., KV-cache eviction) adaptively based on the forget gate values"): during autoregressive generation with a KV cache, tokens whose cumulative forget gate product with respect to the current position has dropped below a threshold (e.g., F_{i, j} < 0.01) contribute negligibly to the attention output and could be evicted from the cache, reducing memory usage and attention computation. Because F_{ij} = \prod_{l=j+1}^i f_l is directly computable from the running cumulative sum c_i, the eviction decision requires no additional forward passes—it is a deterministic function of the forget gate values already being computed. A concrete heuristic: maintain the KV cache entries, and after each new token is processed, compute the effective retention F_{i, j} for all cached positions j; evict any position where F_{i, j} < \epsilon for a threshold \epsilon (e.g., 0.01). The paper does not evaluate this strategy, but it is a natural extension. The potential gain is significant: for long-context generation, the KV cache is the dominant memory consumer, and evicting tokens whose information has been "forgotten" could reduce peak memory by 30–50% for typical forgetting patterns (where decay to <0.01 occurs within a few thousand tokens for most heads). This is more principled than heuristic eviction strategies based on attention scores or position because the forget gate provides a causal, cumulative measure of information retention that is independent of the current query—the eviction decision for token j depends only on forget gate values at positions j+1 through i, not on the specific query being computed.
When to Prefer This Method
The paper positions FoX as a general-purpose replacement for standard softmax attention, demonstrating advantages across language modeling, length extrapolation, and downstream tasks. It does not articulate a sharp tradeoff matrix against named alternatives (e.g., "prefer FoX over Mamba when X, prefer Mamba over FoX when Y"). The closest the paper comes to a conditional recommendation is the finding that FoX's advantage grows with training context length and shrinks with model size (Section 4.5), which implies a decision rule based on the ratio of context length to parameter count, but this is an empirical trend rather than a prescribed algorithm. The paper's stated comparison is primarily against the standard Transformer and secondarily against recurrent models, but the recurrent model comparison is parameter-matched rather than FLOPs-matched, making it insufficient to support a deployment recommendation. A conditional preference matrix would therefore be speculative rather than grounded in the paper's direct evidence—the paper's contribution is demonstrating that forget gates improve softmax attention, not identifying the precise conditions under which softmax attention with forget gates should be chosen over other architectural families. I omit this subsection to avoid fabricating tradeoffs the paper does not establish.
