LLM : 논문리뷰 : Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach

Overall Summary

이 논문은 latent space에서의 recurrent reasoning을 통해 test-time computation을 scaling하는 새로운 language model architecture를 제안하고, 3.5B parameter model을 800B tokens로 pretraining하여 proof-of-concept를 제시한다. Reasoning benchmark에서 test-time compute를 scaling함으로써 performance를 향상시키는 것을 보여주며, latent reasoning이 test-time compute scaling을 위한 promising research direction임을 시사한다.

전반적인 중요성과 의의: Latent reasoning을 통한 test-time computation scaling이라는 새로운 연구 방향을 제시하고, language model의 reasoning capability를 향상시키는 데 기여할 수 있는 잠재력을 보여준다.

쉬운 설명

이 논문은 "생각하는 시간"을 늘려주는 언어 모델을 제안합니다. 기존 언어 모델은 한 번에 한 단계씩 생각하는 반면, 이 모델은 "잠재 공간 (latent space)"이라는 내부 공간에서 여러 번 반복해서 생각 (recurrent) 함으로써 더 복잡한 문제를 해결할 수 있습니다. 마치 사람이 어려운 문제를 풀 때 여러 번 곱씹어 생각하는 것과 비슷합니다.

 

 

 

Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach 학습 노트

Purpose of the Paper

• 기존 연구의 한계: Chain-of-thought reasoning은 bespoke training data가 필요하고, long context window를 요구하며, verbalized intermediate results에 의존하여 reasoning capability를 제한.
• 새로운 접근 방식 제시: Latent space에서 implicitly reasoning하는 recurrent block을 통해 test-time computation을 scaling하는 새로운 language model architecture를 제안.
• 차별점: Specialized training data 없이 standard training data로 학습 가능, small context window에서도 작동, verbal representation으로 표현하기 어려운 reasoning type을 capture.

Key Contributions

• Novel Architecture:
  • Latent depth-recurrent block: Recurrent unit을 사용하여 test-time에 arbitrary depth로 unrolling 가능.
  • Prelude, core recurrent block, coda로 구성: Input embedding, latent space에서의 recurrent computation, output generation을 분리.
  • Path independence: 안정적인 recurrence 학습을 위해 random state initialization과 input injection을 매 step마다 적용.
• Training Objective:
  • Truncated backpropagation: Recurrent unit의 마지막 k iterations에 대해서만 backpropagation을 수행하여 computation and memory cost를 절감.
  • Randomly sampled iteration counts: Test-time scaling을 위해 training 중에 random number of iterations를 할당.
• Scaling & Performance:
  • 3.5B parameter, 800B tokens 규모로 scaling하여 proof-of-concept model을 구축.
  • Reasoning benchmark에서 test-time compute를 scaling함으로써 performance를 (때로는 dramatically) 향상시킴 (computation load equivalent to 50B parameters).

Novelty

• Latent space에서 recurrent reasoning을 통해 test-time computation을 scaling하는 방식을 최초로 제안.
• Recurrent model을 대규모(Large scale)에서 pre-train.
• Train time과 inference time에 다른 Recurrence를 적용.

Experimental Highlights

• Datasets: Generic web text, scientific writing, code를 포함한 data mixture를 사용하여 pretraining.
• Model: 3.5B parameter language model (prelude: 2 layers, core recurrent block: 4 layers, coda: 2 layers, hidden size: 5280).
• Training: AMD cluster (Frontier)에서 800B tokens, 4096 GPUs로 training.
• Results (Figure 1):
  • Test-time compute recurrence를 늘림에 따라 reasoning benchmark (ARC challenge, GSM8K CoT, OpenBookQA)에서 accuracy가 지속적으로 향상.
  • Reasoning이 많이 필요한 task (GSM8k)일수록 더 많은 compute를 활용하여 성능 향상.
• Model Evaluation 결과: Recurrence step을 증가시킴에 따라, 벤치마크 데이터들에서 지속적인 성능 향상.

Limitations and Future Work

• Limitations:
  • Limited compute budget으로 인해 single medium-scale model만 training.
  • Data mixture, learning rate schedule, optimization setup 등이 fully optimized되지 않음.
  • Memorization capacity가 제한적일 수 있음.
