La Softmax escalable es superior para la atención.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Autores: Ken M. Nakanishi

cs.AI

Resumen

El elemento máximo del vector producido por la función Softmax tiende a cero a medida que aumenta el tamaño del vector de entrada. Los modelos de lenguaje basados en transformadores dependen de Softmax para calcular puntuaciones de atención, lo que provoca que la distribución de atención se aplane a medida que crece el tamaño del contexto. Esto reduce la capacidad del modelo para priorizar la información clave de manera efectiva y potencialmente limita su capacidad de generalización de longitud. Para abordar este problema, proponemos Escalable-Softmax (SSMax), que reemplaza Softmax en escenarios donde varía el tamaño del vector de entrada. SSMax puede integrarse fácilmente en arquitecturas basadas en transformadores existentes. Los resultados experimentales en modelado de lenguaje muestran que los modelos que utilizan SSMax no solo logran una reducción más rápida de la pérdida durante el preentrenamiento, sino que también mejoran significativamente el rendimiento en contextos largos y en la recuperación de información clave. Además, un análisis de las puntuaciones de atención revela que SSMax permite al modelo centrar la atención en la información clave incluso en contextos largos. Además, aunque los modelos que utilizan SSMax desde el inicio del preentrenamiento logran una mejor generalización de longitud, aquellos que ya han comenzado el preentrenamiento aún pueden adquirir parte de esta capacidad al reemplazar Softmax en las capas de atención con SSMax, ya sea durante o después del preentrenamiento.

English

The maximum element of the vector output by the Softmax function approaches zero as the input vector size increases. Transformer-based language models rely on Softmax to compute attention scores, causing the attention distribution to flatten as the context size grows. This reduces the model's ability to prioritize key information effectively and potentially limits its length generalization. To address this problem, we propose Scalable-Softmax (SSMax), which replaces Softmax in scenarios where the input vector size varies. SSMax can be seamlessly integrated into existing Transformer-based architectures. Experimental results in language modeling show that models using SSMax not only achieve faster loss reduction during pretraining but also significantly improve performance in long contexts and key information retrieval. Furthermore, an analysis of attention scores reveals that SSMax enables the model to focus attention on key information even in long contexts. Additionally, although models that use SSMax from the beginning of pretraining achieve better length generalization, those that have already started pretraining can still gain some of this ability by replacing Softmax in the attention layers with SSMax, either during or after pretraining.