Thinking into the Future: Latent Lookahead Training for Transformers

1. 문제 정의 (Problem Statement)

현대 대규모 언어 모델 (LLM) 은 **다음 토큰 예측 (Next-Token Prediction, NTP)**과 Teacher-forcing을 기반으로 한 자기회귀 (Autoregressive) 학습 패러다임을 따릅니다. 이 방식은 확장성이 뛰어나지만, 다음과 같은 두 가지 근본적인 한계를 가집니다.

1. 단순한 미래 시야 (Myopic Future Horizon): 모델은 매 단계에서 오직 다음 한 토큰만을 예측하도록 훈련됩니다. 이는 모델이 여러 단계 앞의 결과를 고려하거나, 현재 선택이 미래에 미칠 영향을 명시적으로 추론 (Reasoning) 하는 능력을 제한합니다.
2. 비효율적인 계산 할당 및 샘플링 오류: 모든 토큰 생성에 동일한 계산량이 할당되며, 매 단계에서 잠재 공간 (Latent Space) 의 상태를 즉시 이산적인 (Discrete) 토큰으로 샘플링하여 결정합니다. 불확실성이 높은 단계 (예: 미로 찾기, 스도쿠의 어려운 칸) 에서도 모델은 즉시 결정해야 하므로, 잘못된 경로로 진입하거나 사고의 연쇄 (Chain-of-Thought) 가 끊기는 오류가 발생할 수 있습니다.

인간은 어려운 결정을 내릴 때 내부적으로 여러 시나리오를 시뮬레이션하며 '생각'한 후 행동합니다. 반면, 기존 LLM 은 이러한 '생각' 과정을 거치지 않고 즉각적인 반응에 의존합니다.

2. 방법론 (Methodology: Latent Lookahead)

저자들은 **"Latent Lookahead (잠재적 선점)"**라는 새로운 학습 전략을 제안합니다. 이는 모델이 다음 토큰을 출력하기 전에, 잠재 공간 (Latent Space) 에서 다단계로 미래를 시뮬레이션할 수 있게 합니다.

핵심 메커니즘

• 가시적 토큰 (Visible Tokens) vs 잠재적 토큰 (Latent Tokens):
  • 가시적 토큰: 실제 데이터에 존재하는 표준 토큰입니다.
  • 잠재적 토큰: 모델이 선택된 위치에서 생성하는 보조 예측값으로, 실제 토큰으로 샘플링되지 않고 잠재 상태 (Hidden States) 로만 유지됩니다.
• 작동 원리:
  1. 모델이 다음 가시적 토큰을 출력하기 직전, 선택된 위치에서 $\tau$ (타우) 단계의 잠재적 사고 (Latent Thought) 를 수행합니다.
  2. 모델은 예측된 잠재 상태를 컨텍스트에 다시 피드백하여, $\tau$번의 재귀적 단계를 거칩니다.
  3. 이 과정에서 샘플링 (Sampling) 을 수행하지 않으므로, 모델은 토큰을 확정하기 전에 내부적으로 여러 가능성을 탐색하고 상태를 정제 (Refine) 할 수 있습니다.
  4. $\tau$단계의 사고가 끝난 후, 첫 번째 잠재 상태의 숨겨진 표현을 기반으로 최종 가시적 토큰을 샘플링합니다.

학습 및 추론

• 학습 (Training):
  • $\tau$개의 잠재 토큰 ($z_{i,1}, \dots, z_{i,\tau}$) 은 각각 $i+1, \dots, i+\tau$ 번째의 정답 (Ground-truth) 토큰에 대해 지도 학습 (Supervised) 됩니다.
  • 이는 모델이 "앞으로 $\tau$단계를 내다보며 (Lookahead)" 예측을 정교화하도록 강제합니다.
  • Attention Mask: 잠재 토큰 간에는 양방향 (Bi-directional/Non-causal) 어텐션을 허용하여, 미래의 잠재 상태가 과거의 잠재 상태를 정제할 수 있게 합니다. 이는 병렬 학습을 가능하게 하며 계산 복잡도를 줄입니다.