• Future Work:
  • Finetuning, reinforcement learning 등 post-training scheme을 통해 further enhancement.
  • Recurrence compression, internalizing reasoning from CoT data 등.
  • Relationship to other modern architecture improvements (e.g., linear attention, MoE) 연구.

 

Abstract

우리는 latent space에서 implicitly reasoning 함으로써 test-time computation을 확장할 수 있는 새로운 language model architecture를 연구합니다. 우리 model은 recurrent block을 반복하여 작동하므로 test-time에 임의의 depth로 unrolling 됩니다. 이는 더 많은 token을 생성하여 compute를 확장하는 주류 reasoning model과는 대조적입니다. chain-of-thought 기반 approach와 달리, 우리 approach는 specialized training data가 필요하지 않으며, 작은 context window에서도 작동할 수 있고, 단어로 쉽게 표현되지 않는 유형의 reasoning을 포착할 수 있습니다. 우리는 proof-of-concept model을 35억 parameters와 8,000억 tokens로 scaling 합니다. 그 결과 model이 reasoning benchmark에서 성능을 향상시킬 수 있으며, 때로는 최대 500억 parameters에 해당하는 computation load까지 극적으로 향상될 수 있음을 보여줍니다.

 

 

1. Scaling by Thinking in Continuous Space

인간은 자연스럽게 어떤 문제를 해결하는 데 다른 문제보다 더 많은 정신적 노력을 기울입니다. 인간은 intermediate results를 verbalizing하고 기록함으로써 장시간에 걸쳐 생각할 수 있지만, 답변의 첫 단어가 발화되기 전에 뇌에서 복잡하고 recurrent firing patterns을 통해 상당한 양의 사고가 발생합니다.

 

language models의 성능을 높이려는 초기 시도는 model size를 scaling하는 데 중점을 두었으며, 이는 엄청난 양의 data와 computation을 필요로 합니다. 최근 연구자들은 test time computation을 scaling하여 models의 reasoning 능력을 향상시키는 방법을 모색했습니다. 주류 approach는 long chain-of-thought examples에 대한 post-training을 통해 model이 context window에서 intermediate calculations을 verbalize하고 생각을 externalize하는 능력을 개발하는 것입니다.

 

그러나 값비싼 internal reasoning이 항상 single verbalized next token으로 projection되어야 한다는 제약은 낭비적으로 보입니다. models가 continuous latent space에서 natively "think"할 수 있다면 더 유능할 수 있다는 것이 타당합니다. 이러한 미개척된 추가 compute dimension을 unlock하는 한 가지 방법은 model에 recurrent unit을 추가하는 것입니다. 이 unit은 loop에서 실행되어 hidden state를 반복적으로 processing하고 updating하며 computation을 무기한으로 수행할 수 있도록 합니다. 이것이 현재 dominant paradigm은 아니지만, 이 idea는 machine learning의 기초이며 recurrent neural networks, diffusion models, universal 또는 looped transformers와 같이 매 decade마다 (재)발견되었습니다.

 

본 연구에서, 우리는 depth-recurrent language models가 효과적으로 학습하고, 효율적인 방식으로 trained될 수 있으며, test-time compute의 scaling 하에서 상당한 performance improvements를 보임을 보여줍니다. 우리가 제안하는 transformer architecture는 training 중에 randomly sampled number의 iterations 동안 실행되는 latent depth-recurrent block을 기반으로 합니다. 우리는 이 paradigm이 수십억 parameters와 5천억 개 이상의 pretraining data tokens로 scale될 수 있음을 보여줍니다. Test-time에 model은 latent space에서 recurrent reasoning을 통해 performance를 improve할 수 있으며, 더 많은 parameters와 training data로부터 benefit을 얻는 다른 open-source models와 경쟁할 수 있습니다. 또한, 우리는 recurrent depth models가 inference time에 per-token adaptive compute, (self)-speculative decoding, KV-cache sharing과 같은 non-recurrent models에서 상당한 tuning과 research effort를 필요로 하는 여러 features를 naturally support함을 보여줍니다. 우리는 latent space에서 token trajectories를 tracking하여 연구를 마무리하며, model이 numerical computations을 위해 latent space에서 shapes를 rotating하는 것과 같은 여러 흥미로운 computation behaviors가 scale에 따라 simply emerge함을 보여줍니다.

 

 

