🤖 AI

Trajectory-Refined Distillation

이 논문은 교사의 가이드 하에 학생의 롤아웃(rollout)에 궤적 수준의 수정을 적용함으로써 구조적 "접두사 실패(prefix failure)" 문제를 완화하고, 다양한 벤치마크와 모델 규모 전반에서 추론 정확도와 커버리지를 향상시키는 새로운 온-폴리시(on-policy) 증류 방법인 궤적 정밀 증류(Trajectory-Refined Distillation, TRD)를 소개한다.

원저자: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

게시일 2026-06-09

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 학생(AI)에게 복잡한 수학 문제를 푸는 법을 가르치고 있다고 상상해 보세요. 당신에게는 정답을 알고 있는 아주 똑똑한 선생님(더 발전된 AI)이 있습니다.

기존 논문에서 설명하는 표준 방식인 **온-폴리시 증류(On-Policy Distillation, OPD)**의 과정은 다음과 같습니다:

학생은 스스로 문제를 풀려고 시도합니다.
학생이 단계를 적어 내려갈 때, 선생님은 지켜보며 이렇게 말합니다. "방금 네가 쓴 이 특정 단어 다음에 올 확률은 이 정도란다."
학생은 모든 단어에 대해 선생님의 확률과 일치하도록 자신의 뇌를 조정하려고 노력합니다.

문제점: "접두사 실패(Prefix Failure)" (잘못된 길로 들기)

이 논문은 이 표준 방식의 주요 결함인 **접두사 실패(Prefix Failure)**를 지적합니다.

학생이 문제를 풀기 위해 길을 따라 걷고 있다고 상상해 보세요. 3단계에서 학생이 잘못된 길로 접어들었습니다. 이제 막다른 골목에 다다랐습니다.

선생님의 딜레마: 선생님은 학생이 길을 잃었다는 것을 봅니다. 선생님은 올바른 경로가 "아, 사실은..."으로 시작하여 완전히 다른 방향으로 진행되어야 한다는 것을 알고 있습니다.
혼란: 하지만 선생님은 학생의 현재 위치(막다른 골목)를 반드시 고려해야만 합니다. 그래서 선생님의 조언은 혼란스러운 메시지가 됩니다. "이미 여기 와 있으니 이 막다른 길에 머물러라"라는 말과 "즉시 뒤돌아 나가라"라는 말이 뒤섞이게 됩니다.
결과: 학생은 혼란스러운 신호를 받게 됩니다. 학생은 선생님으로부터 배우려고 노력하지만, 학생이 이미 잘못된 경로에 빠져 있기 때문에 선생님의 조언은 파편화됩니다. 학생은 올바른 경로로 돌아오는 법을 배우는 대신, "길을 잃은 상태에서 잘 지내는 법"을 배우게 됩니다. 논문에서는 이를 "양봉 혼합(bimodal mixture)"이라고 부릅니다. 즉, 선생님이 동시에 두 가지 상반된 것을 가르치려 함으로써 학생을 혼란에 빠뜨리는 것입니다.

기존의 해결책들은 특정 단어에 대한 "성적"(손실 함수)을 미세하게 조정하려고 시도했지만, 학생이 여전히 잘못된 경로에 머물러 있다는 근본적인 문제는 해결하지 못했습니다.

해결책: 궤적 정제 증류 (Trajectory-Refined Distillation, TRD)

저자들은 새로운 방식인 **궤적 정제 증류(TRD)**를 제안합니다. 학생의 엉성한 초안을 단어 단위로 채점하는 대신, 작업 흐름을 바꿉니다:

초안 작성: 학생은 여전히 문제에 대해 첫 번째 시도(원시 롤아웃)를 합니다.
재작성: 채점하기 전에, 선생님은 그 엉성한 초안을 가져와서 다시 씁니다. 선생님은 잘못된 길을 바로잡고, 논리를 교정하며, 학생의 시작점을 바탕으로 깔끔하고 완벽한 버전의 솔루션을 만들어냅니다.
학습: 학생은 엉성한 원본이 아니라, 이 깔끔하게 교정된 버전으로부터 배웁니다.

