M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

이 논문은 TC$^0$ 복잡도 한계를 극복하고 긴 시퀀스에서의 상태 추적 및 언어 모델링 성능을 획기적으로 개선하기 위해 행렬 값 상태와 비선형 전이를 도입한 M$^2$RNN 아키텍처를 제안하고, 이를 하이브리드 모델에 적용하여 기존 선형 어텐션 기반 모델보다 우수한 성능을 입증했습니다.

핵심

🧠 M2RNN: AI 의 '기억력'을 혁신한 새로운 뇌 구조

이 논문은 인공지능 (AI) 이 글을 쓰거나 코드를 작성할 때, 매우 긴 문맥을 기억하고 논리적으로 추론하는 능력을 획기적으로 개선한 새로운 기술인 M2RNN을 소개합니다.

기존의 AI 모델 (트랜스포머) 은 병렬 처리가 빨라 글을 빠르게 읽을 수 있지만, 복잡한 논리나 긴 이야기의 흐름을 기억하는 데는 한계가 있었습니다. 반면, M2RNN 은 이 두 가지 장점을 모두 잡은 '초능력'을 가진 새로운 구조입니다.

이 복잡한 기술을 일상적인 비유로 쉽게 설명해 드리겠습니다.

---

1. 문제: 왜 기존 AI 는 긴 이야기를 잊어버릴까?

기존의 AI 모델 (트랜스포머) 은 도서관 사서와 같습니다.

• 장점: 책장 (데이터) 을 한 번에 훑어보며 필요한 정보를 빠르게 찾을 수 있습니다 (병렬 처리).
• 단점: 하지만 책장 전체를 한 번에 기억할 수 없기 때문에, 책이 너무 많으면 (긴 문맥) 중요한 정보를 놓치거나, "누가 언제 무엇을 했지?" 같은 복잡한 추론 (예: 체스 게임 기록, 코드 실행) 을 하기가 어렵습니다.

한편, 기존에 있던 **RNN(순환 신경망)**은 메모리 노트를 쓰는 학생과 같습니다.

• 장점: 앞선 내용을 하나하나 기억하며 논리를 이어갈 수 있습니다.
• 단점: 노트의 페이지 수가 매우 적어 (상태 크기 제한) 긴 이야기를 다 적어두면 중요한 내용이 지워지거나, 노트를 한 줄씩만 읽을 수 있어 속도가 매우 느립니다.

2. 해결책: M2RNN (행렬 - 행렬 RNN) 의 등장

저자들은 이 문제를 해결하기 위해 M2RNN이라는 새로운 구조를 만들었습니다. 이걸 이해하기 위해 두 가지 핵심 비유를 사용하겠습니다.

비유 1: '단순한 메모지'에서 '거대한 백과사전'으로 (상태 크기 확장)

기존 RNN 은 정보를 한 줄의 메모 (벡터) 로만 저장했습니다. 정보가 쌓이면 이전 내용이 지워지기 일쑤였습니다.
하지만 M2RNN 은 정보를 **거대한 백과사전 (행렬)**으로 저장합니다.

• 비유: 학생이 작은 메모지 대신, 수백 페이지의 두꺼운 공책을 사용하게 된 것입니다.
• 효과: 이제 AI 는 긴 이야기 속의 모든 등장인물, 사건, 인과관계를 한눈에 펼쳐놓고 기억할 수 있게 되었습니다. 덕분에 "누가 언제 무엇을 했는지"를 정확히 추적하는 능력이 비약적으로 향상되었습니다.

비유 2: '수동 정리'에서 '자동 분류 시스템'으로 (하드웨어 효율성)

기존의 메모 방식은 컴퓨터 칩 (GPU) 이 일을 할 때, 불필요한 공간 (패딩) 을 채우는 낭비가 많았습니다. 마치 택배 박스에 빈 공간만 가득 채운 채로 보내는 것과 같습니다.
M2RNN 은 이 공책을 컴퓨터 칩이 가장 좋아하는 형태로 딱 맞게 정리했습니다.

• 효과: 컴퓨터 칩이 일을 할 때 빈 공간 없이 100% 효율적으로 작동합니다. 그래서 속도가 느려질 것 같았던 RNN 이지만, 실제로는 매우 빠르게 작동할 수 있게 되었습니다.

3. M2RNN 의 놀라운 성과

이 논문에서는 M2RNN 을 기존 모델에 섞어 사용했을 때의 결과를 보여줍니다.

• 완벽한 기억력: 훈련할 때 보지 못한 아주 긴 이야기에서도, 등장인물의 관계나 코드 실행 결과를 100% 정확히 기억해냅니다. (기존 모델은 긴 이야기일수록 기억력이 떨어집니다.)
• 혼합 모델의 승리: M2RNN 을 기존 AI 모델 (트랜스포머나 Mamba 등) 에 하나만 섞어도 성능이 크게 좋아집니다.
  • 비유: 거대한 도서관 (기존 AI) 에 **한 명의 슈퍼 사서 (M2RNN)**를 고용하자, 복잡한 질문에도 정답을 찾아내는 능력이 비약적으로 상승했습니다.
