기술 요약: LARK (Learnability-Grounded Trajectory Selection)

문제 정의

본 논문은 더 큰 교사(Teacher) 모델이 생성한 추론 궤적(Reasoning Trajectories)을 모방하도록 더 작은 학생(Student) 거대 언어 모델(LLM)을 미세 조정하는 패러다임인 **추론 증류(Reasoning Distillation)**의 과제를 다룹니다. 기존의 데이터 선택 방법들은 주로 질문 수준에서 작동하거나 궤적의 품질, 정답 여부, 또는 모델의 확신도와 같은 휴리스틱에 의존하지만, 결정적인 요소인 **학습 가능성(Learnability)**을 간과하는 경우가 많습니다.

구체적으로, 단일 질문에 대해 여러 개의 후보 추론 궤적이 존재할 수 있습니다(서로 다른 교사 모델, 샘플링 실행 또는 스타일 등에 의해 생성됨). 이러한 궤적들이 동일한 정답에 도달하더라도, 학생 모델에 제공하는 학습 신호는 매우 다를 수 있습니다. 어떤 궤적이 "고품질"이거나 "자연스럽다"고 해서 반드시 학생 모델이 가장 효율적으로 학습할 수 있는 것은 아닙니다. 본 논문이 제기하는 핵심 질문은 다음과 같습니다: 우리는 증류된 추론 궤적이 학생 모델에 의해 가장 잘 학습될 수 있도록 보장하는 원칙적인 궤적 선택 기준을 설계할 수 있는가?

방법론: LARK

저자들은 정적인 품질 지표 대신 최적화 이론에 기반한 학습 가능성 정의를 바탕으로 궤적을 선택하는 방법론인 LARK(Learnability-grounded Anchor-time Ranking)를 제안합니다.

1. 최적화 목적 함수로서의 학습 가능성

LARK는 궤적 선택을 정책 최적화(Policy Optimization) 문제로 프레임화합니다. 이 방법은 학생의 훈련 손실(Training Loss)이 감소하는 속도를 특징짓는 학습 가능성 계수 $\rho(\theta; q)$를 정의합니다. 이 계수는 다음과 같이 정의됩니다:
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
여기서 $q$는 $K$개의 후보 궤적에 대한 가중치 벡터입니다. Polyak–Łojasiewicz (PL) 조건 하에서, 높은 $\rho$는 더 빠른 수렴을 의미합니다.

2. 앵커 타임 근사 (Anchor-Time Approximation)

전체 미세 조정 과정 동안 시간 통합(Time-integrated) 학습 가능성 목적 함수를 직접 최적화하는 것은 다루기 어렵습니다(Intractable). LARK는 **앵커 상대 조건(Anchor-relative condition)**을 도입합니다. 이는 과잉 매개변수화된(Over-parameterized) 모델이 "레이지 트레이닝(Lazy training)" 영역에 있을 때, 학습 가능성 비율이 초기 앵커 파라미터 $\theta_{ref}$에서의 값에 비례한다는 가설입니다. 이를 통해 목적 함수를 $\rho(\theta_{ref}; q)$를 최대화하는 것으로 근사할 수 있습니다.

3. 국소 선형화 및 순방향 패스 대리물 (Local Linearization and Forward-Pass Proxy)

$\rho(\theta_{ref}; q)$를 직접 최대화하는 것은 $\rho$가 단일 궤적만을 선택하는 퇴화된 "원-핫(One-hot)" 해로 수렴하는 준 볼록(Quasi-convex) 함수이기 때문에 문제가 됩니다. 이를 방지하기 위해 LARK는 다음을 수행합니다:

• 국소 선형화(Local Linearization): 선택이 없는 표준 SFT 상태인 균등 분포 $p$ 주변에서 1차 테일러 전개를 사용하여 $\rho$를 확장합니다.
• 순방향 패스 대리물 ($\hat{g}_k$): 모든 후보에 대해 역전파(Backward pass)를 수행하여 정확한 그래디언트를 계산하는 대신, LARK는 순방향 패스(Forward pass)만을 사용하여 계산 가능한 대리물 $\hat{g}_k$를 도출합니다. 이 대리물은 학생의 토큰 단위 예측 잔차(Residual)와 크로스 엔트로피 손실을 기반으로 학습 가능성 비율의 국소 그래디언트를 추정합니다.
  • 구체적으로, $\hat{g}_k$는 학생이 구조적으로 "확신을 가지고 틀린(Confidently wrong)" 경우를 선호합니다. 즉, 손실 $\ell_k$는 높지만(예측이 좋지 않음을 나타냄), 잔차 질량(Residual mass)이 소수의 "근접 오답(Near-miss)" 토큰들에 집중되어 있어 명확한 교정 신호를 제공하는 경우입니다.

4. $\chi^2$-정규화된 선택 정책

