La Scalable-Softmax è superiore per l'attenzione.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Autori: Ken M. Nakanishi

cs.AI

Abstract

L'elemento massimo del vettore prodotto dalla funzione Softmax si avvicina a zero all'aumentare delle dimensioni del vettore di input. I modelli linguistici basati su Transformer dipendono dalla Softmax per calcolare i punteggi di attenzione, causando l'aplanamento della distribuzione di attenzione all'aumentare delle dimensioni del contesto. Ciò riduce la capacità del modello di dare priorità in modo efficace alle informazioni chiave e potenzialmente limita la sua capacità di generalizzazione della lunghezza. Per affrontare questo problema, proponiamo Scalable-Softmax (SSMax), che sostituisce la Softmax nei casi in cui le dimensioni del vettore di input variano. SSMax può essere integrato senza problemi nelle architetture basate su Transformer esistenti. I risultati sperimentali nella modellizzazione del linguaggio mostrano che i modelli che utilizzano SSMax non solo ottengono una riduzione più rapida della perdita durante la preformazione, ma migliorano significativamente le prestazioni nei contesti lunghi e nel recupero delle informazioni chiave. Inoltre, un'analisi dei punteggi di attenzione rivela che SSMax consente al modello di concentrare l'attenzione sulle informazioni chiave anche in contesti lunghi. Inoltre, sebbene i modelli che utilizzano SSMax fin dall'inizio della preformazione raggiungano una migliore generalizzazione della lunghezza, quelli che hanno già iniziato la preformazione possono comunque acquisire in parte questa capacità sostituendo la Softmax nei livelli di attenzione con SSMax, sia durante che dopo la preformazione.

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.