기술 요약: 잠재 사고 벡터(Latent Thought Vectors)를 이용한 추론 시점의 재사고(Inference-Time Rethinking)

1. 문제 정의

대규모 언어 모델(LLM)의 표준적인 생각의 사슬(Chain-of-Thought, CoT) 추론은 단일 순방향 패스(forward pass)를 통해 작동하며, 이 과정에서 모델은 생성되는 각 토큰에 대해 돌이킬 수 없는 결정을 내립니다. 이러한 구조는 선언적 지식(추론 계획의 내용)과 절차적 실행(토큰의 구문적 생성)을 결착시킵니다. 결과적으로, 모델이 출력을 확정하기 전에 자신의 추론 전략을 되돌아보거나, 비판하거나, 정교화할 수 있는 독립적인 메커니-즘이 없기 때문에 생성 초기의 오류가 검증 없이 전파됩니다.

저자들은 효과적인 추론을 위해서는 "무엇에 대해 추론할 것인가"(선언적)와 "그것을 어떻게 언어화할 것인가"(절차적) 사이의 기능적 분리가 필요하다고 주장하며, 이는 인간 인지에서 의미 처리(예: 베르니케 영역)와 구문 생성(예: 브로카 영역) 사이의 구분을 모사합니다.

2. 방법론

2.1 잠재 사고 벡터($z$)를 가진 생성 모델

제안된 프레임워크는 추론을 두 가지 뚜렷한 구성 요소로 분해합니다:

1. 잠재 사고 벡터 ($z$): 추론 전략에 대한 연속적이고 선언적인 표현("무엇").
2. 절차적 디코더: $z$에 조건화되어 토큰 시퀀스를 생성하는 자기회귀 트랜스포머("어떻게").

결합 모델은 다음과 같이 정의됩니다:
$$p_\theta(x, z) = p_\alpha(z) p_\beta(x|z)$$
여기서 $x$는 지면 토큰 시퀀스(질문 + 추론 흔적)이며, $\theta = (\alpha, \beta)$는 모든 모델 파라미터를 나타냅니다.

• 사전 모델 ($p_\alpha(z)$): 고정된 사전 분포를 사용하는 표준 VAE와 달리, 이 모델은 **학습 가능한 전송 맵(transport map)**을 사용합니다. 인코더($U_\alpha$)는 구조화되지 않은 노이즈 토큰 $z_0 \sim \mathcal{N}(0, I)$를 구조화된 잠재 사고 벡터 $z$로 변환합니다. 이를 통해 유효한 추론 패턴을 반영하는 기하학적 구조를 가진 학습된 매니폴드(manifold)를 생성합니다.
• 사고 가이드 생성기 ($p_\beta(x|z)$): 잠재 벡터 $z$가 모든 레이어에 교차 주의(cross-attention)를 수행하는 트랜스포머 디코더입니다. 결정적으로, 디코더는 짧은 컨텍스트 윈도우($w=64$)를 사용합니다. 이러한 구조적 병목 현상은 모델이 장기 의존성을 $z$를 통해 전달하도록 강제하며, 잠재 벡터가 중복된 로컬 컨텍스트가 아닌 전역적이고 시퀀스 수준의 구조를 인코딩하도록 유도합니다.

2.2 학습: 이중 속도 최적화(Dual-Rate Optimization)

학습에는 빠른 인스턴스별 적응과 느린 전역 학습을 분리하기 위해 이중 속도 체계를 채택합니다:

• 빠른 로컬 업데이트: 각 훈련 시퀀스에 대해, 노이즈 토큰 $z_0$에 대한 변분 사후 확률 파라미터 $(\mu, \sigma^2)$는 높은 학습률을 가진 Adam 경사 하강법을 통해 직접 최적화됩니다. 별도의 추론 네트워크를 학습시키지 않으며, 각 인스턴스는 자신만의 최적화된 잠재 표현을 유지합니다.
• 느린 전역 업데이트: 모델 파라미터 $\theta$(인코더 및 디코더 가중치)는 표준적인 낮은 학습률로 업데이트됩니다.

이 접근 방식은 사후 확률 붕괴(posterior collapse)를 방지하고, 모델이 일반적인 스키마($\theta$)를 축적하는 동시에 잠재 변수를 통해 인스턴스별 정밀도를 유지할 수 있게 합니다.

2.3 추론 시점의 재사고(Inference-Time Rethinking)

테스트 시점에 모델은 추론 전략을 정교화하기 위해 깁스 방식(Gibbs-style)의 반복 절차를 사용합니다. 질문 $x_q$가 주어졌을 때, 목표는 생성과 성찰을 교대로 수행하여 주변 확률 $p_\theta(x_r|x_q)$를 근사하는 것입니다:

