재귀 혼합: 적응적 토큰 수준 계산을 위한 동적 재귀 깊이 학습

초록

언어 모델의 규모 확장은 인상적인 성능을 발휘하지만, 이에 수반되는 계산 및 메모리 요구량으로 인해 학습과 배포 모두 비용이 많이 든다. 기존의 효율성 개선 노력은 일반적으로 매개변수 공유 또는 적응형 계산 중 하나를 대상으로 하여, 두 가지를 동시에 달성하는 방법에 대한 의문을 남겼다. 본 연구에서는 Mixture-of-Recursions(MoR)를 소개한다. MoR는 단일 Recursive Transformer 내에서 두 가지 효율성 축을 결합한 통합 프레임워크로, 재귀 단계 간 공유된 레이어 스택을 재사용하여 매개변수 효율성을 달성한다. 동시에 경량 라우터는 개별 토큰에 서로 다른 재귀 깊이를 동적으로 할당함으로써 적응형 토큰 수준 사고를 가능하게 한다. 이를 통해 MoR는 주어진 재귀 깊이에서 여전히 활성화된 토큰들 간에만 이차적 주의 계산을 집중시키고, 이들의 키-값 쌍만 선택적으로 캐싱함으로써 메모리 접근 효율성을 더욱 개선한다. 이러한 핵심 메커니즘 외에도, 첫 번째 재귀에서의 키-값 쌍을 재사용하는 KV 공유 변형을 제안하여, 특히 프리필 지연 시간과 메모리 사용량을 줄이도록 설계하였다. 135M에서 1.7B 매개변수에 이르는 다양한 모델 규모에서 MoR는 새로운 파레토 최적점을 형성한다: 동일한 학습 FLOPs와 더 작은 모델 크기에서, 검증 perplexity를 크게 낮추고 few-shot 정확도를 향상시키며, 기본 및 기존 재귀 모델 대비 더 높은 처리량을 제공한다. 이러한 성과는 MoR가 대규모 모델의 품질을 유지하면서도 대규모 모델의 비용을 발생시키지 않는 효과적인 접근법임을 입증한다.

English

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 decrease prefill latency and 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: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

초록

Support