행렬 기반 비선형 RNN으로 확장 가능한 언어 모델링

초록

본 논문은 기존 비선형 RNN이 상태 크기 제한으로 언어 모델링에서 뒤처지는 문제를 해결하기 위해, 행렬 형태의 숨겨진 상태와 외적 기반 상태 확장 메커니즘을 도입한 M²RNN을 제안한다. M²RNN은 대규모 언어 모델에 적용했을 때 긴 문맥 추적과 코드 실행 같은 TC⁰를 초월하는 작업에서 우수한 일반화 능력을 보이며, 하이브리드 구조에서도 기존 선형 RNN 대비 0.4~0.5 퍼플렉시티 개선과 8점 이상의 LongBench 성능 향상을 달성한다. 또한 텐서코어 활용을 최적화해 학습 효율성을 크게 높인다.

상세 분석

M²RNN은 기존 비선형 RNN이 벡터 형태의 은닉 상태(hₜ∈ℝᵈ)만을 사용해 파라미터와 메모리 효율성 사이에서 트레이드오프를 겪는 점을 근본적으로 재설계한다. 논문은 상태를 K×V 크기의 행렬 Hₜ∈ℝ^{K×V}로 확장함으로써, 동일한 파라미터 수에서 선형 RNN이 사용하는 행렬 상태와 동등하거나 더 큰 표현 용량을 확보한다. 핵심은 외적(outer‑product) 연산을 이용한 상태 확장 메커니즘이다. 입력 xₜ에 대한 키 φ(kₜ)와 값 vₜ를 외적으로 결합해 ΔHₜ=φ(kₜ)vₜᵀ를 생성하고, 이를 기존 상태에 누적한다(Hₜ=H_{t‑1}+ΔHₜ). 이 과정은 기존 LSTM·GRU와 달리 게이트가 상태와 독립적으로 작동하도록 설계된 ‘forget gate’를 도입해 그래디언트 소실을 방지하면서도 병렬 계산이 가능하도록 만든다.

표현력 측면에서, 저자는 TC⁰에 머무르는 트랜스포머와 선형 RNN이 해결할 수 없는 S₅와 같은 대칭군 퍼뮤테이션 문제를 M²RNN이 정확히 해결함을 실험적으로 입증한다. 이는 비선형 활성화와 행렬 상태가 결합돼 NC¹ 수준의 계산을 수행할 수 있음을 의미한다. 또한 상태 크기가 충분히 크면 비선형 RNN도 선형 RNN과 동등한 언어 모델링 성능을 달성한다는 가설을 검증한다. 실험에서는 상태 차원(K,V)을 늘려 파라미터 수를 일정하게 유지하면서도 퍼플렉시티가 크게 감소하는 것을 확인한다.

하드웨어 효율성도 중요한 기여점이다. 외적 기반 상태 확장은 텐서코어가 최적화된 16×16 매트릭스 연산 형태와 일치해, 기존 플래시RNN이 배치 차원 패딩으로 발생시키던 FLOP 낭비를 최소화한다. 논문은 Triton으로 구현한 전·후방 커널을 제시하고, 토폴로지‑의존적·독립적 두 가지 텐서 병렬(TP) 전략을 통해 다중 GPU 환경에서도 통신 오버헤드 없이 확장 가능함을 보인다.

하이브리드 실험에서는 7B MoE 모델에 M²RNN 레이어를 삽입했을 때, 동일한 파라미터 대비 3배 작은 상태 크기로도 Gated DeltaNet 하이브리드 대비 0.4~0.5 퍼플렉시티 개선을 얻었다. 특히 단일 M²RNN 레이어만 교체해도 전체 모델 성능이 크게 상승했으며, LongBench에서 최대 8점 이상의 장기 문맥 일반화 향상을 기록했다. 이는 M²RNN이 비선형 RNN의 표현력과 선형 RNN의 효율성을 동시에 제공한다는 강력한 증거다.

결론적으로, M²RNN은 비선형 RNN의 오래된 한계인 “작은 상태 → 낮은 성능”을 행렬 상태와 외적 확장으로 극복하고, 하드웨어 친화적인 구현을 통해 대규모 언어 모델에 실용적으로 적용할 수 있는 새로운 설계 패러다임을 제시한다.