Dense Local Dependencies Induce Attention-Logit Explosion and Training Instability During Long-Sequence Transformer Training

기술 요약: 조밀한 국소적 의존성이 어텐션 로짓 폭발 및 학습 불안정성을 유발함

문제 정의
자기회귀적 트랜스포머 언어 모델은 긴 시퀀스를 학습할 때 빈번하게 학습 불안정성을 보이며, 이러한 현상은 저정밀도 산술 연산(예: bfloat16) 환경에서 특히 두드러지게 나타난다. 이러한 불안정성은 종종 갑작스러운 손실 값의 급증이나 발산으로 나타나며, 이는 흔-히 "어텐션 로짓 폭발(attention-logit explosion)"—즉, 어텐션 로짓의 크기가 급격하고 무제한적으로 증가하는 현상—을 동반한다. 기존 연구들은 이러한 문제를 잔차 연결(residual connection)의 증폭, 옵티마이저와의 상호작용, 또는 일반적인 장기 시퀀스 관련 문제로 돌렸으나, 왜 시퀀스 길이가 길어질수록 어텐션 로짓의 성장을 촉진하는지에 대한 구체적인 근본 원인은 여전히 명확히 밝혀지지 않았다. QK-정규화(QK-normalization)나 특수 옵티마이저와 같은 기존의 해결책들은 근본 원인보다는 증상을 완화하는 데 그치는 경우가 많다.

방법론 및 분석 프레임워크
저자들은 **조밀한 국소적 의존성(dense local dependencies)**이 이러한 불안정성의 주요 동인이라고 제안한다. 이들의 분석은 표준 풀 셀프 어텐션(full self-attention) 메커니즘에 내재된 **저계수 병목 현상(low-rank bottleneck)**에 기초한다:

1. 이론적 근거: 표준 트랜스포머에서 어텐션 로짓 행렬 $S = QK^T$의 계수(rank)는 어텐션 차원 $d$에 의해 제한된다. 그러나 조밀한 국소적 의존성(토큰이 주변의 많은 토큰에 의존하는 경우)을 포함하는 태스크는 실질적으로 고계수(high-rank)를 요구하는 어텐션 패턴을 필요로 하며, 이때 계수는 시퀀스 길이 $n$에 따라 선형적으로 증가한다.
2. 불일치: $n \gg d$인 상황에서 모델은 표현상의 불일치에 직면한다. 저계수 인수분해($QK^T$)를 사용하여 고계수의 타겟 어텐션 패턴을 근사하려 할 때, 최적화 과정에서 어텐션 로짓의 크기를 팽창시킬 수밖에 없다.
3. 실험적 검증: 저자들은 다음 두 가지 주요 경로를 통해 이 가설을 검증하였다:
   • 합성 태스크(Synthetic Tasks): 타겟 어텐션 행렬이 조밀한 국소적 의존성(밴드 행렬 형태)을 포함하지만 장거리 구조는 없는 제어된 태스크를 구축하였다. 모델은 이 행렬을 재구성하는 과제를 수행했다.
   • 언어 모델링: PG-19 데이터셋을 사용하여 GPT-2 아키텍처를 이용한 광범위한 실험을 수행하였다. 저자들은 시퀀스 길이, 어텐션 헤드 차원, 그리고 어텐션 메커즘의 구조(풀 어텐션과 슬라이딩 윈도우 로컬 어텐션을 혼합한 "Full-Local Attention" 도입)를 변화시켰다. 또한 마스킹을 통해 국소적 의존성의 밀도를 조절하였다.

주요 결과
실증적 증거는 제안된 메커니즘을 일관되게 뒷받침한다:

• 시퀀스 길이와 로짓 성장: 합성 태스크와 언어 모델링 모두에서, 어텐션 차원 $d$가 고정되어 있을 때 시퀀스 길이 $n$이 증가함에 따라 고정된 재구성 정확도를 달성하기 위해 필요한 최대 어텐션 로짓이 급격히 증가한다.
• 차원 스케일링: 어텐션 헤드 차원 $d$를 늘리면 로짓 성장이 상당히 완화되며, 이는 병목 현상이 계수 제약과 시퀀스 길이 사이의 관계와 관련이 있음을 확인시켜 준다.
• 로컬 어텐션의 효능: 일부 풀 어텐션 헤드를 슬라이딩 윈도우 로컬 어텐션 헤드로 교체하면(이는 저계수 병목 현상 없이 조밀한 국소적 의존성을 자연스럽게 모델링함), 로짓 성장이 극적으로 감소한다. 이러한 감소는 남겨진 풀 어텐션 헤드에서도 관찰되는데, 이는 로컬 헤드가 표현의 부담을 분담(offload)하고 있음을 시사한다.
• 밀도 대 국소성: 불안정성은 단순히 국소적 의존성의 존재가 아니라, 국소적 의존성의 밀도에 의해 발생한다. 국소적 연결을 희소화(90% 마스킹)했을 때 로짓 성장이 억제되었고 학습이 안정화되었다.
• 정밀도 민감도: 로짓 폭발은 FP32와 저정밀도(BF16) 학습 모두에서 발생하지만, 학습 불안정성(손실 발산)은 주로 저정밀도 설정에서 관찰된다. 이는 큰 로짓의 영향을 수치적 오류가 증폭시키기 때문이다.
• 모델 규모: 불안정성의 심각도는 모델 크기에 따라 증가한다. 동일한 조건 하에서 더 큰 모델(예: GPT-2 Large)은 더 작은 모델(예: GPT-2 Small)보다 더 뚜렷한 로짓 성장과 불안정성을 보였으며, 안정성을 유지하기 위해 더 큰 로컬 어텐션 용량을 필요로 했다.

의의 및 주장
본 논문은 안정적이고 효율적이며 확장 가능한 롱 컨텍스트 트랜스포머를 위한 근본적인 설계 원칙을 제시한다: 조밀한 국소적 의존성을 명시적으로 모델링하는 것.

• 근본 원인 식별: 본 연구는 풀 셀프 어텐션의 저계수 병목 현상이 조밀한 국소 구조를 표현하도록 강제될 때 발생하는 어텐션 로짓 폭발의 근본 원인임을 확립하였다.
• 설계 원칙: 로짓 폭발을 단순히 정규화하여 없애야 할 수치적 아티팩트로 취급하기보다, 아키텍처가 조밀한 국소 상호작용을 처리할 전용 용량(예: 로컬 어텐션 헤드)을 할당해야 한다고 주장한다. 이 접근 방식은 저계수 병목 현상을 자연스럽게 완화한다.
• 효율성 및 안정성: 로컬 의 dependency를 로컬 어텐션 헤드로 분담함으로써, 모델은 학습 안정성을 높일 뿐만 아니라 계산 효율성(복잡도를 $O(n^2)$에서 $O(nw)$로 감소)과 추론 시 KV 캐시 메모리 요구량을 줄이는 이점을 얻는다.
• 광범위한 시사점: RoPE, ALiBi 또는 구조적 어텐션(structured attention)을 사용하는 것과 같이 국소성 편향(locality-biased)을 가진 다양한 아키텍처들의 성공은, 순수 풀 어텐션 모델을 괴롭히는 저계수 병목 현상을 피하기 위해 이러한 조밀한 국소적 의존성을 암시적 또는 명시적으로 모델링할 수 있는 능력 덕분일 수 있음을 시사한다.