깊이 재귀적 어텐션 믹스처: 잠재 추론을 위한 새로운 패러다임

초록

본 논문은 파라미터를 깊이마다 공유하면서도, 시퀀스, 깊이, 전문가 3차원 어텐션을 결합한 모듈형 프레임워크 Dreamer를 제안한다. 깊이 재귀(Depth‑Recurrence)를 통해 숨은 상태 크기의 한계를 극복하고, FLOP·파라미터·메모리 조건을 동일하게 맞춘 기존 최첨단 모델 대비 2‑8배 적은 학습 토큰으로 동등하거나 우수한 성능을 달성한다.

상세 분석

Dreamer는 세 가지 핵심 어텐션을 하나의 레이어에 통합한다. 첫째, 전통적인 Transformer와 동일한 시퀀스 어텐션(SA)은 토큰 간 관계를 모델링한다. 둘째, 깊이 어텐션(DA)은 동일 토큰이 여러 재귀 깊이에서 생성된 은닉 상태들을 수직으로 연결해, 깊이 차원에서도 점곱 어텐션을 수행한다. 이는 RNN에서 발생하는 숨은 크기 병목을 완화하고, 깊이 증가에 따라 복잡한 논리 연산을 지속적으로 축적할 수 있게 한다. 셋째, 전문가 어텐션(EA)은 MoE 방식을 차용해 수천 개의 경량 MLP 전문가를 두고, 쿼리·키 기반 라우팅으로 토큰‑깊이 조합에 맞는 전문가를 선택한다. 라우팅은 sigmoid‑top‑K와 bias‑balancing을 이용해 과부하를 방지하고, 공유 전문가를 도입해 top‑1 라우팅의 그라디언트 손실을 보완한다.

Dreamer는 모든 어텐션 가중치를 깊이 전역으로 공유한다. 이를 위해 Residual‑Norm(RMSNorm)으로 잔차 스케일을 정규화하고, 어텐션 프로젝션 자체를 선형 MoE로 전환해 파라미터와 FLOP를 일정하게 유지한다. 또한 KV‑캐시를 깊이 방향에만 적용해 메모리 사용량을 시퀀스 길이에 독립적으로 만든다.

실험에서는 자연어 추론 벤치마크(예: GSM‑8K, ARC‑E, MMLU)에서 FLOP·파라미터·메모리 매칭된 최신 MoE 모델과 비교했으며, 동일 토큰 수에서 2‑8배 적은 학습량으로 동등하거나 더 높은 정확도를 기록했다. 특히 1B 파라미터 규모의 Dreamer는 2배 큰 기존 SOTA 모델과 비슷한 성능을 보였으며, 전문가 선택 다양성 측면에서도 2‑11배 높은 다양성을 나타냈다.

이 논문의 주요 기여는 (1) 깊이‑재귀와 다차원 어텐션을 결합한 통합 프레임워크, (2) 깊이 어텐션을 통한 숨은 크기 병목 해소, (3) FLOP·파라미터·메모리 매칭 실험을 통한 공정 비교, (4) 깊이별 지식 재사용 및 전문가 활용 패턴에 대한 정량적 분석이다. 이러한 설계는 모델을 수평(전문가 수)과 수직(깊이)으로 독립적으로 확장할 수 있게 하여, 향후 대규모 언어 모델에서 효율적인 잠재 추론을 구현하는 새로운 방향을 제시한다.