• 추론 (Inference):
  • 모델은 선택된 위치에서 $\tau$개의 잠재 단계를 거친 후, 첫 번째 잠재 상태로부터 다음 토큰을 생성합니다.
  • 이는 모델이 "생각"한 후 "말하는" 구조를 구현합니다.

3. 주요 기여 (Key Contributions)

1. Latent Lookahead 프레임워크 제안: 가시적 토큰을 생성하기 전에 잠재 공간에서 재귀적으로 사고를 확장하는 새로운 학습 및 추론 프로세스를 상세히 기술했습니다.
2. 효율적인 병렬 학습 설계: $n$개의 사고 위치에서 $\tau$단계를 수행할 때, 순차적 처리 ($O(n\tau)$) 대신 어텐션 마스크를 설계하여 **병렬 처리 ($O(\tau)$)**가 가능하도록 하여 학습 비용을 획기적으로 줄였습니다.
3. 계획 (Planning) 작업에서의 성능 입증: 미로 찾기 (Maze), 스도쿠 (Sudoku), ProsQA 등 **선견지명 (Foresight)**이 필수적인 작업에서 기존 자기회귀 모델 및 비자기회귀 베이스라인을 압도적으로 능가하는 성능을 보였습니다.
4. 잠재 토큰의 해석: 생성된 잠재 토큰이 상태의 중첩 (Superposition of states) 을 인코딩하며, 사고 과정을 통해 불확실성이 줄어들고 정답으로 수렴하는 것을 시각적으로 증명했습니다.

4. 실험 결과 (Results)

저자들은 2 층 GPT-2 유사 모델을 사용하여 다양한 계획 작업에서 실험을 수행했습니다.

• 성능 향상:
  • 9x9 스도쿠: 기존 자기회귀 (NTP) 베이스라인의 정확도 (12.5%) 를 **35.5%**까지 크게 향상시켰습니다. (Pause Token 기법은 12.5% 에 머무름)
  • 4x4 미니 스도쿠: 78.0% (NTP) → 93.5% (Latent Lookahead)
  • ProsQA 및 Maze: 두 작업 모두에서 NTP 및 Pause Token 베이스라인보다 우수한 성능을 기록했습니다.
• 계산량 스케일링 ($\tau$ 증가): 잠재 토큰의 수 ($\tau$) 를 늘릴수록 성능이 선형적으로 증가했습니다. 반면, 기존 Pause Token 기법은 성능이 포화되거나 오히려 하락하는 경향을 보였습니다.
• 깊이 (Depth) vs 잠재 사고: 32 층의 깊은 모델보다 2 층 모델에 Latent Lookahead 를 적용한 것이 더 높은 정확도 (93.5% vs 80.0%) 를 보였습니다. 이는 성능 향상이 단순히 모델 깊이의 증가 때문이 아니라, 계산 자원을 할당하여 사고하는 능력에서 비롯됨을 시사합니다.
• 어텐션 마스크의 중요성: 잠재 토큰 간에 양방향 어텐션 (Non-causal) 을 허용하는 것이, 단방향 (Causal) 어텐션보다 훨씬 좋은 성능을 냈습니다. 이는 미래의 추론이 과거의 추론을 정제하는 과정이 필수적임을 의미합니다.

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

이 논문은 LLM 이 단순한 다음 토큰 예측을 넘어, 복잡한 계획과 추론이 필요한 작업을 수행할 수 있도록 하는 새로운 패러다임을 제시합니다.

• 계산 자원의 동적 할당: 어려운 결정 단계에서 더 많은 계산 (Latent Lookahead) 을 할당하여 모델의 표현력을 극대화합니다.
• 샘플링 오류 감소: 토큰을 확정하기 전에 잠재 공간에서 여러 시나리오를 탐색하고 정제함으로써, 초기 오류가 누적되는 것을 방지합니다.
• 미래 방향: 추후 사전 학습 (Pre-training) 단계에서 이 방식을 적용하거나, 불확실성 지표 (엔트로피 등) 를 기반으로 '어디서 생각할지'를 동적으로 학습하는 방향으로 발전할 수 있습니다.

결론적으로, Latent Lookahead는 모델이 "생각한 후 말하기 (Think before talking)"를 가능하게 하여, 구조화된 계획 작업에서 인간과 유사한 추론 능력을 구현하는 강력한 방법론으로 평가됩니다.