우리는 depth recurrence를 가진 3.5B parameter language model을 train 합니다. Test time에 model은 더 많은 compute를 사용하고 performance를 improve하기 위해 더 오래 iterate할 수 있습니다. Test-time reasoning을 long Chains-of-Thought에서 "verbalizing"하여 scaling하는 대신, model은 latent space에서의 reasoning을 통해 entirely improve됩니다. OpenBookQA와 같이 less reasoning을 require하는 tasks는 더 많은 compute를 effectively make use하는 GSM8k와 같은 tasks보다 더 quickly converge합니다.

 

1. Scaling by Thinking in Continuous Space 정리 노트 (AI 연구자 대상)

핵심: Latent space에서의 recurrent reasoning을 통한 test-time computation scaling

기존 방식과의 차별점:

• Chain-of-Thought (CoT)의 한계 극복: CoT는 intermediate calculations을 verbalize하여 externalize하는 방식이지만, 이는 internal reasoning을 single verbalized next token으로 projection해야 하는 제약이 있음.
• Recurrent Unit 활용: Model에 recurrent unit을 추가하여 hidden state를 iteratively processing and updating. 이를 통해 test-time에 arbitrary depth로 unrolling하여 computation을 무기한 수행 가능.
• No Specialized Training Data: CoT와 달리 specialized training data가 필요 없음.
• Small Context Window: 작은 context window에서도 작동 가능.
• Not Easily Represented in Words Reasoning Capture: 언어로 표현하기 어려운 유형의 reasoning 포착.

제안하는 Transformer Architecture:

• Latent depth-recurrent block 기반.
• Training 중 randomly sampled number의 iterations 동안 block 실행.
• Test-time에 latent space에서 recurrent reasoning을 통해 performance 향상.
• Per-token adaptive compute, (self)-speculative decoding, KV-cache sharing 등 inference time features를 naturally support.
• Latent space에서 "emergent behaviors" 발생 (e.g., numerical computations을 위한 shape rotation).

Scaling:

• 3.5B parameters, 5천억 tokens 이상의 pretraining data로 scaling 가능.
• Test-time reasoning scaling으로, 더 많은 parameters와 training data를 활용하는 open-source models와 경쟁.

쉬운 설명:

기존의 language models는 답변을 생성하기 전에 "말로 표현하는" 방식으로 생각을 정리하는 chain-of-thought 방식을 사용했습니다(사람처럼). 하지만 이 방법은 생각을 한 단어씩 표현해야 하기 때문에 비효율적일 수 있습니다.

이 논문에서는 model이 "머릿속으로"(latent space) 더 깊이 생각할 수 있도록 recurrent unit을 추가하는 새로운 방법을 제안합니다. 이 unit은 정보를 반복적으로 처리하고 업데이트하여 model이 test-time에 필요한 만큼 오랫동안 "생각"할 수 있게 해줍니다. 마치 사람이 어려운 문제를 풀 때 머릿속으로 여러 단계를 거치는 것과 비슷합니다.

이 방법의 장점은:

1. 특별한 training data가 필요 없고,
2. 짧은 문맥에서도 잘 작동하며,
3. 말로 표현하기 어려운 복잡한 추론도 가능합니다.
4. Model이 알아서 똑똑해지는 부가적인 능력도 생긴다.

결론적으로, 이 논문은 language models가 더 효율적이고 강력하게 "생각"할 수 있는 새로운 architecture를 제시합니다.

 

 

2. Why Train Models with Recurrent Depth?

2. Why Train Models with Recurrent Depth?

Recurrent layers는 transformer model이 token을 emit하기 전에 arbitrarily many computations를 perform할 수 있게 합니다. In principle, recurrent mechanisms은 test-time compute scaling을 위한 simple solution을 provide합니다. Long context reasoning (OpenAI, 2024; DeepSeek-AI et al., 2025)의 more standard approach와 compared했을 때, latent recurrent thinking은 several advantages를 가집니다.

