Schalable-Softmax is superieur voor aandacht.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Auteurs: Ken M. Nakanishi

cs.AI

Samenvatting

Het maximale element van de vector die wordt uitgevoerd door de Softmax-functie nadert nul naarmate de grootte van de invoector toeneemt. Op transformer-gebaseerde taalmodellen vertrouwen op Softmax om aandachtscores te berekenen, wat leidt tot een vlakkere aandachtsverdeling naarmate de contextgrootte toeneemt. Dit vermindert het vermogen van het model om effectief belangrijke informatie te prioriteren en beperkt mogelijk de generalisatie van de lengte. Om dit probleem aan te pakken, stellen we Scalable-Softmax (SSMax) voor, die Softmax vervangt in scenario's waar de grootte van de invoector varieert. SSMax kan naadloos worden geïntegreerd in bestaande op transformer gebaseerde architecturen. Experimentele resultaten in taalmodellering tonen aan dat modellen die SSMax gebruiken niet alleen sneller verliesvermindering bereiken tijdens pretraining, maar ook aanzienlijk betere prestaties leveren in lange contexten en bij het ophalen van belangrijke informatie. Bovendien onthult een analyse van aandachtscores dat SSMax het model in staat stelt de aandacht te richten op belangrijke informatie, zelfs in lange contexten. Bovendien kunnen modellen die vanaf het begin van de pretraining SSMax gebruiken betere lengtegeneralisatie bereiken, terwijl diegenen die al zijn begonnen met pretraining nog steeds enigszins deze mogelijkheid kunnen verkrijgen door Softmax in de aandachtlagen te vervangen door SSMax, zowel tijdens als na pretraining.

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.