Reasoning to Learn from Latent Thoughts

Meta info.

• Authors: Yangjun Ruan, Neil Band, Chris J. Maddison, Tatsunori Hashimoto
• Paper: https://arxiv.org/pdf/2503.18866
• Affiliation: Stanford Univ., Univ. of Toronto, Vector Institute.
• Published: March 24, 2025

TL; DR

LLM에 bootstrapping으로 구조화된 internal reasoning representation(여기서는 Token)인 latent thoughts 생성을 학습하여 reasoning ability 향상 가능성 제안

Background

CoT prompting의 일반화로 intermediate reasoning step의 유용성 확인

• 정작 LLM 학습할 때 중간 step으로 학습하는건 아니지 않나 (최종 답변으로 튜닝)

Problem States

LLM한테 Reasoning step 중간에도 supervision을 주면 reasoning 성능이 좋아지지 않을까

Suggestions

BoLT (Bootstrapped latent thought) model 제안

• Latent Thought Sampling > Answer Generation 흐름
• $Z \sim q(Z|X; M_t)$
  • Z: latent thoughts, CoT같은 token sequence.
  • $M_t$: t번째 iteration에서의 모델. EM 알고리즘으로 학습.
  • objective: ELBO
    • 좋은 $Z$를 샘플링하면서

    • $q(Z

      X)$분포가 prior $p(Z

      X)$에 비슷하도록

• Expectation step: Z를 K개 샘플링 > Z와 함께 $Y \sim p(Y|X, Z; M_t)$생성
  • Y에 대해 중요도가 높다고 생각하는 thought Z* 선택

  • (각 $Z_k$마다 $p(Y

    X, Z_k)$ log-likelihood 기반으로 weight w로 활용)

• Maximization step: {Z*, X, Y}로 모델 학습해서 $M_{t+1}$ 업데이트
  • p를 위해 , q를 위해 , Z의 시작과 끝을 알리기 위해 <start/endoflatent> 활용
    • $Z$를 $X$랑 같이 학습하기 위해 데이터 포맷팅에 유의

    • posterior $q(Z

      X)$: Z 예측모델 (문제 $X$를 보면 무슨 생각$Z$가 떠올라?) > 문제는, 얘만 학습하면 Z가 진짜 reasoning하는 Z가 아니라, 보고있는 문제에만 맞는 Y를 맞추는데에 힌트로써만 역할하는 hacky reasoning이 될 우려=overfitting)

    • joint $p(Z,X)$: $Z$랑 $X$를 같이 생성하는 모델 (정답 생각 $Z^*$가 있다면 X같은 문제였겠네 의 흐름) > Z가 정말 괜찮은 생각이라면 Z만 봐도 X가 뭔지 알 수 있지 않나? 같은 제약을 Z에 주면서 단순 Y만 보고 만들지 않도록 의도.

Effects

• Experiments setup
  • Benchmarks: MATH, GSM8K
  • backbone: TinyLlama-1.1B
  • EM iteration은 4회 반복
  • baseline
    • Raw Token Match: latent thought없이 raw-corpus만으로 BoLT만큼 동일한 토큰수 학습
    • Train FLOP Match: 더 작은 corpus를 여러번 돌려서 FLOP 연산량만 BoLT만큼 맞춰 학습
    • Latent Warmstart: M_0(gpt-4o-mini)가 생성한 synthetic latent로 Z를 써서 시작한 모델 (bootstrapping 없음)
• results:
  • latent thought개념을 적용한 BoLT가 그렇지 않은 baseline들에 대해 SOTA tab 1
  • ELBO objective는 iteration 돌 때마다 개선 > 실제 downstream task 성능향상 Fig 8
    • ELBO loss 떨어지면서 MATH 성능 개선 (warm start부터 iter=4까지 )
  • fine-tuning에서도 BoLT가 raw-data보다 유익 Fig 9
    • MATH는 계속 향상되는데에비해 GSM8K는 iter=2이후 plateau
  • bootstrapping을 반복할수록 개선됨 확인 Fig 11

Personal note. 간만에 좋은 논문 같아요. Latent 개념을 쓰긴 하는데 그렇다고 vector로 가져온건 아니라서 눈으로 해석해봄직 하단것도 그렇고(이 연구에서 최초 제안한 듯), CoT는 어쨌든 inference-time의 reasoning이긴 하니까 training에 어떻게 reasoning까지 학습할지를 구체적으로 디자인해서 적용해보고, 모델의 생성 결과 reasoning을 실제 학습에 다시 넣는 self-improvement라는 지점이 이론적으로도 실험적으로도 검증된 점, 모델이 joint랑 posterior 구조 학습하려고 special token을 활용한점이 똑똑했다고 생각합니다. iteration 돌 때마다 비싼거 (그래서 1.1B 모델로 실험할수밖에,,,) EM알고리즘 특성상 Z initialization이 어떻게 되는지에 의존도가 클거라는 점이 구조적 한계인거랑 z는 어쨌든 backbone의 upperbound에 제약이 있긴 하다는 점은 당연히 후속연구가 필요하겠지만…