비유하자면:

기존 방식: 당신이 이야기 중간에 설정 오류(plot hole)가 있는 글을 썼습니다. 편집자는 당신의 글을 읽고 "3문단의 '고양이'라는 단어 대신 '강아지'라고 썼어야 했어요"라고 말합니다. 하지만 당신는 여전히 왜 '고양이'라고 썼는지 고민하며 3문단에 갇혀 있습니다. 매우 혼란스럽죠.
TRD 방식: 당신이 설정 오류가 있는 이야기를 씁니다. 편집자는 당신의 초안을 가져와서 설정 오류를 고치고, 이야기가 완벽하게 흐르도록 중간 부분을 통째로 다시 씁니다. 그런 다음, 당신은 다음에 더 좋은 이야기를 쓰는 법을 배우기 위해 편집자가 재작성한 버전을 공부합니다.

이것이 왜 더 효과적인가

근본 원인 해결: 학생이 배우기 전에 "잘못된 길"을 미리 바로잡음으로써, 선생님의 조언이 더 이상 혼란스럽지 않습니다. 명확하고 단일한 경로가 됩니다.
새로운 경로 탐색: 학생이 첫 번째 시도에서 정답을 맞혔더라도, 선생님은 더 짧거나 더 스마트한 해결 방식을 보여줄 수 있습니다. 이는 학생이 더 유연해지도록 가르칩니다.
효율성: 연구진은 "재작성된" 경로가 학생의 원래 장황한 시도보다 훨씬 짧고 명확한 경우가 많았으며, 이것이 실제로 훈련 속도를 높였다는 것을 발견했습니다.

결과

연구진은 이 방법을 어려운 수학 경시 대회(AIME, HMMT 등)와 코딩 챌린지에 테스트했습니다. 그 결과는 다음과 같습니다:

새로운 방식(TRD)이 기존 방식들을 지속적으로 압도했습니다.
특히 AI가 보통 막히게 되는 가장 어려운 문제들을 해결하는 데 탁월한 효과를 보였습니다.
선생님이 별도의 더 큰 AI이든, 혹은 "특권 정보(privileged information)" 기술을 사용하여 동일한 AI가 스스로의 똑똑한 스승 역할을 하든 상관없이 잘 작동했습니다.

요약하자면, 이 논문은 AI를 효과적으로 가르치기 위해서는 실수를 저지르는 도중에 단어 하나하나를 채점할 것이 아니라, 선생님이 먼저 전체 경로를 바로잡게 한 뒤 학생이 그 교정된 여정을 통해 배우게 해야 한다고 주장합니다.

기술 요약: 궤적 정밀 증류 (Trajectory-Refined Distillation, TRD)

1. 문제 정의: 온-폴리시 증류(On-Policy Distillation)에서의 접두사 실패(Prefix Failure)

본 논문은 대규모 언어 모델(LLM)을 위한 온-폴리시 증류(OPD) 및 온-폴리시 자기 증류(OPSD)에서 발생하는 근본적인 구조적 한계인 **접두사 실패(prefix failure)**를 식별한다.

표준적인 OPD/OPSD에서 학생 모델(student model)은 롤아웃(trajectory) $y_o$ 를 생성하고, 교사 모델(teacher model)은 해당 특정 롤아웃에 기반하여 밀집된 토큰별 감독 신호를 제공한다. 저자들은 학생의 롤아웃이 잘못된 추론 경로(즉, "실패한 접두사")를 따를 경우, 교사의 감독 신호가 구조적으로 손상된다고 주장한다.

