O Scalable-Softmax é Superior para Atenção.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Autores: Ken M. Nakanishi

cs.AI

Resumo

O elemento máximo do vetor produzido pela função Softmax se aproxima de zero à medida que o tamanho do vetor de entrada aumenta. Modelos de linguagem baseados em Transformer dependem do Softmax para calcular os escores de atenção, o que faz com que a distribuição de atenção se torne mais uniforme à medida que o tamanho do contexto cresce. Isso reduz a capacidade do modelo de priorizar informações-chave de forma eficaz e potencialmente limita sua capacidade de generalização de comprimento. Para resolver esse problema, propomos o Softmax Escalonável (SSMax), que substitui o Softmax em cenários onde o tamanho do vetor de entrada varia. O SSMax pode ser integrado de forma transparente em arquiteturas baseadas em Transformer existentes. Resultados experimentais em modelagem de linguagem mostram que modelos que utilizam o SSMax não apenas alcançam uma redução mais rápida na perda durante o pré-treinamento, mas também melhoram significativamente o desempenho em contextos longos e na recuperação de informações-chave. Além disso, uma análise dos escores de atenção revela que o SSMax permite que o modelo concentre a atenção em informações-chave mesmo em contextos longos. Adicionalmente, embora os modelos que utilizam o SSMax desde o início do pré-treinamento alcancem uma melhor generalização de comprimento, aqueles que já iniciaram o pré-treinamento ainda podem adquirir parte dessa capacidade substituindo o Softmax nas camadas de atenção pelo SSMax, seja durante ou após o pré-treinamento.

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.