Pré-entraînement à long contexte avec l'Attention Phare

Résumé

L'entraînement de transformeurs causaux à des longueurs de séquence extrêmes est limité par le temps et la mémoire quadratiques de l'attention par produit scalaire mis à l'échelle (SDPA). Dans ce travail, nous proposons Lighthouse Attention, un algorithme d'attention hiérarchique basé sur une sélection symétrique, dédié à l'entraînement, qui encapsule le SDPA standard et peut être facilement retiré vers la fin de l'entraînement. Notre sélection hiérarchique est également sans gradient, ce qui nous dispense de traiter un noyau de passage arrière complexe et potentiellement inefficace. Notre contribution est triple : (i) une étape de pré- et post-traitement hiérarchique sous-quadratique qui effectue une compression et une décompression adaptatives de la séquence ; (ii) une stratégie de compression symétrique qui regroupe simultanément les requêtes, les clés et les valeurs, tout en préservant la causalité de gauche à droite, ce qui améliore considérablement le parallélisme ; (iii) une approche d'entraînement en deux étapes où nous pré-entraînons pendant la majeure partie du temps avec Lighthouse Attention, puis récupérons un modèle d'attention complet à la fin par un court entraînement. Nous menons des expériences préliminaires de pré-entraînement de LLM à petite échelle qui montrent l'efficacité de notre méthode par rapport à un entraînement par attention complète, tous les autres paramètres étant identiques, où nous obtenons un temps d'entraînement total plus rapide et une perte finale plus faible après la phase de récupération. Le code complet est disponible à l'adresse : https://github.com/ighoshsubho/lighthouse-attention

English

Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention