Tiny Recursive Reasoning with Mamba-2 Attention Hybrid

이 논문은 Mamba-2 상태 공간 모델을 TRM 의 재귀적 추론 구조에 통합하여 파라미터 수를 동등하게 유지한 결과, ARC-AGI-1 벤치마크에서 상위 K 개 정답률 (pass@K) 을 향상시키면서도 1 순위 정확도를 유지함으로써 재귀적 추론을 위한 SSM 기반 연산자의 유효성을 입증했습니다.

핵심

1. 배경: "생각하는 시간"이 더 중요하다?

과거에는 AI 가 똑똑해지려면 "머리 (파라미터)"를 엄청나게 크게 키워야 한다고 믿었습니다. 하지만 최근 연구들은 **"생각하는 과정 (반복적인 사고)"**이 더 중요하다고 말합니다.

• 기존 방식 (TRM): 아주 작은 AI(700 만 개의 파라미터, 즉 아주 작은 뇌) 가 문제를 풀 때, 답을 바로 말하지 않고 숨겨진 공간에서 여러 번 생각을 반복합니다. 마치 퍼즐을 풀 때 종이에 답을 쓰지 않고, 머릿속으로 "아니야, 이건 아니야, 저건 어때?"라고 계속 고민하는 것과 같습니다.
• 문제: 이 "생각하는 과정"을 담당하는 부품을 기존에는 **Transformer(attention)**라고 불리는 복잡한 장치를 썼습니다. 그런데 이 장치는 무겁고 느립니다.

2. 실험: "새로운 엔진 (Mamba-2)"을 달아보자

연구진은 "생각하는 과정"을 담당하는 장치를 Mamba-2라는 새로운, 더 가볍고 빠른 엔진으로 바꿔봤습니다.

• 비유: 같은 차체 (모델 크기) 에, 기존에 달았던 무거운 **V8 엔진 (Transformer)**을 떼어내고, 최신형 **하이브리드 엔진 (Mamba-2)**을 달아본 것입니다.
• 목표: 엔진만 바꿨을 때, 차의 성능 (추리 능력) 이 떨어지지 않고 오히려 좋아질까?

3. 결과: "더 많은 후보를 찾아내는 능력"이 좋아졌다

결과는 놀라웠습니다. 엔진을 바꾼 하이브리드 모델이 기존 모델보다 더 좋은 성적을 냈습니다.

• 핵심 발견 (Coverage vs Selection):

  • 기존 모델 (Transformer): "내가 가장 확신하는 답 하나를 골라내서 정답일 가능성이 높게 맞췄다." (1 등 선정 능력은 비슷함)
  • 새로운 모델 (Mamba-2 Hybrid): "내가 생각한 답들 중에서 정답이 포함될 확률이 훨씬 높은 다양한 후보군을 많이 만들어냈다." (정답을 포함하는 '후보 목록'의 범위가 넓어짐)

• 일상적인 비유:

  • 기존 모델: 친구가 "내 생각엔 이 식당이 최고야!"라고 단 한 군데를 강력 추천합니다. (맞을 수도, 틀릴 수도 있음)
  • 새로운 모델: 친구가 "이 식당, 저 식당, 그리고 그 옆 골목 식당까지 다양하게 10 군데를 추천해 줍니다. 그중에서 정답인 식당이 꼭 들어있습니다."
  • 결과: 정답을 고르는 '1 등' 선정 능력은 비슷하지만, 정답을 '포함'하고 있는 후보 목록의 크기가 훨씬 커서 전체적인 성공률이 높아진 것입니다.

4. 왜 이런 일이 일어났을까?

• Mamba-2 의 특징: 이 엔진은 정보를 순서대로 (한 방향으로) 처리하는 데 매우 능숙합니다.
• 시너지: 복잡한 추리 문제 (ARC-AGI) 는 공간적, 시각적 패턴을 이해해야 합니다. Mamba-2 가 순차적으로 정보를 처리하면서, 기존 엔진과는 **다른 각도에서 문제를 바라보는 '새로운 생각의 경로'**를 만들어냈습니다.
• 결론: 서로 다른 엔진을 섞어쓰면, AI 가 문제를 풀 때 더 다양한 시나리오를 상상할 수 있게 되어, 정답을 놓치지 않을 확률이 높아진 것입니다.