• 실제 적용: 70 억 개의 파라미터를 가진 대형 모델에서, M2RNN 을 섞은 모델은 기존 모델보다 언어 이해도, 긴 문맥 기억, 복잡한 추론에서 모두 더 좋은 점수를 받았습니다.

4. 결론: 왜 이것이 중요한가?

M2RNN 은 **"빠른 속도 (병렬 처리)"와 "깊은 이해 (논리적 추론)"**라는 서로 상충되던 두 가지 목표를 동시에 달성하게 해줍니다.

• 기존: 빠른 모델은 멍청하고, 똑똑한 모델은 느리다.
• M2RNN: 빠르면서도 똑똑한 AI 의 시대를 열었습니다.

이 기술이 발전하면, AI 는 더 긴 소설을 읽고 줄거리를 완벽하게 이해하거나, 복잡한 코드를 작성하고 디버깅하는 등, 인간과 더 자연스럽게 소통하고 복잡한 작업을 수행할 수 있게 될 것입니다.

---

한 줄 요약:

M2RNN 은 AI 에게 '작은 메모지' 대신 '거대한 공책'을 주고, 컴퓨터 칩이 이를 가장 효율적으로 읽을 수 있게 정리해준 기술로, AI 가 긴 이야기를 잊지 않고 논리적으로 추론할 수 있게 만든 혁신입니다.

기술 요약

M2RNN: 확장 가능한 언어 모델을 위한 행렬 값 상태의 비선형 RNN 기술 요약

이 논문은 현대 언어 모델링의 한계를 극복하기 위해 제안된 행렬 값 상태 (Matrix-Valued States) 를 가진 새로운 비선형 순환 신경망 (RNN) 아키텍처인 M2RNN (Matrix-to-Matrix RNN) 을 소개합니다.

1. 문제 정의 (Problem)

현재 언어 모델의 주류인 Transformer는 병렬 처리가 가능하지만, 학습 시 $O(N^2)$의 시간 복잡도와 추론 시 선형적으로 증가하는 메모리 요구사항으로 인해 확장성에 한계가 있습니다. 이를 해결하기 위해 선형 RNN (Linear RNNs, 예: Mamba, Gated DeltaNet) 이 대안으로 등장했으나, 다음과 같은 근본적인 한계를 가지고 있습니다.

• 제한된 상태 추적 (State Tracking) 능력: 선형 RNN 은 $TC_0$ 복잡도 클래스에 속하며, 체스 수 추적, 코드 실행, 엔티티 추적 등 비선형 RNN 만이 해결할 수 있는 복잡한 상태 추적 작업 (NC1 클래스) 을 수행하는 데 이론적, 실용적 한계가 있습니다.
• 낮은 언어 모델링 및 검색 성능: 기존 비선형 RNN (LSTM, GRU 등) 은 표현력은 높지만, 상태 크기 (State Size) 가 벡터 형태로 제한되어 있어 언어 모델링 성능과 문맥 내 검색 (In-context Retrieval) 성능이 선형 RNN 나 Transformer 에 비해 현저히 낮습니다.
• 하드웨어 활용도 저하: 기존 비선형 RNN 은 시퀀스 길이에 따라 순차적으로 계산되어야 하므로 병렬화가 어렵고, 배치 차원 패딩 (Padding) 으로 인해 Tensor Core 의 연산 효율이 떨어지는 문제가 있습니다.

2. 방법론 (Methodology)

저자들은 선형 RNN 의 효율성과 비선형 RNN 의 표현력을 결합한 M2RNN을 제안합니다.

핵심 아키텍처

• 행렬 값 상태 (Matrix-Valued Hidden States): 기존 RNN 의 벡터 상태 ($h_t \in \mathbb{R}^d$) 를 행렬 상태 ($H_t \in \mathbb{R}^{K \times V}$) 로 확장합니다. 이는 외적 (Outer Product) $k_t v_t^\top$을 통해 상태를 업데이트하는 방식으로, 파라미터 수를 크게 늘리지 않으면서 상태 용량을 획기적으로 증가시킵니다.
• 비선형 상태 전이: $Z_t = \tanh(H_{t-1}W + k_t v_t^\top)$와 같은 비선형 전이를 사용하여 복잡한 상태 추적을 가능하게 합니다.
• 병렬 계산 가능한 Forget Gate: LSTM/GRU 와 달리, 잊음 게이트 (Forget Gate) 가 이전 상태에 의존하지 않고 입력에만 의존하여 계산되므로, 비선형성에도 불구하고 게이트 계산은 병렬화가 가능합니다.
• 하이브리드 아키텍처: M2RNN 레이어를 Attention 레이어와 선형 RNN 레이어 (Mamba-2, Gated DeltaNet) 와 혼합하여 사용합니다. 이는 Attention 이 문맥 검색을 담당하고, M2RNN 이 비선형 상태 추적을 담당하도록 설계되었습니다.