• Latent reasoning은 bespoke training data의 construction을 require하지 않습니다. Chain-of-thought reasoning은 model이 domain of interest에서 constructed된 long demonstrations에 대해 trained되도록 require합니다. In contrast, our proposed latent reasoning models는 specialized demonstrations 없이 standard training data를 사용하여 variable compute budget으로 train할 수 있으며, additional compute가 주어지면 test-time에 their abilities를 enhance할 수 있습니다.
• Latent reasoning models는 chain-of-thought reasoning models보다 less memory for training and inference를 require합니다. 후자는 extremely long context windows를 require하기 때문에 token-parallelization (Liu et al., 2023a)과 같은 specialized training methods가 needed될 수 있습니다.
• Recurrent-depth networks는 standard transformers보다 parameter 당 more FLOPs를 perform하여 scale에서 accelerators 간의 communication costs를 significantly reducing합니다. This는 especially slower interconnects로 training할 때 higher device utilization을 enables합니다.
• Compute-heavy하고 parameter count가 small한 architecture를 constructing함으로써, 우리는 problems를 "thinking", 즉 memorizing 대신 meta-strategies, logic and abstraction을 learning함으로써 solve하는 models를 향한 strong prior를 set하기를 hope합니다. Complex algorithms를 learning하기 위한 recurrent priors의 strength는 이미 "deep thinking" literature (Schwarzschild et al., 2021b; Bansal et al., 2022; Schwarzschild et al., 2023)에서 demonstrated되었습니다.

More philosophical note에서, 우리는 latent reasoning이 spatial thinking, physical intuition 또는 (motor) planning과 같이 verbalization을 defy하는 human reasoning의 facets를 captures하기를 hope합니다. Recurrent process의 many iterations에 걸쳐, high-dimensional vector space에서의 reasoning은 linear thinking 대신 multiple directions의 deep exploration을 simultaneously enable하여, novel and complex reasoning behavior를 exhibiting할 수 있는 system으로 leading합니다.

이러한 manner로 compute를 scaling하는 것은 extended (verbalized) inference scaling (Shao et al., 2024) 또는 pretraining에서 parameter counts를 scaling (Kaplan et al., 2020)하는 것과 odds에 있지 않으며, 우리는 이것이 model performance를 scale하는 third axis를 build할 수 있다고 argue합니다.

———————— Table of Contents ————————

• Section 3에서는 our latent recurrent-depth model architecture와 training objective를 introduces합니다.
• Section 4에서는 AMD cluster인 Frontier에서의 our large-scale training run에 대한 data selection과 engineering을 describes합니다.
• Section 5에서는 benchmark results를 reports하여, inference compute를 scaling할 때 model이 how improve하는지 showing합니다.
• Section 6에서는 recurrent models가 LLM usecases를 naturally simplify하는 how를 showing하는 several application examples를 includes합니다.
• Section 7에서는 이 architecture와 training objective로 scale에서 어떤 computation patterns가 emerge하는지 visualizes하여, numerical reasoning을 requiring하는 prompts에 responding할 때 "orbiting"과 같은 context-dependent behaviors가 latent space에서 emerge함을 showing합니다.

 

2. Why Train Models with Recurrent Depth? 정리 노트 (AI 연구자 대상)

핵심 주장: Latent recurrent thinking이 기존의 long context reasoning (e.g., Chain-of-Thought)보다 효율적이고 강력한 대안이다.

Latent Recurrent Thinking의 장점 (vs. Long Context Reasoning):

1. No Bespoke Training Data:
   • CoT: Domain-specific long demonstrations 필요.
   • Latent Recurrent: Standard training data로 variable compute budget을 사용하여 training 가능. Test-time에 compute를 늘려 성능 향상.
2. Less Memory:
   • CoT: Extremely long context windows → specialized training methods (e.g., token-parallelization) 필요.
   • Latent Recurrent: Training 및 inference에 더 적은 memory 요구.
3. Higher FLOPs/Parameter:
   • Parameter 당 더 많은 FLOPs 수행 → scale에서 accelerators 간 communication costs 감소.
   • Slower interconnects 환경에서 training 시 higher device utilization 가능.
4. Strong Prior Towards "Thinking":
   • Compute-heavy, small parameter count architecture → "Thinking" (meta-strategies, logic, abstraction learning)을 통해 problems를 solve하는 models 지향 (memorization 지양).
   • "Deep thinking" literature에서 recurrent priors의 effectiveness 이미 입증.