5. 요약 및 의미

이 연구는 **"AI 의 두뇌 구조를 바꾸면, 아주 작은 모델로도 더 똑똑해질 수 있다"**는 것을 증명했습니다.

• 핵심 메시지: 거대한 AI 를 만드는 것뿐만 아니라, **"어떻게 생각하는가 (구조와 엔진)"**를 바꾸는 것이 더 중요할 수 있습니다.
• 미래: 이제 AI 개발자들은 'Transformer'만 고집하지 않고, 'Mamba' 같은 새로운 엔진을 섞어쓰며 더 효율적이고 똑똑한 AI 를 만들 수 있는 길을 열었습니다.

한 줄 요약:

"작은 AI 에 새로운 엔진을 달으니, 정답을 딱 하나만 고르는 능력은 그대로 유지하면서, 정답을 포함할 수 있는 '생각의 목록'을 훨씬 더 많이 만들어내서 전체적인 성적이 좋아졌습니다!"

기술 요약

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 최근 TRM(Tiny Recursive Model) 과 같은 연구는 7M(백만) 파라미터 규모의 초소형 네트워크가 잠재 공간 (latent space) 에서의 반복적 정제 (latent recursion) 를 통해 추상적 추론 과제를 해결할 수 있음을 입증했습니다. 이는 모델의 규모보다는 '추론 시간 (inference-time computation)'과 '반복 과정'이 추론 능력의 핵심임을 시사합니다.
• 기존 접근법의 한계: 기존 TRM 은 Transformer 블록 (Attention 기반) 을 기본 연산자 (operator) 로 사용합니다. 그러나 Attention 은 계산 비용이 높고, 명시적 사고 과정 (Chain-of-Thought) 은 불필요한 토큰 생성으로 인해 비효율적일 수 있으며, 모델의 실제 계산 과정을 반영하지 못할 수 있습니다.
• 핵심 질문: 반복적 정제의 본질과 유사한 상태 공간 모델 (State Space Models, SSM) 인 Mamba-2를 TRM 의 재귀적 구조에 도입할 경우, 추론 능력이 유지되거나 향상될 수 있는가? 특히 Mamba-2 의 내재적 반복 (recurrence) 이 외부 루프 (outer-loop) 재귀와 어떻게 상호작용하는지 규명하는 것이 목표입니다.

2. 방법론 (Methodology)

저자는 TRM 아키텍처를 유지하면서, 각 단계의 업데이트 연산자를 Transformer 블록에서 Mamba-2 와 Attention 의 하이브리드 구조로 교체했습니다.

• 아키텍처 변경 (TR-mamba2attn):
  • 기존 TRM 의 Attention-only 블록을 Mamba-2 → Mamba-2 → Attention → MLP 파이프라인으로 대체했습니다.
  • 이유: Mamba-2 는 순차적 (causal) 처리에 강점이 있지만, ARC-AGI 나 미로 찾기 같은 공간적 작업에서는 양방향 (bidirectional) 의존성 파악이 필수적입니다. 따라서 Mamba-2 의 효율적인 순차 처리 능력과 Attention(또는 MLP-t) 의 전역적 (all-to-all) 정보 교환 능력을 결합하여 하이브리드 구조를 설계했습니다.
• 파라미터 매칭 (Parameter Parity):
  • 공정한 비교를 위해 기존 TRM(attn) 과 하이브리드 모델의 파라미터 수를 거의 동일하게 맞췄습니다 (기존: 6.83M, 하이브리드: 6.86M).
  • 숨겨진 차원 (Hidden size) 은 512 로 고정했습니다.
• 정규화 전략 (Post-Norm):
  • 재귀적 구조에서 반복 적용 시 발산 (NaN) 을 방지하기 위해 Post-Norm을 사용했습니다. 이는 잔여 연결 (residual connection) 후 정규화를 수행하여 재귀 깊이와 무관하게 숨겨진 상태의 크기를 제한합니다. 이는 재귀 안정성에 필수적입니다.

3. 주요 기여 (Key Contributions)

