PonderLM-2: Pretraining LLM with Latent Thoughts in Continuous Space

이 논문은 생성 단계의 연산량을 늘려 각 토큰 예측 전에 잠재적 사고 (latent thought) 를 생성하는 'PonderLM-2'를 제안함으로써, 추론 비용은 동일하게 유지하면서 모델 크기를 두 배로 늘린 기존 모델보다 더 뛰어난 성능을 달성하는 새로운 사전 학습 방법을 소개합니다.

핵심

🧠 "생각하는 시간"을 늘려주는 AI: PonderLM-2 설명

이 논문은 **"AI 가 말을 하기 전에 잠시 '생각'할 시간을 주면, 훨씬 똑똑해진다"**는 아주 직관적이지만 혁신적인 아이디어를 제안합니다.

기존의 AI 는 질문을 받자마자 바로 대답을 뱉어내야 했지만, 이 새로운 방법 (PonderLM-2) 은 AI 가 대답을 하기 전에 **눈에 보이지 않는 '잠재적 생각 (Latent Thought)'**을 하나 더 거쳐서 더 정확한 답을 찾도록 훈련시킵니다.

일상생활에 비유해서 쉽게 설명해 드릴게요.

---

1. 기존 AI vs 새로운 AI: "속도전" vs "고민"

• 기존 AI (일반적인 언어 모델):

  • 상황: 친구가 "오늘 점심 뭐 먹지?"라고 물었을 때, 0.1 초 만에 "김치찌개!"라고 바로 외칩니다.
  • 문제: 너무 급하게 대답하다 보니, "아, 김치찌개는 어제 먹었지"라는 사실을 놓치거나, "아침에 배탈 났는데..."라는 상황을 고려하지 못해 엉뚱한 답을 할 수 있습니다.
  • 비유: 무작정 달리는 마라토너처럼, 멈추지 않고 계속 앞으로만 나갑니다.

• 새로운 AI (PonderLM-2):

  • 상황: 친구가 물어보면, 잠시 멈춰서 "음... 김치찌개는 어젯밤에 먹었으니 안 되겠고, 비가 오니까 따뜻한 국물이 좋겠네... 아, 된장찌개 어때?"라고 눈에 보이지 않는 머릿속에서 고민을 합니다.
  • 결과: 그 '고민' 과정을 거친 후, 훨씬 더 적절한 "된장찌개!"라고 대답합니다.
  • 비유: 명상하는 요가 강사처럼, 한 걸음을 내딛기 전에 균형을 잡고 생각한 후 움직입니다.

2. 핵심 기술: "보이지 않는 생각 (Latent Thought)"

이 논문에서 말하는 **'잠재적 생각 (Latent Thought)'**이란 무엇일까요?

• 디지털 공간에서의 '초안' 작성:
  우리가 글을 쓸 때, 바로 최종 원고를 쓰는 게 아니라 먼저 **초안 (Draft)**을 쓰거나 메모를 적는 경우가 있죠? PonderLM-2 는 AI 가 실제 단어 (Token) 를 출력하기 전에, **컴퓨터 내부의 '보이지 않는 메모장'**에 그 단어에 대한 초안을 적어보는 과정을 추가합니다.
• 연속된 공간에서의 사고:
  기존 AI 는 단어장 (사전) 에서 단어를 고르는 방식이라서, "김치찌개"와 "된장찌개" 사이에는 명확한 경계가 있습니다. 하지만 이 새로운 방법은 **연속된 공간 (Continuous Space)**에서 생각을 합니다. 마치 "김치찌개"와 "된장찌개" 사이의 중간 지점까지 생각의 흐름을 이어가며 더 정교한 결정을 내리는 것과 같습니다.

3. 어떻게 훈련시킬까? "야코비 (Jacobi) 반복"이라는 마법

여기서 가장 재미있는 점은, AI 가 이렇게 '생각'하는 과정을 훈련할 때 어떻게 빠르게 시켰느냐는 것입니다.