5. Capture Non-Verbalizable Reasoning:
   • Spatial thinking, physical intuition, motor planning 등 verbalization으로 표현하기 어려운 human reasoning의 facets 포착 가능성.
   • High-dimensional vector space에서의 reasoning → multiple directions 동시 탐색 → novel and complex reasoning behavior.

결론: Latent recurrent thinking은 model performance를 scaling하는 세 번째 축 (extended inference scaling, parameter count scaling과 함께)을 구축할 수 있다.

쉬운 설명:

이 섹션에서는 왜 굳이 model에게 "머릿속으로" 깊게 생각하는 능력을 줘야 하는지 설명합니다. 기존의 Chain-of-Thought 방식은 model에게 "말로 설명하면서 풀어봐"라고 하는 것과 같습니다. 하지만 이 방법은 몇 가지 단점이 있습니다.

1. "말로 설명하는" 예시 데이터가 많이 필요: 특정 문제 유형(domain)에 맞게 사람이 직접 만든 긴 설명(demonstration) 데이터가 있어야 잘 작동합니다.
2. 메모리 많이 차지: "말"이 길어지면, 즉 context window가 커지면, model이 한 번에 처리해야 할 정보가 많아져서 메모리가 많이 필요하고, 특수한 training 기법도 필요합니다.
3. 비효율적: Model의 parameter 수에 비해 실제로 하는 "계산"(FLOPs)이 적어서, 여러 장치(accelerators) 간의 통신 비용이 커집니다.
4. 단순 암기: 문제 해결에 필요한 진정한 사고력(thinking)보다는, 주어진 데이터를 "암기"하는 경향이 있습니다.

Latent recurrent thinking은 이러한 단점을 극복합니다.

1. 별도 데이터 불필요: 일반적인 training data로도 충분.
2. 메모리 절약.
3. 효율적 계산.
4. "진짜" 사고 유도: Model이 "머릿속으로"(latent space) 깊게 생각(meta-strategies, logic, abstraction)하도록 유도.
5. 말로 표현 못 할 능력까지: 공간 지각, 직관 등 말로 설명하기 어려운 능력도 학습 가능.

결론적으로, 이 논문은 latent space에서의 recurrent thinking이 model을 더 똑똑하게 만드는 새로운 방법이라고 주장합니다.

 

3. A scalable recurrent architecture

3. A scalable recurrent architecture

이 섹션에서는 latent recurrent depth를 가진 transformer를 위한 제안된 architecture를 설명하고, design choices와 small-scale ablations에 대해 논의합니다. Architecture diagram은 Figure 2에서 찾을 수 있습니다. 우리는 항상 sequence dimension을 n, model의 hidden dimension을 h, 그리고 vocabulary를 set V로 refer합니다.

3.1. Macroscopic Design

Model은 주로 decoder-only transformer blocks (Vaswani et al., 2017; Radford et al., 2019)를 중심으로 structured됩니다. 그러나 이러한 blocks는 세 가지 functional groups로 structured됩니다.

1. Prelude P: Input data를 multiple transformer layers를 사용하여 latent space로 embedding합니다.
2. Core recurrent block R: Recurrent computation의 central unit으로, states s ∈ R^(n×h)를 modifying합니다.
3. Coda C: Latent space에서 un-embedding하고, model의 prediction head를 contains합니다.

Core block은 prelude와 coda blocks 사이에 set되며, core를 looping함으로써 우리는 노래에 indefinite amount of verses를 넣을 수 있습니다.

Recurrent iterations 횟수 r과 input tokens sequence x ∈ V^n이 주어지면, 이러한 groups는 다음과 같은 way로 output probabilities p ∈ R^(n×|V|)를 produce하는 데 used됩니다.

• e = P(x)
• s0 ∼ N (0, σ^2 * I_(n·h))
• si = R(e, si-1) for i ∈ {1, . . . , r}
• p = C(sr)

여기서 σ는 random state를 initializing하기 위한 some standard deviation입니다. 이 process는 Figure 2에 shown되어 있습니다.

Initial random state s0가 주어지면, model은 repeatedly core block R을 applies합니다. R은 latent state s_(i-1)과 embedded input e를 accepts하고 new latent state s_i를 outputs합니다. All iterations를 finishing한 후, coda block은 last state를 processes하고 next token의 probabilities를 produces합니다.