선택기가 원래의 데이터 분포에서 너무 멀어져 일반화 성능을 해치는 방향으로 학습 가능성 기준을 "해킹"하는 것을 방지하기 위해, LARK는 $\chi^2$-정규화( $\chi^2$-regularization) 항을 부과합니다. 이는 다음과 같은 폐쇄형(Closed-form)의 예산 매개변수화된 선택 규칙을 도출합니다:

• 궤적들은 점수 $\hat{g}_k$에 따라 순위가 매겨집니다.
• 상위 $B$개의 궤적이 선택됩니다.
• 단순한 상위 $B$개 절단 방식과 달리, L_ARK는 $(B+1)$ 번째 궤적과의 점수 차이에 비례하는 **소프트 가중치(Soft weights)**를 할당합니다. 이 가중치 규칙은 균등 사전 확률(Uniform prior)으로부터의 $\chi^2$ 발산을 제어하면서 학습 가능성 개선의 하한을 최대화하도록 유도되었습니다.

주요 기여

1. 정식화: 저자들은 추론 궤적 선택을 학습 가능성에 기반한 정책 최적화 문제로 정식화하였으며, 앵커 타임 학습 가능성 비율 $\rho$를 원칙적인 목적 함수로 도입하였습니다.
2. 알고리즘: 순방향 패스 대리물 $\hat{g}_k$와 $\chi^2$-정규화를 사용하는 폐쇄형의 예산 매개변수화된 선택 규칙을 도출하였습니다. 이는 매 궤적마다 비용이 많이 드는 역전파를 피하면서도, 추정 오차와 일반화에 대한 명시적인 경계(Bound)를 이론적으로 보장합니다.
3. 실험적 검증: LARK가 다양한 베이스 모델(Qwen-2.5-7B/1.5B, Llama-3.2-3B)과 추론 벤치마크(AIME, AMC, GPQA, MATH-500) 전반에서 기존 베이스라인(휴리스틱, 품질 기반, 학생 측 스코어링 방식 포함)을 일관되게 능가함을 입증하였습니다.

실험 결과

• 성능: LARK는 테스트된 모든 모델-예산 설정에서 가장 높은 평균 Acc@5를 달 기록했습니다. 예를 들어, $B=3$인 예산 설정에서 Qwen-2.5-7B를 대상으로 했을 때, LARK는 가장 강력한 베이스라인인 RSR보다 성능을 7.07 퍼센트 포인트 향상시켰습니다.
• 효율성: LARK는 순방향 패스만을 사용하는 방식으로, 가장 비용이 적은 학생 측 베이스라인들과 유사한 계산 비용을 가지며, 이는 정확한 그래디언트 기반 방식이 훨씬 더 많은 계산을 요구하는 것과 대조적입니다.
• 진단 분석:
  • 손실 감소(Loss Decay): LARK에 의해 선택된 궤적들은 훈련 중 SFT 손실을 가장 빠르게 감소시켰습니다.
  • 예측 타당성(Predictive Validity): 교사 모델별로 집계했을 때, LARK 점수 $\hat{g}_k$는 다운스트림 정확도와 강한 양의 상관관계(Pearson $r=0.69$)를 보였습니다.
  • 토큰 수준 동작: 선택된 궤적들은 정답 토큰의 확률은 낮지만 "오답 토큰의 집중도"는 높게 나타났습니다. 이는 LARK가 학생이 틀렸지만 그 오류가 구조적이고 교정 가능한 사례를 선택함을 확인시켜 줍니다.

의의 및 주장

본 논문은 **학습 가능성(Learnability)**이 궤적의 품질이나 난이도와는 별개인, 효율적인 추론 증류를 위한 독특하고 효과적인 원리라고 주장합니다. LARK는 선택 과정을 학생 모델의 최적화 역학에 결합함으로써, 가장 효율적인 학습 신호를 제공하는 궤적을 식별합니다.

저자들은 자사의 방법론이 다음과 같음을 강조합니다:

• 일반화 보존: $\chi^2$-정규화를 통해 선택 정책이 좁은 범위의 "쉬운" 또는 "해킹된" 궤적에 과적합되는 것을 방지합니다.
• 이론적 근거: 선택 규칙은 명시적인 오차 경계를 가진 최적화 이론(PL 조건, Gradient Flow)의 기본 원리로부터 도출되었습니다.
• 실용성: 선택 규칙 자체를 위한 하이퍼파라미터 튜닝이 필요 없으며(정규화 강도는 예산 $B$에 의해 결정됨), 역전파의 계산 오버헤드를 피할 수 있습니다.

이 연구는 향후 학습 가능성 기반의 관점을 더 노이즈가 많은 후보 풀과 더 넓은 추론 영역에 적용할 수 있음을 시사하지만, 현재의 범위는 수학적 문제 해결에서의 추론 증류에 엄격히 제한되어 있습니다.