Adaptive Computation Pruning for the Forgetting Transformer

要約

最近提案された忘却トランス（FOX）には、忘れられたゲートがソフトマックスの注意に組み込まれており、標準的なロープベースのトランスと比較して、一貫して優れたパフォーマンスまたはオンパーパフォーマンスを示しています。
特に、FOXの多くの注意ヘッドはすぐに忘れがちで、各タイムステップでの出力により、主にローカルコンテキストに依存します。
この観察結果に基づいて、FOXの適応計算剪定（ACP）を提案します。これは、忘却ゲートによって強く減衰される入出力依存性を含む計算を動的にプルーン化する方法です。
これは、剪定された注意の重みが無視できるようにする動的に設定された剪定しきい値を使用して達成されます。
ACPをFOXで言語モデルの事前削減に適用し、ソフトマックスの注意のフロップ数を一貫して減少させて、異なるモデルサイズとコンテキストの長さにわたって約70％減少し、トレーニングスループットの約10％から35％の改善をもたらします。
さらに、コンテキストの長さが長くなると、計算の節約が大きくなります。
これらの速度の改善はすべて、パフォーマンスの劣化なしに達成されます。
また、いくつかの分析を実行して、剪定パターンを調べたり、さまざまな注意ヘッドにわたるフロップ節約の分布を分析するなど、私たちの方法に関するより深い洞察を提供します。
私たちのコードは、https：//github.com/zhixuan-lin/arctic-foxで入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google