Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

Scaling language models unlocks impressive capabilities, but the accompanying computational and memory demands make both training and deployment expensive. Existing efficiency efforts typically target either parameter sharing or adaptive computation, leaving open the question of how to attain both simultaneously. We introduce Mixture-of-Recursions (MoR), a unified framework that combines the two axes of efficiency inside a single Recursive Transformer. MoR reuses a shared stack of layers across recursion steps to achieve parameter efficiency, while lightweight routers enable adaptive token-level thinking by dynamically assigning different recursion depths to individual tokens. This allows MoR to focus quadratic attention computation only among tokens still active at a given recursion depth, further improving memory access efficiency by selectively caching only their key-value pairs. Beyond these core mechanisms, we also propose a KV sharing variant that reuses KV pairs from the first recursion, specifically designed to further decrease memory footprint. Across model scales ranging from 135M to 1.7B parameters, MoR forms a new Pareto frontier: at equal training FLOPs and smaller model sizes, it significantly lowers validation perplexity and improves few-shot accuracy, while delivering higher throughput compared with vanilla and existing recursive baselines. These gains demonstrate that MoR is an effective path towards large-model quality without incurring large-model cost.

한 줄 요약

Mixture-of-Recursions(MoR)는 재귀 트랜스포머 내에서 파라미터 공유와 적응형 계산을 결합하여 성능과 효율성을 동시에 향상시키는 새로운 프레임워크이다.

핵심 기여도

**Mixture-of-Recursions(MoR)** 프레임워크를 제안하여 파라미터 효율성과 적응형 토큰 수준 계산을 통합.
**가벼운 라우터**를 통해 토큰별 재귀 깊이를 동적으로 할당함으로써 불필요한 FLOPs를 절감.
**재귀별 KV 캐싱**을 통해 메모리 트래픽을 줄이고, **KV 공유 변형**을 도입하여 메모리 용량을 추가로 절감.
135M~1.7B 파라미터 규모에서 베이스라인 대비 **+1.6%~+3.2%의 few-shot 정확도 향상**과 **더 높은 처리량** 달성.

핵심 아이디어

기존 트랜스포머 모델은 파라미터 공유(예: layer tying)와 적응형 계산(예: early-exiting)을 별도로 다루는 경향이 있었으나, MoR은 이 두 효율성 축을 하나의 재귀 트랜스포머 아키텍처 내에서 통합한다. MoR은 공유된 레이어 스택을 재귀 단계에 걸쳐 재사용함으로써 파라미터 수를 줄이고, **가벼운 라우터**를 통해 각 토큰에 필요한 재귀 깊이를 동적으로 할당함으로써 불필요한 계산을 제거한다. 이는 토큰별로 재귀 깊이를 조정하여 **적응형 토큰 수준의 '생각'**을 가능하게 한다. 또한, MoR은 재귀 단계별로 활성화된 토큰의 key-value 쌍만 캐싱함으로써 메모리 트래픽을 줄이고, **KV 공유 변형**을 통해 첫 번째 재귀의 KV를 재사용함으로써 메모리 용량을 추가로 절감한다.

기술적 접근법

**Mixture-of-Recursions(MoR)**: 재귀 트랜스포머 내에서 파라미터 공유와 적응형 재귀 깊이 할당을 결합.
**가벼운 라우터**: 토큰별 재귀 깊이를 동적으로 결정.
**재귀별 KV 캐싱**: 활성화된 토큰의 key-value 쌍만 캐싱하여 메모리 트래픽 최적화.
**KV 공유 변형**: 첫 번째 재귀의 KV를 재사용하여 메모리 용량 절감.
**모델 규모**: 135M, 360M, 730M, 1.7B 파라미터.
**데이터셋**: FineWeb-Edu, LAMBADA, HellaSwag, PIQA, WinoGrande, ARC, MMLU.

주요 결과

**FineWeb-Edu 검증 세트**에서 MoR은 1.7B 파라미터 모델 기준 **NLL 2.16** 달성 (베이스라인 대비 -0.23).
**Few-shot 정확도**에서 360M 파라미터 모델 기준 **+1.6%~+3.2% 개선**.
**1.7B 파라미터 모델**에서 **베이스라인 대비 1.6× 처리량 향상**.
**360M 파라미터 모델**에서 베이스라인 730M 모델과 동일한 성능 달성.

의의 및 한계

MoR은 파라미터 효율성과 적응형 계산을 동시에 달성하는 새로운 아키텍처를 제시하며, 기존 재귀 트랜스포머의 한계를 극복한다. 특히, **동적 재귀 깊이 할당**과 **재귀별 KV 캐싱**은 메모리와 계산 효율성을 동시에 개선하며, **1.7B 파라미터 모델에서도 1.6× 처리량 향상**을 보인다. 그러나, 재귀 트랜스포머 자체는 파라미터 공유로 인해 성능 저하가 발생할 수 있으며, 이는 MoR의 라우터와 KV 캐싱 메커니즘으로 완화된다. 또한, 라우터의 가중치 학습 과정에서 추가적인 트레이닝 복잡성이 발생할 수 있다.

실용적 활용

MoR은 대규모 언어 모델의 성능을 유지하면서도 **메모리와 계산 비용을 30% 이상 절감**할 수 있어, 클라우드 외의 일반 서버나 모바일 기기에서도 효율적으로 배포 가능하다. 특히, **few-shot 학습**이나 **실시간 추론**이 필요한 산업 분야(예: 고객 서비스, 번역, 콘텐츠 생성)에서 유용하게 활용될 수 있다.