• 문제: 보통 AI 가 한 번 생각하면, 그 결과가 다음 생각의 입력이 되어야 합니다. (A → B → C) 이렇게 순서대로 하면 훈련 속도가 매우 느립니다.
• 해결책 (야코비 반복): 연구자들은 마치 거울 앞에 거울을 여러 개 놓고 이미지를 빠르게 반복하듯, AI 가 모든 단계를 한 번에 병렬로 (Parallel) 계산하도록 만들었습니다.
  • 비유: 교실에서 선생님이 "1 번 학생, 답을 말해봐"라고 하면 1 번이 말하고, 그 다음 2 번이 말해야 하는 게 아니라, 모든 학생이 동시에 답을 적어내서 선생님이 한 번에 확인하고 수정하는 방식입니다. 이렇게 하면 훈련 시간이 훨씬 단축됩니다.

4. 왜 이것이 중요한가? "작은 뇌"가 "큰 뇌"를 이긴다

이 방법의 가장 놀라운 성과는 효율성입니다.

• 기존 방식: 더 똑똑한 AI 를 만들려면 **매우 큰 뇌 (파라미터 수)**를 만들어야 했습니다. (예: 28 억 개의 파라미터)
• 새로운 방식: PonderLM-2 는 **작은 뇌 (14 억 개의 파라미터)**를 사용하더라도, '생각하는 시간'을 늘려주면 2 배 큰 뇌보다 더 좋은 성적을 냅니다.
  • 비유: **작은 차 (1.4B)**가 **거대한 트럭 (2.8B)**보다 더 빠르고 똑똑하게 목적지에 도착하는 것입니다. 단, 그 작은 차는 운전사가 **미리 경로를 꼼꼼히 계획 (생각)**했기 때문에 가능합니다.

5. 요약: AI 의 '생각'을 늘리자

이 논문은 **"AI 에게 더 많은 데이터나 더 큰 모델을 주는 것보다, 매 순간 '잠시 멈춰서 생각'하는 시간을 훈련시키는 것이 더 효과적이다"**라고 말합니다.

• 기존: "말하기 전에 생각하라 (Chain-of-Thought)" → 질문할 때만 적용.
• PonderLM-2: "말하기 전에, 매 단어마다 생각하라" → 훈련 단계부터 습관화.

이처럼 AI 가 **눈에 보이지 않는 '생각의 과정'**을 통해 스스로를 다듬는다면, 앞으로 더 적은 비용으로 훨씬 더 똑똑하고 정확한 AI 를 만들 수 있을 것입니다. 마치 우리가 급하게 대답하기보다, 한 번 더 깊이 생각하면 더 좋은 답을 얻는 것과 같은 원리입니다.

기술 요약

1. 문제 제기 (Problem Statement)

대형 언어 모델 (LLM) 의 성능 향상을 위한 기존 접근법은 주로 파라미터 수와 데이터 양의 확장 (Scaling) 에 의존해 왔으나, 데이터 부족과 스케일링 법칙의 한계로 인해 체감 효율이 떨어지고 있습니다. 이에 따라 테스트 시간 확장 (Test-time Scaling), 특히 체인 오브 씽킹 (Chain-of-Thought, CoT) 을 통한 추론 단계 확장이 주목받고 있습니다.

그러나 기존 CoT 기반 방법론들은 다음과 같은 한계를 가집니다:

• 데이터 의존성: 특수한 지시 데이터 (Instruction Data) 나 강화 학습 (RL) 이 필요하여 일반 코퍼스로 사전 학습하기 어렵습니다.
• 이산적 공간의 제약: 사고 과정을 토큰 (Discrete Token) 단위로 표현하므로 표현력이 제한적입니다.
• 모델 능력의 한계: 기본 모델의 능력에 종속적이며, 복잡한 추론을 위해서는 모델 자체의 능력을 먼저 키워야 합니다.

