[논문 리뷰] Tiny Recursive Reasoning with Mamba-2 Attention Hybrid

Wenlong Wang, Fergal Reid|arXiv (Cornell University)|2026. 02. 12.

Advanced Graph Neural Networks인용 수 0

한 줄 요약

본 논문은 작은 재귀 추론 모델(TRM)의 Transformer 블록을 Mamba-2 하이브리드 연산자로 대체하고, ARC-AGI-1 pass@2에서의 성능 향상 및 높은-K 커버리지를 보이면서도 pass@1은 유사하게 유지되어, 아주 작은 모델에서 재귀 연산자로서 Mamba-2의 실용성을 검증한다.

ABSTRACT

Recent work on recursive reasoning models like TRM demonstrates that tiny networks (7M parameters) can achieve strong performance on abstract reasoning tasks through latent recursion -- iterative refinement in hidden representation space without emitting intermediate tokens. This raises a natural question about operator choice: Mamba-2's state space recurrence is itself a form of iterative refinement, making it a natural candidate for recursive reasoning -- but does introducing Mamba-2 into the recursive scaffold preserve reasoning capability? We investigate this by replacing the Transformer blocks in TRM with Mamba-2 hybrid operators while maintaining parameter parity (6.83M vs 6.86M parameters). On ARC-AGI-1, we find that the hybrid improves pass@2 (the official metric) by +2.0\% (45.88\% vs 43.88\%) and consistently outperforms at higher K values (+4.75\% at pass@100), whilst maintaining pass@1 parity. This suggests improved candidate coverage -- the model generates correct solutions more reliably -- with similar top-1 selection. Our results validate that Mamba-2 hybrid operators preserve reasoning capability within the recursive scaffold, establishing SSM-based operators as viable candidates in the recursive operator design space and taking a first step towards understanding the best mixing strategies for recursive reasoning.

연구 동기 및 목표

• Mamba-2 상태 공간 재귀가 능력을 잃지 않으면서 작은 재귀 추론 모델에서 Transformer 블록을 대체할 수 있는지 조사한다.
• Mamba-2 하이브리드가 추상 추론 벤치마크(ARC-AGI-1)와 기타 과제(Sudoku, Maze)에 미치는 영향을 평가한다.
• 잠재 재귀 추론에서 연산자 선택이 후보 커버리지와 상위-1 선택에 어떤 영향을 주는지 특성화한다.

제안 방법

• 같은 잠재 상태 z_H 및 z_L를 유지하면서 3개의 바깥 순환과 4–6개의 안쪽 순환으로 TRM 재귀 구조를 보존한다.
• 매 스텝의 Transformer 블록을 두 가지 변형에서 Mamba-2 하이브리드 스택으로 대체한다: TR-mamba2attn (Mamba-2 → Mamba-2 → Attention → MLP) 및 TR-mamba2mlpt (Mamba-2 → Mamba-2 → MLP-t).
• 연산자 효과를 분리하기 위해 원래 TRM-attn과 매개변수를 맞춘다(대략 6.83M 대 6.86M).
• 재귀 계산을 안정시키기 위해 post-norm(RMSNorm)을 사용한다.
• ARC-AGI-1, Sudoku-Extreme, Maze-30x30-Hard에서 pass@K(K ∈ {1,2,5,10,100,1000})와 해당되는 경우 정확도(Exact accuracy)를 평가한다.

실험 결과

연구 질문

• RQ1 Mamba-2 하이브리드 연산자가 TRM 유사 재귀 스캐폴드에서 추론 능력을 보존할 수 있는가?
• RQ2 Mamba-2 하이브리드가 후보 커버리지(pass@K가 큰 경우)을 개선하면서 상위-1 정확도를 희생하지 않는가?
• RQ3 Attention 기반 TRM과 비교했을 때 Sudoku 및 Maze와 같은 다른 추론 과제에서 Mamba-2 하이브리드의 성능은 어떤가?
• RQ4 재귀 추론에서 Mamba-2 하이브리드 사용 시 커버리지와 선택 간의 트레이드오프는 무엇인가?

주요 결과

모델	매개변수	pass@1	pass@2	pass@5	pass@10	pass@100	pass@1000
TRM-attn	6.83M	40.75	43.88	49.25	52.13	60.50	65.50
TR-mamba2attn	6.86M	40.50	45.88	51.88	54.50	65.25	69.75

• ARC-AGI-1에서 하이브리드가 pass@2를 2.0 포인트 향상시키며(45.88% 대 43.88%),
• 하이브리드는 더 높은 K 값에서 일관되게 우수한 성능을 보이며 pass@100에서 +4.75%에 도달하고 pass@1의 동등성은 유지(−0.25%)한다.
• Sudoku-Extreme은 MLP-t 변형이 유리하며, TRM-mlp-t가 87.4% 정확도이고 TR-mamba2mlpt가 84.2%로 주의 기반 모델보다 높다.
• Maze-30x30-Hard는 불안정한 경향을 보이며 TR-mamba2attn은 80.6%에 도달하는 반면 MLP-t 변형은 실패(0.0%)로 나타나, 과제 의존적으로 하이브리드의 효과가 다름을 시사한다.
• ARC-AGI-1 결과는 Mamba-2 하이브리드의 후보 커버리지 개선을 시사하지만 상위-1 정확도를 저하시지 않는다.
• Post-norm은 비정형 루프의 안정적 재귀에 필수적으로 강조된다(안정적 재귀를 위한 필수 요소).