Élagage adaptatif des calculs pour le Transformer à oubli

Adaptive Computation Pruning for the Forgetting Transformer

April 9, 2025

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

cs.AI

Résumé

Le Transformer à Oubli (FoX) récemment proposé intègre une porte d'oubli dans l'attention softmax et a démontré des performances systématiquement meilleures ou équivalentes par rapport au Transformer standard basé sur RoPE. De manière notable, de nombreuses têtes d'attention dans FoX ont tendance à oublier rapidement, ce qui fait que leur sortie à chaque pas de temps dépend principalement du contexte local. Sur la base de cette observation, nous proposons l'Élagage Adaptatif des Calculs (ACP) pour FoX, une méthode qui élimine dynamiquement les calculs impliquant des dépendances entrée-sortie fortement atténuées par la porte d'oubli. Cela est réalisé en utilisant un seuil d'élagage dynamiquement ajusté qui garantit que les poids d'attention élagués restent négligeables. Nous appliquons ACP à l'apprentissage préalable de modèles de langage avec FoX et montrons qu'il réduit systématiquement le nombre de FLOPs dans l'attention softmax d'environ 70% pour différentes tailles de modèles et longueurs de contexte, entraînant une amélioration du débit d'entraînement d'environ 10% à 35%. De plus, des longueurs de contexte plus longues génèrent des économies de calcul plus importantes. Toutes ces améliorations de vitesse sont obtenues sans aucune dégradation des performances. Nous effectuons également plusieurs analyses pour fournir une compréhension plus approfondie de notre méthode, comme l'examen des motifs d'élagage et l'analyse de la distribution des économies de FLOPs entre les différentes têtes d'attention. Notre code est disponible à l'adresse https://github.com/zhixuan-lin/arctic-fox.

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.