따라서, 사전 학습 (Pretraining) 단계에서 각 토큰 생성 전에 추가적인 계산 단계를 도입하여, 이산적 토큰 공간이 아닌 연속적 잠재 공간 (Continuous Latent Space) 에서 사고를 확장하는 새로운 방법론이 필요합니다.

---

2. 방법론 (Methodology)

저자들은 PonderLM-2를 제안합니다. 이는 언어 모델이 다음 토큰을 예측하기 전에, 해당 위치의 **잠재적 사고 (Latent Thought)**를 생성하도록 사전 학습하는 방법입니다.

핵심 메커니즘

1. 잠재적 사고 생성 (Latent Thought Generation):

   • 표준 Transformer 은 현재 입력 시퀀스를 기반으로 다음 토큰을 직접 예측합니다.
   • PonderLM-2 는 먼저 현재 위치의 **마지막 은닉 상태 (Last Hidden State)**를 "잠재적 사고"로 생성합니다.
   • 이 은닉 상태는 실제 토큰이 아닌 **연속 벡터 (Continuous Vector)**이며, 이를 다음 단계의 입력 임베딩으로 재사용하여 실제 토큰을 예측합니다.
   • 이를 통해 모델은 제약 없는 연속 공간에서 예측을 정제 (Refine) 할 수 있습니다.

2. 수평적 확장 (Horizontal Scaling):

   • 기존 '수직적 확장 (Vertical Scaling, 모델 깊이 증가)'과 달리, 모델 구조를 변경하지 않고 **생성 과정의 길이 (Sequence Length)**를 늘려 계산량을 확장합니다.
   • 각 실제 토큰 앞에 1 개 이상의 잠재적 사고 토큰이 삽입됩니다.

3. 병렬 학습을 위한 자코비 반복 (Jacobi Iteration):

   • 문제: 잠재적 사고를 순차적으로 생성하면 (h1 → h2 → ...), 긴 시퀀스 학습 시 순차적 의존성으로 인해 병렬 학습이 불가능해집니다.
   • 해결: 자코비 반복 (Jacobi Iteration) 기법을 도입하여 병렬 학습을 가능하게 합니다.
     • 초기화: 원본 토큰 임베딩으로 초기 은닉 상태를 계산합니다.
     • 반복: 이전 반복 단계의 은닉 상태와 원본 토큰 임베딩을 번갈아 가며 (Interleaving) 새로운 입력 시퀀스를 구성하고, 이를 모델에 한 번의 패스로 통과시켜 모든 위치의 은닉 상태를 동시에 업데이트합니다.
     • 수렴: 이 과정은 고정점 (Fixed Point) 을 찾아가며, 이론적으로 $T$ (시퀀스 길이) 번 이내, 실제로는 매우 적은 반복 (3~4 회) 만에 순차적 추론 결과와 일치하는 값으로 수렴합니다.

4. 학습 목적 함수:

   • $K$번의 자코비 반복 후, 최종 은닉 상태로부터 다음 실제 토큰을 예측하도록 크로스 엔트로피 손실을 계산합니다.
   • 과적합을 방지하기 위해 반복 횟수 $K$를 학습 시마다 무작위로 샘플링합니다 (예: 2 또는 3).

---

3. 주요 기여 (Key Contributions)

1. 새로운 사전 학습 패러다임: 일반 코퍼스 (General Corpus) 를 사용하여, 특수한 지시 데이터나 강화 학습 없이도 모델이 **연속 공간에서의 사고 (Latent Thoughts)**를 학습하게 하는 최초의 방법론 중 하나입니다.
2. 효율적인 병렬 학습 알고리즘: 순차적 의존성을 가진 '생각' 과정을 자코비 반복을 통해 병렬화하여, 긴 시퀀스 학습의 계산적 비효율성을 해결했습니다.
3. 파라미터 효율성 극대화: 추론 비용 (Inference Cost) 이 동일한 조건에서, 파라미터 수가 2 배인 표준 모델보다 우수한 성능을 달성했습니다.
4. CoT 와의 시너지: 잠재적 사고의 개수를 늘려 '잠재적 사고의 사슬 (Chain of Latent Thoughts)'을 형성하면, CoT 와 유사하게 성능이 지속적으로 향상됨을 입증했습니다.

