기술 요약: 하이브리드 선형 어텐션에 대한 체계적 분석

문제 정의

트랜스포머 아키텍처는 대규모 언어 모델(LLM)을 지배하고 있으나, 시퀀스 길이($L$)가 증가함에 따라 발생하는 이차 복잡도($O(L^2)$)와 선형적으로 증가하는 키-값(KV) 캐시 메모리 요구 사항이라는 문제를 안고 있다. 선형 어텐션 메커니즘은 KV 캐시를 고정된 크기의 은닉 상태로 압축함으로써 $O(L)$의 복잡도를 제공하지만, 일반적으로 풀 어텐션(full attention)에 비해 열등한 회상(recall) 성능을 보인다. 이를 해결하기 위해 "하이브리드 선형 어텐션" 모델들이 등장하였으며, 이는 풀 어텐션 레이어와 선형 어텐션 레이어를 교차하여 배치한다. 그러나 하이브리드 구조 내에서 선형 어텐션 구성 요소를 선택하는 과정은 여전히 상당 부분 임의적이다. 기존 연구들은 특정 선형 구성 요소의 아키텍처적 특성보다는 풀 어텐션 대비 선형 레이어의 비율에 과도하게 집중해 왔다. 또한, 단독(standalone)으로 성능이 우수한 선형 모델이 반드시 우수한 하이브리드 아키텍처로 이어지는지, 그리고 선형 대 풀 어텐션 비율이 언어 모델링(LM) 손실과 장거리 회상 성능에 어떻게 차별적으로 영향을 미치는지도 불분명하다.

방법론

저자들은 두 가지 파라미터 규모에 걸쳐 72개의 모델을 학습 및 평가하는 포괄적인 체계적 분석을 수행하였다:

• 340M 파라미터: 200억 토큰으로 학습.
• 1.3B 파라미터: 1,000억 토큰으로 학습.

평가된 모델들:
본 연구는 세 세대의 선형 어텐션 메커니즘을 다룬다:

1. 1세대 (벡터 재귀): HGRN (게이트형 벡터 상태).
2. 2세대 (외적을 통한 행렬 상태): RetNet (고정 감쇠), GLA (데이터 의존적 게이트), HGRN-2 (계층적/결합 게이트), 그리고 Mamba-2 및 RWKV-6와 같은 변형들.
3. 3세대 (델타 규칙/제어된 망각): DeltaNet 및 Gated DeltaNet.

실험 설계:

• 하이브리드화: 각 선형 변형 모델을 풀 어텐션 레이어와 함께 다섯 가지 특정 선형 대 풀 어텐션 비율(24:1, 12:1, 6:1, 3:1, 그리고 순수 선형 베이스라인)로 테스트하였다. 표준 풀 트랜스포머가 상한선(upper-bound) 베이스라인 역할을 한다.
• 벤치마크:
  • 언어 모델링 (LM): ARC-Challenge, ARC-Easy, HellaSwag, LAMBADA, OpenBookQA, PIQA를 포함한 표준 벤치마크.
  • 회상 (Recall): 검색(retrieval), 멀티홉 추적(multi-hop tracing), 집계(aggregation), 질의응답(question answering)을 평가하는 RULER 벤치마크 제품군.
• 학습: fineweb-edu 데이터셋을 사용하여 flash-linear-attention 라이브러리로 사전 학습되었으며, AdamW와 코사인 학습률 스케줄을 통해 최적화되었다. 평가는 제로샷(zero-shot) 방식으로 수행되었다.

핵심 기여

1. 세대 간 비교 분석: 본 논문은 세 세대의 선형 어텐션 아키텍처를 모두 아우르는 최초의 체계적인 비교를 제공하며, 이들을 단독 모델로서뿐만 아니라 하이브리드 스택으로 통합했을 때도 평가한다.
2. 단독 성능과 하이브리드 성능의 분리: 단독으로 우수한 선형 모델이 반드시 우수한 하이브리드 모델을 보장한다는 가설에 이의를 제기한다. 연구 결과, 풀 어텐션 레이어가 도입될 때 아키텍처의 순위가 크게 변화함을 입증한다.
3. LM vs 회상의 동인 분리: 선형 대 풀 어텐션 비율이 언어 모델링 능력과 회상 성능에 어떻게 각각 다르게 영향을 미치는지 분석하여, 이 두 지표가 아키텍처 변화에 서로 다르게 반응함을 밝혀낸다.
4. 핵심 아키텍처 구성 요소 식별: 하이브리드 성능을 최적화하는 데 필수적인 설계 원칙인 선택적 게이팅(selective gating), 계층적 재귀(hierarchical recurrence), 제어된 망각(controlled forgetting)을 분리하여 규명한다.

주요 결과

1. 단독 vs 하이브리드 성능의 불일치

최고의 단독 성능을 보이는 선형 모델이 최고의 하이브리드 성능을 보장하지는 않는다.

