El proyecto presentado explora una nueva arquitectura de atención llamada 'Atención Híbrida' para modelos de lenguaje, con un enfoque particular en la generación de código Rust. El objetivo principal es acelerar significativamente la inferencia sin comprometer la calidad del modelo. Tradicionalmente, los modelos de lenguaje basados en Transformers utilizan la atención completa (full attention), que tiene una complejidad computacional de O(n²), donde 'n' es la longitud de la secuencia. Esto significa que el tiempo de cálculo aumenta cuadráticamente con la longitud de la secuencia, lo que se vuelve prohibitivo para secuencias largas.
Atención Híbrida aborda este problema combinando tres componentes: una capa de atención lineal inicial, una capa de atención cuadrática intermedia y una capa de atención lineal final. Esta estructura permite reducir la complejidad computacional a O(n·W + n·D), donde 'W' y 'D' representan parámetros relacionados con las ventanas locales y la recurrencia, respectivamente. La clave está en usar una atención local (una ventana de contexto) para capturar relaciones sintácticas a corto plazo y un camino recurrente (similar a un GRU) para mantener un estado comprimido de información a largo plazo. Una puerta (gate) aprende a equilibrar la importancia de estos dos caminos.
En el proyecto, se entrenó un modelo de lenguaje de 25.6 millones de parámetros desde cero, utilizando un corpus de código Rust de 173.5 MB. La arquitectura GPT estándar se modificó para incorporar Atención Híbrida, y se optimizó el proceso de inferencia utilizando kernels Triton y operaciones personalizadas de PyTorch. El resultado fue un aumento de velocidad de 51x en la inferencia, pasando de 5.6 tokens por segundo a 286.6 tokens por segundo, con una pérdida mínima en la calidad (medida por la perplexity). Además, se implementó una estrategia de caché KV (Key-Value) que comprime tokens antiguos a 8 bits, permitiendo un manejo más eficiente de la memoria y contribuyendo a la mejora de la velocidad.
Aunque la arquitectura Atención Híbrida es crucial, el autor destaca que la expansión del corpus de entrenamiento tuvo un impacto aún mayor en el rendimiento. El modelo genera código Rust que se asemeja a la sintaxis correcta, aunque la comprensión semántica aún es limitada. Los próximos pasos incluyen experimentos para comparar Atención Híbrida con variantes puramente locales o recurrentes, evaluar diferentes puntos de control del entrenamiento y explorar el aumento de la longitud del contexto y la transición a una tokenización BPE (Byte Pair Encoding).
En resumen, Atención Híbrida representa una solución prometedora para acelerar la inferencia en modelos de lenguaje, especialmente en el contexto de la generación de código, al reducir la complejidad computacional y optimizar el uso de la memoria.