Adaptive Computation Pruning for the Forgetting Transformer
Authors: Zhixuan Lin
Johan Obando-Ceron
Xu Owen He
Aaron Courville
Venue: COLM 2025 (2025)
Keywords: transformer
forgetting transformer
efficient transformer
sequence modeling
adaptive computation pruning
forget gate
sparse attention
FlashAttention
hardware-aware optimization
TL;DR: We propose a method that adaptively prunes computations in the Forgetting Transformer based on forget gate values.

Abstract: The recently proposed Forgetting Transformer (FoX) incorporates a forget gate into softmax attention and has shown consistently better or on-par performance compared to the standard RoPE-based Transformer. Notably, many attention heads in FoX tend to forget quickly, causing their output at each timestep to rely primarily on local context. Based on this observation, we propose Adaptive Computation Pruning (ACP) for FoX, a method that dynamically prunes computations involving input-output dependencies that are strongly decayed by the forget gate. In particular, our method performs *provably safe* pruning via a dynamically set pruning threshold that guarantees the pruned attention weights are negligible. 
We apply ACP to language model pretraining with FoX and show it consistently reduces the number of FLOPs and memory accesses in softmax attention by around 70\% across different model sizes and context lengths, resulting in a roughly 50\% to 70\% reduction in attention runtime (or a 2--3$\times$ speedup) and a roughly 10\% to 40\% increase in end-to-end training throughput. Furthermore, longer context lengths yield greater computational savings. All these speed improvements are achieved without any performance degradation. Our code is available at https://github.com/zhixuan-lin/forgetting-transformer.

Published as a conference paper at COLM 2025
Adaptive Computation Pruning for the Forgetting Trans- former
Zhixuan Lin
Johan Obando-Ceron Mila Quebec AT Institute
Xu Owen He MakerMaker Al owen. hexu@gmail.com com
Mila Quebec AI Institute Université de Montréal zxlin. cri.inll.........mmmmmmmmm.m...m com
Université de Montréal
jobando@73@gmail.cco
Aaron Courville
Mila Quebec AI Institute Université de Montréal courvila@mile quebec
Abstract
The recently proposed Forgetting Transformer (FoX) incorporates a forget gate into softmax attention and has shown consistently better on-par performance compared to the standard RoPE-based Transformer. Notably, many attention heads in FoX tend to forget quickly, causing their output each timestep to rely primarily on local context. Based on this observation, we propose Adaptive Compulation Pruning (ACP) for FoX, a method that dynamically prunes computations involving input-output dependencies that are strongly decayed by the forget gate. In particular, our method performs prozably safe pruning via dynamically set pruning threshold that guarantees the pruned attention weights are negligible. We apply ACP language model pretraining with FoX and show it consistently reduces the number of FLOPs and memory accesses in softmax attention by around 70% across different model sizes and context lengths, resulting in roughly 50% 70% reduction attention runtime (or 2-3x speedup) and a roughly 10% to 40% increase in end-to-end training throughput. Furthermore, longer context lengths yield greater computational savings. All these speed improvements are achieved without any performance degradation. Our code is available at https: //g:::::iiii thub. cttc httm//hixxann-iiee http:/hhixunnnnnfoogttttiinnnnn -eriiiinng-tsssnreee transformer.
1 Introduction
Transformers (Vaswani et al., 2017) have quadratic time complexity with respect to context length, resulting in significant computational costs over long sequences. The recently pro- posed Forgetting Transformer (FoX) (Lin et al., 2025) features modified softmax attention mechanism with a forget gate, which allows some attention heads to downweight distant dependencies and focus mainly on the local context. FoX has been shown to consistently achieve better or on-par performance compared to the standard RoPE-based (Su et al., 2024) Transformer in various tasks, including long-context language modeling and downstream tasks such as the needle-in-a-hassteel test (Kamradt, 2023). It is also compatible with the FlashAttention (Dao, 2024) algorithm, which allows efficient processing of long sequences.
Lin et al. (2025) show that many attention heads in FoX tend to forget quickly. For these heads, the dependencies between distant input-output pairs are extremely weak and can potentially be ignored. Based on this observation, we propose Adaptive Computation Pruming (ACP) for FoX, method that dynamically prunes computations involving input-output dependencies that are strongly decayed by the forget gate. In particular, our method performs provably safe pruning via a dynamically set threshold that guarantees the total pruned attention weights are bounded by hyperparameter In practice, we set
*Correspondence Zhixuan Lin.

---

Published as a conference paper at COLM 2025
Without
With
Pruning boundary
Query
Query
Key/Value
ValyV
Figure 1: Illustration of Forgetting Attention with and without ACP. Each cell represents block in the FlashAttention algorithm. Darker colors indicate more-negative decay bias values and thus stronger decay. The solid arrows indicate the set of blocks that would be visited (in the indicated order) in the FlashAttention iterations.
0.000045, which is effectively negligible. Furthermore, as shown in Figure 1, the decay structure in FoX induces a Boddnngwindow-liie pruning pattern, enabling an efficient two- stage implementation. First, we identify pruning boundary across the grid of computations in FlashAttention via a linear-time algorithm. Once the pruning boundary is identified, we restrict the FlashAttention iterations to the remaining blocks, avoiding any wasted computation on pruned dependencies.
We apply ACP to language model prefraining with FoX with sizes from 125M to 760M parameters and training context lengths from 4k to 16k tokens. We find that ACP consistently prunes around 70% of the FLOPs and memory accesses in softmax attention across the tested model sizes and context lengths, resulting in a roughly 50% to 70% reduction in attention runtime (or 2-3x speedup) and a roughly 10% to 40% increase in training throughput. In particular, longer context lengths lead to greater computational savings and speedups. These speed improvements are achieved without affecting language modeling loss and downstream task performance. To provide further insight into our method, we conduct series of analyses such as examining the pruning boundaries and analyzing the distribution of computational savings across different attention heads. Notably, our analysis reveals the existence of "local heads" and " "global heads" that are responsible for modeling dependencies of different lengths. Finally, in addition to our current results that focus on applying ACP during pretraining, we also discuss how ACP could be used to reduce computation and memory usage for prefilling and decoding during inference, along with preliminary results.
Preliminaries: Forgetting Transformer
This section gives a brief introduction to the Forgetting Transformer and particular its FlashAttention-based implementation. Throughout this work, we follow Yang et al. (2024) and use notation such as A mj and A [m][n]] to index block of a matrix (or a vector).
For example, for matrix E Rlxx and block sizes Bg and By for the two dimensions of A, A [m][n]] E BBxB would be block of A such that (A (Amm]n]) (m - 1) Bg + and = (n-1) B +
Aij, where =
The Forgetting Transformer features a modified softmax attention mechanism with a for- get gate, called Forgetting Attention. Forgetting Attention takes sequence of input vec- tors (x;) and produces sequence of output vectors (o:)/ In addition to the usual query/ /key value projections 4i, k;, D; = WyR; Wax,, Wox; E Rd at each timestep we also compute scalar forget gate = o(w] x: + by) E R, where is the sigmoid function. The
2

