Scaling Latent Reasoning via Looped Language Models
목차
개요
현재의 LLM은 주로 chain-of-thought(CoT)처럼 명시적인 텍스트 생성을 통해 “사고”하도록 학습된다. 이 방식은 추론 능력의 획득을 post-training 단계로 미루고, 사전학습 데이터를 충분히 활용하지 못한다는 한계가 있다. ByteDance Seed, UC Santa Cruz, Princeton, Mila 등이 공동으로 발표한 이 논문은 재귀를 뜻하는 Ouroboros에서 이름을 딴 Ouro라는 Looped Language Model(LoopLM) 계열을 제안하고 공개한다.
Ouro는 세 가지 축으로 추론 능력을 사전학습 단계에 내장한다. 첫째는 잠재 공간에서의 반복 연산이고, 둘째는 학습된 깊이 할당을 위한 엔트로피 정규화 목적함수이며, 셋째는 7.7T 토큰 규모로의 확장이다. 그 결과 Ouro 1.4B와 2.6B 모델은 광범위한 벤치마크에서 최대 12B급 SOTA 모델에 필적하는 성능을 보인다.
논문의 핵심 주장은 이 이점이 지식 저장 용량의 증가에서 오는 것이 아니라는 점이다. 통제된 실험을 통해 저자들은 루프가 파라미터당 지식 저장량을 늘리지 않으며, 대신 지식 조작(knowledge manipulation) 능력을 크게 향상시킨다는 것을 보인다. 또한 LoopLM이 만들어내는 추론 궤적이 명시적 CoT보다 최종 출력과 더 잘 정렬되어 있음을 실험적으로 확인한다.
저자들은 이를 통해 루프 깊이를 모델 크기와 데이터에 이은 제3의 스케일링 축으로 제시한다.
방법론
LoopLM 아키텍처
표준 언어 모델은 L개의 레이어를 쌓아 한 번 통과시킨다. LoopLM은 동일한 depth-L 레이어 스택을 t번 재사용한다.
1
2
3
4
5
6
비루프 모델:
F(x) = lmhead ∘ M_L ∘ emb(x), M_L = T_L ∘ ... ∘ T_1
루프 모델 (t회 반복):
F^(t)(x) = lmhead ∘ M_L ∘ M_L ∘ ... ∘ M_L ∘ emb(x)
\_____ t회 반복 _____/
t = 1이면 비루프 모델과 동일하다. 각 재귀 스텝 t에서 모델은 언어 모델링 헤드 출력을 생성하며, 스텝별 손실은 다음과 같이 정의된다.
1
2
L^(t) = E_x [ - sum_ℓ log p^(t)(x_{ℓ+1} | x_{1:ℓ}) ]
p^(t)(· | x_{1:ℓ}) = softmax( lmhead(h_ℓ^(t)) )
여기서 h_ℓ^(t)는 t번 루프를 돈 뒤 위치 ℓ의 은닉 상태다. 이 설계의 이점은 세 가지다. 첫째, 학습된 조기 종료 메커니즘으로 적응형 연산이 가능해 연산 깊이가 파라미터 수와 분리된다. 둘째, CoT처럼 출력 시퀀스를 늘리는 대신 내부 계산 그래프를 깊게 만들기 때문에 컨텍스트 길이가 불어나지 않는다. 셋째, 동일 데이터로 학습했을 때 더 큰 표준 트랜스포머를 능가할 수 있다.
적응형 종료 게이트
각 스텝 t에서 LM 헤드와 병렬로 동작하는 exit gate를 추가한다. 게이트는 최종 레이어 은닉 상태로부터 순간 종료 확률을 출력한다.
1
2
3
4
5
6
7
8
9
10
11
λ_t(x) = σ( Linear_φ( h^(t) ) ) ∈ (0, 1)
생존 확률:
S_t(x) = Π_{j=1..t} ( 1 - λ_j(x) ), S_0(x) = 1
스텝 t에서 처음 종료할 비정규화 확률:
p̃_t(x) = λ_t(x) · S_{t-1}(x), t = 1, ..., T_max - 1
유효 분포:
p(t | x) = p̃_t(x) (t = 1, ..., T_max - 1)
p(T_max | x) = S_{T_max - 1}(x)
| 추론 시에는 누적 종료 확률 CDF(n | x)를 계산하고, 임계값 q를 처음 넘는 스텝에서 종료한다. |
1
2
CDF(n | x) = 1 - Π_{j=1..n} ( 1 - λ_j(x) )
t_exit(x) = min { m : CDF(m | x) ≥ q }
q는 연산량과 정확도의 트레이드오프를 조절하는 배포 시점 하이퍼파라미터다. q가 작으면 더 일찍 종료해 연산을 아끼고, 크면 더 깊은 연산을 허용한다. 이 결정론적 분위수 기반 정책은 PALBERT의 Q-exit 기준을 따른다.
Stage I 엔트로피 정규화 목적함수
단순히 다음 토큰 예측 손실만 최적화하면 문제가 생긴다. 깊은 루프일수록 스텝별 손실이 낮아지므로 확률 질량이 후반 스텝으로 이동하고, 후반 스텝이 더 많은 학습 신호를 받아 손실이 더 떨어지며, 다시 질량이 끝으로 몰리는 자기 강화가 일어난다. 결국 분포가 t = T_max에 붕괴한다.
이를 막기 위해 엔트로피 정규화 항을 추가한다.
1
2
L = Σ_t p(t | x) · L^(t) - λ · H( p(· | x) )
H( p(· | x) ) = - Σ_t p(t | x) log p(t | x)
앞의 항은 기대 태스크 손실이고, 뒤의 항이 엔트로피 정규화다. λ가 크면 탐색적인 고엔트로피 분포를 장려하고, 작으면 확신이 있을 때 특정 스텝에 질량을 몰 수 있게 한다.
저자들은 이를 변분 추론 관점에서도 해석한다. 종료 스텝 z를 잠재 변수로 보고 학습된 종료 분포를 변분 사후분포, π를 사전분포로 두면 negative ELBO는 기대 손실과 KL 항의 합이 된다. 균등 사전분포 π_t = 1/T_max를 쓰면 KL 항이 음의 엔트로피에 상수를 더한 형태가 되어, 위 목적함수와 동치가 된다.
균등 사전분포를 선택한 이유도 명시되어 있다. 기하 사전분포나 Poisson-lognormal 사전분포는 조기 종료를 선호하는 편향을 넣는 반면, 균등 사전분포는 깊이에 대해 편향이 없다. 입력 난이도에 따른 종료 결정을 전역 연산 선호도와 분리하기 위한 선택이다.
Stage II 게이트 집중 학습
2단계에서는 LM 파라미터를 동결하고 종료 게이트만 학습한다. detach된 스텝별 손실로부터 실제 개선량을 계산해 종료 여부의 학습 라벨을 만든다.
1
2
I_i^(t) = max( 0, L_{i, stop}^(t-1) - L_{i, stop}^(t) )
w_i^(t) = σ( k · ( I_i^(t) - τ ) ), k = 50.0, τ = 0.005
w가 1에 가까우면 계속 진행, 0에 가까우면 종료를 권장하는 라벨이다. 게이트의 예측 지속 확률과 이 라벨 사이의 이진 교차 엔트로피를 스텝 길이 M에 대해 평균해 적응형 종료 손실을 구성한다.
1
2
L_adaptive^(t) = -(1/M) Σ_i [ w_i^(t) log(1 - λ_i^(t)) + (1 - w_i^(t)) log λ_i^(t) ]
L_adaptive = (1 / T_max) Σ_{t=2..T_max} L_adaptive^(t)
이 구성은 두 가지 실패 모드를 동시에 벌한다. 계속해야 할 때 종료하는 underthinking과, 종료해야 할 때 계속하는 overthinking이다.
학습 파이프라인과 데이터
Ouro 모델은 표준 decoder-only 트랜스포머를 기반으로 하며, 루프 연산 자체를 깨끗하게 검증하기 위해 부가적인 구조 변경을 최소화했다. 각 블록은 RoPE를 쓰는 Multi-Head Attention과 SwiGLU FFN으로 구성되며, 깊은 재귀 연산의 안정성을 위해 attention과 FFN 앞뒤에 RMSNorm을 두는 sandwich normalization을 적용한다. 토크나이저는 SmolLM2의 49,152 토큰 어휘를 사용한다.
모델 구성
| 모델 | 파라미터 | 레이어 | 히든 크기 | Attention | FFN | 위치 인코딩 | 어휘 크기 |
|---|---|---|---|---|---|---|---|
| Ouro 1.4B | 1.4B | 24 | 2048 | MHA | SwiGLU | RoPE | 49,152 |
| Ouro 2.6B | 2.6B | 48 | 2048 | MHA | SwiGLU | RoPE | 49,152 |
전체 사전학습은 7.7T 토큰이며, 4개 스테이지로 나뉜다.
스테이지별 학습 레시피
| 항목 | Stage 1a 사전학습 I | Stage 1b 사전학습 II | Stage 2 CT Annealing | Stage 3 LongCT | Stage 4 Mid-training |
|---|---|---|---|---|---|
| 학습률 | 3.0e-4 | 3.0e-4 | 3.0e-5 | 3.0e-5 | 1.0e-5 |
| 스케줄러 | Constant | Constant | Cosine Decay | Constant | Cosine Decay |
| 배치(토큰) | 4M에서 8M | 8M | 8M | 8M | 8M |
| 시퀀스 길이 | 4K | 4K | 16K | 64K | 32K |
| 학습 토큰 | 3T | 3T | 1.4T | 20B | 300B |
| 재귀 스텝 | 8 | 4 | 4 | 4 | 4 |
| KL 계수 람다 | 0.1 | 0.05 | 0.05 | 0.05 | 0.05 |
| RoPE base | 10K | 10K | 40K | 1M | 1M |
옵티마이저는 전 스테이지에서 AdamW를 사용하며 weight decay 0.1, beta는 (0.9, 0.95), gradient clipping은 1.0이다. 프레임워크는 torchtitan 기반의 flame을 사용했다.
데이터는 전부 오픈소스 코퍼스로 구성해 재현 가능성을 확보했다. Stage 1은 6T 토큰 규모의 데이터셋에서 샘플링하며, Nemotron-CC 73.4%, MAP-CC 13.0%, OpenCoder-pretrain 7.5%, MegaMath-web 4.1%, Ultra-FineWeb-zh 2.0%로 구성된다. 중국어 데이터는 토크나이저에 중국어 어휘가 없어 바이트 단위로 파편화되는 문제 때문에 Stage 2부터 제외했다.
Stage 2의 CT Annealing은 1.4T 토큰 규모이며 Nemotron-CC-high-quality 66.5%, Nemotron-CC-Math-v1 15.0%, Nemotron-SFT-General 6.2%, MegaMath-high-quality 4.6%, Synthetic-Code 3.8%, Nemotron-SFT-Code 3.4%, OpenCoder-Annealing 0.5%로 구성된다. Stage 3은 ProLong의 64K 길이 서브셋 20B 토큰을 사용한다. Stage 4는 20개 이상의 오픈소스 SFT 데이터셋을 통합해 182B 토큰을 만든 뒤 90B를 샘플링하고, 분포 안정화를 위해 Stage 1에서 30B, Stage 2에서 180B를 리플레이해 총 300B 토큰을 구성한다.
SFT는 약 830만 개 예제로 진행했다.
SFT 데이터 구성
| 도메인 | 데이터 소스 | 규모 |
|---|---|---|
| Math | OpenThoughts3, AceReason-1.1-SFT | 3.5M |
| Code | AceReason-1.1-SFT, OpenCodeReasoning, Llama-Nemotron-Post-Training-Dataset, OpenThoughts3 | 3.2M |
| Science | OpenThoughts3, Llama-Nemotron-Post-Training-Dataset | 808K |
| Chat | OO1-Chat-747K, DeepWriting-20K | 767K |
SFT는 LlamaFactory로 2 epoch, 최대 시퀀스 길이 32K, Adam 옵티마이저에 학습률 2e-5, beta (0.9, 0.95), cosine decay 스케줄로 진행했다.
학습 안정성 조정
재귀 아키텍처는 표준 트랜스포머와 다른 최적화 특성을 보였고, 저자들은 안정성을 우선해 여러 조정을 가했다.
Stage 1a에서 재귀 스텝 8로 시작했을 때 loss spike와 gradient oscillation이 발생했다. 여러 재귀 반복을 거치며 gradient가 누적되어 작은 섭동이 증폭되는 것으로 저자들은 추정한다. 따라서 Stage 1b부터 재귀 스텝을 8에서 4로 줄였다.
이 시점에서 모델이 두 갈래로 나뉜다. 1.4B는 원래의 24 레이어를 유지하고, 2.6B는 레이어 복제를 통해 24 레이어를 48 레이어로 upcycling한다. 공유 가중치 구조 덕분에 이 upcycling이 표준 트랜스포머에서 흔히 보이는 불안정성 없이 매끄럽게 진행됐다.
배치 크기는 4M에서 8M 토큰으로 점진적으로 늘렸고, KL 계수 람다는 0.1에서 0.05로 줄였다. 람다 축소는 태스크 손실과 KL 페널티 사이의 gradient 충돌을 줄이는 동시에, 균등 사전분포의 당김을 완화해 모델이 유용한 깊이 패턴을 자유롭게 탐색하도록 한다. 저자들은 또한 재귀 아키텍처가 파라미터 수가 같은 트랜스포머보다 더 작은 학습률을 필요로 한다는 점을 경험적으로 확인했다.
SFT 이후 DAPO와 GRPO를 사용한 RLVR 정렬을 시도했으나 SFT 체크포인트를 넘어서지 못했다. 근본 원인은 동적 조기 종료 메커니즘에 있다. vLLM과 SGLang은 고정된 실행 경로로 빠른 rollout을 제공하는데, LoopLM의 가변 깊이 연산에서는 이 전제가 깨진다. 4스텝 rollout을 모두 생성한 뒤 종료 임계값을 넘는 첫 토큰을 골라 조기 종료를 흉내내는 off-policy 방식과, 고정 4스텝으로 rollout과 업데이트를 모두 수행하는 방식을 시도했지만 둘 다 성능 향상이 없었다.
주요 결과
베이스 모델 벤치마크
평가는 lm-eval-harness와 evalplus 프레임워크를 사용했고, 모든 모델을 동일한 파이프라인으로 평가했다.
Ouro 1.4B와 1~4B 베이스라인 비교
| 벤치마크 | Qwen2.5 1.5B | Qwen3 1.7B | Qwen2.5 3B | Qwen3 4B | Gemma3 4B | Ouro 1.4B R4 |
|---|---|---|---|---|---|---|
| 학습 토큰 | 18T | 36T | 18T | 36T | 4T | 7.7T |
| MMLU | 60.99 | 62.46 | 65.62 | 73.19 | 58.37 | 67.35 |
| MMLU-Pro | 29.11 | 37.27 | 37.87 | 51.40 | 34.61 | 48.62 |
| BBH | 43.66 | 53.51 | 55.37 | 70.95 | 66.32 | 71.02 |
| ARC-C | 54.44 | 55.72 | 55.46 | 63.65 | 60.92 | 60.92 |
| HellaSwag | 67.73 | 67.09 | 74.54 | 75.66 | 75.58 | 74.29 |
| Winogrande | 66.77 | 66.30 | 70.17 | 71.19 | 71.07 | 72.30 |
| GSM8K | 60.73 | 70.28 | 74.60 | 72.86 | 68.69 | 78.92 |
| MATH500 | 17.60 | 25.80 | 42.60 | 59.60 | 68.60 | 82.40 |
| HumanEval+ | 46.30 | 59.80 | 62.20 | 70.70 | 29.30 | 67.40 |
| MBPP+ | 50.00 | 58.50 | 54.20 | 65.90 | 51.10 | 62.70 |
1.4B Ouro는 대부분의 벤치마크에서 4B Qwen3-Base와 비슷한 수준에 도달한다. 특히 BBH에서 71.02 대 70.95, GSM8K에서 78.92 대 72.86, MATH500에서 82.40 대 59.60으로 4B 모델을 앞선다.
Ouro 2.6B와 3~12B 베이스라인 비교
| 벤치마크 | Qwen2.5 7B | Llama3.1 8B | Qwen3 8B | Gemma3 12B | Ouro 2.6B R4 |
|---|---|---|---|---|---|
| 학습 토큰 | 18T | 15T | 36T | 12T | 7.7T |
| MMLU | 74.20 | 73.02 | 76.63 | 72.14 | 74.60 |
| MMLU-Pro | 43.55 | 43.24 | 53.72 | 49.21 | 55.73 |
| BBH | 53.72 | 71.56 | 77.65 | 78.41 | 80.46 |
| ARC-C | 63.65 | 60.75 | 66.10 | 72.44 | 66.40 |
| HellaSwag | 79.98 | 81.97 | 79.60 | 83.68 | 79.69 |
| Winogrande | 76.48 | 77.11 | 76.80 | 77.74 | 75.85 |
| GSM8K | 81.50 | 78.17 | 83.09 | 77.18 | 81.58 |
| MATH500 | 61.20 | 52.90 | 62.30 | 83.20 | 90.85 |
| HumanEval+ | 70.60 | 31.10 | 75.30 | 37.20 | 70.70 |
| MBPP+ | 63.50 | 51.60 | 67.90 | 66.10 | 66.60 |
2.6B Ouro는 추론 집약적 벤치마크에서 8B 모델을 상회한다. MMLU-Pro 55.73, BBH 80.46, MATH500 90.85로 각각 8B Qwen3-Base의 53.72, 77.65, 62.30을 넘어선다. 성능 향상이 가장 두드러지는 곳이 MMLU-Pro, BBH, GSM8K, MATH500이라는 점은 반복 연산이 다단계 추론과 지식 조작을 강화한다는 가설과 부합한다.
추론 모델 평가
SFT를 거친 Ouro-Thinking 모델을 AIME 2024/2025, OlympiadBench, GPQA, SuperGPQA, BeyondAIME, HLE에서 평가했다. 디코딩은 별도 언급이 없는 한 temperature 1.0, top_p 0.7이며, 모든 시스템을 동일한 in-house 하네스와 동일한 프롬프트로 평가하고 고정 rubric의 LLM-as-judge 프로토콜을 적용했다.
고난도 추론 벤치마크 결과
| 모델 | AIME24 pass@1 | AIME24 pass@10 | AIME25 pass@1 | AIME25 pass@10 | Olympiad bench | Beyond AIME | HLE | Super GPQA | GPQA |
|---|---|---|---|---|---|---|---|---|---|
| Ouro-1.4B-Thinking-R4 | 65.0 | 83.3 | 46.3 | 73.3 | 71.6 | 34.0 | 5.21 | 47.4 | 45.5 |
| Ouro-2.6B-Thinking-R4 | 64.7 | 90.0 | 50.3 | 76.7 | 76.4 | 39.0 | 5.58 | 53.7 | 52.7 |
| Qwen3-1.7B | 32.0 | 55.6 | 22.0 | 33.3 | 56.4 | 15.0 | 4.13 | 35.9 | 34.0 |
| Qwen3-4B | 61.3 | 75.0 | 51.3 | 63.3 | 73.2 | 31.0 | 5.21 | 51.9 | 54.5 |
| Qwen3-8B | 73.0 | 86.7 | 66.7 | 81.3 | 75.3 | 38.0 | 2.22 | 48.0 | 59.1 |
| DeepSeek-Distill-Qwen-1.5B | 29.6 | 66.7 | 23.0 | 43.33 | 56.44 | 9.0 | 4.2 | 26.5 | 33.2 |
| DeepSeek-Distill-Qwen-7B | 57.3 | 83.3 | 36.0 | 73.3 | 72.0 | 30.0 | 5.14 | 46.6 | 51.0 |
1.4B 모델은 OlympiadBench 71.55(Qwen3-4B 73.18), BeyondAIME 34.0(Qwen3-4B 31.0)을 기록한다. 2.6B 모델은 OlympiadBench 76.44(Qwen3-8B 75.25), BeyondAIME 39.0(Qwen3-8B 38.0)을 기록한다.
재귀 깊이와 외삽
모델은 최대 재귀 스텝 4로 학습됐다. 스텝 5부터 8은 학습 설정을 벗어난 외삽 구간이다.
Ouro 1.4B 베이스 모델의 재귀 스텝별 성능
| 재귀 스텝 | ARC-C | ARC-E | CommonsenseQA | HellaSwag | MMLU | Winogrande |
|---|---|---|---|---|---|---|
| 1 | 37.63 | 63.85 | 44.64 | 55.24 | 41.21 | 56.99 |
| 2 | 54.86 | 80.30 | 67.98 | 71.15 | 60.43 | 66.69 |
| 3 | 59.47 | 83.33 | 74.37 | 74.07 | 66.71 | 71.35 |
| 4 | 60.92 | 83.96 | 75.43 | 74.29 | 67.45 | 72.30 |
| 5 (외삽) | 58.96 | 82.91 | 75.35 | 73.72 | 66.64 | 70.32 |
| 6 (외삽) | 59.73 | 82.58 | 74.94 | 72.77 | 65.77 | 71.03 |
| 7 (외삽) | 58.96 | 81.99 | 74.28 | 72.35 | 65.28 | 70.09 |
| 8 (외삽) | 58.19 | 82.07 | 73.55 | 71.60 | 64.49 | 69.30 |
Ouro-Thinking 모델의 재귀 스텝별 성능
| 모델 및 벤치마크 | T=1 | T=2 | T=3 | T=4 | T=5 | T=6 | T=7 | T=8 |
|---|---|---|---|---|---|---|---|---|
| 1.4B OlympiadBench | 2.22 | 59.70 | 70.67 | 71.55 | 72.30 | 69.48 | 69.04 | 66.81 |
| 1.4B SuperGPQA | 2.03 | 33.07 | 44.50 | 47.37 | 48.73 | 46.15 | 45.29 | 42.88 |
| 1.4B AIME 2024 | 0.00 | 37.33 | 62.33 | 65.00 | 60.67 | 50.67 | 42.33 | 38.67 |
| 1.4B AIME 2025 | 0.33 | 25.00 | 43.33 | 46.30 | 47.00 | 43.00 | 41.00 | 38.00 |
| 2.6B OlympiadBench | 18.96 | 68.59 | 75.56 | 76.44 | 71.85 | 69.19 | 57.63 | 39.26 |
| 2.6B SuperGPQA | 15.66 | 48.58 | 56.70 | 53.68 | 56.45 | 55.44 | 53.32 | 46.84 |
| 2.6B AIME 2024 | 3.00 | 52.00 | 70.33 | 64.70 | 57.00 | 56.33 | 49.67 | 39.00 |
| 2.6B AIME 2025 | 2.00 | 40.67 | 50.67 | 50.30 | 49.33 | 46.00 | 38.00 | 24.33 |
T=1에서 성능이 매우 낮다는 사실은 복잡한 태스크에서 반복적 정제가 필수적임을 확인해 준다. 1.4B Thinking 모델은 T=4 또는 T=5 부근에서, 2.6B Thinking 모델은 T=3 또는 T=4에서 정점을 찍는다. 두 모델 모두 T=6부터 8 구간의 외삽에서는 성능이 떨어진다.
한편 안전성 지표는 반대 양상을 보인다. 벤치마크 성능은 학습 깊이를 넘어서면 저하되지만, 안전 정렬은 외삽 구간에서도 재귀 스텝이 늘어날수록 개선된다.
조기 종료 전략과 KV 캐시
세 가지 조기 종료 전략을 MMLU에서 비교했다. 정적 종료(고정 스텝), 은닉 상태 차이 임계값(연속 스텝 간 은닉 상태 L2 변화량이 임계값보다 작아지면 종료), 그리고 Q-exit 기준을 쓰는 학습된 게이트다. 학습된 게이트는 다시 사전학습 목적함수만 적용한 설정과, Stage II 적응형 종료 손실까지 적용한 설정으로 나뉜다.
Stage II 학습을 거친 게이트가 모든 연산 예산에서 최고 정확도를 보였다. 평균 종료 라운드 2.5에서 특화 학습 게이트는 66% 정확도를, 표준 게이트는 약 64%를 기록한다. 특화 학습 게이트와 표준 게이트의 격차는 대부분의 동작점에서 2~3%p 수준이다.
은닉 상태 차이 임계값 전략도 예상보다 경쟁력이 있어, 평균 2~3 라운드 구간에서 특화 학습 게이트와 1~2%p 이내로 근접한다. 정적 baseline은 1라운드에서 2라운드로 갈 때 40%에서 60%로 급격히 오르지만, 3라운드에서 4라운드로 갈 때는 67.35%로 한계 이득이 작다. 이 수확 체감이 적응형 방법이 효과적인 이유다.
재귀 구조는 KV 캐시 측면에서 문제를 만든다. 소박하게 구현하면 각 재귀 스텝이 자체 KV 캐시를 유지해야 하므로 4스텝 모델에서 4배의 메모리 오버헤드가 발생한다. prefill 단계에서는 네 스텝이 모두 자체 캐시를 필요로 하며, 캐시를 재사용하면 GSM8K에서 10점 이상 성능이 떨어진다. 반면 디코딩 단계에서는 캐시 재사용이 가능하다.
디코딩 단계 KV 캐시 공유 전략
| 전략 | GSM8K | MATH-500 | 메모리 절감 |
|---|---|---|---|
| Full (4배 캐시) | 78.92 | 82.40 | 1.00배 |
| 첫 스텝만 재사용 | 18.73 | 8.43 | 4.00배 |
| 마지막 스텝만 재사용 | 78.85 | 80.40 | 4.00배 |
| 평균 캐시 재사용 | 78.73 | 78.52 | 4.00배 |
첫 스텝 캐시만 재사용하면 GSM8K가 78.92에서 18.73으로 붕괴한다. 초기 표현만으로는 이후 디코딩 스텝을 감당할 수 없다는 뜻이다. 반면 마지막 스텝 재사용과 평균 재사용은 전체 캐시 대비 GSM8K 기준 0.3점 이내로 근접하면서 메모리를 4배 줄인다. 이 결과는 LoopLM을 동급 파라미터 표준 트랜스포머와 비슷한 메모리 footprint로 배포할 수 있게 한다.
지식 용량과 지식 조작
같은 파라미터 수로 더 좋은 성능을 내는 이유가 무엇인지 확인하기 위해, 저자들은 Physics of Language Models 계열의 통제 가능한 합성 태스크를 사용한다.
첫째는 Capo 태스크 기반 지식 용량 측정이다. 성별, 생년월일, 대학, 전공, 고용주 다섯 가지 속성을 갖는 합성 인물 전기 데이터셋 bioS(N)을 생성하고, 정답 속성 토큰 위치의 교차 엔트로피 손실 합으로 모델이 기억한 정보량을 비트 단위로 추정한다. 1M에서 40M 파라미터 범위의 GPT-2 스타일 모델을 학습시켰고, N은 20K에서 500K 인물 범위에서 1000회 노출로 학습했다.
결과는 명확하다. 루프 모델과 비루프 모델 모두 파라미터당 약 2비트의 용량 비율에 수렴한다. 루프를 늘리는 것만으로는 지식 용량 자체도, 용량의 스케일링 특성도 개선되지 않는다.
둘째는 Mano 태스크 기반 지식 조작 측정이다. 이 태스크는 중간 사고 과정 없이 modulo 23 산술 트리를 직접 계산하는 문제로, 파라미터에 인코딩된 산술 규칙 적용과 이진 트리 구조 파싱을 동시에 요구한다. 최대 표현 길이 L을 10, 16, 24로 두어 난이도를 나눴다.
Mano 태스크 정확도
| 모델 구성 | L = 10 | L = 16 | L = 24 |
|---|---|---|---|
| Base (12 x 1) 기준 모델 | 93.6 | 94.4 | 34.8 |
| Base (2 x 1) | 21.5 | 8.4 | 7.5 |
| Loop (2 x 6) | 98.1 | 96.3 | 78.0 |
| Base (3 x 1) | 75.4 | 29.8 | 11.0 |
| Loop (3 x 4) | 97.9 | 95.8 | 92.2 |
| Base (6 x 1) | 84.7 | 59.5 | 20.0 |
| Loop (6 x 2) | 93.4 | 88.5 | 35.1 |
동일 파라미터 조건에서 루프 모델은 항상 비루프 모델을 앞선다. 게다가 FLOPs가 같은 12 x 1 baseline과 비교해도 종종 더 나은 성능을 보인다. 특히 L = 24의 어려운 설정에서 Loop (3 x 4)는 92.2를 기록해 iso-FLOP baseline의 34.8을 크게 앞선다.
셋째는 자연어 multi-hop QA다. “A의 강사는 B이다”, “B의 선생은 C이다” 같은 합성 사실로부터 “A의 강사의 선생은 누구인가”를 묻는 3-hop 질문을 다룬다. 전체 가능한 QA 쌍의 2.5%에서 25% 범위로 고유 학습 샘플 수를 바꿔가며, 동일한 총 학습 토큰 예산 아래 성능을 비교했다. 루프가 많은 모델일수록 더 적은 샘플로 3-hop QA를 학습하며, 동일한 고유 샘플 수에서도 더 빠르게 학습한다.
이 결론은 실제 벤치마크에서도 재현된다. MMLU의 57개 세부 카테고리를 Loop 1 대비 Loop 4의 상대 향상률로 분석했다.
MMLU 카테고리별 상대 향상률
| 향상이 큰 카테고리 | 향상률 | 향상이 작은 카테고리 | 향상률 |
|---|---|---|---|
| Elementary Mathematics | +155.6% | Moral Scenarios | +7.8% |
| Formal Logic | +143.3% | Global Facts | +8.3% |
| Logical Fallacies | +127.8% | Virology | +13.7% |
| High School Statistics | +126.9% | Anatomy | +21.4% |
논리, 수학, 절차적 추론을 요구하는 카테고리에서 향상이 크고, 암기된 사실 검색에 의존하는 카테고리에서는 향상이 미미하다. global_facts의 거의 평탄한 향상률은 반복 연산이 접근 가능한 지식량을 늘리는 것이 아님을 보여준다.
이론 분석
저자들은 왜 루프가 파라미터 증가 없이 지식 그래프 탐색을 가능하게 하는지 이론적으로 다룬다. 그래프 도달 가능성 문제를 설정하되, 지식 그래프의 일부만 컨텍스트에 주어지고 나머지 관계는 파라미터에 인코딩되어 있다고 가정한다.
정리 1(비형식): 결합된 지식 그래프의 최대 크기를 n으로 고정하면, 컨텍스트 그래프의 인접 행렬과 질의 쌍이 주어졌을 때, 컨텍스트 그래프와 무관한 1레이어 트랜스포머를 O(log^2 D)회 루프시켜 결합 그래프에서 s에서 t로의 경로 존재 여부를 판정할 수 있다. 여기서 D는 결합 그래프의 지름이다.
잠재 추론 방식별 순차 연산 스텝
| 잠재 추론 방식 | 순차 연산 스텝 |
|---|---|
| Discrete CoT | O(n^2) |
| Continuous CoT | O(D) |
| Universal Transformer | O(log D) |
Universal Transformer는 전체 쌍 연결성 탐색의 병렬성을 극대화해 순차 연산 스텝을 O(n^2)에서 O(log D)로 지수적으로 줄인다.
샘플 효율성에 대해서는 별도의 가설을 제시한다. 추론 태스크가 학습된 파라메트릭 지식이나 알고리즘 절차의 반복 적용을 요구한다면, 모델은 깊이가 다른 레이어들에 걸쳐 반복 구조를 학습해야 한다. 가중치를 공유하지 않는 깊은 트랜스포머는 훨씬 큰 함수 클래스를 탐색해야 하지만, 파라미터 공유 레이어는 실현 가능한 가설 클래스를 크게 줄여 학습의 샘플 복잡도를 낮춘다.
안전성과 충실성
안전성은 11개 금지 카테고리 330개 예제로 구성된 HEx-PHI 데이터셋으로 평가했다. GPT-4o를 judge로 사용해 1에서 5까지의 유해성 점수를 매기고, 최고 점수 5를 받은 케이스 비율을 harmfulness rate로 계산한다. Base 모델은 greedy decoding에 max_new_tokens 128, Thinking 모델은 temperature 1.0, top_p 0.7, max_new_tokens 8192로 평가했다.
두 모델 모두 재귀 스텝이 늘어날수록 안전성이 개선되며, 학습되지 않은 5에서 8 스텝 구간에서도 개선이 이어진다. Thinking 체크포인트는 재귀 스텝 4에서 harmful rate를 1.4B는 0.009, 2.6B는 0.003까지 낮춘다. 이는 Qwen3-4B-Thinking의 0.009와 동등하거나 더 낮은 수치다.
원인을 파악하기 위해 상위 레이어 마지막 입력 토큰의 은닉 표현에 PCA를 적용했다. 동일 형식의 benign 질문 100개와 harmful 질문 100개를 사용했다. 재귀 스텝이 늘어날수록 모델이 benign 프롬프트와 harmful 프롬프트를 더 잘 분리하며, 안전하지 않은 응답에 해당하는 점들은 대부분 두 클러스터의 경계 근처에 나타난다. 유해성 판별의 어려움이 안전하지 않은 응답으로 이어지고, 재귀 스텝 증가가 이를 완화한다는 해석이다.
충실성(faithfulness)은 더 흥미로운 결과다. 표준 LLM은 CoT 텍스트를 생성하기 전에 이미 답을 정해두고, 이후 텍스트로 그 결정을 사후 정당화하는 경향이 보고되어 왔다. 저자들은 의미적 동치 판단이 모호한 Quora Question Pairs 데이터셋으로 이를 검증했다.
Qwen3-4B-Thinking에서 최종 토큰 로짓에 대한 단순 선형 프로브는 모델의 최종 답을 0.99 ROC AUC로 예측한다. 사고 과정이 결과에 거의 영향을 주지 않는다는 뜻이다.
Ouro 1.4B는 재귀 스텝당 24 레이어를 사용한다. 레이어 1부터 24i까지의 은닉 상태로 스텝 i의 답을 예측하는 선형 프로브를 학습했다. 같은 재귀 스텝 안에서는 답이 잘 예측되지만, 직전 스텝의 표현(레이어 24(i-1))으로는 스텝 i의 결정을 신뢰성 있게 예측하지 못한다. 새 재귀 패스가 잠정적 선택을 수정하는 추가 연산을 수행한다는 증거다.
1,000개 Quora Question Pairs에 대한 스텝 간 일치 행렬도 같은 결론을 지지한다. 스텝 2와 스텝 4의 라벨 일치는 361건으로 36.1%에 불과하고, 스텝 2와 스텝 3의 일치는 551건으로 55.1%다. 반면 i가 4 이상인 구간에서는 인접 스텝 간 일치가 1000에 가까워진다. 저자들은 이를 학습 깊이를 넘어선 구간에서 재귀적 추론이 학습되지 않았고, 루프 수가 늘수록 답이 고정점으로 수렴하기 때문으로 해석한다.
배포 관점의 파생 이점도 논의된다. 중간 스텝의 LM 헤드를 제안자로, 최종 스텝을 검증자로 두면 외부 draft 모델 없이 speculative decoding을 구성할 수 있다. 같은 분해 위에서 사용자에게 토큰이 노출되기 전에 draft 분포에 안전성 스크리닝을 끼워 넣을 수도 있다. 또한 학습 목적함수가 스텝별 기대 손실의 단조 개선 성질을 보존하므로, 어느 중간 스텝에서든 디코딩을 시작해 스트리밍하면서 이후 스텝이 검증 및 수정을 이어가는 anytime 알고리즘이 된다.
LoopLM 스케일링 법칙
부록에서는 53M, 134M, 374M, 778M, 1.36B의 다섯 가지 모델 크기와 1, 2, 4, 8의 네 가지 재귀 깊이 조합으로 소규모 스케일링 실험을 수행했다. 모두 FineWeb-Edu 코퍼스 20B 토큰으로 학습하고, ARC-Challenge, ARC-Easy, HellaSwag, LAMBADA, OpenBookQA, PIQA 여섯 개 벤치마크로 평가했다.
첫 번째 관찰은 표준 모델과 LoopLM 모두 모델 크기와 재귀 스텝이 커질수록 성능이 향상된다는 점이다. 두 번째 관찰은 더 조심스럽게 읽어야 한다. 연산 예산을 고려하지 않은 동일 조건에서는 표준 모델이 LoopLM보다 항상 점수가 높다.
표준 모델 대비 LoopLM 평균 성능 격차
| 모델 크기 | 재귀 스텝 2 | 재귀 스텝 4 |
|---|---|---|
| 170M | 0.021 | 0.039 |
| 340M | 0.023 | 0.037 |
| 680M | 0.015 | 0.026 |
| 1.3B | 0.017 | 0.025 |
격차는 재귀 스텝이 늘어날수록 커지고, 모델 크기가 커질수록 줄어든다. 재귀 스텝 4에서 이 경향이 특히 일관적이다.
손실 측면에서는 Chinchilla 형태의 멱법칙 적합을 시도했다.
1
2
Total Loss Scaling Law:
L_t = E + A / (N + t1)^alpha + B / (D + t2)^beta + C / (Tm + t3)^gamma
N은 모델 크기, D는 학습 데이터 크기, Tm은 최대 재귀 스텝이며, t1, t2, t3는 변수가 0 근처에서 폭발하지 않도록 하는 평활 항이다. 전체 데이터 포인트로 적합했을 때 결정계수는 0.9596이다.
스텝별 손실에 대해서도 유사한 형태의 Step-wise Loss Scaling Law를 적합했다. 최대 재귀 스텝 2, 4, 8 각각에 대해 결정계수는 0.8898, 0.8146, 0.795다. 적합 파라미터가 양수로 나와, 재귀 스텝이 늘어날수록 스텝별 손실이 감소함을 확인했다. 총 손실보다 스텝별 손실이 본질적으로 덜 안정적이어서 결정계수가 상대적으로 낮다.
사전분포 선택에 대한 ablation도 부록에 있다. 776M LoopLM에 T_max 4를 두고 FineWeb-Edu 20B 토큰, 글로벌 배치 50K 토큰, 약 40K 스텝으로 학습하며 기하 사전분포의 감마를 0.1부터 0.9까지 스윕했다. 균등 사전분포가 일관되게 더 낮은 학습 손실과 깔끔한 수렴을 보였고, 기하 사전분포는 감마가 커질수록 격차가 벌어졌다. 감마가 큰 기하 사전분포는 t=1, 2에 질량을 몰아 깊은 스텝의 credit assignment를 약화시키고, “깊을수록 좋다”는 성질을 훼손한다.
한계와 주의사항
첫째, 학습 깊이를 넘어선 외삽은 대체로 손해다. 안전 정렬은 5에서 8 스텝에서도 개선되지만, 태스크 성능은 T=4를 지나면 저하된다. 2.6B Thinking 모델의 OlympiadBench는 T=4에서 76.44였다가 T=8에서 39.26까지 떨어진다.
둘째, 연산 예산을 무시하면 파라미터를 공유하지 않는 표준 모델이 더 강하다. 성능 격차는 재귀 스텝이 늘어날수록 커진다. LoopLM의 가치는 파라미터 효율이라는 제약 아래에서 성립하는 것이지, 절대 성능 우위를 뜻하지 않는다.
셋째, 학습 안정성 문제가 실재한다. 재귀 스텝 8로 시작한 Stage 1a에서 loss spike와 gradient 진동이 발생해 4로 줄여야 했다. 연산 자원 제약으로 학습률 스윕도 충분히 수행하지 못했고, 안정 수렴을 우선한 보수적 값을 채택했다.
넷째, RLVR 정렬이 실패했다. 동적 조기 종료가 vLLM과 SGLang의 고정 실행 경로 전제를 깨뜨리기 때문이다. 고정 4스텝으로 off-policy 문제를 피한 경우에도 SFT 체크포인트를 넘지 못했으며, 저자들은 광범위한 SFT를 거친 소형 모델의 RL 여지가 제한적일 가능성을 원인으로 추정한다. 흥미롭게도 4라운드로 학습한 모델이 추론 시에는 더 적은 라운드를 사용하는 현상이 관찰됐으나, 이 일반화의 메커니즘은 아직 규명되지 않았다.
다섯째, KV 캐시 절감은 디코딩 단계에 한정된다. prefill 단계에서 캐시를 재사용하면 GSM8K에서 10점 이상 성능이 떨어져, 네 스텝 모두 자체 캐시를 유지해야 한다.
여섯째, 이론 결과에도 제약이 있다. 잠재 공간의 중첩이 최악의 경우 n개 노드 정보를 인코딩해야 하므로 은닉 차원이 n 규모로 요구되며, 저자들은 이를 이론적 한계로 명시한다.
일곱째, 충실성 검증은 관찰적 대리 지표에 기반한다. 잠재 추론 과정을 직접 조작할 수 없기 때문에, 중간 은닉 표현을 읽어내 재귀가 깊어질 때 예측이 바뀌는지를 측정하는 방식으로 우회했다. CoT에 대한 표준적 개입 실험과 동일한 강도의 인과 주장은 아니다.
여덟째, 부록의 multi-hop QA 추가 실험에서 iso-FLOP baseline 결과 일부가 본문의 얕은 버전보다 크게 낫지 않았고, 저자들은 무작위성이나 하이퍼파라미터 튜닝 부족 가능성을 언급하며 후속 실험의 필요성을 인정한다.
결론
Ouro는 반복 연산과 적응형 깊이를 사전학습 단계에 직접 통합해 7.7T 토큰 규모로 학습한 Looped Language Model 계열이다. 1.4B와 2.6B 모델이 각각 4B와 8B 표준 트랜스포머를 지속적으로 따라잡거나 능가하며, 2~3배의 파라미터 효율 향상을 보여준다.
논문의 기여는 성능 수치 자체보다 그 원인을 분리해낸 데 있다. 루프는 파라미터당 약 2비트라는 지식 저장 용량을 늘리지 않는다. 대신 파라미터에 인코딩된 지식을 조합하고 조작하는 능력을 강화하며, 이는 Mano 태스크와 multi-hop QA 같은 합성 실험과 MMLU 카테고리별 향상 패턴 양쪽에서 일관되게 확인된다.
균등 사전분포 위의 엔트로피 정규화라는 실용적 목적함수와, 마지막 스텝 KV 캐시 재사용으로 메모리를 4배 줄이는 배포 전략도 함께 제시된다. 반복적 정제 과정이 최종 출력과 인과적으로 결합된 추론 궤적을 제공하고, 안전 정렬이 재귀 스텝 증가에 따라 개선된다는 점은 표준 CoT에서 보기 어려운 성질이다.
저자들은 잠재 반복 연산을 파라미터와 데이터에 이은 제3의 스케일링 축으로 규정하고, 더 깊은 깊이에서의 성능 외삽 개선과 더 복잡한 재귀 메커니즘 탐색을 후속 과제로 제시한다.