기술 요약: SSA: 더 나은 점수 함수로 성능 향상

문제 제기

Transformer 모델은 강력한 문맥 학습 (ICL) 능력을 보여주지만, 테스트 시 입력이나 알고리즘 구조가 훈련 분포와 다른 분포 이동 (distribution shifts) 상황에 일반화할 때는 체계적인 실패를 보입니다. 최근 문헌은 대규모 언어 모델 (LLM) 이 "저확률 설정"에서 어려움을 겪으며, 기본 규칙이 변하지 않음에도 불구하고 일반화에 실패한다고 지적했습니다.

저자들은 통제된 ICL 작업을 통해 이러한 실패의 구조적 기원을 조사했습니다. 그들은 어텐션 메커니즘 내의 Softmax 점수 함수가 이러한 일반화 실패의 주요 원인임을 확인했습니다. 구체적으로, 입력 시퀀스에 훈련 중 보았던 값보다 훨씬 큰 크기의 값이 포함될 때 Softmax 는 **어텐션 포화 (attention saturation)**를 유발합니다. 이 현상은 어텐션 가중치가 최대 값을 가진 단일 토큰으로 거의 완전히 붕괴되게 하여, 작업에 필수적인 모든 다른 문맥 토큰을 효과적으로 무시하게 만듭니다. 이러한 "어텐션 붕괴"는 모델이 시퀀스 전반에 걸쳐 정보를 통합하는 것을 방해하여, 논리적 집계나 기능적 추론이 필요한 작업에서 체계적인 오류를 초래합니다.

방법론

통제된 실험 설정

사전 훈련 아티팩트와 구조적 효과를 분리하기 위해, 저자들은 두 가지 합성 ICL 작업에 대해 Transformer 모델을 처음부터 훈련했습니다:

1. 정량화 작업 (Quantification Task): 숫자 시퀀스에 대한 논리 양화사 ("모든"과 "일부") 를 해석하는 모델의 능력을 평가합니다.
2. 선형 함수 작업 (Linear Function Task): 문맥 내 예시로부터 아핀 함수 $f(x) = ax + b$를 추론하는 모델의 능력을 평가합니다.

모델들은 표준 정규 분포 $N(0, 1)$에서 추출된 데이터로 훈련되었으며, 더 긴 시퀀스 길이와 더 높은 분산을 가진 분포 ($N(0, \sigma)$, 여기서 $\sigma \ge 1$) 에서 추출된 입력을 포함한 분포 이동 하에서 테스트되었습니다.

제안된 해결책: 스케일드 부호 평균 (Scaled Signed Averaging, SSA)

저자들은 Softmax 의 학습 가능한 대안으로 **스케일드 부호 평균 (SSA)**을 제안합니다. SSA 는 지수 점수 함수를 요소별 (element-wise) 로 적용되는 매개변수화 함수로 대체합니다:
$$x \mapsto (1 + b|x|)^{\text{sgn}(x)n}$$
여기서 $b > 0$이고 $n \ge 1$은 학습 가능한 스칼라 매개변수입니다.

주요 이론적 특성:

• 다항식 대 지수 성장: 로그이트 차이가 증가할 때 "하드맥스 (Hardmax)" 분포로 빠르게 포화되는 Softmax 와 달리, SSA 는 지수적으로 성장하지 않고 다항식적으로 성장합니다.
• 이론적 보장: 저자들은 정리 1 과 정리 2 를 통해 전역 로그이트 확대 또는 내부 불균형 (한 토큰의 점수가 다른 것보다 훨씬 크게 초과하는 경우) 하에서 Softmax 가승수가 승자에 대해서는 1 로, 나머지에 대해서는 0 으로 지수적으로 빠르게 수렴함을 증명합니다. 반면, SSA 는 다항식적으로만 수렴하여 2 차 토큰에 대해 0 이 아닌 "문맥 질량 (contextual mass)"을 유지합니다. 이는 극단적인 값이 존재할 때도 모델이 문맥에 대한 균형 잡힌 시각을 유지할 수 있게 합니다.

평가 범위

이 연구는 SSA 를 세 가지 다른 설정에서 평가했습니다:

