Mixture-of-Recursions: 토큰 단위 적응적 재귀 깊이 학습
목차
개요
Transformer를 수천억 파라미터 규모로 확장하면서 few-shot 일반화와 추론 능력이 크게 향상되었지만, 그에 따르는 메모리 사용량과 연산 요구량 때문에 학습과 배포 비용이 함께 증가했다. 이 비용 문제를 완화하는 대표적인 축은 두 가지다. 하나는 가중치를 줄이거나 공유하는 파라미터 효율(parameter efficiency)이고, 다른 하나는 필요할 때만 연산을 더 쓰는 적응적 연산(adaptive computation)이다.
파라미터 효율의 대표적 접근은 레이어 타잉(layer tying)으로, 공유된 가중치 집합을 여러 레이어에서 재사용하는 방식이다. 적응적 연산의 대표적 접근은 early-exiting으로, 예측이 쉬운 토큰에 대해서는 네트워크 중간에서 빠져나가 연산을 절약한다. 그러나 두 축을 하나의 아키텍처 안에서 실질적으로 결합한 설계는 아직 부재했다. Recursive Transformer는 동일한 공유 레이어 집합을 반복 적용하므로 가중치 공유가 내장되어 있어 좋은 출발점이지만, 기존의 동적 재귀 시도는 별도의 특수한 학습 절차를 요구하거나 실제 배포 효율을 확보하기 어려웠다. 그 결과 대부분의 구현은 모든 토큰에 동일한 연산량을 적용하는 고정 깊이 재귀에 머물렀다.
이 논문이 제안하는 Mixture-of-Recursions(MoR)는 경량 라우터를 종단간(end-to-end)으로 학습시켜 토큰별 재귀 깊이를 할당하는 통합 프레임워크다. 공유 파라미터 블록을 각 토큰에 몇 번 적용할지를 라우터가 결정하므로, 연산이 실제로 필요한 토큰에 집중된다. 또한 이 토큰 단위 동적 재귀는 재귀 단계별 KV 캐싱과 자연스럽게 결합되어, 각 토큰이 할당받은 재귀 깊이에 해당하는 key-value 쌍만 선택적으로 저장하고 조회한다.
정리하면 MoR는 하나의 아키텍처 안에서 다음 세 가지를 동시에 달성한다.
| 축 | 메커니즘 | 절감 대상 |
|---|---|---|
| 파라미터 효율 | 레이어 가중치 공유(재귀 블록) | 고유 파라미터 수 |
| 적응적 연산 | 토큰 단위 재귀 깊이 라우팅 | 불필요한 FLOPs |
| 메모리 효율 | 재귀 단위 KV 캐싱 | KV 캐시 메모리와 IO |
개념적으로 MoR는 잠재 공간 추론(latent space reasoning)을 위한 사전학습 프레임워크에 해당한다. 증강된 연속 프롬프트를 생성 이전에 숙고하는 방식과 달리, MoR는 각 토큰을 디코딩하는 도중에 잠재적 사고를 수행한다. 또한 라우팅 메커니즘 덕분에 모델의 수직축(깊이축)을 따라 적응적 추론이 가능해지며, 기존 연구에서 흔했던 균일하고 고정된 사고 깊이를 넘어선다.
방법론
Recursive Transformer와 파라미터 공유 전략
표준 Transformer는 self-attention과 feed-forward network로 구성된 L개의 고유 레이어 스택으로 토큰 표현을 만든다. 은닉 상태는 다음과 같이 갱신된다.
1
h_{l+1} = f(h_l ; Phi_l), l = 0, ..., L-1
여기서 Phi_l은 l번째 레이어의 파라미터다. Recursive Transformer는 L개의 서로 다른 가중치 대신 모델을 Nr개의 재귀 블록으로 분할하고, 각 블록이 공유 파라미터 풀 Phi’를 사용한다. 이 설계는 파라미터 크기를 늘리지 않으면서 유효 네트워크 깊이를 늘린다.
논문은 네 가지 파라미터 공유 전략을 검토한다. 9개 레이어 모델을 3회 재귀로 바꾸는 예를 들면 다음과 같다.
| 전략 | 레이어 전개 방식 | 특징 |
|---|---|---|
| Cycle | (0,1,2), (0,1,2), (0,1,2) | 동일 블록을 순환 재사용하여 반복적 정제를 유도 |
| Sequence | (0,0,0), (1,1,1), (2,2,2) | 같은 레이어를 연속 적용한 뒤 다음 레이어로 이동 |
| Middle-Cycle | 첫 레이어와 마지막 레이어는 고유, 중간만 Cycle 공유 | 입출력 변환의 표현력을 유지 |
| Middle-Sequence | 첫 레이어와 마지막 레이어는 고유, 중간만 Sequence 공유 | 입출력 변환의 표현력을 유지 |
Middle-Cycle 전략의 갱신식은 첫 레이어에서 f(h_0 ; Phi_0), 마지막 레이어에서 f(h_{L-1} ; Phi_{L-1})를 사용하고, 중간 레이어에서는 인덱스 ((l-1) mod ((L-2)/Nr)) + 1에 해당하는 공유 파라미터를 사용한다.
파라미터 공유는 학습 측면에서도 이점이 있다. FSDP(Fully Sharded Data Parallel) 환경에서 기존에는 한 번의 all-gather가 한 번의 반복만 지원했지만, 재귀 모델은 모아 온 파라미터를 모든 재귀 단계에서 재사용하므로 gather당 Nr번의 반복이 가능하다. 추론 측면에서는 연속 깊이별 배칭(continuous depth-wise batching)이 가능해진다. 서로 다른 단계에 있는 토큰들이 동일한 파라미터 블록을 사용하므로 하나의 배치로 묶을 수 있고, 다른 샘플의 완료를 기다리며 생기는 버블을 제거할 수 있다.
기존 연구의 한계는 명확하다. 파라미터는 공유되더라도 KV 캐시는 깊이별로 별도 유지되는 것이 일반적이어서 캐시 크기가 줄지 않고, 높은 조회 지연이 추론 병목으로 남는다. 또한 대부분의 재귀 모델은 모든 토큰에 고정 재귀 깊이를 적용해 토큰별 난이도 차이를 무시한다. early-exit 같은 사후(post-hoc) 방법은 별도 학습 단계를 요구하며 성능 저하를 동반하는 경우가 많다. 반대로 동적 경로를 도입하면 조기 종료된 토큰의 KV 쌍이 이후 재귀 깊이에서 결측되는 새로운 문제가 생기며, 이를 해결하려면 병렬 디코딩 같은 복잡한 엔지니어링이 필요하다.
라우팅 전략: Expert-choice와 Token-choice
MoR는 두 가지 라우팅 방식을 설계하고 비교한다.
Expert-choice 라우팅은 MoD(Mixture-of-Depths)의 top-k 게이팅에서 착안한 방식으로, 각 재귀 깊이를 하나의 전문가(expert)로 보고 그 전문가가 선호하는 top-k 토큰을 선택한다. 재귀 단계 r에서 라우터는 은닉 상태 h와 라우팅 파라미터 theta로부터 스칼라 점수를 계산한다.
1
2
3
g_r^t = G(theta_r^T * h_r^t)
h_r^{t+1} = f(h_r^t, Phi') + h_r^t, if g_r^t > P_beta(G_r)
h_r^{t+1} = h_r^t, otherwise
여기서 G는 sigmoid나 tanh 같은 활성 함수이고, P_beta(G_r)는 재귀 단계 r의 전체 점수에 대한 beta 백분위 임계값이다. 단계 간 일관된 진행을 위해 계층적 필터링(hierarchical filtering)을 적용하여, 재귀 단계 r에서 선택된 토큰만 r+1에서 재평가된다. 이는 처음부터 학습하면서도 early-exit 동작을 모사하는 효과를 낸다.
Token-choice 라우팅은 각 토큰이 처음부터 전체 재귀 경로를 확정하는 방식이다. 초기 은닉 상태로부터 라우터가 전문가별 점수를 softmax 또는 sigmoid로 계산하고, top-1 게이팅으로 전문가 i를 선택하면 해당 토큰은 재귀를 i번 순차 적용받는다.
두 방식을 동일 연산 조건에서 비교하기 위해, expert-choice의 토큰 용량(top-k)을 완벽하게 부하 균형이 맞은 token-choice의 기대 토큰 분포에 맞춘다. 완전 균형 상태에서 각 토큰이 재귀 깊이 1부터 Nr까지 균등 확률 1/Nr로 배정되면, 재귀 단계 r은 전체 토큰의 (Nr - r + 1)/Nr 비율을 처리한다. 예를 들어 Nr이 3이면 재귀 단계 1, 2, 3은 각각 토큰의 3/3, 2/3, 1/3을 처리한다.
| 라우팅 | 장점 | 단점 | 완화 기법 |
|---|---|---|---|
| Expert-choice | 정적 연산 예산 확보, 완벽한 부하 균형 | 인과성 위배(정보 누출) | 보조 라우터, 보조 손실 |
| Token-choice | 정보 누출 없음 | 전문가 간 부하 불균형 | 밸런싱 손실, Loss-free 바이어스 |
Expert-choice의 인과성 위배는 top-k 선택이 시퀀스 뒤쪽 토큰 정보를 필요로 하기 때문에 발생한다. 보조 라우터는 별도 경량 네트워크를 함께 학습하되 추론 시에만 사용하며, 메인 라우터의 top-k 선택을 타깃으로 이진 교차 엔트로피로 학습한다. 그래디언트 차단으로 메인 학습 목적과 분리된다. 보조 손실은 메인 라우터 자체에 이진 교차 엔트로피를 적용하여 top-k 토큰의 점수를 1로, 나머지를 0으로 밀어낸다.
Token-choice의 부하 불균형은 밸런싱 손실로 완화한다. 전문가 수 Nr, 시퀀스 길이 T에 대해 다음과 같이 계산된다.
1
2
3
L_Balance = Nr * sum_{i=1..Nr} f_i * P_i
f_i = (1/T) * sum_{t=1..T} I(Token t selects Expert i)
P_i = (1/T) * sum_{t=1..T} g_i^t
Loss-free 방식은 명시적 정규화 손실 대신 전문가별 바이어스 b_i를 조정한다. 각 배치에서 전문가별 배정 토큰 수 c_i를 세고 부하 위반 오차 e_i = c_avg - c_i를 계산한 뒤, b_i를 b_i + u * sign(e_i)로 갱신한다. 이 바이어스는 top-k 선택에만 영향을 주고 최종 라우팅 가중치 값 자체에는 반영되지 않는다.
학습 안정화를 위해 라우터 z-loss도 사용한다. 게이팅 네트워크가 만들어내는 큰 로짓을 억제하여 수치적 불안정을 줄이는 정규화 항이다.
1
L_z(x) = (1/B) * sum_{i=1..B} ( log sum_{j=1..Nr} exp(x_ij) )^2
KV 캐싱 전략: Recursion-wise Caching과 Recursive Sharing
동적 깊이 모델은 자기회귀 디코딩 과정에서 KV 캐시 일관성 문제를 겪는다. 토큰이 조기 종료하면 깊은 레이어에 해당하는 key와 value가 존재하지 않게 되고, 이는 이후 토큰에 필요한 정보일 수 있다. MoR는 이를 위해 두 가지 전략을 설계한다.
Recursion-wise KV 캐싱은 특정 재귀 단계로 라우팅된 토큰만 그 단계에 자신의 key-value 항목을 저장한다. 각 재귀 깊이의 KV 캐시 크기는 expert-choice의 용량 계수 또는 token-choice의 실제 균형 비율로 정확히 결정된다. 어텐션은 해당 단계에 캐시된 토큰들로 제한되어 블록 지역적(block-local) 연산이 되고, 메모리 효율과 IO 요구량이 개선된다.
Recursive KV 공유는 모든 토큰이 최소한 첫 번째 재귀 블록은 통과한다는 점을 활용한다. 첫 단계에서만 KV 쌍을 캐시하고 이후 모든 재귀에서 재사용한다. 쿼리 길이는 선택 용량에 따라 재귀 깊이마다 짧아질 수 있지만, key와 value 길이는 항상 전체 시퀀스를 유지한다. 따라서 분포 불일치가 있더라도 모든 토큰이 재계산 없이 과거 문맥에 접근할 수 있다.
바닐라 Transformer를 1로 정규화한 상대 비용은 다음과 같다. 여기서 Nr은 재귀 횟수, k는 레이어당 선택 토큰 수(전체 문맥 길이 ctx보다 작음)이며, 재귀 블록만 비교한 값이다.
| 전략 | KV 메모리 | KV 캐시 IO | 어텐션 FLOPs (레이어당) |
|---|---|---|---|
| Recursion-wise Caching | (Nr + 1) / (2 * Nr) | (Nr + 1) / (2 * Nr) | (k / ctx)의 제곱 |
| Recursive Sharing | 1 / Nr | 1 | k / ctx |
Recursion-wise 캐싱은 KV 메모리와 IO를 모델 전체 기준으로 약 (Nr+1)/(2Nr)배로 줄이고, 레이어당 어텐션 FLOPs를 (k/ctx)의 제곱 배로 낮춘다. Recursive 공유는 문맥을 전역 재사용하므로 메모리 절감폭이 가장 크고, 공유 깊이에서 KV 프로젝션과 prefill을 건너뛰어 추가 속도 향상을 얻을 수 있다(Cycle 전략과만 호환). 다만 어텐션 FLOPs 감소는 k/ctx 수준에 그치고, 여전히 많은 KV IO가 디코딩 병목으로 남는다.
실험 셋업
학습 설정은 다음과 같다.
| 항목 | 값 |
|---|---|
| 아키텍처 | Llama 기반 Transformer, SmolLM 오픈소스 설정 참조 |
| 데이터셋 | SmolLM-Corpus 내 FineWeb-Edu 중복 제거 서브셋 (교육 자료 기반 220B 토큰) |
| 하드웨어 | H100 또는 A100 GPU 4장 |
| 학습률 스케줄러 (메인, isoFLOP) | Trapezoid (warmup 약 5%, stable, cooldown 20%) |
| 학습률 스케줄러 (그 외 실험) | Cosine annealing |
| 어휘 크기 | 49K |
| 최대 입력 시퀀스 길이 | 2K 토큰 |
Trapezoid 스케줄러는 stable 구간의 중간 체크포인트를 저장한 뒤 예산별로 cooldown 단계만 추가 실행하는 방식으로 재사용된다. 덕분에 FLOPs 예산마다 모델을 처음부터 다시 학습할 필요가 없다. warmup은 가장 작은 예산(2e18)의 전체 학습 스텝의 5%로 설정하고, cooldown은 각 예산별 전체 학습 스텝의 20%로 설정한다.
평가는 Language Model Evaluation Harness로 여섯 개 few-shot 벤치마크에서 수행한다. LAMBADA(LD), HellaSwag(HS), PIQA(PQ), WinoGrande(WG), ARC(Easy와 Challenge), MMLU가 사용된다. LAMBADA, WinoGrande, MMLU를 제외한 데이터셋은 타깃 문자열의 바이트 길이로 정확도를 정규화한다. MMLU는 단순화를 위해 continuation 태스크를 사용한다. 평가는 H100 또는 A100 GPU 한 장에서 측정한다.
기반이 되는 네 가지 바닐라 모델 사양은 다음과 같다. 작은 세 모델은 Grouped-Query Attention을 사용하여 KV 헤드 수를 줄였다.
| 모델 | 비임베딩 파라미터 | 임베딩 파라미터 | 레이어 수 | 히든 차원 | 어텐션 헤드 | KV 헤드 | 헤드 차원 | FFN 차원 |
|---|---|---|---|---|---|---|---|---|
| Vanilla 135M | 106M | 28M | 30 | 576 | 9 | 3 | 64 | 1536 |
| Vanilla 360M | 315M | 47M | 32 | 960 | 15 | 5 | 64 | 2560 |
| Vanilla 730M | 654M | 75M | 26 | 1536 | 24 | 8 | 64 | 4096 |
| Vanilla 1.7B | 1.61B | 101M | 24 | 2048 | 32 | 32 | 64 | 8192 |
FLOPs 계산은 Kaplan 등의 근사를 따른다. 역전파 FLOPs가 순전파의 약 두 배이므로 순전파만 계산하며, 어텐션을 제외한 선형 프로젝션 연산은 파라미터 수의 두 배로 계산한다. 어텐션은 쿼리-키 내적과 softmax 가중 value 스케일링만 반영하고, causal 마스킹으로 인해 실제 손실에 기여하지 않는 상삼각 부분은 제외한다. FlashAttention, 정규화, 비선형 연산 비용은 생략한다. MoR는 용량 계수와 캐싱 메커니즘으로 결정되는 각 재귀 깊이의 시퀀스 길이를 기준으로 FLOPs를 계산하며, token-choice의 경우 완전 균형을 가정해 근사한다. 일부 MoR 모델은 유효 깊이를 재귀 횟수로 나누어떨어지게 하려고 레이어를 추가했고, 그만큼 학습 스텝을 줄여 FLOPs 예산을 맞췄다.
주요 결과
동일 FLOPs 및 동일 토큰 조건 비교
16.5e18 FLOPs 예산에서 360M 기반 모델들을 비교한 결과는 다음과 같다. NLL은 FineWeb-Edu 검증셋 기준이며 낮을수록 좋고, 정확도는 여섯 개 벤치마크 평균이다. 파라미터는 비임베딩 기준이다.
| 모델 | 라우팅 / KV | Nr | 파라미터 | FLOPs (e18) | 학습 토큰 | NLL | LD | HS | PQ | WG | ARC | MMLU | 평균 |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Vanilla | - | - | 315M | 16.5 | 20B | 2.7824 | 32.0 | 37.8 | 65.6 | 50.5 | 39.6 | 28.0 | 42.3 |
| Recursive | - | 2 | 167M | 16.5 | 20B | 2.8079 | 31.0 | 37.1 | 66.7 | 52.3 | 40.8 | 27.5 | 42.6 |
| Recursive | - | 3 | 118M | 16.5 | 20B | 2.8466 | 29.8 | 35.9 | 65.0 | 52.3 | 39.0 | 27.2 | 41.5 |
| Recursive | - | 4 | 98M | 16.5 | 19B | 2.8781 | 28.2 | 35.4 | 65.5 | 52.5 | 38.0 | 26.8 | 41.0 |
| MoR | Expert / Cache | 2 | 167M | 16.5 | 27B | 2.7511 | 34.4 | 39.3 | 65.7 | 51.2 | 39.6 | 28.1 | 43.1 |
| MoR | Expert / Cache | 3 | 118M | 16.5 | 30B | 2.7925 | 33.1 | 37.9 | 66.9 | 52.1 | 38.3 | 27.4 | 42.6 |
| MoR | Expert / Cache | 4 | 98M | 16.5 | 30B | 2.8204 | 30.1 | 37.3 | 65.0 | 51.1 | 38.9 | 27.4 | 41.6 |
| MoR | Expert / Cache | 2 | 167M | 12.3 | 20B | 2.7749 | 33.2 | 38.3 | 65.2 | 52.6 | 40.1 | 28.1 | 42.9 |
| MoR | Expert / Cache | 3 | 118M | 11.0 | 20B | 2.8246 | 31.9 | 37.0 | 65.7 | 50.5 | 38.3 | 27.4 | 41.8 |
| MoR | Expert / Cache | 4 | 98M | 11.0 | 20B | 2.8519 | 30.2 | 36.5 | 64.3 | 52.3 | 38.6 | 27.2 | 41.5 |
| MoR | Token / Cache | 3 | 118M | 16.5 | 30B | 2.9163 | 27.6 | 34.1 | 63.8 | 50.6 | 37.4 | 26.8 | 40.0 |
| MoR | Expert / Share | 3 | 118M | 16.5 | 31B | 2.7983 | 31.7 | 37.2 | 65.1 | 51.0 | 39.0 | 27.1 | 41.9 |
동일 학습 예산에서 expert-choice 라우터와 2회 재귀를 사용하는 MoR는 검증 손실이 더 낮고 평균 few-shot 정확도도 바닐라를 앞선다(43.1% 대 42.3%). 파라미터는 거의 절반 수준이다. 연산 효율이 높아 동일 FLOPs 예산 안에서 더 많은 학습 토큰(27B 대 20B)을 처리할 수 있기 때문이다. Nr을 3이나 4로 늘려도 재귀 베이스라인을 일관되게 앞서면서 바닐라와 근소한 차이를 유지한다.
학습 토큰을 20B로 고정해 아키텍처 차이만 분리하면, Nr이 2인 MoR는 학습 FLOPs를 25% 적게 쓰면서도 바닐라와 재귀 베이스라인보다 낮은 검증 손실과 높은 정확도를 달성한다. 실제 학습 시간은 바닐라 대비 19% 감소했고 최대 메모리 사용량은 25% 감소했다. 계층적 필터링과 재귀 단위 어텐션이 시퀀스 길이를 줄이기 때문이다.
설계 변형 간 차이도 뚜렷하다. Nr이 3일 때 token-choice 라우팅은 40.0%로 expert-choice의 42.6%보다 낮았으며, 라우팅 세분성이 성능에 결정적임을 보여준다. KV 공유를 적용하면 독립 캐싱 대비 성능이 소폭 하락하는 대신 메모리 효율이 개선된다.
1.7B 규모에서 68.5e18 FLOPs 예산으로 검증한 결과는 다음과 같다.
| 모델 | 라우팅 / KV | Nr | 파라미터 | 학습 토큰 | LD | HS | PQ | WG | ARC | MMLU | 평균 |
|---|---|---|---|---|---|---|---|---|---|---|---|
| Vanilla | - | - | 1.61B | 20B | 40.8 | 49.4 | 70.6 | 54.8 | 47.4 | 30.2 | 48.9 |
| Recursive | - | 2 | 0.87B | 18B | 37.3 | 46.5 | 68.9 | 52.6 | 44.2 | 29.6 | 46.5 |
| Recursive | - | 3 | 0.67B | 20B | 36.4 | 45.3 | 69.5 | 52.7 | 43.9 | 29.1 | 46.2 |
| MoR | Expert / Cache | 2 | 0.87B | 26B | 41.1 | 47.5 | 70.0 | 55.6 | 46.0 | 30.3 | 48.4 |
| MoR | Expert / Cache | 3 | 0.67B | 27B | 37.2 | 46.6 | 69.1 | 53.7 | 44.1 | 29.7 | 46.7 |
| MoR | Token / Cache | 3 | 0.67B | 30B | 35.6 | 43.2 | 68.1 | 53.0 | 43.4 | 29.0 | 45.4 |
이 규모에서는 expert-choice MoR가 재귀 베이스라인을 크게 앞서지만 바닐라가 근소하게 우세하다(48.9% 대 48.4%). 논문은 각 모델의 최적 스케일링 정책에 따른 최적 토큰 수 차이 때문일 수도 있고, 현재 MoR 설계가 스케일업에 완전히 적합하지 않다는 신호일 수도 있다고 해석한다.
고유 파라미터를 118M으로 고정하고 재귀 횟수만 늘린 실험도 있다. 10B 토큰 학습, Middle-Cycle 기준이다.
| 모델 | Nr | 학습 NLL | 검증 NLL | 평균 정확도 |
|---|---|---|---|---|
| Vanilla | 1 | 2.9182 | 2.9678 | 39.48 |
| MoR | 2 | 2.8767 | 2.9205 | 39.96 |
| MoR | 3 | 2.8667 | 2.9111 | 40.10 |
모델 크기를 고정한 채 재귀 단계만 늘려도 학습과 검증 NLL이 모두 감소하고 평균 정확도가 개선된다.
IsoFLOP 스케일링 분석
135M, 360M, 730M, 1.7B 네 규모에서 재귀 횟수를 3으로 고정해(고유 파라미터는 약 3분의 1) 2e18, 5e18, 16.5e18 세 FLOPs 예산으로 학습했다. 16.5e18 예산에서의 검증 NLL과 평균 few-shot 정확도는 다음과 같다.
| 기반 모델 크기 | Vanilla NLL / 평균 | Recursive NLL / 평균 | MoR NLL / 평균 |
|---|---|---|---|
| 135M | 2.8432 / 42.15 | 2.9552 / 39.79 | 2.9490 / 38.11 |
| 360M | 2.7824 / 42.49 | 2.8466 / 41.39 | 2.7924 / 42.68 |
| 730M | 2.7048 / 43.43 | 2.7886 / 42.21 | 2.7438 / 43.55 |
| 1.7B | 2.8658 / 40.54 | 3.0042 / 38.87 | 2.8316 / 40.97 |
MoR는 모든 모델 크기와 연산 예산에서 재귀 베이스라인을 일관되게 앞선다. 가장 작은 135M 규모에서는 재귀에 따른 용량 병목으로 바닐라보다 낮지만, 이 격차는 규모가 커지면서 빠르게 사라진다. 360M 이상에서는 파라미터의 약 3분의 1만 사용하면서도 바닐라와 대등하거나 이를 상회하며, 특히 저예산 및 중간 예산 구간에서 우위가 두드러진다.
추론 처리량 평가
파라미터 공유 아키텍처인 MoR는 연속 깊이별 배칭을 활용할 수 있다. 디코딩 중 완료된 시퀀스 자리를 대기 중인 새 쿼리로 즉시 채워 GPU 활용률을 높게 유지하며, MoR의 early-exit 메커니즘이 배치 내 버블을 추가로 제거한다.
측정 조건은 다음과 같다. 360M 기반, 재귀 깊이 2·3·4, 16.5e18 FLOPs 예산으로 학습한 모델을 사용한다. FineWeb-Edu 검증셋 1K 샘플을 큐에 넣고 동적으로 스케줄링하며, 입력 프리픽스 없이 평균 256을 갖는 정규분포에서 샘플링한 길이만큼 토큰을 생성한다. 배칭은 고정 배치 32와 상대 최대 배치(바닐라와 MoR의 최대 배치 크기 비율을 32에 곱한 값) 두 가지를 사용한다. H100 VRAM 기준 최대 배치는 MoR-2가 42, MoR-3이 48, MoR-4가 51이다. 가변 길이 KV 캐시 지원을 위해 FlashAttention 2를 사용하고, torch.compile 호환을 위해 정적 크기 캐시를 채택했다. 안정적 측정을 위해 100회 warmup을 거친 뒤 측정했으며, KV 저장과 갱신에 드는 시간은 구현 편의상 제외했다.
| 모델 | 배치 32 기준 처리량 | 최대 배치 기준 처리량 |
|---|---|---|
| MoR-2 | 1.23배 | 1.60배 |
| MoR-3 | 1.33배 | 1.95배 |
| MoR-4 | 1.42배 | 2.06배 |
두 배칭 설정 모두에서 MoR 변형들은 연속 시퀀스 배칭을 사용하는 바닐라 베이스라인을 상회한다. 재귀 깊이를 늘릴수록 조기 종료 토큰이 늘고 KV 캐시 사용량이 더 줄어 처리량이 크게 증가한다. 성능은 소폭 하락하지만 처리량 이득을 감안하면 유리한 트레이드오프로 평가된다.
Ablation: 파라미터 공유, 라우터, KV 캐시
파라미터 공유 전략은 135M과 360M 기반 재귀 Transformer에서 10B 토큰으로 학습해 비교했다. 검증 NLL 값은 다음과 같다.
| 기반 모델 / Nr | Cycle | Sequence | Middle-Cycle | Middle-Sequence |
|---|---|---|---|---|
| 135M / 2 | 3.0071 | 3.1093 | 3.0330 | 3.0991 |
| 135M / 3 | 3.1154 | 3.1637 | 3.1048 | 3.1602 |
| 360M / 2 | 2.8487 | 2.9467 | 2.8295 | 2.9303 |
| 360M / 3 | 2.9363 | 3.0245 | 2.8760 | 2.9753 |
360M 모델에서는 Middle-Cycle이 두 재귀 깊이 모두에서 가장 낮은 NLL을 기록한다. 135M에서는 재귀 2일 때 Cycle이 근소하게 앞서지만(3.0071 대 3.0330), 재귀 깊이가 3으로 늘면 Middle-Cycle이 역전한다(3.1048 대 3.1154). 순수 Sequence 공유는 네 설정 모두에서 최악의 NLL을 기록하며 재귀 깊이가 늘수록 정확도 격차가 벌어진다. 고유 파라미터 수를 동일하게 맞춘 730M 비교에서도 Middle-Cycle이 유효 레이어 수가 더 적음에도(26 대 30) 더 나은 성능을 보인다(NLL 2.7552 대 2.7573, 평균 정확도 42.32 대 41.90). 이에 따라 논문의 모든 후속 MoR 및 재귀 모델은 Middle-Cycle을 채택한다.
라우터 설계 ablation은 360M 기반, 재귀 3회, recursion-wise 캐싱 조건에서 수행했다. Expert-choice에서는 인과성 위배 완화 방법(보조 라우터 대 보조 손실), 정규화 함수(sigmoid 대 tanh), 라우터 구조(Linear, MLP, Wide-MLP), z-loss 유무를 비교했다. 평가 지표로는 마지막 재귀에서 한 번도 선택되지 않는 dead 토큰 비율을 함께 측정했다(2K 길이 500 샘플 기준).
| 설정 | Dead 토큰 비율 (%) | NLL | 평균 정확도 |
|---|---|---|---|
| 보조 라우터 + sigmoid + MLP | 0.0 | 2.8893 | 39.4 |
| 보조 라우터 + tanh + MLP | 66.7 | 2.8720 | 36.2 |
| 보조 손실 + sigmoid + MLP | 0.0 | 2.8816 | 40.0 |
| 보조 손실 + tanh + MLP | 0.0 | 2.9933 | 38.8 |
| 보조 손실 + sigmoid + Linear | 0.1 | 2.8667 | 40.1 |
| 보조 손실 + sigmoid + Wide-MLP | 0.4 | 2.8716 | 39.4 |
| 보조 손실 + sigmoid + Linear + z-loss | 0.0 | 2.8824 | 40.0 |
추론 시 동작 측면에서 보조 손실이 별도 보조 라우터 학습보다 효과적이다. sigmoid 정규화와 단순 Linear 라우터 구조가 가장 좋은 성능을 낸다. z-loss는 정확도에는 거의 영향이 없지만 dead 토큰 비율을 소폭 줄인다. 일부 보조 라우터 변형에서 dead 토큰 비율이 66.7%까지 치솟는데, 이는 라우터가 입력과 무관하게 항상 같은 위치의 토큰을 선택하는 위치 편향을 뜻한다.
Token-choice에서는 z-loss를 기본 적용하고 밸런싱 손실과 Loss-free 방식을 비교했다. 부하 불균형은 MaxVio로 측정한다.
1
MaxVio = (max_i Load_i - Load_expected) / Load_expected
| 설정 | MaxVio | NLL | 평균 정확도 |
|---|---|---|---|
| 밸런싱 손실 (계수 0.01) + softmax + MLP | 0.682 | 2.9118 | 39.4 |
| Loss-free + softmax + MLP | 0.852 | 2.9081 | 39.4 |
| Loss-free + sigmoid + MLP | 1.281 | 3.0188 | 37.6 |
| 최종 채택 설정 (밸런싱 손실 + softmax) | 0.266 | 2.9358 | 39.1 |
두 방식의 로그 확률과 few-shot 정확도는 비슷하지만, 명시적 밸런싱 손실이 MaxVio를 뚜렷하게 낮춰 라우팅 안정성 면에서 선호된다. 그럼에도 모델은 학습 스텝의 거의 절반 동안 이질적인 전문가들 사이의 부하 균형에 어려움을 겪는다. 재귀 깊이가 다른 연산 블록을 전문가로 사용하기 때문에 생기는 이질성이 부하 균형을 어렵게 만드는 것으로 해석된다.
KV 캐시 공유 ablation에서는 바닐라와 재귀 Transformer 모두 처음부터 학습할 경우 KV 공유로 인한 성능 손실이 크지 않음을 확인했다. 재귀 Transformer는 자유도가 낮음에도 비교적 견고했으며, 논문은 각 재귀 깊이의 KV 쌍을 크기(magnitude)와 방향으로 분해해 그 근거를 제시한다. 파라미터를 공유하는 깊이들은 크기 패턴이 매우 일관되고 코사인 유사도가 높은 대각 패턴을 보인다. 다만 value 상태 공유가 key 상태 공유보다 어려운 것으로 나타났다. MoR 프레임워크 안에서는 recursive KV 공유가 expert-choice 라우팅에서 고정 토큰 조건 기준 상당한 성능 저하를 유발했다. 해당 재귀 깊이에서 활성인 토큰만 갱신하고 참조하는 편이 더 유리하다는 뜻이다. 반대로 token-choice 라우팅은 KV 공유의 이득을 볼 수 있는데, 부정확한 라우팅 결정을 공유 KV가 제공하는 추가 문맥 정보가 보완하기 때문이다.
추가 분석: 컴퓨트 최적 스케일링과 테스트타임 스케일링
isoFLOPs 제약 하에서 MoR의 컴퓨트 최적 스케일링 거동은 베이스라인과 구별된다. 최적 지점을 잇는 경로의 기울기가 더 완만하며, 이는 MoR가 파라미터 수 증가로부터 더 큰 이득을 얻는다는(즉 데이터 요구가 상대적으로 적다는) 뜻이다. 공유 파라미터 블록 자체의 성능이 추가 데이터 투입보다 중요해지기 때문으로 해석된다. 따라서 MoR의 최적 스케일링 정책은 더 큰 모델을 더 짧게 학습시키는 쪽으로 자원을 배분하는 것이다.
라우팅 분석에서는 토큰별 재귀 깊이가 주어진 문맥에서 다음 토큰을 얼마나 쉽게 예측할 수 있는지를 반영함을 보인다. 예를 들어 한 단어의 두 번째 서브워드 부분은 예측이 쉬워 적은 단계만 필요하다. 반면 “—”, “(“, “.”, “,” 같은 기능어는 여는 부분의 생성은 쉬운 편이지만, 닫는 부분이나 여는 기호 바로 뒤 첫 토큰의 예측은 더 어려운 경향을 보인다.
보조 손실을 사용하는 expert-choice 라우터는 선택 토큰과 미선택 토큰을 완전히 분리한다. 모든 재귀 단계에서 선택된 토큰의 라우팅 점수는 1.0 근처에 밀집하고 미선택 토큰은 0.0 근처에 모인다.
테스트타임 스케일링 분석에서는 Nr이 2, 3, 4인 MoR 모델에서 최대 사고(재귀) 깊이를 점진적으로 늘릴 때 로그 우도가 개선되는 것을 확인했다(500 샘플 기준). 더 깊은 재귀는 추가 연산을 제공할 뿐 아니라 각 단계가 해당 깊이에서 토큰 표현을 정제하는 데 특화되도록 만든다. 다만 360M 모델 기준으로 Nr을 키울수록 MoR의 고유 파라미터 수는 줄어들어 전체 성능은 점진적으로 하락한다.
한계와 주의사항
논문이 명시한 한계와 후속 과제는 다음과 같다.
추론(reasoning) 특화 모델로의 확장이 남아 있다. MoR 프레임워크는 토큰별 필요 재귀 깊이를 적응적으로 정하므로 잠재 추론을 본질적으로 지원하지만, 실제 추론 데이터셋으로 후속 학습했을 때 라우터가 chain-of-thought 체인의 필요성에 맞춰 학습되는지는 아직 검증되지 않았다.
모델 패밀리의 추가 확장이 필요하다. 연산 자원 제약으로 실험은 17억 파라미터까지로 제한되었다. 다음 단계는 30억 파라미터를 넘는 규모에서 훨씬 큰 코퍼스로 학습하는 것이다. 확장성을 높이려면 비공유 블록의 크기를 키우거나, 깊이별 LoRA 또는 전문가(MoE 원리)를 도입하고 전문가 병렬화를 활용하는 방안이 제시된다. 사전학습 비용을 줄이기 위해 기존 바닐라 LLM 체크포인트에서 출발하는 continued pre-training(uptraining)도 후속 과제다.
적응적 용량 제어가 어렵다. Expert-choice 라우팅은 사전 정의된 용량 계수로 완벽한 부하 균형을 보장하지만, 추론 시 다른 용량을 할당하고 싶을 때 제약이 생긴다. 보조 손실을 사용할 경우 선택 토큰과 미선택 토큰의 라우터 출력이 거의 완벽히 분리되기 때문에, 학습 이후 top-k 값을 조정하기 어렵다. 학습과 추론 양쪽에서 서로 다른 용량을 활용할 수 있는 설계가 필요하다.
희소(sparse) 알고리즘과의 결합은 아직 탐색되지 않았다. 구조적 희소성을 통합하면 토큰과 레이어 수준에서 불필요한 연산을 동적으로 가지치기할 수 있으며, 프루닝이나 양자화 같은 기법도 MoR와 상보적일 것으로 기대된다.
멀티모달 및 비텍스트 도메인 확장도 후속 과제다. MoR의 재귀 블록은 모달리티에 무관하므로 비전, 음성, 통합 멀티모달 아키텍처로 확장 가능하지만, 실제 검증은 이루어지지 않았다.
이 외에 실험 설정에서 유의할 점이 있다. 처리량 측정에서 KV 쌍의 캐싱과 갱신에 드는 시간은 구현 편의상 제외되었으며, 논문은 이를 반영한 보다 정밀한 속도 비교를 후속 과제로 남겼다. Token-choice 라우팅의 FLOPs는 완전한 부하 균형을 가정한 근사치다. 또한 135M 규모에서는 MoR가 바닐라보다 낮은 성능을 보이므로, 작은 모델에서는 재귀 용량 병목을 고려해야 한다.
결론
Mixture-of-Recursions는 파라미터 공유, 적응적 재귀 깊이, 효율적인 KV 캐싱을 하나의 Transformer 아키텍처 안에서 동시에 활용하면서 모델 품질을 유지한다. 경량 라우터로 토큰에 재귀 깊이를 동적으로 할당하고, 선택된 토큰에 대해서만 key-value 상태를 선택적으로 캐싱함으로써 이차 어텐션 연산과 중복 메모리 접근 비용을 함께 줄인다. 135M부터 1.7B까지의 규모에서 MoR는 동일 학습 FLOPs와 더 작은 모델 크기 조건에서 검증 perplexity를 낮추고 few-shot 정확도를 높이면서, 바닐라 및 기존 재귀 베이스라인보다 높은 처리량을 제공한다. 360M 이상 규모에서 파라미터의 약 3분의 1만으로 바닐라와 대등하거나 이를 상회하는 결과는 파라미터 효율과 적응적 연산을 통합하는 방향이 실질적인 성과를 낼 수 있음을 보여준다.