확장 가능한 소프트맥스가 어텐션에 뛰어나다.

Scalable-Softmax Is Superior for Attention

January 31, 2025

저자: Ken M. Nakanishi

cs.AI

초록

소프트맥스 함수에 의해 출력된 벡터의 최대 요소는 입력 벡터 크기가 증가함에 따라 접근하는 값이 제로에 가까워집니다. 트랜스포머 기반 언어 모델은 어텐션 점수를 계산하기 위해 소프트맥스를 활용하는데, 이로 인해 컨텍스트 크기가 커질수록 어텐션 분포가 평탄해지게 됩니다. 이는 모델이 핵심 정보를 효과적으로 우선순위를 매기는 능력을 감소시키고 길이 일반화를 제한할 수 있습니다. 이 문제를 해결하기 위해 우리는 입력 벡터 크기가 다양한 경우 소프트맥스를 대체하는 Scalable-Softmax (SSMax)를 제안합니다. SSMax는 기존의 트랜스포머 기반 아키텍처에 매끄럽게 통합될 수 있습니다. 언어 모델링에서의 실험 결과는 SSMax를 사용하는 모델이 사전 훈련 중 손실 감소 속도가 빨라지는 것뿐만 아니라 긴 컨텍스트와 핵심 정보 검색에서 성능이 크게 향상되는 것을 보여줍니다. 또한 어텐션 점수 분석 결과, SSMax를 사용하면 모델이 긴 컨텍스트에서도 핵심 정보에 집중할 수 있게 됩니다. 게다가, 사전 훈련 시작부터 SSMax를 사용하는 모델이 더 나은 길이 일반화를 달성하지만, 이미 사전 훈련을 시작한 모델은 사전 훈련 중이나 이후에 어텐션 레이어에서 소프트맥스를 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.