Parallel Test-Time Scaling for Latent Reasoning Models

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 대규모 언어 모델 (LLM) 의 성능을 향상시키기 위해 테스트 시간 확장 (Test-Time Scaling, TTS) 기법이 주목받고 있습니다. 특히, 병렬 TTS는 여러 추론 경로를 생성하여 투표 (voting) 나 검색 (search) 을 통해 결과를 집계하는 방식으로, 명시적 체인 오브 씽킹 (Explicit Chain-of-Thought, CoT) 모델에서 매우 효과적입니다.
• 잠재적 추론 (Latent Reasoning) 의 등장: 최근 연구들은 추론 과정을 자연어 토큰이 아닌 **연속적인 벡터 공간 (Continuous Vector Space)**에서 수행하는 '잠재적 추론'을 제안했습니다. 이는 토큰 기반 CoT 보다 컴퓨팅 효율이 높고 추상적인 패턴을 포착할 수 있다는 장점이 있습니다.
• 핵심 문제: 기존 토큰 기반 모델의 병렬 TTS 기법을 잠재적 추론 모델에 적용하는 것은 두 가지 근본적인 장벽으로 인해 불가능했습니다.
  1. 샘플링 (Sampling) 부재: 토큰 모델은 확률 분포 (logits) 를 기반으로 다양한 토큰을 샘플링할 수 있지만, 잠재적 추론 모델은 연속 벡터를 다루며 명시적인 확률 분포가 없어, 어떻게 여러 개의 다른 추론 경로를 확률적으로 생성할지 (stochasticity) 에 대한 메커니즘이 없습니다.
  2. 집계 (Aggregation) 의 어려움: 토큰 기반 모델은 로그 가능도 (log-likelihood) 를 기반으로 경로를 평가하고 순위 매길 수 있지만, 잠재적 추론 모델은 중간 단계의 벡터에 대한 확률적 신호가 없어 생성된 경로들의 품질을 평가하거나 최선의 경로를 선택하기 어렵습니다.

2. 제안된 방법론 (Methodology)

이 논문은 위 두 가지 문제를 해결하기 위해 샘플링 전략과 집계 메커니즘을 재설계한 병렬 TTS 프레임워크를 제안합니다.

A. 잠재 공간에서의 확률적 샘플링 (Stochastic Sampling)

연속 공간에서 다양한 추론 경로를 생성하기 위해 불확실성 (Uncertainty) 이론을 기반으로 두 가지 전략을 도입했습니다.

1. 몬테카를로 드롭아웃 (Monte Carlo Dropout, MC-dropout):
   • 원리: 추론 시 드롭아웃을 활성화하여 모델 가중치에 무작위 마스크를 적용합니다.
   • 목적: **인지적 불확실성 (Epistemic Uncertainty)**을 포착합니다. 즉, 모델이 학습 데이터에서 얻은 지식이 제한적임을 반영하여, 서로 다른 가중치 구성을 가진 서브네트워크를 통해 구조화되고 방향성 있는 추론 확장을 유도합니다.
2. 가산 가우시안 노이즈 (Additive Gaussian Noise, AGN):
   • 원리: 각 추론 단계의 잠재 벡터에 직접 가우시안 노이즈를 추가합니다.
   • 목적: **알레atoric 불확실성 (Aleatoric Uncertainty)**을 시뮬레이션합니다. 입력 데이터의 내재적 노이즈를 모방하여 결정론적 중심을 중심으로 광범위하고 등방성 (isotropic) 인 탐색을 유도합니다.

B. 잠재 경로 집계 (Latent Trajectory Aggregation)

생성된 여러 잠재적 추론 경로 중 가장 좋은 경로를 선택하기 위해 **잠재 보상 모델 (Latent Reward Model, LatentRM)**을 설계했습니다.