이 architecture는 deep thinking literature에 based하며, 여기서 latent inputs e를 every step (Bansal et al., 2022)에 injecting하고 latent vector를 random state로 initializing하는 것이 recurrence를 stabilizes하고 initialization에 independent한 steady state, 즉 path independence (Anil et al., 2022)로의 convergence를 promotes한다는 것이 shown되었습니다.

Motivation for this Design: 이 recurrent design은 stable iterative operators를 learn하는 데 required되는 minimal setup입니다. 좋은 example은 function E(x, y)의 gradient descent입니다. 여기서 x는 variable of interest이고 y는 data일 수 있습니다. 이 function의 gradient descent는 initial random state, 여기서 x0에서 starts하고, previous state xk와 data y에 depends하는 simple operation (optimizes하는 function의 gradient)을 repeatedly applies합니다. Our function을 actually optimize하려면 every step에서 y를 use해야 한다는 점에 note하십시오. Similarly 우리는 recurrence의 every step에서 data e를 repeatedly inject합니다. If e was provided only at the start, e.g. via s0 = e, then the iterative process would not be stable, as its solution would depend only on its boundary conditions.

Input tokens을 hidden latent space로 embed하기 위해 several layers를 using하는 structure는 standard fixed-depth transformers (Skean et al., 2024; Sun et al., 2024; Kaplan et al., 2024)를 analyzing하는 empirical results에 based합니다. 이 body of research는 LLMs의 initial and the end layers가 noticeably different한 반면, middle layers는 interchangeable하고 permutable하다는 것을 shows합니다. For example, Kaplan et al. (2024)은 few layers 내에서 standard models가 이미 sub-word tokens을 latent space의 single concepts로 embed하고, model이 then operates한다는 것을 show합니다.

Remark 3.1 (Is this a Diffusion Model?). 이 iterative architecture는 다른 modern iterative modeling paradigm, diffusion models (Song and Ermon, 2019), especially latent diffusion models (Rombach et al., 2022)와 familiar하게 look할 것입니다. 우리는 si = R(e, si-1) + n where n ∼ N (0, σi * I_(n·h))와 같이 diffusion models와 even more similar한 iterative schemes로 several ablations를 ran했지만, noise의 injection이 our preliminary experiments에서 help하지 않는다는 것을 find했으며, 이는 possibly our training objective와 connected되어 있습니다. 우리는 또한 si = Ri(e, si-1), 즉 current step을 input (Peebles and Xie, 2023)으로 takes하는 core block을 evaluated했지만, 이것이 path independence와 badly interact하여 models가 extrapolate할 수 없다는 것을 find했습니다.

3.2. Microscopic Design

Each group 내에서, 우리는 broadly standard transformer layer design을 follow합니다. Each block은 multiple layers를 contains하고, each layer는 base가 50000인 RoPE (Su et al., 2021)를 사용하는 standard, causal self-attention block과 gated SiLU MLP (Shazeer, 2020)를 contains합니다. 우리는 RMSNorm (Zhang and Sennrich, 2019)을 our normalization function으로 use합니다. Model은 queries와 keys에 learnable biases를 가지며, nowhere else입니다.

Recurrence를 stabilize하기 위해, 우리는 norm layers n_i를 사용하여 다음과 같은 "sandwich" format으로 all layers를 order합니다. 이는 (Ding et al., 2021; Team Gemma et al., 2024)의 similar strategies와 related되지만, identical하지는 않습니다.

• xˆl = n2 (xl-1 + Attn(n1(xl-1)))
• xl = n4 (xˆl + MLP(n3(xˆl)))

Small scales에서는 pre-norm, post-norm 등 most normalization strategies가 almost equally well work하지만, 우리는 이러한 options를 ablate하고 이 normalization이 scale에서 recurrence를 train하는 데 required된다는 것을 find했습니다.

Embedding matrix E와 embedding scale γ가 주어지면, prelude block은 first input tokens x를 γE(x)로 embeds하고, then 위에 described된 layout으로 lP many prelude layers를 applies합니다.

Our core recurrent block R은 si와 e의 concatenation을 hidden dimension h로 mapping하는 adapter matrix A : R^(2h) → R^h로 starts합니다 (Bansal et al., 2022). Smaller models에서는 addition을 통한 initial embedding features의 re-incorporation이 equally well work하지만, 우리는 concatenation이 scale에서 best work한다는 것을 find했습니다. This는 then lR transformer layers로 fed됩니다. Core block의 end에서 output은 again RMSNorm nc로 rescaled됩니다.

