Adaptive Computation Pruning for the Forgetting Transformer
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 the 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. This is achieved using a dynamically set pruning threshold
that ensures that the pruned attention weights remain negligible. We apply ACP
to language model pretraining with FoX and show it consistently reduces the
number of FLOPs in softmax attention by around 70% across different model sizes
and context lengths, resulting in a roughly 10% to 35% improvement in training
throughput. Furthermore, longer context lengths yield greater computational
savings. All these speed improvements are achieved without any performance
degradation. We also perform several analyses to provide deeper insights into
our method, such as examining the pruning patterns and analyzing the
distribution of FLOP savings across different attention heads. Our code is
available at https://github.com/zhixuan-lin/arctic-fox.
67ff12ea58ed263257af79b5
Zhixuan Lin
claimed_verified
2025-04-18T09:40:25.937Z
6694cc1009326cb83f2d11bb
/avatars/1ddaaed70a16ac475a9404848aef5d48.svg
Zhixuan Lin
user
zhixuan-lin
67ff12ea58ed263257af79b6
Johan Obando-Ceron
67ff12ea58ed263257af79b7
Xu Owen He
extracted_pending
2025-04-16T02:16:11.020Z
66906c4e37eadb9c577984d3
/avatars/b81765472942fdf94c0ee885ca62df2d.svg
Owen He
user
littleowen
67ff12ea58ed263257af79b8
Aaron Courville
Adaptive Computation Pruning for the Forgetting Transformer
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 the 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. This is achieved using a dynamically set pruning threshold
that ensures that the pruned attention weights remain negligible. We apply ACP
to language model pretraining with FoX and show it consistently reduces the
number of FLOPs in softmax attention by around 70% across different model sizes
and context lengths, resulting in a roughly 10% to 35% improvement in training
throughput. Furthermore, longer context lengths yield greater computational
savings. All these speed improvements are achieved without any performance
degradation. We also perform several analyses to provide deeper insights into
our method, such as examining the pruning patterns and analyzing the
distribution of FLOP savings across different attention heads. Our code is
available at https://github.com/zhixuan-lin/arctic-fox.
Adaptive Computation Pruning (ACP) applied to Forgetting Transformer (FoX) reduces the number of FLOPs in softmax attention by about 70% without degrading performance, improving training throughput by 10% to 35%.
Forgetting Transformer FoX forget gate softmax attention Adaptive Computation Pruning ACP input-output dependencies pruning threshold FLOPs training throughput context lengths pruning patterns distribution of FLOP savings attention heads