AtnRes: Nueva técnica optimiza modelos de lenguaje

Fuentes: GitHub - MoonshotAI/Attention-Residuals

Attention Residuals (AttnRes) es una innovadora técnica que optimiza las conexiones residuales en los Transformers, una arquitectura fundamental en modelos de lenguaje grandes (LLMs) como GPT. Tradicionalmente, las conexiones residuales simplemente suman la salida de cada capa con una ponderación uniforme. Si bien esto facilita el entrenamiento, a medida que la profundidad de la red aumenta (más capas), la contribución de cada capa se diluye, un problema conocido como 'dilución de PreNorm', y los valores de los estados ocultos tienden a crecer sin límite. AttnRes aborda este problema permitiendo que cada capa agregue selectivamente las representaciones de las capas anteriores mediante un mecanismo de atención ponderada, dependiente de la entrada.

En esencia, AttnRes reemplaza la suma simple con una operación de softmax attention sobre las salidas de las capas precedentes. Esto significa que cada capa aprende a dar más peso a las representaciones de las capas anteriores que son más relevantes para la tarea actual. La implementación original (Full AttnRes) tiene una complejidad de memoria de O(Ld), donde L es la profundidad de la red y d es la dimensión de la representación, lo que puede ser prohibitivo para modelos muy profundos. Para solucionar esto, se introduce Block AttnRes, una variante que divide la red en bloques y aplica la atención solo entre los bloques, reduciendo la complejidad de memoria a O(Nd), donde N es el número de bloques. Esto mantiene la mayoría de los beneficios de Full AttnRes con una sobrecarga computacional mínima.

¿Para qué sirve? AttnRes es útil para entrenar modelos Transformer más profundos y eficientes, especialmente aquellos que requieren razonamiento complejo o generación de código. Los investigadores y desarrolladores de LLMs son los principales usuarios potenciales. Se ha demostrado que mejora el rendimiento en una variedad de benchmarks, incluyendo razonamiento multi-paso (GPQA-Diamond) y generación de código (HumanEval), a menudo superando a los modelos base con un menor costo computacional.

Consideraciones: Aunque Block AttnRes reduce la complejidad de memoria, aún introduce una sobrecarga computacional, aunque marginal. La elección del tamaño del bloque (N) es un hiperparámetro que puede requerir ajuste. Además, como con cualquier técnica de atención, la interpretabilidad puede ser un desafío. Alternativas a AttnRes incluyen otras técnicas de normalización y optimización de conexiones residuales, pero AttnRes ofrece una combinación única de rendimiento y eficiencia.