GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

1. 문제 정의 (Problem Statement)

배경:
대규모 언어 모델 (LLM) 의 추론 성능을 향상시키기 위해 추론 시간 확장 (Inference-Time Scaling, ITS) 기법이 주목받고 있습니다. 이는 모델 파라미터를 수정하지 않고 테스트 시 추가적인 계산 자원을 할당하여 여러 추론 경로를 생성하고 그중 최선의 답을 선택하는 방식입니다.

현재의 한계:

• 이산적 (Discrete) 모델: 토큰 기반 모델에서는 다음 토큰에 대한 명시적인 확률 분포를 기반으로 샘플링 (예: Temperature Scaling, Nucleus Sampling) 이 가능합니다.
• 연속적 잠재 공간 (Continuous Latent Space) 모델: 최근 등장한 잠재 추론 모델 (Latent Reasoning Models) 은 중간 단계를 텍스트 토큰이 아닌 숨겨진 상태 (Hidden State) 공간에서 처리합니다. 이 경우 명시적인 토큰 확률 분포가 없어, 기존 ITS 기법인 휴리스틱한 교란 (Heuristic Perturbations) 에 의존합니다. (예: Dropout, 고정된 가우시안 노이즈 주입)

핵심 문제:
휴리스틱한 교란 방식은 무작위성 (Stochasticity) 을 주입하지만, 명시적인 조건부 샘플링 분포를 정의하지 않습니다.

• 교란의 강도 (Strength) 를 단순히 증가시킨다고 해서 더 좋은 샘플링 품질이 보장되지 않습니다.
• 오히려 과도한 교란은 분포의 큰 이동 (Distribution Shift) 을 일으키지만, 정답을 지지하는 영역으로의 탐색을 방해하거나 (Over-exploration), 반대로 탐색이 너무 제한적 (Under-exploration) 이 되어 성능이 저하될 수 있습니다.
• 즉, "얼마나 많은 무작위성을 주입할 것인가"가 아니라 "어떻게 탐색을 효과적으로 제어하고 최적화할 것인가" 가 핵심 과제입니다.

---

2. 제안 방법: 가우시안 사고 샘플러 (Gaussian Thought Sampler, GTS)

저자들은 잠재 공간에서의 탐색을 학습 가능한 조건부 분포 (Learnable Conditional Distribution) 를 통한 샘플링 문제로 재정의했습니다.

핵심 구성 요소:

1. 조건부 가우시안 정책 (Context-Conditioned Gaussian Policy):

   • 고정된 노이즈 대신, 각 추론 단계의 숨겨진 상태 (Hidden State) 에 따라 평균 ($\mu$) 과 분산 ($\sigma$) 을 예측하는 경량 모듈 (MLP) 을 도입합니다.
   • 잠재 상태 $h_k$에 대해 교란 변수 $z_k$를 다음과 같이 정의합니다:
     $$\tilde{h}_k = h_k^{det} + z_k$$
     여기서 $z_k \sim \mathcal{N}(\mu_\phi(c_k), \text{diag}(\sigma_\phi^2(c_k)))$이며, $c_k$는 컨텍스트 (주로 $h_k^{det}$) 입니다.
   • 이를 통해 휴리스틱한 교란을 명시적이고 최적화 가능한 확률적 샘플링 정책으로 전환합니다.

2. 학습 방법 (GRPO 스타일 정책 최적화):

   • 백본 (Backbone) 모델의 파라미터는 고정 (Frozen) 하고, GTS 모듈만 학습합니다.
   • 강화 학습 (Reinforcement Learning) 접근법을 사용합니다. 각 입력에 대해 $N$개의 추론 경로를 생성 (Rollout) 하고, 정답 여부 (Correctness) 와 신뢰도 (Confidence) 를 기반으로 보상 (Reward) 을 계산합니다.
   • GRPO (Group Relative Policy Optimization) 를 사용하여 정책 업데이트를 안정화합니다.
     • 보상 설계: 정답 여부에 대한 대칭적 신호와, 정답/오답 그룹 내에서의 신뢰도 기반 스케이핑 (Shaping) 항을 결합합니다.
     • 목적 함수: 정책 그래디언트 손실과 참조 정책 (Reference Policy) 과의 KL 발산 패널티를 결합하여, 과도한 편향을 방지하면서 탐색을 최적화합니다.