1. 생성: 현재의 잠재 사고 $z^{(t-1)}$에 조건화된 후보 추론 흔적 $x_r^{(t)}$를 샘플링합니다.
   $$x_r^{(t)} \sim p_\beta(x_r | U_\alpha(z_0^{(t-1)}), x_q)$$
2. 성찰: 생성된 흔적의 가능도(likelihood)를 최대화하도록 잠재 노이즈 벡터 $z_0$를 최적화합니다(이는 변분 사후 확률과 실제 사후 확률 사이의 KL 발산을 최소화하는 것과 같습니다).
   $$z_0^{(t)} = \arg \min_{z_0} D_{KL}(q(z_0|x_r^{(t)}, x_q) \parallel p_\theta(z_0|x_r^{(t)}, x_q))$$

이 루프를 통해 모델은 "재사고"할 수 있습니다. 즉, 생성된 흔적의 오류가 정교화된 잠재 벡터에 정보를 제공하고, 이 벡터는 다음 반복에서 수정된 흔적을 가이드합니다. 이 과정은 변분 하한(ELBO)에 대한 **좌표 상승 최적화(coordinate-ascent optimization)**로서 이론적 근거를 가지며, 결합 사후 확률의 근사를 단조적으로 개선함을 보장합니다.

3. 주요 기여

1. 생성 프레임워크: 선언적 내용(잠재 벡터)과 절차적 생성(토큰)을 분리하여, 작은 규모의 모델이 거대한 단일 구조의 모델보다 뛰어난 성능을 낼 수 있도록 하는 새로운 아키텍처를 제시했습니다.
2. 추론 시점의 재사고: 테스트 시점에 잠재 사고를 반복적으로 정교화하는 깁스 방식의 절차를 통해, 추가적인 학습 데이터나 모델 재매개변수화 없이 자가 수정(self-correction)을 가능하게 했습니다.
3. 경험적 검증: 추론 시점의 연산이 파라미터 수와는 독립적인 매우 효과적인 스케일링 축임을 입증했습니다.

4. 실험 결과

저자들은 GSM8K-Aug 데이터셋(385K 예시)으로 0.2B 파라미터 모델을 처음부터 학습시켰으며, 세 가지 벤치마크(GSM8K, SVAMP, MultiArith)에서 평가했습니다.

• 성능: 30회의 재사고 반복을 수행한 Rethink-30 모델은 다음과 같은 성과를 거두었습니다:
  • GSM8K에서 31.54% 달성 (3B 규모의 CoT-SFT 모델의 22.70% 및 2B 규모의 MARCoS의 24.11%를 상회).
  • SVAMP에서 51.50% 달성 (2B MARCoS의 27.77% 대비).
  • MultiArith에서 68.00% 달성 (2B MARCoS의 42.33% 대비).
• 효율성: 단 한 번의 패스(Rethink-1)만으로도 10~15배 더 많은 파라미터를 가진 대부분의 베이스라인 모델을 능가했습니다.
• 일반화: 잠재 벡터가 표면적인 토큰 가변성을 추상화하는 능력 덕분에, 도메인 외(out-of-domain) 벤치마크에서도 상당한 견고함을 보여주었습니다.

5. 의의 및 주장

본 논문은 효과적인 수학적 추론이 단순히 거대한 파라미터 수에서만 오는 것이 아니라, 정교한 추론 시점의 연산으로부터 발생할 수 있음을 주장합니다.

• 직교하는 스케일링(Orthogonal Scaling): 결과는 "더 오래 생각하는 것"(반복적 정교화에 연산을 투자하는 것)이 모델 크기를 키우는 것을 보완하는 강력한 스케일링 축임을 시사합니다.
• 인지적 정렬: 선언적 계획과 절차적 실행을 분리함으로써, 모델은 인간의 인지 메커니즘을 모사하여 확정 전 단계에서 사고를 검사하고 수정할 수 있습니다.
• 대리 지표로서의 가능도: 저자들은 전문가의 시연(demonstration)을 바탕로 학습될 때, 가능도 지형(likelihood landscape)이 자연스럽게 유효한 논리 경로를 선호한다고 가정합니다. 따라서 생성된 흔적의 가능도를 최대화하도록 잠재 벡터를 최적화하는 것은 올바른 추론 전략으로 모델을 유도하는 효과적인 방법이 됩니다.

저자들은 한계점 또한 언급했습니다. 가능도를 정답의 대리 지표로 사용하는 것은 양질의 훈련 데이터가 존재한다는 가정에 의존하므로, 노이즈가 많은 환경에서는 이 상관관계가 약해질 수 있으며, 향후 연구에서 외부 검증 신호가 필요할 수 있다고 밝혔습니다.