Skalierbares Softmax ist überlegen für Aufmerksamkeit.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Autoren: Ken M. Nakanishi

cs.AI

Zusammenfassung

Das maximale Element des Vektors, der von der Softmax-Funktion ausgegeben wird, nähert sich null an, wenn die Größe des Eingabevektors zunimmt. Transformer-basierte Sprachmodelle verlassen sich auf Softmax, um Aufmerksamkeitswerte zu berechnen, was dazu führt, dass die Aufmerksamkeitsverteilung abflacht, wenn die Kontextgröße wächst. Dies verringert die Fähigkeit des Modells, Schlüsselinformationen effektiv zu priorisieren und begrenzt möglicherweise seine Längenverallgemeinerung. Um dieses Problem zu lösen, schlagen wir Scalable-Softmax (SSMax) vor, das Softmax in Szenarien ersetzt, in denen die Größe des Eingabevektors variiert. SSMax kann nahtlos in bestehende Transformer-basierte Architekturen integriert werden. Experimentelle Ergebnisse im Sprachmodellieren zeigen, dass Modelle, die SSMax verwenden, nicht nur eine schnellere Reduzierung des Verlusts während des Pretrainings erreichen, sondern auch die Leistung in langen Kontexten und bei der Schlüsselinformationsabfrage signifikant verbessern. Darüber hinaus zeigt eine Analyse der Aufmerksamkeitswerte, dass SSMax dem Modell ermöglicht, die Aufmerksamkeit auch in langen Kontexten auf Schlüsselinformationen zu richten. Zusätzlich können Modelle, die von Anfang an SSMax verwenden, eine bessere Längenverallgemeinerung erzielen, während solche, die bereits mit dem Pretraining begonnen haben, durch den Ersatz von Softmax in den Aufmerksamkeitsschichten durch SSMax, entweder während oder nach dem Pretraining, einige dieser Fähigkeiten erlangen.

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.