Coda는 lC layers, nc에 의한 normalization, 그리고 tied embeddings E^T를 사용하여 vocabulary로의 projection을 contains합니다.

Summary에서, 우리는 architecture를 각 stage의 layers 수를 describing하는 triplet (lP, lR, lC)와 each forward pass에서 vary할 수 있는 recurrences 수 r로 summarize할 수 있습니다. 우리는 shape (1, 4, 1)과 hidden size h = 1024를 가진 a number of small-scale models와 shape (2, 4, 2)와 h = 5280을 가진 large model을 train합니다. 이 model은 only 8 "real" layers를 가지지만, recurrent block이 iterated될 때, e.g. 32 times, it unfolds to an effective depth of 2 + 4r + 2 = 132 layers, constructing computation chains that can be deeper than even the largest fixed-depth transformers (Levine et al., 2021; Merrill et al., 2022).

3.3. Training Objective

Training Recurrent Models through Unrolling. Model이 test-time에 recurrent iterations를 scale up할 때 function할 수 있도록 ensure하기 위해, 우리는 training 중에 iteration counts를 randomly sample하여 every input sequence (Schwarzschild et al., 2021b)에 random number of iterations r을 assigning합니다. 우리는 distribution X에서 random samples x와 distribution Λ에서 random iteration counts r에 대한 loss function L의 expectation을 optimize합니다.