• 관찰: 340M 규모에서 GatedDeltaNet이 가장 높은 단독 정확도를 달 기록했다. 그러나 하이브리드 구성에서는 HGRN-2(6:1 비율)가 GatedDeltaNet을 능가하였으며, 트랜스포머 베이스라인과 단독 리더 모델보다 1.2 퍼센트 포인트 더 높은 성능을 보였다.
• 시사점: 단독 벤치마크는 하이브리드 백본을 선택하는 데 있어 신뢰할 수 없는 예측 지표이다.

2. 선형 대 풀 어텐션 비율의 영향

• 언어 모델링 (LM): 모든 비율(24:1에서 3:1까지)에서 성능이 대체로 안정적이다. 대부분의 아키텍처가 평균 LM 점수 0.55~0.57 근처에 밀집되어 있으며, 풀 어텐션 레이어의 비율에 매우 적은 민감도를 보인다.
• 회상 성능: 회상은 비율에 매우 민감하다. 풀 어텐션 레이어의 비율이 높아질수록 성능이 유의미하게 향상된다.
  • 순수 선형 모델은 낮은 회상 성능(RULER 점수 ~0.1–0.35)을 보인다.
  • 3:1 비율에서 회상 성능이 트랜스포머 베이스라인에 근접하거나 이를 초과한다.
  • DeltaNet 및 GatedDeltaNet과 같은 모델들은 3:1 구성에서 회상 성능이 정점에 도달한다.

3. 성공의 아키텍처 결정 요인

효과적인 하이브리드 모델을 위한 세 가지 상호 보완적인 특성이 식별되었다:

• 선택적 게이팅 (Selective Gating): 학습된 토큰 단위 게이트를 가진 모델(예: GatedDeltaNet, HGRN-2)은 고정된 감쇠를 가진 모델(예: RetNet)보다 일관되게 우수한 성능을 보이며, 고정 감쇠 모델은 장거리 단서를 보호하는 데 실패한다.
• 계층적 재귀 (Hierarchical Recurrence): HGRN-2의 이중 경로(느린 조대 메모리 + 빠른 토큰 상세 정보)는 단일 경로 절제 모델(HGRN)에 비해 회상 성능을 두 배로 높였으며, 이는 넓게 배치된 풀 어텐션 레이어가 정보를 "고정(latch)"할 수 있는 재귀적 계층 구조의 이점을 얻음을 시사한다.
• 제어된 망각 (Controlled Forgetting): 델타 규칙이나 게이트형 대각 감쇠를 통해 오래된 콘텐츠를 명시적으로 삭제하는 메커니즘은 상태 혼잡(state crowding)을 방지한다. 이러한 기능이 결여된 모델(예: 순수 가산 업데이트 방식인 GLA)은 비율과 관계없이 트랜스포머 수준의 회상 성능에 도달하는 데 어려움을 겪는다.

4. 실무적 권장 사항

본 연구는 HGRN-2 또는 GatedDeltaNet과 같은 아키텍처를 3:1에서 6:1 사이의 선형 대 풀 어텐션 비율로 사용할 것을 권장한다. 이 구성은 KV 캐시 메모리 점유율을 4배에서 7배까지 줄이면서도 트랜스포머 수준의 회상 능력을 달성한다.

의의 및 주장

본 논문은 하이브리드 프레임워크 내에서 세 세대의 선형 어텐션 메커니즘을 체계적으로 비교한 첫 번째 연구라고 주장한다. 주요 의의는 경험적인 비율 탐색에서 원칙적인 구성 요소 선택으로 초점을 전환했다는 점에 있다.

저자들은 다음과 같이 단언한다:

• 하이브리드의 품질은 단독 성능으로부터 추론될 수 없으며, 선형 레이어와 풀 어텐션 레이어 사이의 상호작용은 비자명(non-trivial)하다.
• 퍼플렉시티(LM 손실)가 아닌 회상이 최적의 선형 대 풀 어텐션 조합을 결정하는 결정적 요인이다.
• 메모리가 제한된 실무자는 선택적 게이팅, 계층적 재귀, 제어된 망각을 갖춘 백본을 선택하고, 약 4분의 1의 레이어를 풀 어텐션에 할당함으로써 KV 캐시 비용의 일부만으로 트랜스포머 수준의 회상 성능을 얻을 수 있다.

저자들은 한계점으로, 본 연구가 최대 1.3B 파라미터와 4,096 토큰 컨텍스트 창에 국한되어 있어 10B 이상의 파라미터 규모 및 더 긴 컨텍스트(예: 128k 토큰)로의 확장성에 대한 질문이 남아 있음을 인정한다. 그러나 저자들은 기존의 프로덕션 모델(예: MiniMax-01)을 예로 들어 하이브리드 모델이 효과적으로 확장될 수 있음을 보여준다.