기술 요약: 짧은 윈도우 어텐션이 장기 기억을 가능하게 함

문제 제기

소프트맥스 어텐션에 기반한 현대의 대규모 언어 모델(LLM)은 시퀀스 길이에 선형적으로 증가하는 키-밸류 (KV) 캐시에 의존하며, 이로 인해 긴 컨텍스트 작업에서 계산 및 메모리 비용이 prohibitive 해집니다. 반면, 선형 순환 신경망 (RNN) 과 상태 공간 모델 (SSM) 은 일정한 메모리 및 계산 비용을 유지하지만, 트랜스포머에 비해 열등한 회상 능력을 보이는 경우가 많습니다.

최근의 하이브리드 아키텍처는 효율성과 성능을 균형 있게 맞추기 위해 로컬 슬라이딩 윈도우 어텐션 (SWA) 과 글로벌 어텐션 또는 선형 RNN 계층을 결합하려는 시도를 하고 있습니다. 그러나 슬라이딩 윈도우 길이가 로컬 기억과 장기 기억 간의 상호작용에 미치는 영향은 아직 충분히 연구되지 않았습니다. 구체적으로, 선형 RNN 이 존재할 때 윈도우 크기가 모델의 장기 의존성 학습 능력에 어떻게 영향을 미치는지, 그리고 더 큰 윈도우가 본질적으로 성능을 향상시키는지 아니면 더 긴 컨텍스트로 외삽하는 모델의 능력을 방해하는지는 명확하지 않습니다.

방법론

저자들은 슬라이딩 윈도우 어텐션 (SWA) 계층과 선형 RNN 계층을 번갈아 배치하는 하이브리드 아키텍처를 조사하며, 장기 기억 구성 요소로 xLSTM(행렬 메모리를 갖춘 mLSTM 의 변형) 을 구체적으로 활용합니다. 제안된 아키텍처인 SWAX는 SWA 계층과 xLSTM 계층을 1:1 비율로 번갈아 배치합니다.

본 연구는 두 가지 주요 방법론적 기여에 초점을 맞춥니다:

1. 윈도우 크기 분석: 저자들은 128 에서 2048 토큰까지의 고정된 슬라이딩 윈도우 크기를 가진 SWAX 모델을 훈련합니다. 그리고 이러한 모델을 단거리 컨텍스트 추론 벤치마크 (예: ARC, Hellaswag, MBPP) 와 장거리 컨텍스트 검색 작업 (예: RULER Needle-In-A-Haystack) 에서 평가합니다.
2. 확률적 윈도우 훈련: 짧은 컨텍스트 정밀도 (더 큰 윈도우가 선호됨) 와 긴 컨텍스트 외삽 (더 작은 윈도우가 선호됨) 간의 트레이드오프를 해결하기 위해, 저자들은 각 배치마다 슬라이딩 윈도우 크기를 확률적으로 샘플링하는 훈련 전략을 제안합니다. 실험에서 윈도우 크기는 짧은 크기 (128) 와 긴 크기 (2048) 사이를 특정 확률 (1.4B 모델의 경우 0.5, 7B 모델의 경우 0.75) 로 번갈아 가며 변경됩니다. 또한, 짧은 컨텍스트 성능을 안정화하기 위해 훈련 마지막 10% 기간에는 확률적 샘플링을 제거하는 어닐링 스케줄을 적용합니다.

실험은 1.4B 와 7B 파라미터를 가진 모델에서 수행되었으며, 사전 훈련 중 고정된 시퀀스 길이 16k 로 1500 억 토큰을 기반으로 훈련되었습니다.

주요 기여