시스템 최적화

• 하드웨어 효율성: 행렬 값 상태의 외적 구조를 활용하여 배치 크기 (Batch Size) 와 관계없이 Tensor Core 를 효율적으로 사용할 수 있도록 설계했습니다. 이는 기존 FlashRNN 의 패딩 오버헤드를 제거합니다.
• 분산 학습 (Tensor Parallelism): M2RNN 레이어를 위한 두 가지 텐서 병렬화 전략 (토폴로지 인식형 및 토폴로지 무관형) 을 제안하여 대규모 모델 학습을 지원합니다.

3. 주요 기여 (Key Contributions)

1. 완벽한 상태 추적 일반화: M2RNN 은 훈련 시 보지 못한 시퀀스 길이에서도 완벽한 상태 추적 (State Tracking) 능력을 보여주며, 선형 RNN 이 실패하는 복잡한 작업 (예: $S_3$ 순열 그룹) 을 해결합니다.
2. 상태 크기의 중요성 규명: 비선형 RNN 의 낮은 성능이 비선형성 자체 때문이 아니라 상태 크기의 부족 때문임을 실험적으로 증명했습니다. 행렬 상태 확장이 언어 모델링 및 검색 성능 향상의 핵심 요소임을 입증했습니다.
3. 하이브리드 모델의 성능 우위: M2RNN 레이어를 기존 하이브리드 모델 (Mamba-2 또는 Gated DeltaNet 기반) 에 소량 (예: 레이어 1 개) 만 추가해도, 전체 레이어를 M2RNN 으로 교체한 것과 유사한 성능 향상을 얻으면서도 학습 처리량 (Throughput) 은 거의 유지됩니다.
4. 효율적인 하드웨어 활용: 패딩 (Padding) 으로 인한 불필요한 연산 (FLOPs) 을 제거하고 Tensor Core 를 효율적으로 활용하는 커널을 구현했습니다.

4. 실험 결과 (Results)

• 언어 모델링 (Language Modeling):
  • 410M 및 7B MoE 모델에서 M2RNN 은 Mamba-2 및 Gated DeltaNet 과 유사하거나 더 나은 퍼플렉시티 (Perplexity) 를 기록했습니다.
  • 하이브리드 설정 (Hybrid M2RNN) 에서 Gated DeltaNet 하이브리드 모델보다 0.4~0.5 포인트 더 낮은 퍼플렉시티를 달성했습니다.
• 문맥 내 검색 (In-Context Retrieval):
  • LongBench 및 RULER 벤치마크에서 M2RNN 을 포함한 하이브리드 모델은 기존 선형 RNN 및 Transformer 기반 모델보다 최대 8 포인트 더 높은 정확도를 보였습니다.
  • 특히, 훈련 시 보지 못한 긴 컨텍스트 (Unseen Context Lengths) 에서도 뛰어난 일반화 성능을 발휘했습니다.
• 실제 데이터 검색:
  • SQuAD, NQ 등 실제 데이터셋 기반 검색 작업에서 Hybrid M2RNN 은 Transformer++ 보다도 우수한 성능을 보이며, 기존 순환 모델 중 가장 높은 점수를 기록했습니다.
• 학습 효율성:
  • M2RNN 레이어 하나만 추가한 하이브리드 모델은 Gated DeltaNet 하이브리드 모델 대비 학습 처리량이 6% 이내로 감소하는 반면, 모든 벤치마크에서 정확도가 크게 향상되었습니다.

5. 의의 및 결론 (Significance)

이 논문은 비선형 RNN 레이어가 효율적이고 확장 가능한 언어 모델의 핵심 구성 요소가 될 수 있음을 입증했습니다.

• 이론적 기여: 선형 RNN 의 표현력 한계를 극복하면서도 Transformer 의 병렬성 장점을 일부 유지하는 새로운 아키텍처를 제시했습니다.
• 실용적 가치: M2RNN 레이어를 기존 모델에 소량만 도입해도 (예: 레이어 1 개) 성능을 획기적으로 개선할 수 있어, 실제 대규모 언어 모델 (LLM) 에 적용하기 매우 실용적입니다.
• 시스템적 혁신: 행렬 값 상태와 외적 연산을 활용한 설계는 GPU Tensor Core 를 효율적으로 활용하게 하여, 비선형 RNN 의 학습 비용을 크게 낮추었습니다.

결론적으로, M2RNN 은 복잡한 상태 추적 작업이 필요한 코드 생성, 논리적 추론, 긴 문맥 이해 등 차세대 언어 모델의 필수적인 블록으로 자리매김할 잠재력을 가지고 있습니다.