title: Adaptive Computation Pruning for the Forgetting Transformer
summary: 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.
paper_authors._id: 67ff12ea58ed263257af79b5
paper_authors.name: Zhixuan Lin
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-04-18T09:40:25.937Z
paper_authors._id: 6694cc1009326cb83f2d11bb
paper_authors.avatarUrl: /avatars/1ddaaed70a16ac475a9404848aef5d48.svg
paper_authors.fullname: Zhixuan Lin
paper_authors.type: user
paper_authors.user: zhixuan-lin
paper_authors._id: 67ff12ea58ed263257af79b6
paper_authors.name: Johan Obando-Ceron
paper_authors._id: 67ff12ea58ed263257af79b7
paper_authors.name: Xu Owen He
paper_authors.status: extracted_pending
paper_authors.statusLastChangedAt: 2025-04-16T02:16:11.020Z
paper_authors._id: 66906c4e37eadb9c577984d3
paper_authors.avatarUrl: /avatars/b81765472942fdf94c0ee885ca62df2d.svg
paper_authors.fullname: Owen He
paper_authors.type: user
paper_authors.user: littleowen
paper_authors._id: 67ff12ea58ed263257af79b8
paper_authors.name: Aaron Courville
paper_title: Adaptive Computation Pruning for the Forgetting Transformer
paper_summary: 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.
paper_ai_summary: 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%.
paper_ai_keywords: 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