Training Large Language Models To Reason In Parallel With Global Forking Tokens

이 논문은 다양한 추론 경로를 보존하고 전역 분기 토큰을 유도하기 위해 쌍대 매칭을 기반으로 한 집합 기반 손실 함수를 도입한 'Set Supervised Fine-Tuning (SSFT)'과 이를 활용한 'Global Forking Policy Optimization (GFPO)'을 제안하여, 수학 추론 및 코드 생성 벤치마크에서 기존 모델보다 우수한 성능을 입증했습니다.

핵심

이 논문은 **"거대 언어 모델 (LLM) 이 복잡한 문제를 풀 때, 어떻게 하면 더 똑똑하고 다양한 방식으로 생각할 수 있게 할까?"**라는 질문에 대한 해답을 제시합니다.

기존의 방식은 모델에게 "생각해 봐"라고만 하면, 모델이 혼자서 길게 생각하다가 (Sequential Thinking) 실수를 하거나, 너무 길게 생각하다가 지쳐서 엉뚱한 답을 내놓는 경우가 많았습니다. 이 논문은 이를 해결하기 위해 **'병렬 사고 (Parallel Thinking)'**와 **'글로벌 포크링 토큰 (Global Forking Tokens)'**이라는 새로운 개념을 도입했습니다.

이 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 문제: "생각이 너무 깊어지면 망한다" (Overthinking)

지금까지 AI 는 문제를 풀 때, 한 가지 길 (한 가지 생각의 흐름) 을 따라 계속 앞으로 나아갔습니다. 마치 미로에서 한 번 길을 잘못 들면, 그 길만 계속 따라가다가 결국 막다른 길에 도달하는 것과 같습니다.

• 기존 방식: "생각해 봐"라고 하면, AI 는 무작위로 다양한 생각을 시도합니다. 하지만 중요한 분기점 (어떤 생각을 선택할지) 에서 올바른 선택을 하기가 어렵습니다.
• 문제점: 다양한 생각을 하려고 온도를 높이면 (Temperature Scaling), 엉뚱한 소리를 할 확률도 함께 늘어납니다. 즉, '다양성'과 '정확성'을 동시에 잡기 힘든 딜레마가 있었습니다.

2. 해결책: "여러 명의 전문가를 한 번에 부른다" (Set of Next-Token Prediction)

이 논문은 AI 가 문제를 풀 때, 한 번에 여러 개의 서로 다른 사고방식 (Reasoning Modes) 을 동시에 가동하도록 훈련시켰습니다.

🌟 핵심 비유: "마법 지팡이와 다양한 길"

이 연구에서는 **'글로벌 포크링 토큰 (Global Forking Tokens)'**이라는 특별한 마법 지팡이를 도입했습니다.

• 이 지팡이는 <think 1>, <think 2>, <think 3> 같은 특수한 태그입니다.
• 기존 방식: AI 가 "어떤 생각을 할까?"라고 고민하다가 우연히 좋은 생각이 떠오르면 그걸로 가는데, 그걸 찾기가 너무 어렵습니다.
• 이 연구의 방식: 우리가 직접 <think 1>을 입력하면 AI 는 A 형 사고방식으로, <think 2>를 입력하면 B 형 사고방식으로 문제를 풉니다. 마치 "이 지팡이를 흔들면 물이 나오고, 저 지팡이를 흔들면 불이 나오는" 것처럼, 각각의 지팡이가 특정한 사고 경로를 정확히 불러일으키는 것입니다.

3. 훈련 방법: "SSFT (세트 감독 미세 조정)"

그렇다면 AI 가 어떻게 이 지팡이들과 사고방식을 연결하는 법을 배울까요? 바로 **'SSFT (Set Supervised Fine-Tuning)'**라는 훈련 방식을 썼습니다.

• 상황: 수학 문제 하나에 대해, 4 가지의 서로 다른 정답 풀이 과정 (A, B, C, D) 이 있다고 가정해 봅시다. 그리고 6 개의 마법 지팡이 (<think 1>~`<think 6>`) 가 있습니다.
• 기존 훈련 (SFT): 그냥 A, B, C, D 풀이를 모두 섞어서 AI 에게 보여줍니다. 그런데 AI 는 이걸 다 섞어서 "그냥 중간쯤 되는 평균적인 생각"만 하게 됩니다. (이걸 **'모드 붕괴 (Mode Collapse)'**라고 합니다. 다양한 생각이 하나로 뭉개지는 현상입니다.)
• 이 연구의 훈련 (SSFT):
  1. AI 가 <think 1>을 봤을 때 A 풀이를 잘 풀고, <think 2>를 봤을 때 B 풀이를 잘 풀도록 **최적의 짝짓기 (Bipartite Matching)**를 찾아냅니다. (헝가리안 알고리즘 사용)
  2. "아, <think 1>은 A 와 짝이 맞고, <think 2>는 B 와 짝이 맞구나!"라고 AI 가 스스로 깨닫게 합니다.
  3. 이렇게 각 지팡이가 고유한 사고방식과 딱 맞게 연결되도록 훈련시킵니다.

