[논문리뷰] Training Large Language Models to Reason in a Continuous Latent Space

dongcho·2025년 11월 13일

논문리뷰

목록 보기
6/8
  • Key words

Chain-of-Thought(CoT): 자연어로 중간 사고 과정을 단계별로 출력하는 방식

Latent Space: hidden representations, 쉽게 말해 말 대신 느낌이나 의미 덩어리로 사고하는 공간

Breadth-First Search(BFS): 여러 경로를 동시에 탐색하며 정답을 찾아나가는 방법. 갈림길에서 하나만 고르는 게 아니라 여러 후보를 동시에 탐색.

선정 이유

자연어로 출력되는 CoT의 이점은 쉬운 해석 가능성도 큰 비중을 차지한다고 생각했는데, Latent space에서의 reasoning을 어떻게 적용시켰을지 궁금해서 읽어보게 되었다.


Background

CoT는 LLM의 사고력을 끌어올렸지만, 말로 모든 걸 설명해야 하는 비효율성이 있었음. 인간도 모든 생각을 말로 설명하지 않듯, 모델도 '생각 공간'에서 사고하는 게 더 자연스러울 수 있음.

+) Fedorenko et al., 2024: '인간의 언어는 추론보다 의사소통에 최적화되어있다.' → 언어가 아닌 다른 공간에서 추론하는 것이 더 효과적일 수 있음.

++) 그 외에도 여러 neuroimaging 연구에서 인간 뇌의 언어 네트워크가 추론 과정에서 거의 비활성화되어있음이 나타남.

LLM이 추론을 위해 언어를 사용할 때의 문제점

  • 각 특정 추론 토큰에 필요한 추론의 양이 다양하지만, 현재 LLM의 구조는 모든 토큰에 대해 거의 동일한 컴퓨팅 예산을 할당하고 있음.
  • CoT 대부분의 토큰은 유창성을 위해서만 사용됨.
  • 언어 토큰을 이용한 한 방향 추론(결정론적인 추론)에서는 초기에 잘못된 판단을 내릴 경우 매우 불리해짐.→ 사람이 글을 쓸 때 문법 오류를 고치지 않고 계속 쓰면 전체 문맥 자체가 어색해지는 원리
  • Pause Token: 모델을 사전 훈련할 때 학습 가능한 pause 토큰을 훈련 코퍼스에 무작위로 삽입하여 LLM의 성능을 향상시키는 방법을 제안.
  • internalization(iCoT): explicit CoT 추론을 위해 훈련된 모델에서 시작하여 점진적으로 중간 단계를 제거하고 모델을 fine-tuning하여 CoT 과정을 내재화시키려 함.
  • +) Knowledge distillation을 통해 Student 모델은 Teacher 모델의 확률분포 또는 hidden state를 모방하도록 학습됨. 이 과정에서 Student 모델은 Teacher 모델의 추론 과정을 암묵적으로 학습하게 되는 원리.
    • 저자는 주로 여기서 언급된 방법론을 사용하되 language space를 Latent space로 전환함.
    • iCoT 방법론: 기존의 CoT 방식을 사용하되 multi-stage 방식으로 초반의 토큰을 점점 잘라내 answer-only 훈련

위와 같은 시도가 있었지만, 여전히 언어 기반이라는 한계가 있음.

Purpose

Coconut 방식을 통해 언어로 사고하지 않고 continuous latent space에서 추론하도록 하는 새로운 패러다임을 제시하여 자연어로 표현하는 비용 없이 더 깊이 있는 사고를 하게 만들고자 함. 이를 통해 BFS와 유사한 프로세스가 가능해짐!

+) 즉, 모델이 초기에 "올바른 결정"을 내리지 못하더라도, 여러 '가능한 옵션'들을 continuous thought 안에 '유지'하면서 추론을 진행할 수 있음.


Method

[Continuous Thought]

Cocounut에서는 LLM의 마지막 hidden state(h_t)를 추론 형태의 표현으로 사용함.