---

Published as a conference paper at COLM 2025
output of the attention is then
0 = Ej=1 Fij expla k///d);; Σj= ///Vd Dijoo kj/ v/v + Di;)
(1)
Σj-1 Fij exp(q) k Vd)
where Fij = II -+1 fi and Dij = log Fij = E -i=j+1 log with Fg = 1 and D; = 0 for any i.
This can be written in matrix form:
O = softmax /VV DVV E Rld,
(2)
where E R is the decay bias matrix containing the Dij factors as its lower triangular entries and -oo above its main diagonal. Q,K, V, O E RLxd are matrices containing qi, ki Dü 0;, E [I, as the rows. For multi-head attention with H heads, we maintain ( ] 1=11 and compute the forget gate
H instances of forget gate parameters 1
==== and
values ( 1 separately for each head. We will omit the (h) superscript throughout this work and assume d represents the dimension of each head.
D is coordinate-wise monotone The matrix D has the following property: for any indices x, y such that 2 x and j و y,, we have Dij I Dxy. This is visualized in Figure 1, where darker colors indicate more-negative Dij values. This property, which we call coordinate-wise monotonicity, crucial for developing an efficient pruning algorithm.
FlashAttention implementation of Forgetting Attention The D matrix can be computed as D cl 1c where E R- contains the cumulative sums c; = Σ log E [1,...,L] L. and 1 E R² is a vector of all ones. This makes it possible to implement Forgetting Attention with a simple modification to the FlashAttention algorithm.
We briefly describe the forward pass. In FlashAttention, queries are divided into M blocks {2
[m] E BB xd m==1 with block size Bq L/M. The keys and values are similarly divided
into N blocks [K,, Vnn E BR J==1 with block size B₃ = L/N. All the computations are then conceptually organized into M x N grid, as shown in Figure 1. In standard softmax attention without forget gates, FlashAttention computes the attention logit blocks
[[m[[]] = Q [m] [ /v in the shared memory (SMEM) of the GPU sequentially across the
key/ /value block dimension N and in parallel across the query dimension M (see Figure left). To implement Forgetting Attention, we only need to additionally load [mm] and Cnn]
into SMEM, construct D ml = Cmm) 1 1e// and compute the modified attention logits
[m][[]] = im K / / + D [n n]] The rest of the forward pass remains the same as in
standard FlashAttention The backward pass is implemented similarly.
3 Adaptive Computation Pruning
We now introduce our method, Adaptive Computation Pruning (ACP). Conceptually, ACP aims to prune all the computations in the term exp(q, kj/vá + Dijio if Dij 5, where 6 < 0 is dynamically set threshold (explained later). The attention outputs after pruning are given by:
0, E 1(Drj 2 8) exp(q, k kj/Va+ Dij)oj 15-1(01) 2 8) exp(q, kj Va Díj)
(3)
where 1(.) is the indicator function that takes if the inner proposition is true and otherwise. The intuition of ACP is as follows. Let sy = 9, /d and u be an upper bound of (1syl)iicc Lee U 2 max [sh]-- Since by definition D 0 for any
3

---

Published as a conference paper at COLM 2025
i, if for some j, Dij is much smaller than -2U, then the corresponding attention weight explsy+Dy) exp(sij s/i + Dij) S exp(2U - -D;)) would be very Aij = E exp(sx+ I exp(sx+Da) small, making the contribution of oj to or negligible. And thus the related computations can be safely skipped.
Safe pruning via dynamically set threshold In practice, we set the threshold 5 dynam- ically based on an upper bound U of {[shp] (1..(,,....)) and the sequence length so that the total pruned attention weights Σ [Dj; < Aij for any i would be bounded by small number > 0. Concretely, we set ā 22 log + log e, which achieves the above guarantee (see Appendix for a proof). We set e = e 10 = 0.000045 throughout this work to ensure that the impact of ACP on attention outputs is negligible.'
Setting& dynamically requires us to know an upper bound U of (/st)))fe (1....)- Since
|sip ||a-|a||1;11e we can set U = PaPk vd where Py and Ps are upper bounds of the L2-norms of
vi
the queries and keys respectively. We can either obtain Pa and Pk by explicitly computing
the L2-norms of the queries and keys, or directly derive them from the corresponding normalization parameters if QK-norm (Dehghani et al., 2023) is used (see Appendix E).
Block-level pruning In Fltentin,,,oon, computations are performed in blocks. Conceptu- ally, these blocks of computation are organized into an M N grid shown in Figure 1, where M is the number of query blocks and N is the number of key and value blocks. Therefore, in practice, ACP operates the block level and we prune a computation block (m, n) if and only if all entries in D [m][[]] are below ö, or equivalently, if the maximum entry of D mm [n] is below 8. Since D is coordinate-wise monotone, the maximum entry of [[m[[n] E IR x B (denoted as max(D m] mi) in the following) is simply its top-right entry (D [m][[[]] 1.B; = (m(m)) (cm) D.. Therefore, we only need to check this entry to determine whether a block should be pruned.
Two-stage implementation Since D is coordinate-wise monotone, it is easy to show that S
if max [m][]] < ô then max(D(xly)) < å for any 2 m and y S n. This means that the set of computation blocks to be pruned constitutes a consecutive region on the lower-left part of the M x N grid, as shown in Figure 1 (right). In addition, this region is separated from the rest of the grid by a pruning boundary that connects the top-left corner and the bottom-right corner of the grid, yielding sliding-windowwlik, pruning pattern. Based on this observation, we can perform ACP in two stages. First, we identify the pruning boundary. Specifically, for each row m, we determine the first computation block (m, nm) on the right of the pruning boundary on row m. In Figure (right), these correspond to the blocks at the start of each arrow. After this is done, for each row m, we start the FlashAttention iterations from block (m, nm) (instead of block (m, 1)), and therefore no computations would be wasted on the pruned blocks. Note that if a block is pruned, the kernel also skips the corresponding memory accesses. For example, in the forward pass, block (m, n) is pruned, the thread block corresponding to query block Qim) will not load K1n) and Vnj- Therefore ACP reduces both the number of FLOPs and memory accesses.
1 For reference, the typical relative rounding error for the popular 6 precision is on the order of 0.001.
21 D |m]jj lies on the diagonal of the grid, not pruned by default as it contains an entry Dj for some which by definition is zero.
4

---

Published as a conference paper at COLM 2025
Algorithm 1 Index search for boundary blocks
Require: Cumsum of log forget gates RL, threshold 6, number of query blocks M
Ensure: Nm be the column index of the boundary block on row in for each !!! E {1,..., M]
1: [
2: for m from 1 to M do
3: Dmax &-
4:
while Dmax < å do
5:
Dmax = (c m²] 1 (cn)),, (This is the top-right and the maximum entry of D m m=(())
6:
I 1 1+1 (We loop until (m,1) is a boundary block)
7:
end while
8: Set 1m =
9: end for
Identifying boundary block indices The final missing piece of ACP is an algorithm to identify the column index nı of the boundary block on each row m. Since D is coordinate- wise monotone, for any two such boundary blocks (m,nm) and (x, yx) we have m 2 x < Am 2 yx- This makes it possible to use an efficient linear-time algorithm to identify the boundary block indices, shown in Algorithm 1. This algorithm has a linear complexity of B }2), compared to the O(1-d) quadratic complexity of standard full attention. In practice, we find that boundary index search accounts for only minimal portion of the total attention kernel runtime (around 2% to 6%; see Appendix D).
O(max( L
4 Experiments
In principle, ACP applies to both pretraining and inference (prefilling and decoding). Be- cause pretraining typically saturates available GPUs, reductions in FLOPs or memory accesses achieved by ACP translate directly into shorter wall-clock time. During inference, ACP should deliver similar reductions in FLOPs and memory accesses; however, realizing comparable end-to-end speedups may require additional optimizations to remove other bottlenecks such as kernel-launch overheads. Therefore, this work focuses on the computa- tional savings and wall-clock improvements of ACP for pretraining. We defer a discussion of inference-time ACP and some preliminary results to Appendix F.
4.1 Experimental setup
Throughout this work, we use the FoX (Pro) architecture introduced in Lin et al. (2025). Following Lin et al. (2025), we do not use RoPE (Su et al., 2024). The Pro architecture enhances the basic LLaMA (Touvron et al., 2023) architecture by incorporating some com- mon components in recurrent sequence models such as QK-norm (Dehghani et al., 2023), output gate, output normalization, and data-dependent token-shift (Peng et al., 2024). For completeness, we also provide results for the FoX (LLaMA) architecture in Appendix C.2, which are similar to the results for FoX (Pro) that we present below.
We train FoX (Pro) models with and without ACP on LongCraw164 (Buckman, 2024) using the standard language modeling objective. We adopt the three training configurations used in the analysis experiments Lin et al. (2025), specified as combinations of the number of model parameters and the number of training tokens: (00M-paramterr166--1 116--tkeen, 360M- parameter/ /7.5B-tokenn and 125M 355M-aaaaeter/ /..B-tokee. For each scale, we train the models with three training context lengths: 4k, 8k, and 16k tokens. The rest of the hyperparameters are the same as those in Lin et al. (2025) and are described in detail in Appendix B.
We use the official Forgetting Transformer repository for the implementation. We imple- ment ACP, including the boundary index search algorithm, on top of the official Forgetting Attention kernel in Triton (OpenAl, 2021).
5

---

Published as a conference paper at COLM 2025
PLOP memory--ccces reduction
80%
75%
77%
73%
72%
76%78%
Attention runtime roduction
62%
62%
67%
70%
60%
60%
51%
2,0002 155
39.59
01
40%
5300 12
4.1/15500
9,74.99
3::3 3,052
1x000 13
20%
20%
0%
0%
125M
360M
760M
125M
360M
760M
Scale (number of parameters)
Scale (number parameters)
Training context length
Training context length
4096
8192
16384
4096
8192
16384
45%
44%
43%
Throughput improvement
40%
30%
24%
2000 196
24%
23%
20%
12%
2000
12%
0%
125M
360M
760M
Scale (number of parameters)
Training context length
4096
8192
16384
Figure 2: (left) Percentage reduction in FLOPs and memory accesses in the attention op- eration due to ACP. (right) P'eccetaage reduction in attention kernel runtime due to ACP. Within each bar we also show the actual runtime with and without ACP milliseconds. The runtime covers one forward and backward pass on batch of 0.5M tokens. (bottom) Percentage training throughput improvement due to ACP. Within each bar we also show the actual values of training throughput with and without ACP. Throughput is measured in tokens per second. Both the attention kernel runtime and throughput are measured on NVIDIA L40S GPUs.
In the following, training throughputs are measured using the final checkpoints on subset of the heldout set of LongCraw164 on 4 NVIDIA L40S GPUs. We find that training throughput typically decreases for a short period the beginning of training and then plateaus, so our reported numbers using the final checkpoints reflect the throughput during the plateau period. When ignoring sub-leading terms, the percentage reduction FLOP' and memory accesses in the attention operation number pruned blocks to the total number of blocks in the FlashAttention grid. We compute this ratio subset of the heldout set of LongCraw164 More details can be found in Appendix
be approximated by the ratio of the
4.2 Computational savings and speedups
In Figure we show the percentage reduction FLOPs and memory accesses in the at- tention operation, the percentage reduction in attention kernel runtime, and the percentage improvement in training throughput due to ACP, across different model sizes and training context lengths. As shown in Figure 2, ACP consistently prunes around 70% of the FLOPs

---

Published as a conference paper at COLM 2025
750M parameters/ 168 triining tokeus
FoX (Pro) ACP, L 166
FoX (Pro) 1166
1.80
0.3
1000
100
1.75
70.9
Depih Percent
30.0
50.9
50.0
1.70
+0.0
A01
90.0
N0.0
X010
1.53
100,3
100.0
512
1024
2048
index
4096
16384
11660
Model
Training contedt length
Document Length
Document Length
FoX (Pto) wi ACP
4096
16384
Score
F6X (Pto) wio ACP
8192
10
Figure 3: (left) Per-token loss given different training context lengths for the 760M- parameter/16B3 116-tokkn setting. This is measured on 2B-token validation set of the LongCraw164. At each token index i, we report the averaged loss over a window of 101 centered at (right) Easy-mode needle-in-a-haystee models with training context length of L = 16k tokens.
results for the 760M-parameter
Model
Wiki ppl]
LMB.
LMB.
Hella. acc-n?
Wine. acc
ARC-e
ARC-c acc-n]
COPA
OBQA
SeQQ
BoolQ
Arg
ppl.
a0c
acc"
3007
acc"
acc-n1
acc1
acc]
FoX (Pro) w/ ACT,
29.66
22.12
37.57
63.11
33.39
48.91
24.66
29.20
79.90
57.16
49.45
FoX (Pro) ACP,
29.98
22.52
37.65
62.84
33.38
52.72
47.94
25.60
67.00
29.00
79.70
54.22
69.06
Fox (Pro) w/ ACP, 1..
28.04
23.20
38.15
60.94
33.46
51.70
48.82
24.66
67.00
28.00
80.00
60.12
29.34
Fox (Pro)//o ACP, 1.
8k
28.07
22.53
38.31
61.11
33,83
50.67
49.28
24.83
69.00
27 40
80.50
61.59
49.75
FoX (Pro) w/ ACT, L=
16k.
27.96
25.16
35.77
62.35
79
50.83
48 02
24.23
69.00
28 20
79.30
58.93
49.06
Fox (Pro) w/o ACR
162
28.04
36.66
62.35
33.32
48.86
25.51
71.00
27.20
82.20
56.76
Table 1: Evaluation results on LM-eval-hamess. All models have roughly 760M non-
embedding parameters and are trained on roughly 16B tokens on LongCrawl64. "acc-n" means length-normaiized accuracy. I. is the training context length.
and memory accesses in softmax attention in all cases, resulting in roughly 50% to 70% reduction in attention runtime (or a 2-3x speedup). These translate into a roughly 10% to 40% increase in end-to-ent training throughput. Note that ACP only affects the speed of the attention kernel, whereas training throughput also depends on the latency of other components such as MLPs and RMSNorms. In particular, longer training context lengths lead to greater throughput improvements, because the proportion of FLOPs and memory accesses in softmax attention increases relative to the rest of the network as context length grows. For example, for 100M-paramete) model with context length of 4k tokens, the attention operation accounts for roughly 16% of the total FLOPs of the model, while for context length of 16k tokens, it accounts for around 45%.
ACP does not damage performance In Figure 3 (left) we show the language modeling loss at different token positions for the 760M-parameter FoX (Pro) models with different training context lengths, with and without ACP. Figure (right) shows the needle-in-a-haystack retrieval results of the 16k-context-length model in Figure 3 (left), following the "easy mode" setup used in Lin et al. (2025) that is suitable for small models without instruction-tuning, Table shows the evaluation results on various downstream tasks from Language Model Evaluation Harness (Gao et al., 2024a) for the models in Figure 3 (left). Additional results can be found in Appendix C.
As shown in these results, the per-token language modeling loss curves with and without ACP almost match exactly (the slight difference within the expected variance across runs). ACP also does not damage long-context retrieval performance, and the downstream task performances of models with and without ACP are similar. Note that is well known that evaluation results runs (Madaan et al., 2024), so it is impossible to obtain exactly the same results even when training the same model with different seeds.
downstream tasks can exhibit high variance across training
7

---

Published as a conference paper at COLM 2025
50%
Percentage of
40%
30%
20%
10%
0%
0%
20%
40%
60%
80%
100%
FLOP & memory-accesss reduction
Figure Distribution of per-head computational savings in a 760M-parameter FoX (Pro) model with 4k training context length. Specifically, we divide percentage computational savings into 20 bins 5%), 55,, 10%), 95%, 100%], and for each bin we count the number of heads in the model whose percentage of pruned attention FLOPs and memory accesses falls into that bin. The counts are then normalized to obtain a distribution.
F1,0P memorry-aneee -eeeoe reluction
100%
80%
40%
20%
0%
Layer
Laye:
Layer
Laye
Layer
Layer
Layer
Layer
Layer
Tayer
Tayer
Toyer 13
Tayer
Tayer
Layer
Layer
Toyer
Layer
Layer 20
Layer 21
Tuaye 20
Tayer 23
Layer 24
Figure 5: Distribution of per-head computational savings in each layer. Each column can be seen as a 90%-rotated (and flipped) version of Figure 4, except the distribution is calculated within each layer. The x-axis of each column is the percentage heads the corresponding layer whose percentage of pruned FLOPs and memory accesses falls within a specific bin. The range of the x-axis of each column from 0% 100%.
4.3 Analyses
In this section, we perform series of analyses to provide deeper insight into method. First, we show the distribution of computational savings across different attention heads. Second, visualize the pruning boundaries in some heads. Finally, investigate how computational savings and model performance vary with e, the hyperparameter that bounds the total pruned attention weights.
Distribution of per-head computational savings In Figure 4, we show the distribution of per-head computational savings in 760M-parameter FoX (Pro) model with a context length of 4k tokens, over the set of all attention heads the model. Figure 4 shows a clear bimodal pattern, and most attention heads are either "local heads" (most computations are pruned) or "global heads" (only small proportion or none of the computations are pruned). Furthermore, a majority of the heads are local heads, consistent with the significant FLOP and memory-access savings shown in Figure 2. In Figure we also show the distribution of per-head savings within each layer. In general, the distribution for each layer matches the distribution for the entire model, except for the first two layers where all the heads are local.
Visualization of pruning boundaries In Figure 6 we show the decay bias matrices D and the attention weight matrices A from three heads in different layers. We also show the
8