4. 결과: "생각의 다채로움이 살아났다"

이 훈련을 받은 모델은 놀라운 능력을 보여줍니다.

• 다양한 사고: <think 1>을 주면 기하학적으로 문제를 풀고, <think 2>를 주면 대수학적으로 풉니다. 각기 다른 전문가가 문제를 푸는 것처럼 서로 다른 전략을 사용합니다.
• 정확도 향상: 단순히 무작위로 생각하게 하는 게 아니라, 올바른 사고방식을 선택할 수 있는 지팡이를 가지고 있으므로, 정답을 맞힐 확률 (Pass@1) 이 크게 올라갑니다.
• GFPO (글로벌 포크링 정책 최적화): 훈련이 끝난 후, 어떤 지팡이를 쓸지 AI 가 스스로 가장 좋은 것을 고르도록 추가 훈련을 시키면, 성능이 더 좋아집니다.

5. 요약: 왜 이 연구가 중요한가요?

이 논문은 **"AI 에게 다양한 생각을 하라고 강요하는 게 아니라, AI 가 다양한 생각을 할 수 있는 '도구 (지팡이)'를 만들어주고, 그 도구들이 각각의 역할을 잘 수행하도록 훈련시켰다"**는 점이 핵심입니다.

• 과거: "생각해 봐!" (AI: "어... 뭐라고 생각하지?") -> 실수하거나 너무 길게 생각함.
• 이제: "이 지팡이 (<think 1>) 를 써봐!" (AI: "알겠습니다, 저는 기하학적으로 접근하겠습니다!") -> 정확하고 다양한 해결책 도출.

이 방법은 수학 문제뿐만 아니라 코딩, 논리 추론 등 복잡한 문제를 풀 때 AI 가 더 똑똑하고 유연하게 작동하게 만들어, 앞으로의 AI 에이전트 개발에 큰 기여를 할 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 은 테스트 시간 (test-time) 에 더 많은 계산 자원을 할당하여 추론 경로를 생성함으로써 성능을 향상시킬 수 있습니다. 그러나 기존 방식에는 두 가지 주요 한계가 존재합니다.

• 과도한 추론 (Overthinking): 순차적으로 토큰을 늘리는 방식은 일정 길이 이상에서 성능이 저하되는 '과도한 생각' 현상을 유발할 수 있습니다.
• 다양성과 정확성의 트레이드오프: 병렬 샘플링 (Parallel Sampling) 을 통해 다양한 추론 경로를 생성하려 할 때, 일반적으로 '온도 (Temperature)' 스케일링을 사용합니다. 하지만 온도를 높이면 다양성은 증가하지만 정확도가 떨어지는 문제가 발생합니다. 특히, 복잡한 문제에서 올바른 추론 모드를 시작하는 '포크 (forking) 토큰'은 샘플링 트리의 깊은 곳에 위치하여 무작위 샘플링으로 찾기 어렵습니다.
• 모드 붕괴 (Mode Collapse): 기존의 감독 미세 조정 (SFT) 이 여러 개의 다양한 추론 궤적 (traces) 을 학습할 때, 모델이 이러한 다양한 모드를 하나의 단일 모드로 압축 (collapse) 하여 다양성을 잃는 문제가 발생합니다.

2. 제안 방법론 (Methodology)

이 논문은 병렬 추론을 '다음 토큰 예측의 집합 (Set of Next-Token Prediction)' 문제로 재정의하고, 이를 해결하기 위해 집합 기반 전역 손실 (Set-based Global Loss) 을 도입한 Set Supervised Fine-Tuning (SSFT) 프레임워크를 제안합니다.

핵심 구성 요소:

1. 글로벌 포크 토큰 (Global Forking Tokens):

   • <think 1>, <think 2> 와 같은 예약된 특수 토큰 집합을 입력 프롬프트에 포함시킵니다.
   • 각 포크 토큰은 서로 다른 고유한 추론 모드 (reasoning mode) 를 시작하는 역할을 합니다.

2. 이분 매칭 (Bipartite Matching) 기반 손실 함수:

   • 입력 질문 $x$에 대해 $M$개의 정답 추론 궤적 (Ground-truth traces) 과 $N$개의 글로벌 포크 토큰 ($N \ge M$) 이 주어집니다.
   • 비용 행렬 (Cost Matrix) 생성: 각 포크 토큰 $g(i)$와 각 궤적 $r(j)$의 조합에 대해 다음 토큰 예측 (NTP) 손실을 계산합니다.
   • 최적 매칭 (Optimal Matching): 할당 문제 (Assignment Problem) 를 해결하여 총 손실을 최소화하는 이분 매칭 $\hat{\sigma}$를 찾습니다 (헝가리안 알고리즘 사용).
   • 손실 계산: 최적의 매칭에 따라 각 궤적이 해당 포크 토큰에 조건부로 생성되도록 손실을 역전파합니다.
   • 효과: 이 방식은 모델이 특정 포크 토큰이 특정 추론 궤적과 매핑되도록 학습하게 하여, 다양한 모드가 붕괴되지 않고 유지되도록 합니다.

