💬 NLP

Pass the Baton: Trajectory-Relayed On-Policy Distillation

본 논문은 학생의 궤적 이탈을 교정하기 위해 교사의 개입을 동적으로 트리거함으로써 프리픽스 실패(prefix failure)를 완화하는 새로운 온-폴리시 증류 방법인 Relay-OPD를 소개하며, 이를 통해 수학적 추론 벤치마크에서 최첨단 성능을 달성하는 동시에 훈련 궤적 길이를 크게 단축한다.

원저자: Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen

게시일 2026-07-29

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Haolei Xu, Xiaowen Xu, Haiwen Hong, Zixuan Ni, Hongxing Li, Yiwen Qiu, Weiming Lu, Yongliang Shen

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

거대하고 똑똑한 컴퓨터(대규모 언어 모델, LLM이라 불리는)가 수학 문제나 논리 퍼즐 같은 복잡한 문제를 해결하는 법을 배우는 세상을 상상해 보십시오. 이들을 가르치기 위해 우리는 흔-히 '교사-학생(teacher-student)' 설정을 사용합니다. 교사는 정답을 알고 있는 거대하고 강력한 모델이며, 학생은 그로부터 배우려고 노력하는 더 작고 빠른 모델입니다. 보통 교사는 완벽한 해답을 그대로 써 내려가고, 학생은 그것을 복사합니다. 하지만 더 똑똑한 방법인 **온-폴리시 증류(On-Policy Distillation)**라는 것이 있습니다. 학생이 스스로 해답을 써 내려가려고 시도할 때, 교사가 단순히 완성된 에세이를 복사하게 하는 대신, 학생이 막히거나 잘못된 길로 빠질 때마다 그 즉시 다음 단어를 안내하며 도움을 주는 방식입니다. 이는 학생에게 단순히 무엇을 암기하는 것이 아니라 어떻게 생각하는지를 가르쳐주기 때문에 매우 효과적입니다.

하지만 이 방법에는 까다로운 결함이 있습니다. 만약 학생이 초기에 실수한다면—마치 미로에서 길을 잘못 든 것처럼—교사는 남은 여정을 안내하려고 노력할 것입니다. 하지만 학생은 이미 길을 잃었기 때문에, 교사의 안내는 종종 혼란스러워지거나 신뢰할 수 없게 되며, 학생은 경로를 벗어나 더욱 방황하게 됩니다. 이는 누군가가 이미 잘못된 방향으로 운전하고 있을 때 그 사람에게 길을 알려주려는 것과 같습니다. 그들이 더 멀리 운전할수록, 상황은 더 혼란스러워집니다. 이 논문은 바로 이 구체적인 문제, 즉 학생이 교사의 도움이 무용지물이 될 정도로 완전히 길을 잃기 전에 어떻게 막을 것인가를 다룹니다.

추론의 릴레이 경주

Qwen 제품군의 모델들을 활용하여 연구한 이 논문의 저자들은, 학생 모델이 궤도를 벗어나기 시작할 때 교사와 학생이 다음에 무엇을 할지에 대해 실제로 의견 차이를 보인다는 점을 깨달았습니다. 자신의 잘못된 생각에 고집스럽게 매달려 있는 학생은 계속 직진하고 싶어 합니다. 반면, 오류를 포착한 교사는 멈추고, 다시 생각하고, 방향을 바꾸고 싶어 합니다.

그들은 이 의견 차이가 발생하는 순간을 **"핸드오프 트리거(handoff trigger, 바통 터치 지점)"**라고 부릅니다. 이것을 릴레이 경주라고 생각해 보십시오. 일반적인 경주에서는 주자(학생)가 발이 걸려 잘못된 방향으로 뛰기 시작하면, 코치(교사)가 사이드라인에서 지시를 외칠 수는 있지만 주자는 계속 비틀거리며 나아갑니다. 이 새로운 방법인 **릴레이-OPD(Relay-OPD)**에서는, 코치가 주자가 절벽을 향해 달려가는 것을 보는 순간, 코치가 개입하여 바통을 잡고 올바른 방향으로 몇 걸음 달려가서 주자의 경로를 재설정합니다. 그런 다음 코치는 바통을 다시 넘겨주고, 학생은 이제 올바른 궤도에 올라 다시 경주를 계속합니다.

작동 원리: "하지만"과 "잠깐" 토큰