1. 최초의 Mamba-하이브리드 잠재 재귀 모델: 재귀적 잠재 추론을 위한 Mamba-2 기반 하이브리드 연산자를 최초로 제안하고 구현했습니다.
2. ARC-AGI-1 에서의 성능 향상: 공식 지표인 pass@2에서 기존 모델 대비 +2.0% 향상 (43.88% → 45.88%) 을 달성했습니다. 또한 pass@100 등 높은 K 값에서 더 큰 격차 (+4.75%) 를 보였습니다.
3. 커버리지 vs 선택 (Coverage vs Selection) 트레이드오프 분석: 하이브리드 모델이 정답을 생성할 확률 (커버리지) 은 높이지만, 최상위 1 순위 선택 (top-1 selection) 은 기존 모델과 유사함을 규명했습니다. 이는 Mamba-2 가 후보 집단의 다양성을 증가시키는 역할을 함을 시사합니다.

4. 실험 결과 (Results)

실험은 ARC-AGI-1, Sudoku-Extreme, Maze-30x30-Hard 세 가지 벤치마크에서 수행되었습니다.

• ARC-AGI-1 (추상 추론):

  • Pass@2 (공식 지표): 하이브리드 모델이 +2.0% 우위.
  • Pass@K 패턴: K 가 커질수록 (Pass@100, Pass@1000) 하이브리드 모델의 우위가 더욱 커짐. 이는 모델이 다양한 후보군 중 정답을 포함할 확률이 높음을 의미합니다.
  • 통계적 분석: 하이브리드 모델은 더 많은 고유한 후보 (Unique Candidates) 를 생성하고 투표 엔트로피가 높았으며, 반면 기존 모델은 상위 1 순위 후보에 투표가 더 집중되었습니다. 이는 하이브리드가 **탐색 (Exploration)**에, 기존 모델이 **선정 (Exploitation)**에 강점이 있음을 보여줍니다.
  • 난이도별 분석: 어려운 문제 (Hard) 에서는 하이브리드 모델이 더 큰 이점을 보였고, 쉬운 문제에서는 기존 모델이 약간 우세했습니다. 이는 두 아키텍처가 서로 다른 문제 유형에 보완적인 강점을 가진다는 것을 시사합니다.

• Sudoku (제약 충족):

  • MLP-t(Attention 대체) 기반 모델이 가장 높은 정확도를 보였으며, Mamba-2 하이브리드도 Attention-only 모델보다 성능이 좋았습니다. 이는 작은 격자 (9x9) 에서는 밀집된 전역 통신이 중요함을 보여줍니다.

• Maze (공간 추론):

  • 30x30 미로에서는 MLP-t 기반 모델이 실패 (0.0%) 했지만, Mamba-2 + Attention 하이브리드는 80.6% 의 높은 정확도를 기록했습니다. 이는 대규모 공간 추론 작업에서는 Mamba-2 의 순차적 처리 능력이 핵심임을 보여줍니다.

5. 의의 및 결론 (Significance & Conclusion)

• 재귀 연산자 설계 공간의 확장: 이 연구는 Mamba-2 와 같은 SSM 기반 연산자가 재귀적 추론 프레임워크에서도 유효하며, 오히려 특정 지표 (커버리지) 에서 Transformer 를 능가할 수 있음을 입증했습니다.
• 효율성과 성능의 균형: Mamba-2 의 선형 복잡도 (linear complexity) 특성을 재귀 구조에 접목함으로써, 추론 시간 계산 (test-time compute) 을 확장하면서도 효율성을 유지할 수 있는 가능성을 제시했습니다.
• 향후 방향: 재귀 루프를 외부 반복이 아닌 SSM 의 내재적 상태 업데이트 (inner recurrence) 로 통합할 수 있는지, 그리고 다양한 혼합 전략 (mixing strategies) 이 재귀 추론에 어떻게 기여하는지에 대한 심층 연구의 기초를 마련했습니다.

요약하자면, 이 논문은 초소형 재귀 모델에 Mamba-2 를 도입하여 추론 능력을 유지하면서도 후보 생성의 다양성과 정답 포함률 (coverage) 을 크게 향상시켰으며, 이는 추론 효율성과 성능을 동시에 개선할 수 있는 새로운 아키텍처 설계 패러다임을 제시합니다.