Native Hybrid Attention for Efficient Sequence Modeling

논문 개요: Native Hybrid Attention (NHA)

이 논문은 시퀀스 모델링 분야에서 Transformer 의 높은 계산 복잡도와 선형 (Linear) 어텐션의 장기 기억 손실이라는 상충되는 문제를 해결하기 위해 제안된 새로운 하이브리드 아키텍처인 **NHA(Native Hybrid Attention)**를 소개합니다. NHA 는 선형 어텐션과 풀 (Full) 어텐션을 단일 계층 설계에 통합하여, 긴 문맥을 효율적으로 처리하면서도 높은 정확도를 유지합니다.

---

1. 문제 제기 (Problem)

• Transformer 의 한계: 자기 어텐션 (Self-Attention) 메커니즘은 장기 의존성을 포착하는 데 탁월하지만, 시퀀스 길이 $n$에 대해 $O(n^2)$의 이차 복잡도를 가집니다. 이는 긴 문서 분석이나 생정보학 등 긴 시퀀스 처리에 큰 병목이 됩니다.
• 선형 모델의 한계: 선형 어텐션 (Linear Attention) 이나 상태 공간 모델 (SSM) 은 $O(n)$의 효율성을 제공하지만, 고정된 크기의 상태 (State) 로 전체 시퀀스를 압축하는 과정에서 정확한 토큰 정보 (특히 장기 기억) 가 손실되어 회상 (Recall) 과 추론 성능이 저하됩니다.
• 기존 하이브리드 모델의 비효율성:
  • 레이어 내 (Intra-layer) 하이브리드: 선형 어텐션과 로컬 슬라이딩 윈도우 어텐션을 별도로 계산한 후 가중 합으로 결합하는 방식은 추가 파라미터가 필요하며, 두 메모리 간의 그라디언트 흐름이 단절될 수 있습니다.
  • 레이어 간 (Inter-layer) 하이브리드: 풀 어텐션 레이어와 선형 레이어를 교대로 쌓는 방식은 이질적인 블록을 관리해야 하며, 최적의 배치 (Placement) 를 찾기 위해 재학습이 필요합니다.

---

2. 방법론 (Methodology)

NHA 는 **레이어 내 (Intra-layer)**와 레이어 간 (Inter-layer) 하이브리드화를 하나의 통합된 설계로 자연스럽게 결합합니다.

가. 레이어 내 하이브리드: 통합된 소프트맥스 어텐션

기존 방식처럼 선형과 로컬 어텐션을 분리하여 계산하는 것이 아니라, 다음과 같은 단일 흐름을 가집니다.

1. 장기 기억 (Long-term Memory): 선형 RNN (Gated Linear RNN) 을 사용하여 과거 토큰들을 고정된 크기의 **키 - 밸브 슬롯 (Key-Value Slots)**으로 압축합니다.
2. 단기 기억 (Short-term Memory): 슬라이딩 윈도우 (Sliding Window) 내에 있는 최근 토큰들을 정밀하게 유지합니다.
3. 통합 어텐션 (Unified Attention): 압축된 장기 기억 슬롯과 정밀한 단기 기억 토큰을 **연결 (Concatenate)**한 후, 단 하나의 소프트맥스 어텐션 연산을 적용합니다.
   • 핵심 이점: 별도의 융합 파라미터 (Fusion Parameters) 없이, 쿼리 (Query) 와 모든 키 (Key) 간의 유사도 점수를 기반으로 자동적으로 장기/단기 기억 간 가중치를 동적으로 할당합니다. 이는 문맥 의존적 (Context-dependent) 인 가중치를 제공하며, 두 메모리 간의 그라디언트 흐름을 자연스럽게 연결합니다.

나. 레이어 간 하이브리드: 윈도우 크기 조절

NHA 는 모든 레이어가 동일한 구조를 공유합니다.

• 윈도우 크기 ($w$) 조절: 슬라이딩 윈도우 크기를 조정함으로써 각 레이어의 동작을 유연하게 제어할 수 있습니다.
  • $w = 0$: 순수 선형 RNN 레이어.
  • $w = \text{전체 시퀀스 길이}$: 풀 어텐션 (Transformer) 레이어.
  • $0 < w < \text{전체 길이}$: NHA 하이브리드 레이어.
• 장점: 모델 구조를 변경하지 않고 윈도우 크기만 조정하여, 각 레이어마다 필요한 정보의 정밀도와 효율성 간의 균형을 맞출 수 있습니다. 이는 추론 시 재학습 없이 최적의 하이브리드 구성을 탐색할 수 있게 합니다.