---

3. 주요 기여 (Key Contributions)

1. 휴리스틱 ITS 의 한계 규명: 단순한 무작위성 주입이 아니라, 샘플링 품질을 제어하는 것이 잠재 공간 ITS 의 핵심임을 분석을 통해 입증했습니다.
2. GTS 제안: 휴리스틱 교란을 명시적이고 최적화 가능한 조건부 샘플링 정책으로 재구성한 경량 모듈을 제안했습니다.
3. 실험적 검증: 다양한 백본 (COCONUT, CODI) 과 벤치마크 (GSM8K 등) 에서 휴리스틱 기반 (Dropout, 고정 가우시안) 방법보다 일관되게 우수한 성능을 보였습니다.

---

4. 실험 결과 (Results)

실험 설정:

• 모델: COCONUT (GPT-2 기반), CODI (LLaMA-3.2-1B 기반).
• 데이터: GSM8K (수학 문제), MultiArith, SVAMP, GSM8K-Hard (OOD 평가).
• 비교 대상: Dropout (p=0.1, 0.5), 표준 가우시안 노이즈 (StandardG).

주요 결과:

1. 스케일링 성능 (Pass@N):
   • 샘플링 예산 ($N$) 이 증가함에 따라 GTS 는 다른 모든 방법보다 더 강력하고 안정적인 성능 향상을 보였습니다.
   • 특히 $N$이 커질수록 GTS 의 우위가 두드러졌으며, 이는 학습된 샘플러가 유한한 예산 하에서 더 유용한 경로 다양성을 생성함을 의미합니다.
2. 샘플링 품질 분석 (Sampling Quality):
   • Sampling Gain (SG): GTS 는 정답에 대한 확신도를 높이는 경로를 더 많이 생성했습니다 (양수 SG). 반면, 강한 Dropout(p=0.5) 은 오히려 SG 를 음수로 만들었습니다.
   • 분포 이동 (Distribution Shift): 휴리스틱 방법은 분포 이동 (JS Divergence) 을 크게 일으켰지만, 이것이 성능 향상으로 이어지지 않았습니다. GTS 는 적절한 수준의 분포 이동을 유지하며 정답 영역을 효과적으로 탐색했습니다.
3. OOD (Out-of-Distribution) 일반화:
   • 훈련 데이터 (GSM8K) 와 다른 벤치마크에서도 GTS 는 휴리스틱 방법들보다 일관된 성능 향상을 보였으며, 이는 학습된 탐색 정책이 문제 유형에 따라 적응적으로 작동함을 시사합니다.
4. 신호 - 잡음비 (SNR) 분석:
   • GTS 는 모델의 내부 추론 역학에 맞춰 각 단계에서 결정론적 제어와 확률적 탐색의 균형을 자동으로 조절했습니다 (예: COCONUT 의 경우 단계별 Saw-tooth 패턴, CODI 의 경우 초기 집중형 패턴).

---

5. 의의 및 결론 (Significance & Conclusion)

• 패러다임 전환: 잠재 공간 추론에서의 ITS 를 단순한 "노이즈 주입"이 아닌 "학습 가능한 조건부 샘플링 정책" 으로 접근해야 함을 제시했습니다.
• 제어 가능성: 휴리스틱 방법의 민감성 (Backbone, Task, Latent Geometry 에 따른 민감도) 을 해결하고, 탐색의 강도와 방향을 명시적으로 제어할 수 있는 프레임워크를 제공합니다.
• 효율성: 백본 모델을 재학습하지 않고도, 경량화된 GTS 모듈만 학습하여 추론 시간 계산 자원을 효율적으로 활용할 수 있습니다.

결론적으로, 효과적인 잠재 공간 추론 시간 확장을 위해서는 단순히 무작위성을 증가시키는 것이 아니라, 맥락에 따라 적응적으로 제어되고 최적화 가능한 샘플링 메커니즘이 필수적임을 이 논문은 입증했습니다.