기술 요약: 수학적 추론 확장을 위한 TrajFusion

문제 정의

대규모 언어 모델(LLM)은 고품질의 솔루션 궤적(trajectory)에 대한 지도 미세 조정(Supervised Fine-Tuning, SFT)을 통해 수학적 추론 분야에서 상당한 진전을 이루었습니다. 이 데이터 구축의 지배적인 패러다임은 **거절 샘플링 미세 조정(Rejection Sampling Fine-Tuning, RFT)**입니다. RFT에서 교사 모델은 주어진 문제에 대해 여러 개의 후보 추론 경로를 생성하며, 최종 정답이 옳은 것들만 훈련을 위해 유지하고, 틀린 모든 궤적은 폐기합니다.

이 방식은 효과적이긴 하지만, 감독(supervision)을 이진 필터(binary filter)로 취급합니다. 이는 교사가 생성한 오류를 체계적으로 배제함으로써, 훈련 과정에서 추론 실패가 어떻게 모델링되는지에 대한 간극을 만듭니다. 저자들은 틀린 궤적들이 흔한 계산 실수, 누락된 가정, 또는 논리적으로는 그럴듯하지만 유효하지 않은 논증 경로와 같은 진단적 정보를 포함하고 있으며, 이러한 정보가 모델이 특정 실패 모드를 인식하고 피하는 데 도움이 될 수 있다고 주장합니다. 현재의 방법들은 이러한 오류를 활용 가능한 구조적 감독이 아닌, 폐기해야 할 노이즈로 취급합니다.

방법론: TrajFusion

본 논문은 RFT를 단순한 이진 필터링 메커니즘에서 구조화된 감독 구축 프로세스로 재정의하는 미세 조정 전략인 TrajFusion을 제안합니다. TrajFusion은 선택된 오답 추론 궤적을 정답 궤적과 통합하여 단일한 융합 훈련 샘플을 만듦으로써, 명시적으로 시행착오(trial-and-error) 추론 과정을 모델링합니다.

핵심 구성 요소

1. 샘플링 및 분석: 각 문제 $x$에 대해, 교사 모델은 $K$개의 궤적을 샘플링합니다. 이들은 자동 검증기를 사용하여 정답($Y^+$) 집합과 오답($Y^-$) 집합으로 분류됩니다.
2. 적응형 선택 (Adaptive Selection): 이 방법은 다음 두 가지 통계치를 기반으로 포함할 오답 궤적의 수를 적응적으로 결정합니다:
   • 오류율 ($r(x)$): 잘못된 시도의 빈도.
   • 오류 다양성 ($u(x)$): 오답 시도들 사이의 고유한 최종 정답 수 (Shannon 엔트로피 또는 고유 정답 수로 측정).
     보존될 오답 궤적의 수 $k(x)$는 다음과 같이 계산됩니다:
     $$k(x) = \min(k_{max}, \lfloor \alpha \cdot r(x) \cdot u(x) \rfloor)$$
     이를 통해 다양하고 빈번한 오류가 발생하는 문제에는 더 풍부한 감독을 제공하고, 희소하거나 정보 가치가 낮은 오류가 발생하는 문제는 자연스럽게 기존의 RFT 방식으로 회귀하도록 합니다.
3. 융합 샘플 구축: 선택된 오답 궤적들은 고정된 성찰 프롬프트(예: "잠깐... 이게 아닌 것 같다. 다시 해보자.")와 함께 정답 궤적 사이에 삽입됩니다. 최종 구조는 다음과 같습니다:
   $$T(x) = [y_{e1}, \rho_1, \dots, y_{ek}, \rho_k, y^*]$$
   여기서 $y_{ei}$는 오답 경로, $\rho_i$는 성찰 프롬프트, $y^*$는 정답 솔루션입니다.
4. 훈련: 융합된 시퀀스들은 모델 아키텍처나 훈련 목적 함수의 변경 없이 표준적인 다음 토큰 예측(next-token prediction, SFT) 방식으로 훈련됩니다.

주요 기여

