하이브리드 선형 어텐션 혁신: HALO와 HypeNet으로 초장문 효율 극대화
초록
본 논문은 대규모 사전학습 없이도 기존 Transformer 모델을 RNN‑Attention 하이브리드 구조로 변환할 수 있는 HALO 파이프라인을 제안한다. 핵심은 효율적인 어텐션 레이어 선택, 새로운 하이브리드 위치 인코딩 HyPE, 그리고 단계별 정렬·지식증류·미세조정이다. 이를 기반으로 만든 HypeNet은 Qwen‑3 시리즈와 동등한 성능을 유지하면서 1M 토큰까지의 초장문 처리 시 메모리·속도 면에서 크게 우수하다. 변환에 필요한 토큰 양은 2.3 B로 전체 사전학습 데이터의 0.01 %에 불과하다.
상세 분석
본 연구는 두 가지 큰 문제점을 동시에 해결한다. 첫째, 기존 하이브리드 모델은 수십~수백 억 토큰 규모의 데이터가 필요해 학계 진입 장벽이 높았다. 둘째, 변환 후 하이브리드 모델은 긴 컨텍스트에서 성능이 급격히 떨어지는 경향이 있었다. 이를 극복하기 위해 저자들은 HALO(Hybrid Attention via Layer Optimization)라는 3단계 파이프라인을 설계하였다.
1️⃣ **초기화 단계(Attention Weight Transfer)**에서는 기존 Transformer의 모든 어텐션 레이어를 동일한 차원의 RNN 레이어로 매핑한다. 여기서 Q, K, V, O 가중치를 그대로 재활용하고, RNN 특유의 전이 행렬(Fₜ) 등은 경험적 초기값으로 채운다. 이렇게 하면 사전학습된 파라미터를 거의 손실 없이 하이브리드 구조에 옮길 수 있다.
2️⃣ Stage 1 – Hidden State Alignment에서는 각 RNN 레이어를 독립적으로 학습시켜, 해당 어텐션 레이어가 생성한 은닉 상태와 최소 평균제곱오차(MSE)를 갖도록 한다. 이 과정은 어텐션‑RNN 간 표현 차이를 최소화해 이후 전체 모델 학습 시 급격한 성능 저하를 방지한다.
3️⃣ Attention Layer Selection은 HALO의 핵심 혁신이다. 저자들은 “Recall 성능 저하가 크고 CSR(Commonsense Reasoning) 저하가 작은” 레이어를 보존 대상으로 선정한다. 구체적으로 각 레이어 i를 RNN으로 교체했을 때의 Recall R(M(i))와 CSR C(M(i)) 차이를 계산하고,
(s_i = \frac{\max(R_{orig} - R(M(i)))}{\max(C_{orig} - C(M(i))) + \epsilon})
와 같은 스코어를 부여한다. 최상위 25% 레이어를 어텐션으로 유지하고 나머지는 RNN으로 교체한다. 이는 기존 연구가 단순히 층 번호나 KL 발산만을 기준으로 삼던 것과 달리, 실제 작업 성능을 직접 반영한다는 점에서 의미가 크다.
4️⃣ Stage 2 – Knowledge Distillation에서는 선택된 어텐션 레이어와 RNN 레이어가 혼합된 최종 하이브리드 모델을 학생 모델로 삼아, 원본 Transformer를 교사 모델로 두고 KL‑다이버전스 손실을 최소화한다. 1 B 토큰, 코사인 LR 스케줄링을 사용해 효율적인 정규화와 일반화가 이루어진다.
5️⃣ Stage 3 – Fine‑tuning은 긴 컨텍스트(예: 1M 토큰) 환경에 맞춰 추가 학습한다. 여기서는 학습률을 낮추고, 동일한 1 B 토큰을 사용해 모델이 장거리 의존성을 더 잘 포착하도록 한다.
이 파이프라인 위에 저자들은 **HyPE(Hybrid Positional Encoding)**를 도입한다. 기존 하이브리드 모델은 어텐션에 RoPE, RNN에 NoPE(위치 정보 없음)를 적용했지만, HyPE는 RNN에 RoPE, 어텐션에 NoPE를 반대로 적용한다. RoPE는 위치‑의존적 회전을 제공해 언어 모델링 성능을 높이고, NoPE는 훈련 없이도 길이 일반화 능력을 제공한다. 또한, 어텐션 스케일링 메커니즘을 추가해 토큰 수가 늘어나도 값이 폭발하지 않도록 설계했다.
HypeNet은 위의 HyPE와 함께 여러 구조적 개선을 포함한다. (①) 어텐션‑RNN 교체 시 발생하는 차원을 맞추기 위한 선형 변환 레이어, (②) 각 레이어마다 독립적인 정규화와 스케일 파라미터, (③) RNN 믹서에 최신 Mamba·DeltaNet 스타일 전이 행렬을 적용해 연산 효율성을 유지하면서도 표현력을 강화한다. 이러한 설계는 모델 크기(≈1 B~7 B 파라미터)와 관계없이 일관된 성능‑효율 트레이드오프를 제공한다.
실험 결과, Qwen‑3(1.7 B, 2.7 B, 6 B) 모델을 HALO로 변환한 HypeNet은 동일 토큰 수·파라미터에서 BLEU, LM‑Eval, NLG‑Bench 등 다양한 벤치마크에서 원본과 거의 차이가 없으며, 128 K1 M 컨텍스트에서 메모리 사용량은 3050 % 감소하고, 토큰당 처리 속도는 2~3배 가량 향상된다. 특히, 장거리 검색(NIAH)과 같은 Recall‑중심 작업에서 기존 하이브리드(distill) 모델이 보였던 급격한 성능 저하가 거의 사라졌다.
핵심 기여는 (1) 2.3 B 토큰이라는 극소량 데이터로 고성능 하이브리드 모델을 만들 수 있음을 증명, (2) 어텐션 레이어 선택을 성능‑중심 스코어로 정량화, (3) RoPE와 NoPE를 반대로 적용한 HyPE로 길이 일반화와 언어 모델링을 동시에 최적화, (4) 이러한 요소들을 통합한 HypeNet 아키텍처를 제시함으로써, 대규모 사전학습 비용 없이도 초장문 처리에 강력한 모델을 손쉽게 구축할 수 있게 되었다.
댓글 및 학술 토론
Loading comments...
의견 남기기