Langkontext-Vortraining mit Leuchtturm-Aufmerksamkeit

Zusammenfassung

Das Training kausaler Transformatoren mit extrem langen Sequenzen wird durch die quadratische Zeit- und Speicherkomplexität der skalierten Punktprodukt-Aufmerksamkeit (SDPA) begrenzt. In dieser Arbeit schlagen wir Lighthouse Attention vor, einen rein trainingsbasierten, symmetrischen, selektionsgesteuerten hierarchischen Aufmerksamkeitsalgorithmus, der die gewöhnliche SDPA umhüllt und gegen Ende des Trainings leicht entfernt werden kann. Unsere hierarchische Selektion ist zudem gradientenfrei, was uns den Umgang mit einem komplexen und potenziell ineffizienten Backward-Pass-Kernel erspart. Unser Beitrag ist dreifach: (i) Ein subquadratischer hierarchischer Vor- und Nachverarbeitungsschritt, der eine adaptive Kompression und Dekompression der Sequenz durchführt. (ii) Eine symmetrische Kompressionsstrategie, die gleichzeitig Queries, Keys und Values zusammenfasst und dabei die Left-to-Right-Kausalität bewahrt, was die Parallelität erheblich verbessert. (iii) Ein zweistufiger Trainingsansatz, bei dem wir die meiste Zeit mit Lighthouse Attention vortrainieren und am Ende mit einem kurzen Training ein vollständiges Aufmerksamkeitsmodell wiederherstellen. Wir führen vorläufige klein angelegte LLM-Vortrainings-Experimente durch, die die Wirksamkeit unserer Methode im Vergleich zum vollständigen Aufmerksamkeitstraining unter ansonsten identischen Bedingungen zeigen. Dabei erzielen wir eine kürzere Gesamttrainingszeit und einen niedrigeren endgültigen Loss nach der Wiederherstellungsphase. Der vollständige Code ist verfügbar unter: 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