1. 통제된 ICL 작업: 처음부터 훈련된 소형 Transformer(12 레이어, 8 헤더).
2. 디코더 전용 벤치마크: FineWeb 코퍼스로 처음부터 훈련된 1 억 1 천 4 백만 파라미터 Nemotron 스타일 모델.
3. 인코더 전용 아키텍처: 아동 지향 언어 (AO-CHILDES 코퍼스) 로 훈련된 BabyBERTa 변형체와 문법적 프로빙 (probing) 스위트에서 평가됨.

저자들은 또한 Sparsemax, Entmax, Temperature-scaled Softmax, CosFormer, SA-Softmax 를 포함한 Softmax 의 기존 대안들을 체계적으로 테스트하여, 특정 작업에서 표준 Softmax 보다 일관된 개선을 제공하는 것은 아무것도 없음을 발견했습니다.

주요 결과

1. 통제된 ICL 작업

• 이동 하의 일반화: SSA 는 분포 외 (OOD) 데이터로 테스트되었을 때 정량화 및 선형 함수 작업 모두에서 Softmax 보다 훨씬 뛰어난 성능을 발휘했습니다.
• 크기에 대한 강건성: 정량화 작업에서 Softmax 모델은 시퀀스에 단일 큰 값이 나타날 때 실패하여 오직 그 값만을 기반으로 예측합니다. SSA 모델은 전체 시퀀스에 걸쳐 정보를 올바르게 집계하여 극단적인 입력 값이 있더라도 낮은 오차율을 유지합니다.
• 선형 함수 추론: SSA 모델은 이동된 분포에서 추출된 계수를 가진 함수를 추론할 때 Softmax 기준 모델에 비해 훨씬 낮은 평균 제곱 오차 (MSE) 를 달성합니다.

2. 디코더 전용 벤치마크

• NLP 성능: SSA 로 훈련된 1 억 1 천 4 백만 파라미터 디코더 전용 모델은 분포 내 (FineWeb) 및 분포 외 (Wikipedia) 텍스트 모두에서 Softmax 대응 모델보다 낮은 퍼플렉시티를 달성했습니다.
• 벤치마크 개선: 표준 NLP 벤치마크 (예: ARC, BoolQ, HellaSwag) 에서 SSA 기반 모델은 제로샷 평가 전반에 걸쳐 Softmax 기준 모델을 일관되게 능가했으며, arc_easy 및 boolq와 같은 작업에서 주목할 만한 개선을 보였습니다.

3. 인코더 전용 모델

• 문법적 프로빙: SSA 기반 BabyBERTa 모델은 Softmax 기반 모델에 비해 여러 문법적 프로빙 작업 (예: 주어 - 동사 일치, 논항 구조) 에서 향상된 성능을 보였으며, 이는 이점이 마스킹 언어 모델링 아키텍처에도 확장됨을 입증했습니다.

중요성 및 주장

이 논문은 Softmax 포화가 큰 크기의 입력에 직면했을 때 어텐션 메커니즘이 작업과 관련이 있는지 여부와 상관없이 단일 토큰으로 붕괴되게 하는 근본적인 구조적 한계라고 주장합니다. 이러한 실패 양상은 분포 이동 하의 ICL 에서 낮은 일반화의 근본 원인으로 식별되었습니다.

이 연구의 중요성은 다음과 같습니다:

1. 구체적인 메커니즘 식별: 모델 규모나 데이터 복잡성이 아닌 Softmax 점수 함수를 특정 ICL 시나리오에서 일반화 실패의 주요 동인으로 지목합니다.
2. 간단하고 효과적인 대안: 어텐션 헤더당 두 개의 추가 스칼라 매개변수만 필요한 Softmax 의 "드롭인 (drop-in)" 대체제인 SSA 를 소개합니다.
3. 광범위한 적용 가능성: SSA 가 디코더 전용 및 인코더 전용 아키텍처뿐만 아니라 합성 및 실제 NLP 작업 전반에 걸쳐 성능을 향상시킨다는 것을 입증합니다.

저자들은 겸손한 입장을 유지하며, SSA 가 Softmax 의 경직성으로 인한 실패를 완화하지만, 입력과 함수 분포 모두 훈련과 크게 다른 극단적인 분포 이동이 포함된 모든 일반화 과제를 해결하는 것은 아니라고 인정합니다. 그들은 토큰의 크기와 관련성을 혼동하는 어텐션의 근본적인 수학적 구조가 향후 연구의 열린 영역으로 남아 있다고 결론지었습니다.