L(θ) = Ex∈X Er∼Λ L(mθ(x, r), x')

여기서 m은 model output을 represents하고, x'는 sequence x를 left로 shifted한 것, 즉 sequence x의 next tokens입니다. 우리는 Λ를 log-normal Poisson distribution으로 choose합니다. Targeted mean recurrence r¯ + 1과 variance σ = 1/2가 주어지면, 우리는 다음을 통해 이 distribution에서 sample할 수 있습니다.

• τ ∼ N (log(¯r) - 1/2 * σ^2, σ) (1)
• r ∼ P(e^τ) + 1, (2)

Normal distribution N과 Poisson distribution P가 주어지면, Figure 3을 see하십시오. Distribution은 most often r¯보다 less than values를 samples하지만, significantly more iterations가 taken되는 occasional events의 heavy tail을 contains합니다.

Truncated Backpropagation. Train time에 computation과 memory를 low하게 keep하기 위해, 우리는 recurrent unit의 last k iterations만을 through backpropagate합니다. This는 maximum activation memory와 backward compute가 now r에 independent하므로, heavy-tailed Poisson distribution Λ로 train할 수 있게 합니다. 우리는 our main experiments에서 k = 8로 fix합니다. Small scale에서는 이것이 k를 uniformly sampling하는 것만큼 well as work하지만, set fixed를 사용하면 training의 each step에서 overall memory usage가 equal합니다. Prelude block은 every step에서 its output e가 injected되므로, still every step에서 gradient updates를 receives한다는 점에 note하십시오. 이 setup은 RNNs에서 commonly done되는 truncated backpropagation through time과 resembles하지만, our setup은 time이 아닌 depth에서 recurrent합니다 (Williams and Peng, 1990; Mikolov et al., 2011).

Section 3에서 described된 대로, Architecture의 visualization입니다. Each block은 a number of sub-layers로 consists됩니다. Blue prelude block은 inputs를 latent space로 embeds하고, 여기서 green shared recurrent block은 final latent state를 compute하기 위해 repeated되는 layers의 block이며, 이는 red coda block의 layers에 의해 decoded됩니다.

 

 

우리는 each training step에 대한 recurrent iterations의 number를 sample하기 위해 log-normal Poisson Distribution을 use합니다.

 

 

3. A Scalable Recurrent Architecture 정리 노트 (AI 연구자 대상)

핵심: Latent space에서 recurrent computation을 수행하는 scalable transformer architecture 제안.

3.1. Macroscopic Design (전반적인 구조):

• Decoder-only Transformer 기반:
• 세 가지 Functional Groups:
  1. Prelude (P): Input을 latent space로 embedding. Multiple transformer layers 사용.
  2. Core Recurrent Block (R): 핵심. Latent state (s ∈ R^(n×h))를 반복적으로 update.
  3. Coda (C): Latent space에서 un-embedding (final output 생성). Prediction head 포함.
• Recurrent Iteration: Core block (R)을 반복하여 "indefinite amount of computation" 수행.
• Mathematical Representation:
  • e = P(x) (embedding)
  • s0 ∼ N(0, σ^2 * I_(n·h)) (initial random state)
  • si = R(e, si-1) (recurrent update)
  • p = C(sr) (final output probabilities)
• Deep Thinking Literature 기반:
  • Every step에서 latent input (e) injection: Recurrence 안정화, path independence 촉진.
  • Initial/end layers는 다르고, middle layers는 interchangeable/permutable하다는 empirical results 반영.

3.2. Microscopic Design (세부 구조):

• Standard Transformer Layer Design 기반:
  • Causal self-attention (with RoPE, base 50000)
  • Gated SiLU MLP
  • RMSNorm
  • Learnable biases on queries and keys
• "Sandwich" Format Normalization: Recurrence 안정화를 위해.
  • xˆl = n2(xl-1 + Attn(n1(xl-1)))
  • xl = n4(xˆl + MLP(n3(xˆl)))
• Core Recurrent Block (R) Details:
  • Adapter matrix (A): si와 e concatenation을 hidden dimension (h)으로 mapping.
  • lR transformer layers.
  • RMSNorm (nc)으로 output rescaling.
• Coda Details:
  • lC layers
  • nc로 normalization
  • Tied embeddings (ET) 사용하여 vocabulary로 projection.
• Architecture Summary: Triplet (lP, lR, lC) (각 stage의 layer 수), recurrence 횟수 r.
• Effective Depth: Recurrent block iteration으로 very deep computation chains 생성 (e.g., 8 "real" layers → 132 effective layers).

3.3. Training Objective:

• Unrolling for Training: Test-time scaling을 위해 training 중 iteration count (r) randomly sampling.
• Loss Function:

  L(θ) = Ex∈X Er∼Λ L(mθ(x, r), x')
  

• Log-normal Poisson Distribution (Λ):
  • Targeted mean recurrence (r¯ + 1).
  • Heavy tail: Occasional large iteration counts.
  • Sampling equations provided.
• Truncated Backpropagation:
  • Last k iterations만 backpropagation (computation/memory 절약).
  • Prelude block은 every step에서 gradient update (because of e injection).
  • Resembles truncated backpropagation through time (RNNs), but recurrent in depth.

쉬운 설명:

이 섹션에서는 "머릿속으로" 깊게 생각하는 model의 architecture와 training 방법을 자세히 설명합니다.

전반적인 구조 (3.1):

• 세 부분:
  1. Prelude: Input을 "머릿속"(latent space)으로 가져오는 역할.
  2. Core Recurrent Block: 핵심! "머릿속"에서 정보를 반복적으로 처리하며 "생각"하는 부분.
  3. Coda: "머릿속"의 최종 결과를 사람이 이해할 수 있는 형태(final output)로 변환.
• 핵심 아이디어: Core block을 여러 번 반복하여 model이 필요한 만큼 오랫동안 "생각"할 수 있게 함.

세부 구조 (3.2):

• Transformer layer를 기본으로 사용하지만, 몇 가지 중요한 수정 사항이 있음:
  • "Sandwich" Normalization: Model이 안정적으로 "생각"할 수 있도록 돕는 특별한 layer 순서.
  • Adapter Matrix: Core block에서 "머릿속" 정보(si)와 input 정보(e)를 합치는 역할.
• 각 부분(Prelude, Core, Coda)이 몇 개의 layer로 구성되는지, Core block을 몇 번 반복할지를 숫자로 표현 ((lP, lR, lC), r).

학습 방법 (3.3):

• Unrolling: Training 중에 Core block을 몇 번 반복할지 random으로 결정 (model이 test-time에 다양한 반복 횟수에 적응하도록).
• Truncated Backpropagation: 계산량을 줄이기 위해 Core block의 마지막 몇 번의 반복에 대해서만 backpropagation.

결론: 이 섹션에서는 latent space에서 recurrent computation을 수행하는 효율적이고 scalable한 transformer architecture와 그 training 방법을 제시합니다.