양봉적 교사 혼합(Bimodal Teacher Mixture): 접두사가 잘못되었을 때, 교사의 분포는 두 가지 모드의 혼합이 된다. 하나는 시퀀스 일관성을 위해 실패한 접우사를 계속 이어가는 것이고, 다른 하나는 올바른 솔루션으로 방향을 전환하는 것이다.
파편화된 그래디언트(Fragmented Gradients): 표준적인 토큰 레벨 손실 함수(KL 발산)는 고정된 실패 궤적을 따라 감독을 수행한다. 결과적으로 그래디언트 업데이트는 파편화된다. 교사는 올바른 수정 경로를 따라 조건화되는 대신, 잘못된 컨텍스트(실패한 접두사)에 조건화되어 반복적으로 수정 토큰(예: "Wait" 또는 "Actually")을 권장하게 된다.
토큰 레벨 수정의 무효성: 기존의 해결책들(예: 손실 클리핑, top-K 절단, 가중치 재설정)은 토큰 손실 수준에서 작동한다. 이들은 "정보가 있는" 토큰을 선택하거나 "노이즈가 있는" 토큰을 억제하려고 시도하지만, 근본 원인인 실패한 접두사 자체가 구조적으로 온전히 남아 있다는 점을 해결하지 못한다. 전달되는 감독 쌍은 이상적인 수정 경로와 괴리되며, 모델이 오류로부터 효과적으로 회복하는 법을 배우는 것을 방해한다.

2. 방법론: 궤적 정밀 증류 (Trajectory-Refined Distillation, TRD)

접두사 실패를 완화하기 위해 저자들은 개입의 단위를 토큰 레벨에서 궤적 레벨로 전환하는 **궤적 정밀 증류(TRD)**를 제안한다.

핵심 메커니즘

TRD는 원본 학생 롤아웃 $y_o$ 를 직접 증류에 사용하는 대신, 손실 계산 전에 정밀화(refinement) 단계를 도입한다.

샘플링: 학생 모델 $\pi_\theta$ 가 원본 온-폴리시 롤아웃 $y_o \sim \pi_\theta(\cdot | x)$ 를 생성한다.
정밀화: 교사 모델 $\pi_T$ $π_{T}$ (또는 정답 $y^*$ $y^{*}$ 와 같은 특권 정보에 조건화된 학생 자신인 OPSD의 경우)가 정밀화된 궤적 $y_r$ $y_{r}$ 을 생성한다.
- OPD의 경우: $y_r \sim \pi_T(\cdot | x, y_o)$ .
- OPSD의 경우: $y_r \sim \text{sg}[\pi_\theta(\cdot | x, y_o, y^*)]$ .
- 결정적으로, 교사는 참조 솔루션(reference solution)에 의해 유도되어 학생의 롤아웃을 수정하도록 프롬프트되며, 이는 학생의 추론 패턴이 유효한 곳은 유지하면서 잘못된 부분의 접두사를 효과적으로 "다시 쓰는" 과정이다.
증류: 정밀화된 궤적 $y_r$ 은 표준 분포 매칭 손실(전체 어휘를 포함한 Forward KL)의 감독 타겟으로 사용된다.

이론적 근거

TRD는 감독 컨텍스트가 실패한 접두사가 아닌 수정 경로를 따라 성장하도록 보장함으로써 "파편화된 그래디언트" 문제를 해결한다.

이상적인 그래디언트 회복: $y_r$ 을 따라 감독함으로써 각 토큰의 컨텍스트에 수정된 접두사가 포함되게 하며, 이를 통해 교사가 반복적이고 컨텍스트가 불일치하는 신호( $g_{frag}$ ) 대신 일관되고 다단계적인 수정 신호( $g_{ideal}$ )를 제공할 수 있게 한다.
탐색(Exploration): $y_o$ 가 올바른 경우에도, TRD는 교사가 제안하는 대안적인 유효 유도 과정을 학생에게 노출시켜 학생이 학습하는 올로 바른 추론 경로의 분포를 넓힌다.

알고리즘

과정은 알고리즘 1로 요약된다:

$y_o \sim \pi_\theta(\cdot | x)$ 를 샘플링한다.
$y_r \sim \pi_T(\cdot | x, y_o)$ 를 생성한다.
$\pi_\theta$ 와 $y_r$ 에 대한 교사의 분포 사이의 KL 발산을 사용하여 $\theta$ 를 업데이트한다.

3. 주요 기여