1. 윈도우 크기에 대한 반직관적 발견: SWA 와 선형 RNN 을 결합한 하이브리드 아키텍처의 경우, 더 큰 슬라이딩 윈도우가 장기 컨텍스트 성능을 저하시킬 수 있음을 논문은 입증합니다. 더 큰 윈도우가 더 많은 컨텍스트를 제공한다는 신념과 달리, 저자들은 큰 윈도우 (예: 2048) 가 모델로 하여금 정확한 로컬 소프트맥스 어텐션에 지나치게 의존하게 만들어, 선형 RNN 계층이 장기 의존성에 대해 충분히 훈련되지 못하게 한다고 발견했습니다.
2. 짧은 윈도우의 이점 메커니즘: 더 짧은 윈도우 (예: 128) 는 모델이 로컬 윈도우 밖의 정보 검색을 위해 선형 RNN 계층에 의존하도록 강제합니다. 이 "감독 신호"는 선형 계층이 장기 의존성 모델링에 특화되도록 하여, 훈련 윈도우를 훨씬 초과하는 시퀀스 길이로 외삽하는 모델의 능력을 향상시킵니다.
3. 확률적 훈련 전략: 저자들은 윈도우 크기를 확률적으로 변화시키는 훈련 절차를 도입합니다. 이 접근법은 모델로 하여금 윈도우의 단기 정밀도와 RNN 의 장기 기억을 모두 활용하도록 강제하며, 어텐션 메커니즘에 대한 드롭아웃의 한 형태로 작용합니다.

결과

• 장기 컨텍스트 성능: 고정된 큰 윈도우 (2048) 로 훈련된 모델은 윈도우보다 긴 시퀀스로 테스트될 때 (예: RULER NIAH) 장기 컨텍스트 검색 작업에서 치명적인 성능 저하를 보입니다. 반면, 더 짧은 윈도우 (128) 로 훈련된 모델은 131k 토큰 길이의 시퀀스에서도 높은 정확도를 유지합니다. 확률적 훈련 방법은 이러한 장기 컨텍스트 작업에서 짧은 윈도우 전용 모델과 동등하거나 더 나은 성능을 달성합니다.
• 단기 컨텍스트 성능: 짧은 윈도우 (128) 만으로는 더 큰 윈도우에 비해 단기 컨텍스트 추론 벤치마크에서 때때로 성능이 떨어지지만, 확률적 훈련 방법은 이러한 성능을 회복시킵니다. 1.4B 와 7B 규모 모두에서 확률적 훈련은 고정된 2048 윈도우로 훈련된 모델과 동등하거나 경우에 따라 더 나은 단기 컨텍스트 성능을 제공합니다.
• 일반화: 이러한 발견은 다양한 선형 어텐션 백본 (xLSTM 및 게이트드 델타넷) 과 다양한 하이브리드 구성 (전체 어텐션 계층과 SWA 번갈아 배치) 에서도 유효합니다. 확률적 훈련 전략은 단기 및 장기 컨텍스트 벤치마크 모두에서 고정 윈도우 베이스라인을 일관되게 능가합니다.

중요성 및 주장

이 논문은 하이브리드 아키텍처에서 슬라이딩 윈도우 크기 선택이 단순히 계산 비용과 성능 간의 트레이드오프가 아니라, 모델이 기억 책임을 어떻게 분배하도록 학습하는지에 대한 결정적 요소라고 주장합니다.

저자들은 다음과 같이 주장합니다:

• 단순한 윈도우는 장기 기억 훈련에 필수적입니다: 로컬 어텐션 필드를 제한함으로써 모델은 장기 의존성을 위해 선형 RNN 계층을 활용하도록 강제되며, 이는 더 나은 길이 외삽으로 이어집니다.
• 확률적 훈련은 트레이드오프를 해결합니다: 훈련 중 동적으로 윈도우 크기를 변경함으로써 모델은 짧은 컨텍스트를 위한 고정밀 로컬 어텐션과 긴 컨텍스트를 위한 글로벌 선형 기억을 모두 활용하는 법을 학습합니다.
• 로컬 어텐션에 대한 과도한 의존은 해롭습니다: 모델이 큰 윈도우로 훈련되면, 그 윈도우 너머의 정보를 검색하는 데 필요한 메커니즘을 학습하지 못해 더 긴 시퀀스에서의 일반화가 저조해집니다.

이 연구는 하이브리드 아키텍처가 무제한 계산 비용 없이 최첨단 장기 컨텍스트 성능을 달성하려면, 단기 작업에 더 정밀할지라도 로컬 어텐션 메커니즘에 과도하게 의존하지 않도록 훈련 절차가 적극적으로 개입해야 함을 시사합니다.