• 감독의 간극 공식화: 저자들은 RFT에서 교사가 생성한 오답 궤적이 체계적으로 배제되는 현상을 식별하고 이를 공식화하였으며, 이러한 오류가 추론 실패를 모델링하기 위한 진단적 가치를 제공한다고 주장합니다.
• TrajFusion 프레임워크: 오답과 정답 궤적을 적응적으로 융합하는 구조화된 감독 프레임워크를 제안합니다. 이 방식은 오류 다양성을 명시적인 제어 신호로 취급하여, 보조 모델, 비판 생성기(critique generator) 또는 외부 난이도 레이블 없이도 다양한 난이도의 문제에 걸쳐 견고하게 확장될 수 있게 합니다.
• 실증적 검증: 이 방법은 여러 베이스 모델(LLaMA3-8B, DeepSeekMath-7B), 데이터 규모(15K ~ 100K 예시), 그리고 컨텍스트 길이(최대 32K)에 대해 기존 RFT보다 일관된 성능 향상을 입증했습니다.

실험 결과

저자들은 여섯 가지 수학적 추론 벤치마크(GSM8K, MATH, CollegeMath, DeepMind-Mathematics, OlympiadBench, TheoremQA)와 경시 대회 수준의 데이터셋(AIME24, AIME25, MATH-500)에서 TrajFusion을 평가했습니다.

• 성능 향상: TrajFusion은 기존 RFT보다 일관되게 우수한 성능을 보였습니다. LLaMA3-8B (100K 데이터)를 사용한 MATH 데이터셋에서 TrajFusion은 **51.4%**의 정확도를 달 기록하여, RFT의 **48.5%**를 상회했습니다. DeepSeekMath-7B의 경우, 52.3% 대 **50.2%**를 기록했습니다.
• 복잡성에 따른 영향: 성능 향상은 추론 경로가 길고 실패 모드가 다양한 도전적인 장문형 추론 벤치마크(MATH, OlympiadBench, TheoremQA)에서 가장 두드러졌습니다. GSM8K와 같은 비교적 단순한 데이터셋에서의 향상은 상대적으로 작았으나 여전히 양의 값을 가졌습니다.
• 데이터 효율성: TrajFusion은 15K개의 훈련 예시만으로도 훨씬 더 큰 데이터셋(예: 2.3M 예시를 사용한 MMIQC)을 사용하는 베이스라인보다 경쟁력 있는 결과를 달성했습니다.
• 긴 컨텍스트로의 확장성: DeepSeek-R1에서 증류된 장문형 데이터를 사용하여 Qwen2.5-Math-7B의 32K 컨텍스트 버전에 적용했을 때, TrajFusion은 AIME24의 Pass@1을 RFT의 40.0%에서 **43.3%**로 개선했습니다.
• 절제 연구 (Ablation Studies):
  • 융합되는 오류의 수($k_{max}$)를 늘리면 성능이 일정 지점까지 향상되었으나, 한계 효용이 관찰되었습니다 ($k_{max}=6$에서).
  • **다양성(Diversity)**을 기준으로 오류를 선택하는 것이(HighDiv) 단순히 빈도만을 기준으로 하는 것(LowDiv)보다 더 나은 결과를 냈으며, 이는 구조적으로 구별되는 실패 모드가 더 유익하다는 것을 시사합니다.
  • TrajFusion으로 훈련된 모델은 초기에 더 긴 시퀀스를 생성했으나, 점차 더 간결해지면서 성찰 토큰의 빈도가 높아졌습니다. 이는 구조적 패턴을 더 효율적으로 사용하고 있음을 나타냅니다.

의의 및 주장

본 논문은 TrajFusion이 일반적으로 버려지는 감독 신호를 활용함으로써 수학적 추론을 확장하는 단순하면서도 효과적인 방법을 제공한다고 주장합니다. 거절 샘플링을 구조화된 시행착오 궤적을 구축하는 과정으로 재정의함으로써, 이 방법은 복잡한 강화 학습 파이프라인이나 아키텍처 수정 없이도 더 풍부한 감독을 제공합니다.

저자들은 TrajFusion이 특히 도전적이고 긴 형태의 추론 문제에서 매우 가치 있다고 강조하며, 여기서 유효한 추론과 결함이 있는 추론을 구분하는 것이 필수적이기 때문입니다. 또한, 이 방법의 적응적 특성이 견고함을 보장한다고 주장합니다. 즉, 오류가 유익할 때는 감독을 풍부하게 하고, 그렇지 않을 때는 안전하게 표준 RFT 동작으로 회귀합니다. 이 접근 방식은 이진적인 정답 필터링과 추론 실패의 미묘한 현실 사이의 간극을 메워, 모델이 지도 학습 환경 내에서 자신의 실수를 통해 학습할 수 있도록 합니다.