접두사 실패의 식별: 본 논문은 "접두사 실패"를 온-폴리시 증류의 성능 저하를 일으키는 구조적 원인으로 공식화하며, 왜 토큰 레벨의 개입이 추론 오류로부터의 회복에 실패하는지 설명한다.
궤적 레벨의 수정: TRD는 과도한 계산 비용 없이 온-폴리시 서포트 제약 조건을 준수하면서 실패한 접두사의 회복을 명시적으로 목표로 하는 최초의 궤적 레벨 최적화 설계를 제안한다.
이중 체제에서의 개선: TRD는 두 가지 체제 모두에서 성능을 향상시킨다:
- 오류 수정: 증류 전 문제가 되는 접두사를 수정함으로써 접두사 실패를 완화한다.
- 탐색: 원본 롤아웃이 올바른 경우에도 대안적인 유효 유도 과정을 노출함으로써 추론 커버리지를 확장한다.
자기 증류에 대한 적용 가능성: 이 방법은 OPSD로 자연스럽게 확장되어, 별도의 더 큰 교사 모델 없이도 특권 정보(ground truth)에 조건화된 학생이 스스로의 추론 경로를 효과적으로 정밀화할 수 있음을 보여준다.

4. 실험 결과

저자들은 Qwen3 모델(1.7B, 4B, 8B)을 사용하여 5개의 경시 수학 벤치마크(AIME24/25, HMMT25, BeyondAIME, AMOBench)와 코드 생성 벤치마크(HumanEval+, MBPP+, LiveCodeBench)에서 TRD를 평가했다.

주요 결과

우수한 성능: TRD는 OPD 및 OPSD 설정 모두에서 기존 베이스라인(클리핑/절단이 포함되거나 포함되지 않은 Forward/Reverse KL)을 지속적으로 능가한다.
어려운 과제에서의 상당한 이득: 개선 효과는 가장 어려운 벤치마크에서 가장 두드러졌다.
- AMOBench: OPSD 설정에서 Qwen3-8B의 경우 TRD는 베이스 모델 대비 Pass@16에서 약 50%의 상대적 개선을 달성했다.
- AIME24/25: TRD는 특히 작은 모델(예: Qwen3-1.7B의 AIME24에서 +4.6%)에서 일관된 이득을 보였다.
안정성: 특정 벤치마크에서 성능이 하락하는 일부 베이스라인(예: AIME24에서 성능이 떨어지는 Reverse KL)과 달리, TRD는 베이스 모델의 성능 미만으로 떨어진 적이 없다.
효율성: OPSD 설정에서 정밀화된 궤적 $y_r$ 은 원본 롤아웃 $y_o$ 보다 현저히 짧았으며(중앙값 길이가 약 9배 압축됨), 추가적인 샘플링 단계에도 불구하고 훈련 실제 시간(wall-clock time)을 약 60% 단축했다.
코드 생성: TRD는 HumanEval+ 및 MBPP+에서 베이스라인과 대등하거나 이를 상회했으나, LiveCodeBench에서의 이득은 제한적이었으며, 이는 교사의 정밀화 능력이 과제의 난이도에 따라 제한될 수 있음을 시사한다.

5. 의의 및 주장

본 논문은 TRD가 LLM의 사후 학습(post-training) 전략에 있어 필수적인 변화를 나타낸다고 주장한다. 저자들의 주장은 다음과 같다:

토큰 레벨 개입은 불충분하다: 접두사 실패의 구조적 특성은 궤적 레벨의 수정을 요구한다.
온-폴리시 서포트가 보존된다: TRD는 고정된 전문가 정책(expert policy)을 사용하는 것과 달리, 온-폴리시 제약을 위반하지 않으면서 이러한 수정을 달성하며, 이는 계산 효율적인 학습 패러다임과 호환된다.
자기 개선(Self-Improvement)이 가능하다: OPSD 설정에서의 TRD 성공은 모델이 별도의 더 큰 교사 모델 없이도 특권 정보(ground truth)를 사용하여 자신의 추론 경로를 효과적으로 정밀화할 수 있음을 입증하며, 이는 경량화된 자기 개선 경로를 제공한다.

이 연구는 향후 증류 방법론이 개별 토큰의 통계적 정렬뿐만 아니라 추론 궤적의 구조적 무결성을 고려해야 함을 시사한다.