---

Published as a conference paper at COLM 2025
Layer 24, Head 1,
Layer 16, Head 4,
Layer 8, Head 4, D
0
0k
0
Query Posiion
Query Posiion
25
Query Position
25
2k
2k
2k
75
75
75
4k
100
4k
100
4k
100
0k
2k
2k Key Position
Key Position
Key Position
Layer 24, Head 1,
Layer 16, Head 4,
Layer 8, Head 4
0k
0k
1.00
0k
Query Position
-0.55
Query Position
0.75
Query Position
0.75
2k
-0.50
2k
0.50
2k
0.50
0.25
0.25
0.25
0.00
0.00
4k
0.00
0k
2k
4k
0k
2k Key Position
2k
4k
Key Position
Key Position
Figure Visualization of the decay matrices D (top row) and the corresponding attention weight matrices A (bottom row) from three heads in different layers. The orange line shows the pruning boundary. Since A is very sparse, we only show entries with scores larger than 0.1. These results use 760M-aara FoX (Pro) model with context length of 4k tokens.
reduction
80%
loss: 2.286
Joss: 2.285
Joss: 2.286
aneenury-aecess access
60%
loss: 2.286
FLOP & memory
40%
20%
loss: 2.285
1000
10000
Value of
Figure 7: Impact of : on FLOP savings for 225--parmmeter model with a training context length of 16k tokens. For each data point we also label the corresponding validation loss.
pruning boundaries on the matrices. The heads on the left and middle are local heads with strong decay, and most off-diagonal blocks are pruned. The rightmost head is typical global head where no blocks are pruned.
Effect of varying In Figure show the impact of : the hyperparameter controlling the maximum total attention weights that can be pruned on computational savings and language modeling loss. As expected, with smaller the computational savings decrease. On the other hand, there is only marginal gain one uses a larger (e.g., e Rs 0.37, which might be unsafe) than our default value -11 0.000045. Therefore, we recommend future work to adopt our default e = -10 as it ensures safe pruning while achieving near-optimal computational savings.

