Longitudes Varias, Velocidad Constante: Modelado de Lenguaje Eficiente con Lightning Attention

Presentamos Lightning Attention, la primera implementación de atención lineal que mantiene una velocidad de entrenamiento constante para diversas longitudes de secuencia bajo un consumo de memoria fijo. Debido al problema con las operaciones de suma acumulativa (cumsum), las implementaciones anteriores de atención lineal no pueden lograr su ventaja teórica en un entorno casual. Sin embargo, este problema puede resolverse eficazmente utilizando diferentes estrategias de cálculo de atención para computar las distintas partes de la atención. Específicamente, dividimos el cálculo de atención en intra-bloques e inter-bloques y usamos el cálculo convencional de atención para intra-bloques y trucos del núcleo de atención lineal para inter-bloques. Esto elimina la necesidad de cumsum en el cálculo de atención lineal. Además, se adopta una técnica de mosaico tanto en los procedimientos hacia adelante como hacia atrás para aprovechar al máximo el hardware GPU. Para mejorar la precisión manteniendo la eficacia, introducimos TransNormerLLM (TNL), una nueva arquitectura adaptada a nuestra lightning attention. Realizamos pruebas rigurosas en conjuntos de datos estándar y autocompilados con tamaños de modelo y longitudes de secuencia variables. TNL es notablemente más eficiente que otros modelos de lenguaje. Además, los resultados de referencia indican que TNL tiene un desempeño comparable con los LLM de última generación que utilizan estructuras transformadoras convencionales. El código fuente está disponible en github.com/OpenNLPLab/TransnormerLLM.