Масштабируемый Softmax превосходит по вниманию.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Авторы: Ken M. Nakanishi

cs.AI

Аннотация

Максимальный элемент вектора, выводимого функцией Softmax, приближается к нулю по мере увеличения размера входного вектора. Языковые модели на основе трансформера полагаются на Softmax для вычисления оценок внимания, что приводит к уплощению распределения внимания по мере увеличения размера контекста. Это снижает способность модели эффективно определять ключевую информацию и потенциально ограничивает ее обобщение на длинные тексты. Для решения этой проблемы мы предлагаем Scalable-Softmax (SSMax), который заменяет Softmax в ситуациях, когда размер входного вектора изменяется. SSMax может быть легко интегрирован в существующие архитектуры на основе трансформера. Экспериментальные результаты в языковом моделировании показывают, что модели, использующие SSMax, не только достигают более быстрого снижения потерь во время предварительного обучения, но также значительно улучшают производительность в длинных контекстах и поиске ключевой информации. Кроме того, анализ оценок внимания показывает, что SSMax позволяет модели сосредотачивать внимание на ключевой информации даже в длинных контекстах. Кроме того, хотя модели, использующие SSMax с самого начала предварительного обучения, демонстрируют лучшее обобщение на длинные тексты, те модели, которые уже начали предварительное обучение, могут приобрести некоторую способность к обобщению, заменив Softmax на SSMax в слоях внимания во время или после предварительного обучения.

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.