다. 효율적인 구현

• Chunkwise Parallel Form: 시퀀스를 청크 (Chunk) 단위로 나누어 병렬 처리할 수 있도록 설계되어, GPU 에서 효율적인 Triton 커널 구현이 가능합니다.

---

3. 주요 기여 (Key Contributions)

1. 통합된 하이브리드 아키텍처 제안: 레이어 내 (단일 소프트맥스) 와 레이어 간 (윈도우 크기 조절) 하이브리드화를 통합한 최초의 설계 중 하나로, 추가 파라미터 없이 동적인 메모리 할당을 가능하게 합니다.
2. 효율적인 GPU 커널 개발: NHA 를 위한 Chunkwise-parallel Triton 커널을 개발하여 대규모 모델 학습 및 추론의 효율성을 확보했습니다.
3. 예비 학습된 LLM 구조적 하이브리드화: 기존 Transformer 기반 LLM (Llama-3, Qwen 등) 을 NHA 구조로 변환하여 짧은 파인튜닝만으로도 경쟁력 있는 성능을 유지하면서 추론 속도와 메모리 효율을 크게 개선했습니다.

---

4. 실험 결과 (Results)

NHA 는 340M 및 1.3B 파라미터 모델부터 8B, 30B 규모의 사전 학습된 모델까지 광범위하게 평가되었습니다.

• 회상 (Recall) 및 추론 작업:
  • Recall-intensive Tasks (FDA, SQuAD, NQ 등): NHA 는 기존 Transformer 및 다른 하이브리드 모델 (Mamba, GLA, GSA 등) 보다 일관되게 최고의 성능을 기록했습니다. 특히 장기 문맥에서의 정보 회상 능력이 뛰어났습니다.
  • Commonsense Reasoning (ARC, HellaSwag 등): 일반 세계 지식과 추론 능력에서도 최상위권의 평균 점수를 달성했습니다.
• 긴 문맥 (Long-Context) 벤치마크:
  • RULER 벤치마크 (NIAH 등) 에서 2K 길이의 문맥으로 학습된 모델이 8K 까지 외삽 (Extrapolation) 되었을 때, NHA 는 Transformer 보다 더 강력한 외삽 능력을 보여주었습니다.
• 효율성 (Efficiency):
  • 연산 속도: 긴 시퀀스에서 FlashAttention (Transformer) 의 계산 시간이 급격히 증가하는 반면, NHA 는 선형에 가까운 스케일링을 유지하며 GSA 와 유사한 속도를 보입니다.
  • 메모리 사용량: NHA-Llama3-8B 는 긴 입력 시퀀스에서도 GPU 메모리 사용량과 지연 시간 (Latency) 이 기존 Llama3-8B 보다 훨씬 느리게 증가하여 확장성이 뛰어납니다.
• 사전 학습 모델 적용:
  • Llama-3-8B 와 Qwen2.5-7B 를 NHA 로 변환한 결과, 풀 어텐션 레이어를 4 개만 유지하더라도 (전체 레이어의 약 12.5%) 원본 Transformer 과 유사한 정확도를 유지하면서 효율성이 크게 향상되었습니다.

---

5. 의의 및 결론 (Significance)

• 이론적 통찰: NHA 는 "로컬 + 글로벌" 정보를 결합할 때, 별도의 가중 합 (Weighted Sum) 이 아닌 단일 소프트맥스 어텐션을 통해 두 메모리 간의 그라디언트 커플링 (Gradient Coupling) 을 자연스럽게 이루어낸다는 점을 증명했습니다. 이는 모델이 토큰별, 헤드별로 문맥에 따라 장기/단기 기억의 중요도를 동적으로 학습할 수 있게 합니다.
• 실용적 가치:
  • 유연한 배포: 추론 시 윈도우 크기만 변경하여 동일한 모델이 고속 RNN 모드와 고정밀 풀 어텐션 모드 사이를 전환할 수 있어, 하드웨어 제약에 따른 동적 배포가 가능합니다.
  • 비용 효율성: 기존 Transformer 모델을 재학습 (Pretraining) 없이도 구조적 변환과 가벼운 파인튜닝만으로 효율적인 하이브리드 모델로 전환할 수 있어, 산업계 적용 가능성이 높습니다.

결론적으로, NHA는 Transformer 의 정확성과 선형 모델의 효율성을 모두 달성하는 이상적인 균형을 제공하며, 긴 문맥을 다루는 차세대 대규모 언어 모델 (LLM) 아키텍처로서 중요한 진전을 이룩했습니다.