---

Published as a conference paper at COLM 2025
5 Related work
Dynamic locality-based computation pruning The most similar methods to ours are context pruning in Selective Attention (Leviathan et al., 2024) and conditional computation in stick-breaking attention (Ian et al., 2024). Similar to FoX, both Selective Attention and stick-breaking attention learn some forms of data-dependent decay, and thus dynamic pruning similar to our ACP is possible. For Selective Attention, this is done at inference time by maintaining mixed memory budget and dropping KV-cache entries that have the strongest decay. However, it is unclear how this can be adapted for training, like what we do in this work with ACP. For stick-breaking attention, this is done by early stopping the stick-breaking process for each query until all attention weights have been assigned. Although for ttick-breaking attention conditional computation can also be used in training, Tan et al. (2024) only investigate applying it during inference, so it is unclear how much speed improvement can be obtained when it is applied during training.
Sliding-window-based computation pruning Methods such as StreamingLLM (Xiao et al., 2024c), LM-Infinite (Han al., 2023), MoA (Fu al., 2024a), and DuoAttention (Xiao et al., 2024b) apply sliding window mask to pretrained models at inference time to reduce computational costs. This approach is also frequently used in KV-cache eviction methods, and is often combined with some importance-based eviction policy (Zhang et al., 2023; Liu et al., 2023; Ge et al., 2024; Oren et al., 2024; Fu et al., 2024b). With ACP, local heads behave similarly to sliding-window attention. However, unlike these related methods where the window size is typically fixed or based on profiling on some dataset, the "window size" of a local head in ACP is determined by the decay bias matrix and the dynamically set threshold, which guarantees that the total attention weights beyond the local window are negligible.
Sparse attention Another category of computation pruning methods exploits the sparsity of softmax attention. These methods mainly differ in how they evaluate the importance of different KV-cache entries based on queries. Most sparse attention methods divide the KV cache into blocks, calculate a summary of each block, and then compute the impor- tance scores using these block summaries (Tang et al., 2024; Xiao et al., 2024a; Gao et al., 2024b; Yuan et al., 2025; Lu et al., 2025). There also exist token-level methods (Desai et al., al., 2023) and more sophisticated methods such as cluster-based method (Liu et al., 2024) or a mixture of different sparse attention methods (Jiang et al., 2024). These are orthogonal to our locality-based approach, and is likely that they can be combined with ACP.
2024; Anagnostidis
6 Conclusion
We propose Adaptive Computation Pruning (ACP) for the Forgetting Transformer (FoX), method that dynamically prunes computations involving input-output dependencies that are strongly decayed by the forget gate in FoX, based on a dynamically set threshold value that ensures negligible impact on the attention output. We apply ACP to language model pretraining and find it leads to significant computational savings and speedups, without sacrificing model performance.
Even though this work primarily focuses on applying ACP during pretraining, we also discuss its potential for inference and present promising preliminary results in the appendix. In particular, for decoding, KV-cache entries could be dynamically evicted based on the pruning boundary, reducing both memory consumption and memory accesses. A more thorough investigation of inference-time ACP is left to future work.
Acknowledgments
ZL thanks Shawn Tan and Songlin Yang for their helpful discussion. AC acknowledges funding from Microsoft Research. This research was enabled in part by the compute resources, software, and technical help provided by Mila (mila. quebec) and the Digital Research Alliance of Canada (al (alliance. waa.innnnccccc... ca).
10

---

Published as a conference paper at COLM 2025
References
Sotiris Anagnostidis, Dario Pavllo, Luca Biggio, Lorenzo Noci, Aurelien Lucchi, and Thomas Hofmann. Dynamic context pruning for efficient and interpretable autoregressive trans- formers. Advances in Neural Information Processing Systems, 36:65202-65223,, 2023.
Jacob Buckman. Longcraw164: A Long-Context Natural-Language Dataset, 2024. URL httts::/mmaifestaaal htw.:/estaa.commaaiiiiiicccoao www.///aaiileeeonnerrrrreeee
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, In The Twelfth International Conference on Learning Representations, 2024.
Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pp- 7480-7512. PMLR, 2023.
Aditya Desai, Shuo Yang, Alejandro Cuadron, Ana Klimovic, Matei Zaharia, Joseph E Gonzalez, and Ion Stoica. Hashattention: Semantic sparsity for faster inference. arviv preprint ar Xiv:2412. 14468, 2024.
Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, et al. Moa: Mixture of sparse attention for automatic large language model compression. CoRR, 2024a.
Yu Fu, Zefan Cai, Abedelkadir Asi, Wayne Xiong, Yue Dong, and Wen Xiao. Not all heads matter: A head-level kv cache compression method with integrated retrieval and reasoning. ar preprint ar ar::::2410 2024b.
Leo Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. framework for few-shot language model evaluation, 07 2024a. URL https://zenodo wrg/ 2628602.
Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. Secrattention: Learning intrinsic sparse attention in your llms. arXiv preprint ar arXiv:2410.13276, 2024b.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. In The Twelfth International Conference on Learning Representations, 2024.
Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Zero-shot extreme length generalization for large language models. arviv preprint XrX:::230..16137 2023.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin-Yew Yew Lin, et al. Minference 1.0: Acceler- ating pre-filling for long-context lims via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-52515, 2024.
Gregory Kamradt, 2023. URL https: com/ htt.://wwaattttt/ttteeeeeeeeeeea blob/main/README, md.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Selective attention improves transformer. ariiv preprint ar arXiv:2410.02703, 2024.
Zhixuan Lin, Evgenii Nikishin, Xu He, and Aaron Courville. Forgetting transformer: Softmax attention with forget gate. In The Thirteenth International Conference on Learning Representations, 2025. URL
11

---

Published as a conference paper at COLM 2025
Guangda Liu, Chengwei Li, Jieru Zhao, Chenqi Zhang, and Minyi Guo. Clusterkv: Ma- nipulating Ilm kv cache in semantic space for recallable compression. arXio preprint (rXiv:2412.03213,, 2024.
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anas- tasios Kyrilliidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for Ilm kv cache compression test time. Advances in Neural Information Processing Systems, 36:52342 36:52342-52364, 52364, 2023.
I Loshchilov. Decoupled weight decay regularization. arviv preprint arX:::1711.01101, 2017.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, et al. Moba: Mixture of block attention for long-context llms. Xiv preprint arXiv:2502.1318) 2025.
Lovish Madaan, Aaditya Singh, Rylan Schaeffer, Andrew Poulton, Sanmi Koyejo, Pontus Stenetorp, Sharan Narang, and Dieuwke Hupkes. Quantifying variance in evaluation benchmarks. arviv preprint arXiv:2406.10223) 2024.
OpenAT, 2021. URL htw.t:
Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, and Roy Schwartz. Transformers multi-state rnns. arviv preprint arXiv:2401.06104, 2024.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemyslaw Kazienko, al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. In First Conference on Language Modeling, 2024.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Nenrocomputingg 568:127063, 2024.
Shawn Tan, Yikang Shen, Songlin Yang, Aaron Courville, and Rameswar Panda. Stick- breaking attention. ar preprint XrX:v:2410.1770 179900,,, 2024.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context 1lm inference. International Conference on Machine Learning, pp. 47901-47911. PMLR, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXio:2302.13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL https://proceedings. neurips. pdf.
file/8ffee243347deee11 1
fbd053c1 (4003aaa-11444aaaaaaaaa
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. InfLLM: Training-free long-context extrapolation for LLMs with an efficient context memory. In The Thirth- Annual Conference on Neural In- formation Processing Systems, 2024a. URL https: / httppr//////peevvieeeeee net fert/foormm? bb rqhASS.
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context 1lm inference with retrieval and streaming heads. arXiv preprint ar xix::2410..088/ 2024b.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient stream- ing language models with attention sinks. The Twelfth International Conference on Learning Representations, 2024c.
12

---

Published as a conference paper at COLM 2025
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear trans- formers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: aareware- and natively trainable sparse attention. arXiv preprint arXiv:2502.11089 2025.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-34710, 10, 2023.
13

---

Published as a conference paper at COLM 2025
A Proof of upper bound of total pruned attention weights
In this section we prove that when the threshold is properly set, the total pruned attention weights _ 1 < Af would be bounded by small number g.
Let sij = 9: Va and u be an upper bound of (|s (L..(.....). i.. u 2 maxi,e Isyl-
Let L. be the sequence length. If we set the threshold to = -2U - log I. + log e, then for any i and j such that Dij < 6, we have that (note that D;: = by definition):
exp(sij + Dij)
exp(si + Dij)
Ai E 1 exp Sit D exp(s; + Dj)
exp(si: (sij - S/ + D1j))
(4)
$ exp(Isij -Sa + Dij) I exp(2U Dij)) S exp(2U 2u - log L + loge)
(5)
(6)
L
Therefore, we have D(j < {.jAi < f. for any and and E/- 1(D;; < 83A <e.
B Experimental details
Configuration
layers
model
head
Peak learning rate
760M params / 16B tokens
24
1536
64
1 x 10-3
360M params 7.5B tokens
24
1024
64
2 x 10-3
125M params / 2.7B tokens
12
768
64
2 10-3
Table 2: Hyperparameters for different configurations. Mayer counts the number of blocks, where each block contains an attention layer and an SwiGLU layer..
Our pretraining hyperparameters follow the setup used in the analysis experiments in Lin al. (2025). We list the hyperparameters for different training configurations used in this work in Table 2. All models are trained with Adam (Loshchilov, 2017) with (B1, (0.9, 0.95), with linear learning rate warmup from the peak learning rate for the first 256 22 tokens and then a cosine decay 0. Each training batch contains 0.5 22 tokens. All models use weight decay of 0.1 and gradient clipping of 1.0. We follow the HuggingFace LLaMA initialization and initialize parameters with (0, 0.02²). ). We do not share the parameters between the embedding layer and the output layer. Weight decay is not applied to the RMSNorm parameters and bias terms in linear layers (only the forget gate projection has a bias term). We use bfloat mixed-precision training for all models.
linear layer weights and embedding
FLOP and memory-access reductions are measured on a 128M-token subset of the LongCraw164 heldout set. This is calculated as the ratio of the number of pruned blocks to the total number of blocks in the FlashAttention grid. In principle, the FLOP and memory- access savings for the forward pass and the backward pass are different due to different FlashAttention block sizes and thus different grid granularity. However, in practice, the block sizes (smaller than 128) are much smaller than the sequence length (larger than 4k), so the difference is negligible. Therefore, throughout this work, we report the FLOP and memory-access savings for the forward pass.
A Attention runtime and training throughput are measured on a 32M-token subset of the same heldout set. The reported attention runtime includes single forward pass and single backward pass, measured using torch. corch.cuda.... Events. .uda.Events. When ACP is used, the runtime includes the time for boundary index search. When measuring throughput, we use gradient checkpointing and gradient accumulation. Each gradient accumulation step processes 32k tokens. Throughput is measured on NVIDIA L40S GPUs with fully sharded data parallel. The power limit of these GPUs is set to 325W.
14

---

Published as a conference paper at COLM 2025
125M parameters/2.7B training tokens
350M parameters77.5 training tokens
2.40
2.00
Loss L.0)
2.35
Loss L())
1.95
2.30
1.90
2.25
1.85
512
1024
2048
4096
8192
16384
512
1024
2048
4096
8192
16384
index
Token index
Model
Training context length
Model
Training context length
FoX (Pro) w/ACP
4096
16384
FoX (Pro) w/ W/CP
4096
16384
FoX (Pro) w/o ACP
8192
FoX (Pro) w/ ACP
8192
Figure 8: (left) Per-token loss given different training context lengths for the 125M- parameter/
2.7-tok 1anameter/2.B--token 2B-token validation set of the LongCraw164. At each token index i, we report the averaged
and 360M-parameter/7.5. 77.5ttttr/r.Bttokenn setting, This is measured on
loss over a window 101 centered at
FoX (Pro) w/o ACP, L 4k
FoX (Pro) w/ ACP, L = 4k
10.0
26.0
20.0
Dopth Pecoent
40.0
40.0
50.0
40.0
40.0
$0.0
$0.0
$0.0
100.0
100.0
Document Length
Document Length
Score
10
FoX (Pro) wo ACP, L= 8k
FoX (Pro) w/ ACP, L =8
3.9
10.0
100
Depih Percent
20.0
30.0
Dopth Percent
30.0
60.0
60.0
80.0
80.1
100.0
1020
Document Length
Document Length
Score
2
10
Figure 9: Easy-mode aaalle-in-a-haystade results for the 760M-parameter models with training context lengths of 4k and 8k tokens.
C Additional results
C.1 Additional FoX (Pro) results
In Figure 8 we show the per-token loss for the 12MM-prrameter/27Btoken and 360M- parameter/7.5B-ttken settings for FoX (Pro) with and without ACP, in addition to the 760M-paaraeeter/6633 166-ttkkrr setting in Figure 3 (left). In Figure 9 we show the mode needle-in-a-haystacd results for models trained with context lengths of 4k and 8k tokens, respectively, in addition to the 16k-context-length results in Figure 3 (right).
15

---

Published as a conference paper at COLM 2025
PLOP memory-acccs reduction
80%
72%
78%
65%
69%
73%
65%
71%
74%
Attention runtime roduction
60%
64%
58%
64%
60%
66%
55%
60%
50%
52%
40%
6.8522200
10.90/5033
(6V22-4--22
40%
198.3
34
+2101
6,994433
10.50220.5
3.7
333
20%
20%
0%
0%
125M
360M
760M
125M
360M
760M
Scale (number of parameters)
Scale (number parameters)
Training context length
Training context length
4096
8192
16384
4096
8192
16384
50%
Throughput improvement
44%
40%
40%
30%
28%
23/3/ 1896
23%
1004-71
21%
20%
13%
11%
11%
10%
0%
125M
360M
760M
Scale (number of parameters)
Training context length
4096
8192
16384
Figure 10: Computational saving and speedup results for FoX (LLaMA). (left) Percentage reduction in FLOPs and memory accesses in the attention operation due to ACP. (right) Percentage reduction in attention kernel runtime due to ACP. Within each bar we also show the actual runtime with and without ACP in milliseconds. The runtime covers one forward and backward pass on batch of 0.5M tokens. (bottom) Percentage training throughput improvement due to ACP. Within each bar we also show the actual values of training throughput with and without ACP. Throughput measured in tokens per second. Both the attention kernel runtime and throughput are measured on 4 NVIDIA L40S GPUs.
C.2 FoX (LLaMA) results
In this section, we present results for the FoX (LLaMA) architecture, in addition the FoX (Pro) results in the main text.
In Figure 10, we show the percentage reduction in FLOPs and memory accesses in the atlention operation, the percentage reduction in attention kernel runtime, and the percentage improvement in training throughput due to ACP, across different model sizes and training context lengths, using the FoX (LLaMA) architecture.
In Figure 11 (left) we show the language modeling loss at different token positions for the 760M-parameter FoX (LLaMA) models with different training context lengths, with and without ACP, Figure 11 (right) shows the -eedle-in---haystack retrieval results of the 16k- context-length model in Figure 11 (left), following the "easy mode" setup used in Lin al. (2025). Table 1 shows the evaluation results on various downstream tasks from Language Model Evaluation Harness (Gao et al., 2024a) for the models in Figure 11 (left).
16

---

Published as a conference paper at COLM 2025
FoX (LLaMA) w/o ACP, 16k
FoX (LLaMA) -= 164
parameters
10.0
200
Depth Peroent
30.0
128
1.50
50.0
50.0
60.0
1.75
70,0
930
80.0
1,70
1030
Training context lengh
Document Length
Document Length
FoX (LLaMA) ACP
4096
16384
FoX (LLaMA) W/OACP
9192
10
Figure 11: FoX (LLaMA) evlauation results. (left) Per-token loss given different training context lengths for the 760M-parameter Saeaanntterrr/--ootle 11B-ttkeen setting. This is measured on 2B-token validation set of the LongCraw164. At each token index i, we report the averaged loss over window of 101 centered at (right) Easy-mode needle-in-a-haystae results for the (60M-parameter models with a training context length = 16k tokens.
Model
Wiki
1ME. ppli
1MB.
PIQAA
Hells.
Winn.
ARC-
ARC-
COPA
OBQA
SriOA
BoolQ
Avg
ppli
3027
a3c
occT
002"
acc-n
s001
30:.0"
7
fot-llama-too---ms fiit---------------------------- te---end-----ee- 4k-77Mmm
32.00
20 42
33-36
51.37
37.46
48.78
47.23
24.6
66:33
77.90
61:22
lo
760m
50.91
21.00
63.00
hn-lllam--hot---me-- bec--end---ee--- fot
58-5mmmmmmmmmmm
len 700m
27 0666
50.43
23.95
65.00
56.88
08.15
int-llamu--ii--in 10-ers ind-emn--000m -enl--0amm- 180m
31.16
28.67
$5.09
37
31.76
50336
47.25
24.40
66.00
28.60
72.00
69.31
48.25
lot nit--llaaaaaa hid oos lo and d len: 1ek: Potmm
31.43
28.41
34,89
61.21
32.27
90.51
66.65
21.06
67.00
2960
77.30
62.07
08.35
Table 3: FoX (LLaMA) evaluation results on LM-eval-harness. All models have roughly 760M non-embedding parameters and are trained roughly 16B tokens on LongCraw164. "acc-n" means length-normalized accuracy. L is the training context length.
D Computational costs of the boundary search algorithm
As discussed in Section 3, the boundary index search algorithm has a linear complexity of O(mxx( L ' Ti 1,),, compared to the O(L'd) quadratic complexity of standard full attention. Note that even though this algorithm runs sequentially, in practice it still has negligible wall-clock time compared to actual attention computations, mainly because FlashAttention also runs sequentially within each thread block with similar number of iteration steps as this algorithm.
In Table we report the percentage of wall clock time spent on boundary index search within the attention kernel, i.e., boundary search time divided by the total runtime of the attention kernel. Note the total runtime of the attention kernel includes the time for boundary index search. As shown in this table, the computational costs of boundary index search are minimal.
125M
350M
760M
L 4096
4%
3%
2%
L 8192
5%
4%
3%
L 16384
6%
4%
3%
Table 4: Percentage of wall-clock time spent on boundary index search within the attention kernel given different model sizes and sequence lengths L.
17

---

Published as a conference paper at COLM 2025
E
Obtaining an attention logit upper bound from QK-norm parameters
When QK-norm is used (assuming we use RMSNorm (Zhang & Sennrich, 2019), as in the FoX (Pro) architecture), the L2-norms of queries and keys are bounded by yLvā and y v respectively, where 7 (L..ad ln is the maximum magnitude of the key RMSNorm
scaling parameters {1) ld i=1 and 7 is defined similarly. Therefore |styl _ ||eilla|kkylle I vdva vd
and thus we can set = 7/ Vd.
F Inference-time ACP
In this section we discuss how ACP can be used at inference time and present some prelimi- nary results. Applying ACP prefilling is straightforward, we mainly discuss ACP for decoding.
For decoding, due to the monotonicity of the pruning boundary, we could maintain pruning boundary index j for each head and update it in an online fashion. Specifically, whenever Dij < 6- where i is the current timestep we increment until Djj 2 Since j never decreases, we can discard any KV-cache entries beyond the pruning boundary, thus reducing memory consumption and memory accesses.
In our preliminary results with naive implementation that does not perform explicit KV cache eviction (but still skips loading pruned blocks to shared memory), applying ACP during inference achieves the same level (around 70%) reduction in memory accesses and FLOPs. This is expected as these savings are opplementation-agnosttt our analysis shows that our implementation overheads during decoding, and wall-clock time improvement is most obvious with long prefilling lengths and large batch sizes. Specifically, with a sufficiently large prefilling length and batch size, ACP reduces the per-step attention kernel runtime by 50% to 60%, which is similar to the 50% to 70% reduction we see during pretraining. However, even in this setup, do not see clear improvement in end-to-end decoding throughput. Our analysis shows this is likely because the kernel execution overlaps with the significant Triton kernel launch overheads (many components in our implementation such RMSNorm are in Triton). Therefore, reducing the attention kernel runtime has little effect on end-to-end throughput, as the throughput would still be bottlenecked by these kernel launch overheads.
However, likely bottlenecked by the kernel launch
We comment that kernel launch overheads could be largely mitigated with optimization (e.g., CUDA graph), so we expect that with properly optimized implementation, ACP should be able to bring the same level of speedup to LLM decoding as it does to pretraining, especially given that the percentage reductions in memory accesses brought by ACP are similar during pretraining and decoding.
18


Reviews:
[{"review_id": "falN86nxwZ", "reviewer": "Reviewer_KGU7", "summary": "The paper introduce Adaptive Computation Pruning (ACP) for the Forgetting Transformer, using the decay bias D_ij to skip entire query–key blocks whose maximum value falls below a theoretically derived threshold delta = −2U − log L + log epsilon, which guarantees that no more than an epsilon fraction of attention mass is discarded.\n\nACP is implemented as a two-stage routine that first finds a monotone block boundary in O(MN) time and then runs FlashAttention only on the retained region; experiments on 125 M–760 M-parameter models and context lengths up to 16 k show 70–79 percent attention-FLOP savings and 10–40 percent wall-clock speed-ups without much loss on language-modeling and retrieval benchmarks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The pruning threshold is mathematically solid, with a compact proof that bounds the discarded attention mass and thus preserves model fidelity; ACP operates at the FlashAttention block granularity, aligning with GPU memory-coalescing patterns and requiring only minimal Triton changes for deployment. Experiments demonstrate large training-time throughput gains without accuracy loss and provide insightful visual diagnostics that explain why certain head become almost completely local.", "reasons_to_reject": "The method is tightly coupled to the FoX decay bias and its generalizability to standard RoPE transformers or other positional encodings is not very clear; all reported results concern training, leaving potential inference-time benefits during autoregressive decoding—where KV-cache bandwidth dominates. scalability is only tested up to 760 M parameters on a single 4 × L40S node, so behavior on multi-billion-parameter models and distributed settings remains unknown.", "questions_to_authors": "Can you quantify end-to-end speed-ups and memory savings during pre-fill and autoregressive decoding when ACP is applied to KV-cache eviction, and how do these gains compare with the training-time numbers? \n\nWhat fraction of total attention runtime is spent on the boundary-search algorithm at sequence length 16 k on L40S versus A100 GPUs", "extra_scores": {}}, {"review_id": "b6kVyw3YaC", "reviewer": "Reviewer_tEnv", "summary": "Adaptive Computation Pruning (ACP) introduces a block‐wise attention pruning mechanism that dynamically discards negligible attention scores during training to reduce computation without sacrificing model quality. By partitioning the attention matrix into fixed-size blocks, ACP computes for each block a tight bound on its maximum contribution (using precomputed norms U) and derives a threshold δ that guarantees any value below it has a provably negligible impact on the softmax output. During each forward pass, a simple boundary‐search algorithm identifies and skips entire blocks whose scores fall below δ, yielding substantial speedups. ACP integrates seamlessly into FoX‐style transformer pretraining, provides theoretical guarantees on output fidelity, and demonstrates significant efficiency gains with minimal accuracy loss across standard benchmarks.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Clearly motivated problem with provable guarantees.", "reasons_to_reject": "Although the authors claim that identifying boundary indices is “negligible” compared to attention computation, there is no quantitative complexity analysis or profiling to substantiate this .\n\n- The threshold \\delta = –2U – log L + log \\epsilon seems to derive directly from bounding softmax tails by exp(2U + Dij)≤\\epsilon/L, a technique common in sparse-attention analyses.\n\n- Although the authors generalize the method to block-level algorithms, the approach remains a straightforward application of existing techniques.\n\n- While Fig 7 shows pruning percentages for varying \\epsilon, there is no guidance on how to choose \\epsilon in novel tasks or architectures, leaving unclear about trade-offs between speed and accuracy .\n\n- There is no discussion of how choices of block sizes (Bq, Bk) and resulting grid dimensions (M, N) affect the granularity of pruning versus the overhead of boundary detection .", "questions_to_authors": "Could you discuss how the choice of block sizes (Bq, Bk) and the resulting grid dimensions (M, N) influence the trade-off between pruning granularity and the overhead associated with boundary detection?", "extra_scores": {}}, {"review_id": "GVOmcYHy3y", "reviewer": "Reviewer_s3o6", "summary": "This paper observes that many attention heads in the FoX (Forgetting Transformer) architecture tend to forget information rapidly, resulting in outputs that primarily depend on local context at each timestep. Motivated by this, the authors propose Adaptive Computation Pruning (ACP), a method that dynamically prunes input-output dependencies in attention computation when they are significantly suppressed by the forget gate. This is done using a dynamic pruning threshold to ensure that the removed attention weights remain negligible.\n\nACP is applied to FoX-based language model pretraining and demonstrates consistent improvements: it reduces softmax attention FLOPs by over 70% across varying model sizes and context lengths, leading to a 10%–40% increase in training throughput. The benefit becomes more pronounced as the context length increases.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper is clearly written and well-organized.\n2. The idea of adaptive pruning for *Forgetting Transformer* is interesting and practically motivated.", "reasons_to_reject": "1. The proposed method essentially combines the Forgetting Transformer's attention mechanism with an adaptive sliding-window strategy. This idea is fairly common and lacks strong novelty.\n2. The technique appears to be generally applicable to standard multi-head self-attention as well, not necessarily tied to the Forgetting Transformer. It is unclear whether the authors have conducted any experiments to explore this applicability.\n3. The computational overhead of the method is not thoroughly discussed. Since pruning is performed at the block level and requires determining block-wise indices, it may not be computationally lightweight. Moreover, it does not appear to guarantee globally optimal pruning, raising concerns about its efficiency compared to other pruning strategies.\n4. The experimental section lacks comparative baselines. The results only compare the proposed method to unpruned FoX, making it difficult to assess the overall competitiveness or practical advantage of ACP. It's possible that its performance remains below that of existing pruning methods.\n5. Some reported results in this paper appear questionable. The FoX (Pro) result on HellaSwag is around 33 in this paper, while the original FoX paper reports a score of 38.39. It's unclear why a different dataset size was used. Furthermore, the original FoX paper provided evaluation on LongBench, which is not included here, missing the way for a more comprehensive and fair comparison.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for improving the training efficiency of the Forgetting Transformer (FoX). The method leverages the special properties of the FoX decay bias matrix and dynamically prunes computations involving input-output dependencies that are strongly decayed by the forget gate. Experiments on language model pre-training with FoX demonstrate a significant reduction in FLOPs (up to 760M parameters and 16B tokens). The approach is technically sound, with provable guarantees, and the empirical results are solid—I do not consider the lack of larger-scale experiments a significant concern in this case. The authors have done a thorough job during the rebuttal, especially in addressing boundary index search and attention runtime, and all reviewers remained positive afterward. Therefore, I recommend acceptance of the paper.

Finally, I’m surprised that no reviewer has commented on the title of the paper. I believe it would be clearer to include terms like pre-training or training in the title, as the paper focuses on improving training efficiency. The word "pruning" may suggest an inference-time technique otherwise.
Decision: Accept