Preentrenamiento de Contexto Largo con Atención Faro

Resumen

El entrenamiento de transformadores causales en longitudes de secuencia extremas está limitado por el tiempo y la memoria cuadráticos de la atención de producto punto escalado (SDPA). En este trabajo, proponemos la Atención Faro, un algoritmo de atención jerárquica basado en selección simétrica exclusivo para entrenamiento que envuelve al SDPA ordinario y puede eliminarse fácilmente hacia el final del entrenamiento. Nuestra selección jerárquica también es libre de gradiente, lo que nos exime de tener que lidiar con un kernel de paso hacia atrás complicado y potencialmente ineficiente. Nuestra contribución es triple: (i) Un paso de preprocesamiento y posprocesamiento jerárquico subcuadrático que realiza compresión y descompresión adaptativa de la secuencia. (ii) Una estrategia de compresión simétrica que agrupa consultas, claves y valores al mismo tiempo, preservando la causalidad de izquierda a derecha, lo que mejora enormemente el paralelismo. (iii) Un enfoque de entrenamiento en dos etapas en el que preentrenamos durante la mayor parte del tiempo con Atención Faro y recuperamos un modelo de atención completa al final con un entrenamiento corto. Realizamos experimentos preliminares de preentrenamiento de LLM a pequeña escala que muestran la efectividad de nuestro método en comparación con el entrenamiento de atención completa con todas las demás configuraciones igualadas, donde logramos un tiempo total de entrenamiento más rápido y una pérdida final más baja después de la fase de recuperación. El código completo está disponible en: 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