Adaptive Berechnungsreduzierung für den Vergessens-Transformer

Adaptive Computation Pruning for the Forgetting Transformer

April 9, 2025

Autoren: Zhixuan Lin, Johan Obando-Ceron, Xu Owen He, Aaron Courville

cs.AI

Zusammenfassung

Der kürzlich vorgeschlagene Forgetting Transformer (FoX) integriert ein Vergessens-Gate in die Softmax-Attention und hat durchweg bessere oder gleichwertige Leistungen im Vergleich zum standardmäßigen RoPE-basierten Transformer gezeigt. Bemerkenswerterweise neigen viele Aufmerksamkeitsköpfe in FoX dazu, schnell zu vergessen, wodurch ihre Ausgabe zu jedem Zeitpunkt hauptsächlich auf dem lokalen Kontext basiert. Basierend auf dieser Beobachtung schlagen wir Adaptive Computation Pruning (ACP) für FoX vor, eine Methode, die Berechnungen dynamisch beschneidet, die Eingabe-Ausgabe-Abhängigkeiten betreffen, die durch das Vergessens-Gate stark abgeschwächt wurden. Dies wird durch einen dynamisch festgelegten Beschneidungsschwellenwert erreicht, der sicherstellt, dass die beschnittenen Aufmerksamkeitsgewichte vernachlässigbar bleiben. Wir wenden ACP auf das Pretraining von Sprachmodellen mit FoX an und zeigen, dass es die Anzahl der FLOPs in der Softmax-Attention durchweg um etwa 70 % über verschiedene Modellgrößen und Kontextlängen reduziert, was zu einer ungefähr 10 % bis 35 % Verbesserung im Trainingsdurchsatz führt. Darüber hinaus führen längere Kontextlängen zu größeren Einsparungen bei den Berechnungen. Alle diese Geschwindigkeitsverbesserungen werden ohne jegliche Leistungseinbußen erreicht. Wir führen auch mehrere Analysen durch, um tiefere Einblicke in unsere Methode zu bieten, wie z. B. die Untersuchung der Beschneidungsmuster und die Analyse der Verteilung der FLOP-Einsparungen über verschiedene Aufmerksamkeitsköpfe hinweg. Unser Code ist unter https://github.com/zhixuan-lin/arctic-fox verfügbar.

English

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.