위의 식은 주어진 입력 시퀀스 (x ≤ t)에 대해 다음 토큰 (x_{t+1})의 확률 분포를 나타내는 LLM의 예측 함수. latent space에서의 추론은 continuous thought가 언어 공간으로 다시 매핑되는 것을 의도하지 않기 때문에, 이 함수는 정의되지 않음.

+) 즉, latent mode에서는 다음 토큰을 예측하는 대신, continuous thought를 다음 단계의 입력으로 직접 사용하는 데 집중하는 것!

[Latent와 Language mode의 전환]

와 토큰 사용하여 시작과 끝 표시.

latent mode 구간(가 j라고 하면 i부터 j-1까지)에는 continuous thought인 마지막 hidden state h가 입력 임베딩으로 사용됨. latent mode가 끝나면 모델은 다시 일반적인 언어 모델처럼 토큰 임베딩을 입력으로 받아서 작동함.

[Multi-stage training]

초기 단계에서는 일반 CoT instance로 학습. 이후 서서히 latent mode로 대체함.

  • k번째 단계에서 CoT의 처음 k개의 추론 단계를 k*c의 continuous thought으로 대체함.c: latent thought 수 지정하는 hyper param. (Fig 2에서 c=1)
  • 학습 과정에서 표준 negative log-likelihood loss를 최적화하지만, 질문과 latent thought에 대한 loss는 masking.

[Inference]

표준 언어 모델 decoding과 유사. 그러나 latent mode에서는 마지막 hidden state → 다음 입력 embedding.

eot 토큰을 언제 생성할지는 1) binary classifier 학습 또는 2) latent thought의 길이를 상수로 padding 하여 결정.

+) 실험에서는 2번의 방법을 사용함. binary classifier은 continuous thought을 입력으로 받아 종료 시점을 판단하는데, '1번 2번 둘 다 잘 작동한다'고 밝힌 뒤 자세한 설명은 생략됨.

++) 정해진 상수 값으로 두고 남는 토큰을 전부 처리하는 것이 구현 측면에서 더 간단해서 그런듯?

Experiment

3개의 데이터셋에 대한 실험을 진행하여 모델이 생성한 답변과 ground truth를 비교하여 정확도를 평가. (pre-trained GPT-2 사용.)

+) GPT-2를 사용한 이유: 선행연구인 iCoT와 공정한 비교를 위해서

Inference 시에는 greedy decoding을 선택하여 복잡한 sampling 없이 가장 가능성 높은 출력을 선택.

[Math Reasoning]

GSM8k: 초등학교 수준의 수학문제, iCoT 연구에서 사용한 합성데이터셋 사용.

기본적으로 c=2로 설정하고 step 3까지 훈련.

  • 결과: 기존의 CoT보다 적은 토큰으로 더 높은 정확도 달성. 정확도만 따진다면 iCoT 또한 능가함.

  • 하이퍼파라미터 실험: 추가적으로 하이퍼파라미터 c를 조정해봤을 때 0, 1, 2로 증가시킬수록 정확도가 향상됨. → CoT와 유사한 chaining 효과가 latent space에서도 나타남을 시사.
  • Case Study - continuous thought

첫 번째 continuous thought를 언어 토큰으로 디코딩한 결과, "180", " 180" (공백 포함), "9"와 같은 토큰이 나타남. 이 문제의 올바른 추론 경로는 "3 × 3 × 60 = 9 × 60 = 540" 또는 "3 × 3 × 60 = 3 × 180 = 540".

→ 첫 번째 continuous thought의 해석 결과인 "180"과 "9"는 실제 계산 과정에서 나타나는 중간 변수이므로 모델이 latent space 안에 실제 문제 해결에 필요한 핵심 정보를 담고 있음 + 두 가지의 가능한 계산 경로를 고려하고 있음!! (BFS와 유사한 프로세스임)을 증명

[Logical Reasoning]

ProntoQA: 5단계(5-hop) 추론을 요구하는 질문들로 구성. 각 질문에는 트리 구조의 지식 그래프가 제공되는데, 추론에 도움이 되지 않는 방해 가지들도 존재함.

ProsQA: 저자들이 ProntoQA보다 더 복잡한 문제를 만들기 위해 DAG를 사용해 고도화시킴.