---

4. 실험 결과 (Results)

저자들은 Pythia, LLaMA 아키텍처를 기반으로 대규모 실험을 수행했습니다.

• 스케일링 효율성 (Scaling Efficiency):

  • 파라미터 효율: PonderLM-2-Pythia-1.4B 모델은 300B 토큰으로 학습되었으며, 파라미터 수가 55% 적은 (1.4B vs 2.8B) 조건에서도 공식 Pythia-2.8B 모델보다 언어 모델링 (Perplexity) 과 다양한 다운스트림 태스크에서 더 높은 성능을 보였습니다.
  • 데이터 효율: 동일한 1.4B 모델이 표준 Pythia-1.4B 와 동일한 최종 성능에 도달하는 데 62% 적은 학습 토큰이 필요했습니다.

• 다운스트림 태스크 성능:

  • Lambada, ARC, WinoGrande, RACE 등 9 가지 일반 태스크에서 표준 Pythia 모델 및 TinyLlama(1.1B, 3T 토큰 학습) 를 능가했습니다.
  • 특히 TinyLlama 는 학습 데이터가 10 배 많음에도 불구하고 PonderLM-2-1.4B 에게 성능에서 밀렸습니다.

• 추론 비용 대비 성능:

  • 추론 FLOPs 가 2 배 증가한 조건 (잠재적 사고 1 회 추가) 에서, 파라미터가 2 배인 모델 (LLaMA-2.8B) 이나 다른 계산 확장 기법 (Looped Transformer, Pause Token 등) 보다 낮은 퍼플렉시티와 높은 정확도를 기록했습니다.

• 기존 모델 적용 (Continual Pretraining):

  • 기존 LLaMA-3-3B 모델을 기반으로 추가 학습 (Continual Pretraining) 시, 기존 방법보다 더 빠른 수렴과 성능 향상을 보여주어 '플러그 앤 플레이' 가능성을 입증했습니다.

• 테스트 시간 확장 (Test-Time Scaling) 과의 호환성:

  • PonderLM-2 로 학습된 모델은 테스트 시간의 Majority Voting, Best-of-N, CoT 프롬프팅과 결합 시 추가적인 성능 향상을 보였습니다. 이는 두 기법이 상호 보완적임을 의미합니다.

---

5. 의의 및 결론 (Significance)

• 새로운 스케일링 차원 제시: 모델의 크기 (파라미터) 나 데이터 양을 늘리는 전통적인 방식 외에, **생성 과정 내의 계산 단계 (Computational Steps)**를 확장하여 모델 능력을 향상시킬 수 있음을 증명했습니다.
• 연속 공간 사고의 가능성: 토큰 단위의 이산적 사고를 넘어, 은닉 상태 (Hidden State) 를 활용한 연속 공간에서의 사고가 언어 모델의 추론 능력을 근본적으로 향상시킬 수 있음을 보여줍니다.
• 실용성: 복잡한 강화 학습이나 특수 데이터 없이도 일반 텍스트 코퍼스로 학습 가능하며, 기존 오픈소스 모델에 적용하여 성능을 개선할 수 있어 실제 적용 가능성이 높습니다.

결론적으로, PonderLM-2 는 사전 학습 단계에서 모델이 '생각하는 시간'을 가질 수 있도록 하여, 더 적은 파라미터와 데이터로도 더 강력한 추론 능력을 갖춘 언어 모델을 구축하는 새로운 방향성을 제시합니다.