3. 글로벌 포크 정책 최적화 (GFPO, Global Forking Policy Optimization):

   • SSFT 로 학습된 모델을 기반으로, 특정 질문에 대해 가장 적합한 포크 토큰을 선택하는 정책을 강화학습 (RL) 단계에서 미세 조정합니다.
   • 전체 토큰이 아닌 포크 토큰 선택에 대한 정책 경량만 업데이트하여 효율성을 극대화합니다.

3. 주요 기여 (Key Contributions)

• SSFT 프레임워크 도입: 다양한 추론 궤적과 제어 토큰 간의 이분 매칭을 통해 집합 기반 손실을 SFT 에 통합했습니다. 이는 단일 전역 포크 토큰에서 다양한 추론 모드를 전역적으로 조종 (steer) 할 수 있게 합니다.
• 모드 붕괴 방지 및 새로운 토큰 생성: 기존 SFT 는 다양한 궤적을 학습해도 단일 모드로 수렴하는 반면, SSFT 는 학습된 글로벌 포크 토큰들이 서로 다른 추론 전략을 명확하게 시작하도록 합니다.
• 성능 향상: 수학 추론 및 코드 생성 벤치마크에서 SSFT 기반 모델은 기존 SFT 모델과 GRPO(Generalized Reinforcement Policy Optimization) 를 적용한 모델보다 Pass@1(단일 경로 정확도) 및 Cons@k(병렬 생성 시 다수결 투표 정확도) 에서 일관되게 우수한 성능을 보입니다.

4. 실험 결과 (Results)

실험은 Qwen2.5-32B 모델을 기반으로 AIME, MATH-500, GPQA-Diamond, LiveCodeBench 등 다양한 벤치마크에서 수행되었습니다.

• Pass@1 및 Cons@k 향상:
  • SSFT-32B 는 AIME 2024 에서 64.06%, AIME 2025 에서 58.13% 의 Pass@1 정확도를 기록하여 기존 SFT 기반 모델 (SFT-mixed-distill-32B) 보다 각각 8.33%, 6.57% 향상되었습니다.
  • 병렬 테스트 시간 계산 (Cons@6, Cons@32) 에서도 일관된 개선을 보였습니다.
• 다양성 검증:
  • 서로 다른 포크 토큰 (<think 1> ~ <think 6>) 을 입력했을 때, 모델은 서로 다른 길이와 전략을 가진 추론 경로를 생성했습니다 (그림 4 참조).
  • 반면, 무작위 매칭을 사용한 SSFT (random $\sigma$) 나 기존 SFT 는 포크 토큰에 따라 추론 패턴이 뚜렷하게 달라지지 않았습니다 (그림 5 참조).
• 범용성 및 강건성:
  • 수학 데이터로 학습된 모델이 코드 생성 (LiveCodeBench) 과 같은 다른 도메인에서도 일반화 성능이 우수함을 확인했습니다.
  • 다양한 모델 크기 (4B, 8B, 32B) 에서도 SSFT 의 효과가 입증되었습니다.

5. 의의 및 결론 (Significance)

이 연구는 LLM 의 추론 능력을 향상시키기 위해 단순히 토큰 길이를 늘리는 순차적 확장 (Sequential Scaling) 을 넘어, 다양한 추론 경로를 병렬로 생성하고 제어할 수 있는 새로운 패러다임을 제시합니다.

• 학습된 제어 메커니즘: 모델이 온도에 의존하지 않고, 학습된 '글로벌 포크 토큰'을 통해 의도적으로 다양한 추론 모드를 활성화할 수 있게 되었습니다.
• 효율성: 복잡한 검색 알고리즘 (MCTS 등) 없이도, 학습된 토큰을 통해 효율적으로 다양한 해결책을 탐색할 수 있습니다.
• 미래 전망: 이 방법은 에이전트 AI 및 복잡한 문제 해결을 위한 LLM 의 테스트 시간 계산 (Test-time Compute) 전략을 혁신할 수 있는 기반을 마련했습니다.

요약하자면, 이 논문은 SSFT를 통해 LLM 이 다양한 추론 모드를 붕괴 없이 학습하고, 이를 글로벌 포크 토큰을 통해 제어할 수 있게 함으로써, 병렬 추론의 정확성과 다양성을 동시에 극대화하는 방법을 증명했습니다.