• 아키텍처: 잠재 추론 백본에 스코어링 헤드를 추가하여, 입력 프롬프트와 현재까지 생성된 잠재 경로 ($h_{1:t}$) 를 받아 스칼라 점수 ($r_t$) 를 출력합니다.
• 학습 목표 (Step-wise Contrastive Objective):
  • 기존 이진 분류 (정답/오답) 방식 대신, **단계별 대비 학습 (Step-wise Contrastive Learning)**을 사용합니다.
  • 각 추론 단계 $t$에서 $N$개의 후보 경로에 대해 Softmax 확률을 계산하고, 정답으로 이어지는 경로에 높은 점수를 부여하도록 학습합니다.
  • 이는 각 단계에서 여러 후보 간의 상대적 품질을 비교하여 더 정교한 스코어링을 가능하게 합니다.
• 집계 전략: 학습된 LatentRM 을 사용하여 생성된 모든 경로의 누적 점수 (Sum of logits) 를 계산하고, 이를 기반으로 Best-of-N 또는 Beam Search를 수행하여 최종 답안을 도출합니다.

3. 주요 기여 (Key Contributions)

1. 잠재적 추론 모델로의 병렬 TTS 확장: 토큰 기반 추론에만 국한되었던 병렬 테스트 시간 확장 능력을 연속 공간 (Latent Space) 모델에도 적용할 수 있게 했습니다.
2. 연속 공간 샘플링 문제 해결: MC-dropout 과 AGN 두 가지 상보적인 전략을 제안하여, 연속 벡터 공간에서도 제어 가능하고 정보량이 풍부한 확률적 추론을 가능하게 했습니다.
3. Latent Reward Model (LatentRM) 설계: 단계별 대비 학습을 통해 잠재적 추론 경로를 평가하고 안내하는 전용 보상 모델을 개발하여, 토큰 확률 없이도 효과적인 경로 선택 (집계) 을 가능하게 했습니다.

4. 실험 결과 (Results)

• 샘플링 확장성: GSM8K, MultiArith 등 수학 추론 벤치마크에서 제안된 두 샘플링 전략 (MC-dropout, AGN) 모두 샘플 수 ($N$) 가 증가함에 따라 해결률 (Coverage) 이 monotonically 증가함을 확인했습니다.
• 샘플링 전략 비교:
  • MC-dropout: 대부분의 $N$에서 더 높은 커버리지를 보였으며, 특히 어려운 문제에서 비표준적인 해결책으로의 구조화된 확장을 유도했습니다.
  • AGN: 높은 다양성 (Diversity) 을 유지하면서도 성능 저하가 적어, 고도 다양성 탐색에 강건한 선택지임을 보였습니다.
  • 시각화: t-SNE 분석 결과, MC-dropout 은 특정 방향으로 밀집된 이동을, AGN 은 방사형으로 퍼지는 '화약' 패턴을 보여주어 서로 다른 탐색 역학을 가짐을 확인했습니다.
• 집계 성능: LatentRM 을 사용한 Best-of-N과 Beam Search는 단순 다수결 투표 (Majority Voting) 보다 일관되게 높은 정확도를 기록했습니다. 특히 Best-of-N 이 가장 효과적인 집계 전략으로 나타났습니다.
• 확장성: GPT-2 기반 모델부터 Llama-3.2-1B, Qwen3-VL(2B/4B) 까지 다양한 백본 모델에서 일관된 성능 향상을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 연속 벡터 공간에서의 추론이 단순한 효율성 향상을 넘어, **병렬 테스트 시간 확장 (Parallel TTS)**을 통해 토큰 기반 모델과 동등하거나 그 이상의 성능 확장 능력을 가질 수 있음을 입증했습니다.

• 효율성: 명시적 CoT 보다 짧은 추론 길이로 더 높은 정확도를 달성할 수 있어, 추론 비용 대비 효율이 뛰어납니다.
• 새로운 방향성: 연속 공간에서의 확률적 샘플링과 보상 기반 집합 학습을 결합함으로써, 추론 모델의 확장성 (Scalability) 에 대한 새로운 패러다임을 제시했습니다.
• 미래 전망: 강화 학습 (RL) 프레임워크와의 통합이나 앙상블 기법의 적용 등을 통해 잠재적 추론의 성능을 더욱 고도화할 수 있는 길을 열었습니다.

요약하자면, 이 연구는 불확실성 기반 샘플링과 대비 학습 기반 보상 모델을 통해 잠재적 추론 모델이 병렬 계산을 통해 지능을 확장할 수 있는 길을 개척한 획기적인 작업입니다.