+) Directed Acyclic Graph(DAG): 방향성 비순환 그래프. 트리 구조보다 더 복잡한 구조를 가질 수 있음.

기본적으로 c=1로 설정하고, step 6까지 훈련.

  • 결과: 기존의 CoT와 비교했을 때 성능 향상. iCoT와는 비슷한 성능 보임.
  • k의 변화에 따른 정확도 변화

  • Reasoning 과정을 그래프 경로로 해석하고, 다음과 같이 분류해봄.
    • Correct Path: 정답으로 이어지는 최단 경로 중 하나.
    • Longer Path: 유효하지만 최단 경로보다 긴 경로.
    • Hallucination: 존재하지 않는 엣지(edge)를 포함하거나 연결이 끊긴 경로.
    • Wrong Target: 유효하지만 목표 노드가 아닌 다른 노드에 도달하는 경로.
    • Correct Label / Incorrect Label: 경로 없이 최종 답변만 출력하는 경우.

[Ablation Study]

  • without curriculum: multi-stage 커리큘럼 없이 바로 continuous thought만으로 훈련 → 성능 급하락.
  • without thought: continuous thought 없이 바로 CoT만 줄여가며 훈련 → 일부 개선되지만 Coconut보다 떨어짐.
  • pause as thought: continuous thought 대신 pause token 삽입→ 생각하는 시간(pause)를 주는 것만으로는 충분하지 않고, 진짜 structured reasoning이 필요하다는 걸 보여줌

[latent space의 이점]

Fig 10에서 나타난대로 height(노드 높이)가 낮을수록 정답과 오답 구분이 쉽고, 높을수록 판단이 어려워짐.

+) Search Tree의 관점으로 보면, leaf node에 해당하는 게 낮은 노드. 하위에 가능성이 없는 leaf node는 금방 평가할 수 있는 것.

이러한 문제에서 Latent space는 일찍 확정하지 않고, 마지막까지 가능한 경로를 유지하면서 정답과 오답을 구분할 수 있게 해줌.→ 섣불리 잘못된 가지로 빠지지 않게 해줌

따라서, planning이 필요한 복잡한 문제에서 훨씬 강력함.


Key point

  • CoT를 latent space로 대체한 시도!
  • planning-heavy 문제에서 성능이 좋다. (ProsQA) → latent space의 이점
  • continuous thought chaining으로 reasoning depth를 늘릴 수 있음
  • BFS 적인 사고 → 여러 선택지를 한꺼번에 유지하면서 점진적으로 잘못된 길을 제거하는 능력. 섣불리 단정짓지 않음

Limitations and Suggestions

  • 완전히 latent reasoning만으로 훈련하면 학습이 잘 안됨.
  • 그냥 답만 맞추기로는 좋은 생각 흐름을 학습하지 못해서 중간 사고과정을 완전히 없앨 수는 없음.
  • training efficiency가 떨어짐. (continuous thought을 여러 번 계산해야 해서 시간이 오래 걸려 병렬화가 어려움.)

Personal Insight

  • LLM이 인간처럼 사고하는 방식으로 진화할 수 있는 가능성을 제안한 것으로 보임.+) iCoT의 방식이 CoT 예시들을 통한 단어 예측 확률 분포 최적화처럼 느껴지는 반면, Coconut은 여러 가능성을 고차원적으로 사고하여 문제를 해결하는 것에 초점을 맞춘 것으로 보임.
  • 각 task별로 k와 c를 어떻게 최적화할 수 있는지 더 많은 예시를 보고싶음.
  • 토큰의 binary classification에 대한 설명이 부족해서 아쉬움.
  • ablation study에서 저자는 pause를 주는 것만으로는 부족하다 라고 언급했지만 내가 느낄 때는 생각보다 성능이 좋아서... pause token 논문도 읽어보고 싶음.

Overall...

iCoT의 방법론을 가져가면서도 공간만을 옮겨보는 발상의 전환이 참신하고, latent space 활용의 잠재력을 보여줌! 노드 서치 알고리즘이나 value function 수정 등의 시도가 이어졌을 것 같은데 더 찾아보고싶다.

profile
비틀비틀짝짜쿵

0개의 댓글