전역 컨텍스트를 필요할 때만 보는 효율적 주의 메커니즘

초록

책을 읽을 때 사람은 현재 페이지에 집중하고, 필요할 때만 이전 페이지를 되돌아본다. 이와 유사하게 대형 언어 모델(LLM)도 언제 전역 컨텍스트에 주의를 기울여야 하는지를 동적으로 학습할 수 있다. 우리는 각 어텐션 헤드에 이진 라우터를 두어 토큰마다 전체 어텐션과 로컬 슬라이딩 윈도우 어텐션을 전환하도록 하는 All-or-Here Attention(AHA)을 제안한다. 윈도우 크기 256 토큰일 때 원본 전체 어텐션 연산의 93%를 슬라이딩 윈도우 어텐션으로 대체해도 성능 저하가 없음을 보였다. 다양한 윈도우 크기로 AHA를 평가한 결과, 컨텍스트 의존성은 긴 꼬리 분포를 보이며 로컬 윈도우가 커질수록 전체 어텐션 필요성이 급격히 감소한다. 로컬 처리와 전역 접근을 분리함으로써 전체 어텐션이 대부분 불필요하고, 효율적인 추론은 전역 컨텍스트에 대한 필요 시점에만 접근하면 된다는 결론에 도달했다.

상세 요약

All-or-Here Attention(AHA)은 기존 트랜스포머의 전역 어텐션이 모든 토큰 쌍에 대해 연산 비용을 발생시키는 문제를 근본적으로 재구성한다. 각 어텐션 헤드에 배치된 이진 라우터는 입력 토큰마다 “전체 어텐션을 사용할지” 혹은 “고정된 크기의 로컬 슬라이딩 윈도우만 사용할지”를 판단한다. 라우터는 토큰의 위치, 현재 레이어의 표현, 그리고 간단한 통계적 신호(예: 토큰의 정보량, 이전 레이어의 어텐션 가중치 분포)를 기반으로 학습되며, 훈련 과정에서 라우터의 결정이 미분 가능하도록 스테이블 스위치(stable switch)와 같은 기법을 적용한다. 결과적으로 모델은 실제로 장거리 의존성이 필요한 경우에만 전역 어텐션을 활성화하고, 대부분의 경우는 비용이 O(window size × seq len)인 로컬 어텐션으로 대체한다.

실험에서는 윈도우 크기 256을 기준으로 전체 어텐션 연산의 93%를 로컬 어텐션으로 교체했음에도 불구하고, 언어 모델링 퍼플렉시티, 독해 정확도, 코드 생성 등 다양한 벤치마크에서 기존 전역 어텐션 모델과 통계적으로 유의미한 차이가 없었다. 윈도우 크기를 128, 512 등으로 변동시킨 추가 실험에서는 컨텍스트 의존성이 ‘긴 꼬리’ 형태를 띠는 것을 확인했다. 즉, 대부분의 토큰은 짧은 거리 내에서 충분히 정보를 얻지만, 소수의 토큰은 매우 긴 거리까지 의존한다. 이러한 분포는 인간이 책을 읽을 때 현재 페이지에 집중하고 필요할 때만 뒤 페이지를 찾아보는 행동과 일맥상통한다.

AHA의 장점은 두 가지로 요약될 수 있다. 첫째, 메모리 사용량과 연산량이 크게 감소해 실제 서비스 환경에서의 추론 비용을 현저히 낮춘다. 둘째, 라우터가 학습 과정에서 전역 어텐션의 필요성을 스스로 파악하므로, 사전에 복잡한 하이퍼파라미터 튜닝 없이도 다양한 작업에 적용 가능하다. 다만 라우터의 이진 결정이 과도하게 보수적이면 전역 어텐션이 충분히 활용되지 않아 성능 저하가 발생할 수 있다. 이를 완화하기 위해 라우터의 스무딩 파라미터를 조정하거나, 일정 비율의 토큰에 대해 무작위로 전역 어텐션을 강제 적용하는 전략이 제안된다. 전반적으로 AHA는 “필요할 때만 전역을 본다”는 인간 독서 메커니즘을 모델에 도입함으로써, 대규모 언어 모델의 효율성을 크게 향상시키는 혁신적인 접근법이라 할 수 있다.