이 시스템은 인간이 언제 실수가 발생하는지 알려줄 필요가 없다는 점에서 매우 영리합니다. 시스템은 확률을 관찰합니다. 만약 교사가 "하지만(But)", "잠깐(Wait)", 또는 **"그러나(However)"**와 같이 사고의 전환을 알리는 '성찰' 단어를 말할 확률이 높은데, 학생은 "그래서(So)"나 "이제(Now)"와 같이 현재의 생각을 이어가려는 단어를 말할 확률이 높다면, 시스템은 다음과 같이 인지합니다: 우리는 경로를 벗어났다!

바로 그 순간, 교사가 여정의 짧은 "구간(leg)"을 맡게 됩니다. 교사는 추론을 바로잡기 위해 짧은 문단을 작성한 후, 제어권을 다시 학생에게 넘깁니다. 그러면 학생은 수정된 경로를 바탕으로 문제를 마무리하며 학습하게 됩니다.

연구진은 이러한 개입이 길 필요가 없다는 것을 발견했습니다. 사실, 교사는 전체 단어의 단 **0.35%**만을 작성함으로써 엄청난 차이를 만들어낼 수 있습니다. 이는 마치 코치가 러너를 위해 전체 경주를 대신 뛰어주는 것이 아니라, 러너의 자세를 바로잡기 위해 단 한 번의 날카로운 넛지(nudge, 가벼운 자극)를 주는 것과 같습니다.

결과: 더 똑똑하고 더 빠르게

연구팀은 다양한 크기(0.6B 및 1.7B 파라미터)의 학생 모델과 4B 파라미터의 교사 모델을 사용하여 8가지 서로 다른 수학 벤치마크에서 이를 테스트했습니다. 결과는 인상적이었습니다:

더 높은 점수: 릴레이-OPD 학생들은 표준적인 방법들을 사용하는 모델들보다 유의미하게 높은 점수를 기록했습니다. 1.7B 파라미터 모델의 경우, 표준 방식 대비 평균 점수가 5.73% 향상되었으며, 그다음으로 우수한 경쟁 모델보다 1.49% 높은 성적을 거두었습니다.
낭비되는 시간 감소: 시스템이 학생이 막다른 길로 헤매는 것을 막아주기 때문에, 훈련 과정이 훨씬 효율적이었습니다. 평균 훈련 경로의 길이가 50% 이상 단축되었습니다. 학생들은 절반의 시간 만에 동일한 양을 학습했습니다.
"구간"의 중요성: 연구진은 교사의 짧은 교정(즉, "교사 구간")이 결정적이라는 것을 입증했습니다. 단순히 학생이 실수를 했을 때 (교사가 먼저 수정하지 않고) 중단하는 것만으로는 효과가 충분하지 않았습니다. 교사는 학생에게 어떻게 다시 궤도로 돌아올 수 있는지를 반드시 보여주어야 했습니다.

이 기술이 아닌 것

이 논문은 이 방법이 무엇이 아닌지를 명확히 밝히고 있습니다. 이는 단순히 답변이 길어지는 것을 중간에 끊는 것(다른 방법들이 시도했던 방식)이 아니며, 학생의 에세이 전체를 다시 쓰는 것도 아닙니다. 이것은 실시간의, "그 순간의" 수정 작업입니다. 또한 교사가 경주 전체를 도맡는 것도 아닙니다. 학생이 여전히 대부분의 달리기를 수행합니다. 교사는 학생이 영구적인 실수를 저지르기 직전, 그 결정적이고 짧은 순간에만 개입합니다.

요약하자면, 릴레이-OPD는 떨어지는 학생을 바닥에 닿기 직전에 잡아내어, 부드럽게 방향을 재설정해주고, 다시 달릴 수 있게 해주는 안전망과 같습니다. 이는 잠재적인 재앙을 학습의 순간으로 바꾸어 놓으며, AI 모델을 더 똑똑하고, 더 빠르고, 자신의 생각 속에서 길을 잃을 가능성이 적게 만듭니다.

기술 요약: Relay-OPD (바톤 터치)

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

온폴리시 증류(OPD)는 강력한 교사(Teacher) 모델로부터 약한 학생(Student) 모델로 지식을 전이할 때, 학생이 생성한 궤적(Trajectory)에 토큰 단위의 감독을 결합함으로써 성능을 개선합니다. 그러나 이 방식은 **접두사 실패(Prefix Failure)**라고 불리는 치명적인 결함을 가집니다. 긴 추론 체인 작업에서 학생 모델이 생성 초기 단계에서 잘못된 추론 방향을 선택하면, 이후의 모든 토큰이 이 편차를 바탕으로 구축됩니다. 이는 다음과 같은 문제를 야기합니다:

교사로부터 신뢰할 수 없거나 잠재적으로 유해한 감독을 유도함.
수정이 불가능한 궤적에 상당한 학습 연산량을 낭비함.

기존의 해결책들은 구조적 한계를 가집니다:

고정 길이 절단 (예: FastOPD, ESR): 실제 추론 실패가 발생하는 지점과 상관없이 경직된 위치에서 롤아웃(Rollout)을 끊어버립니다.
오프라인 재작성 (예: TRD): 궤적이 완료된 후에야 경로를 수정하며, 종종 눈에 보이는 흔적(Artifacts)을 남깁니다.
토큰 수준 혼합 (예: SKD): 명시적인 추론 실패 신호가 아닌 일반적인 분포 불일치에 기반하여 스위칭을 수행합니다.

온라인상에서 실패가 발생하는 즉시 이를 수정하고, 추론 상태 자체에 기반하여 개입하는 메커니즘이 부족한 실정입니다.

2. 방법론: Relay On-Policy Distillation (Relay-OPD)

저자들은 Relay-OPD를 제안하며, 이는 감지된 실패 지점에서 학생의 생성과 짧고 집중적인 교사의 개입을 교차시키는 방식입니다.

2.1 핵심 통찰: 교사-학생 지속의 비대칭성 (Teacher-Student Continuation Asymmetry)

이 방법은 실패한 접두사에서 관찰되는 비대칭성에 기초합니다:

학생: 현재의 (잘못된) 방향을 따라 계속 진행하려는 경향이 있습니다.
교사: 멈추고, 재검토하며, 추론의 방향을 재설정하는 경향이 있습니다 (주로 But, Wait, However와 같은 "성찰(Reflection)" 토큰을 사용).

이러한 차이는 레이블이 필요 없는 핸드오프 트리거(Label-free handoff trigger) 역할을 합니다. 교사의 가장 확률 높은 다음 토큰이 성찰 토큰인 반면, 학생의 상위- $K$ 지원 집합(Support set)에 해당 토큰이 포함되지 않을 때 핸드오프가 트리거됩니다.

2.2 릴레이 궤적 구축 (Relay Trajectory Construction)

트리거가 발생하면, 교사가 잠시 주도권을 가져와 **교사 레그(Teacher leg)**를 생성한 후 학생이 다시 재개합니다.

핸드오프 트리거 ( $\phi(h)$ ): $\phi(h) = 1[a_T(h) \in \mathcal{R}] \cdot 1[K_S(h) \cap \mathcal{R} = \emptyset]$ 로 정의됩니다. 여기서 $\mathcal{R}$ 은 성찰 토큰의 집합, $a_T$ 는 교사의 상위 토큰, $K_S$ 는 학생의 상위- $K$ 집합입니다.
릴레이 예산 ( $M, L$ ): 궤적이 학생의 정책에서 너무 멀어지는 것을 방지하기 위해 개입을 제한하는 예산을 설정합니다:
- $M$ : 최대 교사 개입 횟수.
- $L$ : 성찰 토큰 이후 교사가 생성하는 추가 문단의 수.
프로세스: 학생이 자기회귀적으로 생성합니다. 트리거가 감지되면 (예산이 남아 있는 경우), 교사가 성찰 토큰으로 시작하는 $L$ 개의 문단을 생성합니다. 그 후 학생은 교정된 접두사로부터 생성을 재개합니다.

2.3 최적화 목적 함수

Relay-OPD는 결과물인 릴레이 궤적 $z$ 에 대해 학생 정책을 최적화합니다. 학생의 원래 토큰을 사용하는 표준 OPD와 달리, Relay-OPD는 릴레이 궤적에서 실제로 생성된 토큰 $z_t$ (교사가 생성한 토큰 포함)를 사용합니다.
목적 함수는 역-KL 스타일의 단일 샘플 어드밴티지를 사용합니다:
$A_t^{Relay} = \log \pi_T(z_t | h_t^z) - \log \bar{\pi}_\theta(z_t | h_t^z)$
이는 클리핑된 PPO 스타일의 손실 함수를 통해 최적화되며, 이를 통해 학생이 교사의 분포에 완전히 적합되는 대신 교사의 교정 신호를 선택적으로 흡수할 수 있게 합니다.

2.4 효율적인 구현

두 가지 생성 파이프라인 사이의 전환 오버헤드를 피하기 위해, Relay-OPD는 단일 투기적 디코딩 엔진(Single speculative decoding engine) 내에서 프로세스를 통합합니다:

학생은 드래프트(Draft) 모델 역할을 합니다.
교사는 타겟(Target) 모델 역할을 합니다.
엔진은 현재 상태(학생 레그 또는 교사 레그)에 따라 타겟 정책( $\pi_{tgt}$ )을 학생과 교사 사이에서 동적으로 전환합니다.
교사 레그는 교사를 대상으로 표준 투기적 디코딩을 수행하며, 학생 레그는 모든 드래프트를 수락합니다. 이를 통해 릴레이 궤적의 분포적 정체성이 정확하게 보존됩니다.

3. 주요 기여

비대칭성 식별: 논문은 실패한 접두사에서의 교사-학생 발산이 개입을 위한 신뢰할 수 있는 레이블 프리 신호임을 밝혀냈습니다. 실험 결과, 접두사 실패를 수정하는 데는 초기 단계의 국소적 개입만 필요하며, 교사 토큰은 전체 생성 토큰의 **0.35%**에 불과하다는 것을 보여주었습니다.
Relay-OPD 프레임워크: 이러한 비대칭성을 예산이 할당된 릴레이 롤아웃으로 변환하고, 단일 투기적 디코딩 엔진 내에서 학생과 교사의 생성을 통합하는 새로운 방법을 제시했습니다.
실험적 우위: 8개의 수학적 추론 벤치마크와 두 가지 규모(0.6B 및 1.7B)의 학생 모델에 대해 검증되었습니다.

4. 실험 결과

저자들은 Qwen3-4B-Instruct-2507을 교사로, Qwen3-0.6B/1.7B-Non-Thinking을 학생으로 사용하여 8개의 벤치마크(AIME 2024/2025/2026, MATH500, AMC 2023, OlympiadBench, HMMT Feb/Nov)에서 Relay-OPD를 평가했습니다.

성능:
- 1.7B 학생: 모든 벤치마크에서 최고 또는 차석 결과를 달성했습니다.
  - 평균 정확도: 46.96% (표준 OPD의 41.23% 대비 +5.73% 향상).
  - 가장 강력한 베이스라인인 FastOPD를 +1.49% 차이로 앞질렀습니다.
- 0.6B 학생: 일관된 이득을 보였으며, 표준 OPD 대비 평균 정확도를 +3.01% 개선했습니다.
효율성:
- 학습 길이: 평균 학습 궤적 길이를 50% 이상 단축했습니다 (1.7B 기준 4,658개에서 2,296개 토큰으로 감소).
- 수렴: 더 적은 스텝에서 최적의 체크포인트에 도달했습니다 (OPD의 Step 55 대비 Step 35).
- 추론: FastOPD 대비 평균 응답 길이를 14~28% 줄이면서도 정확도를 높였습니다.
절제 연구 (Ablation Studies):
- 교사 레그의 필요성: 단순히 트리거에서 멈추는 방식(Trigger-stop)은 43.48%의 정확도를 보였으나, 교사 레그( $L=3$ )를 추가했을 때 46.25%로 상승하여 교정된 컨텍스트의 가치를 입증했습니다.
- 목적 함수: 생성된 토큰에 대한 역-KL 스타일의 목적 함수가 "학생 드래프트 토큰" 목적 함수와 "교사 FKL" 목적 함수보다 우수한 성능을 보였으며, 이는 교정 신호의 선택적 흡수의 이점을 확인시켜 줍니다.

5. 의의 및 주장

본 논문은 접두사 실패가 기존의 절단 또는 재작성 방식이 동적으로 해결하지 못하는 온폴리시 증류의 근본적인 병목 현상이라고 주장합니다. Relay-OPD는 다음과 같이 이를 해결합니다:

온라인 작동: 사후 처리가 아닌, 실패가 발생하는 즉시 교정합니다.
최소한의 개입: 교사가 오직 방향을 재설정하기 위해 짧게 개입하는 "릴레이" 접근 방식을 사용하여 궤적을 학생의 정책 분포에 가깝게 유지합니다.
효율성: 궤적을 단축하고 회복 불가능한 오류에 대한 계산 낭비를 방지함으로써 학습 비용을 크게 줄입니다.

저자들은 결론적으로 Relay-OPD가 추론 작업에서 강력한 모델에서 약한 모델로의 전이를 위한 견고하고 레이블이 필요 없는 메커니즘을 제공하며, 학습 비용을 획기적으로 줄이면서도 최첨단(SOTA) 결과를 달성한다고 밝히고 있습니다.