💻 computer science

When Are Teacher Tokens Reliable? Position-Weighted On-Policy Self-Distillation for Reasoning

본 논문은 엔트로피가 아닌 토큰 위치에 의해 가장 잘 예측되는 궤적 구조 패턴을 따르는 교사 토큰의 신뢰도를 식별함으로써 추가적인 교사 계산 없이 표적 증류를 가능하게 하여 추론 모델의 성능을 향상시키는 방법인 위치 가중 온-정책 자기 증류 (PW-OPSD) 를 소개한다.

원저자: Xiaogeng Liu, Xinyan Wang, Yingzi Ma, Yechao Zhang, Chaowei Xiao

게시일 2026-05-22

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Xiaogeng Liu, Xinyan Wang, Yingzi Ma, Yechao Zhang, Chaowei Xiao

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

"교사 토큰은 언제 신뢰할 수 있는가?"라는 논문에 대한 설명을 쉬운 언어와 창의적인 비유로 풀어냅니다.

큰 그림: 학생이 스스로 문제를 풀며 가르치는 것

어려운 수학 문제를 푸는 법을 학생에게 가르치고 있다고 상상해 보세요. 답을 알고 있는 "교사"(초지능 AI) 와 배움을 시도하는 "학생"(조금 덜 똑똑한 AI) 이 있습니다.

온-정책 자기 증류 (On-Policy Self-Distillation) 라는 방법에서 과정은 다음과 같습니다:

학생이 스스로 문제를 풀며 생각을 단계별로 적어냅니다.
교사는 학생이 지금까지 쓴 내용을 정확히 보고, "좋아, 방금 쓴 내용을 바탕으로 다음에 적을 가장 좋은 단어는..."라고 말합니다.
학생은 이 피드백을 통해 배웁니다.

문제점:
이 작업을 수행하는 표준적인 방법은 교사가 제안하는 단어 하나하나를 모두 동등하게 중요하고 동등하게 정확하다고 취급합니다. 이는 러너가 곧 바위 위에 걸려 넘어질지도 모르는 상황인데도, 코치가 "네가 밟는 모든 발걸음이 완벽하다"고 말하는 것과 같습니다.

이 논문의 저자들은 때때로 교사가 혼란을 겪는다는 점을 깨달았습니다. 복잡한 추론 과정에서 교사는 모두 그럴듯해 보이지만 실제로는 오직 하나만이 올바른 최종 답으로 이어지는 여러 가지 다른 "다음 단계"를 제안할 수 있습니다. 학생이 "그럴듯하지만 틀린" 단계를 맹목적으로 따르면 막히게 됩니다.

발견: "분기 생존성 (Branch Viability)" 테스트

연구자들은 궁금했습니다: 교사는 언제 실제로 신뢰할 수 있고, 언제 단순히 추측을 할까요?

이를 알아내기 위해 그들은 "분기 생존성 (Branch Viability)" 이라는 진단 테스트를 고안했습니다. 등산 비유를 사용하여 작동 방식을 설명하면 다음과 같습니다:

설정: 학생이 산을 오르는 상황 (문제를 풀고 있는 상황) 을 상상해 보세요. 그들은 안전하고 올바른 길 위에 있습니다.
테스트: 다양한 지점에서 교사는 "hey, 대신 이 다른 길을 갈 수도 있어"라고 제안합니다.
실험: 연구자들은 학생이 그 제안된 "다른 길"을 가도록 강요하지만, 교사의 도움 (요령) 없이 진행합니다.
결과:
- 시나리오 A (유해하지 않은 다양성): 학생이 새로운 길을 따라 계속 산을 오르고, 여전히 정상에 도달합니다. 교사의 제안은 문제를 해결하는 다른 유효한 방법일 뿐입니다. 이는 안전합니다.
- 시나리오 B (실제 불확실성): 학생이 새로운 길을 따라가다가 길을 잃고 정상에 도달하지 못합니다. 교사의 제안은 함정이었습니다. 이는 신뢰할 수 없습니다.

놀라운 발견: "얼마나 혼란스러운가"가 아니라 "언제"가 중요하다

연구자들은 교사가 "혼란스러워 보일 때"(높은 엔트로피, 즉 많은 옵션을 가질 때) 신뢰할 수 없으리라고 예상했습니다.

하지만 그들은 틀렸습니다.

그들은 교사의 혼란이 어느 시점에서 발생했는지에 비해 덜 중요하다는 사실을 발견했습니다.

등산 초기: 교사가 일찍 우회를 제안한다면, 그것은 거의 확실히 막다른 길입니다. 학생은 주된 길에 머무러야 합니다.
등산 후기: 교사가 끝부분에서 우회를 제안한다면, 그것은 보통 일을 마무리하는 다른 방법일 뿐입니다. 학생은 이를 안전하게 탐색할 수 있습니다.

그들은 문장에서 단어의 "위치"를 살펴봄으로써 이를 테스트했습니다. 그들은 교사가 얼마나 "혼란스러워 보이는지"를 확인하는 것보다, 학생이 문제에서 얼마나 진척되었는지를 아는 것이 신뢰성을 예측하는 훨씬 더 좋은 지표임을 발견했습니다.

해결책: PW-OPSD (위치 가중 학습)

이를 바탕으로 그들은 PW-OPSD라는 새로운 학습 방법을 고안했습니다.

수업이 진행됨에 따라 변하는 교사의 목소리 볼륨 조절기라고 생각하세요:

문제 시작 시: 볼륨이 낮아집니다. 학생은 교사의 말을 듣지만 맹목적으로 신뢰하지는 않습니다. "교사가 여기서 잘못된 우회를 제안할 수도 있으니 조심해"라고 알려줍니다.
문제 끝부분: 볼륨이 높아집니다. 학생은 교사를 완전히 신뢰합니다. "교사는 이 일을 마무리하는 방법을 정확히 알고 있으니 그들의 지시를 따르세요."

이 방법은 교사가 추가 작업을 하거나 추가 테스트를 실행할 필요가 없습니다. 단순히 학생이 얼마나 진척되었는지에 따라 조언을 어떻게 가중할지 변경할 뿐입니다.

결과: 더 똑똑한 수학 풀이기

그들은 이 새로운 방법을 어려운 수학 대회 (AIME 및 HMMT 등) 에 적용하여 테스트했습니다.

결과: "위치 가중" 방법으로 훈련된 학생들은 "모든 것을 동등하게 신뢰"하는 구식으로 훈련된 학생들보다 훨씬 더 많은 문제를 정확하게 풀었습니다.
교훈: 복잡한 추론에서 타이밍이 중요합니다. 교사의 조언에 대한 신뢰도는 무작위가 아니라 패턴을 따릅니다. 그 패턴을 존중함으로써 더 강력한 컴퓨터가 필요 없이 더 똑똑한 AI 를 구축할 수 있습니다.

한 문장으로 요약한 내용

이 논문은 AI 수학 추론에서 교사의 조언은 문제 시작 부분에서는 종종 위험하지만 끝부분에서는 안전하다는 것을 발견했습니다. 따라서 학생을 훈련시키는 가장 좋은 방법은 모든 조언을 동등하게 완벽하게 취급하는 것이 아니라, 시작 부분에서는 교사를 덜 신뢰하고 끝부분에서는 더 신뢰하는 것입니다.

기술 요약: 추론을 위한 위치 가중 온-정책 자기 증류

문제 정의

온-정책 자기 증류 (OPSD) 는 학생 모델이privileged teacher(정답에 조건부인 모델의 복사본) 가 토큰 수준의 타겟을 제공하는 방식으로, 모델이 생성한 롤아웃을 기반으로 수학 추론을 증류하는 실용적인 방법입니다. 그러나 표준 OPSD 목적 함수는 생성된 모든 토큰에 균일한 순방향 KL(Kullback-Leibler) 감독 가중치를 적용합니다. 이는 학생이 방문하는 모든 접두어에서privileged teacher의 타겟이 동등하게 신뢰할 수 있다는 것을 암묵적으로 가정합니다.

최근의 적응형 증류 방법들은 teacher 엔트로피, teacher-학생 확률 비율, 또는 합의와 같은 국소 신호에 기반하여 감독을 조절함으로써 이러한 균일성을 완화하려 시도합니다. 저자들은 이러한 국소 모호성 신호들이 teacher 신뢰도를 위한 충분한 대리 변수가 아니라고 주장합니다. 높은 teacher 엔트로피는 두 가지 구별되는 현상을 나타낼 수 있습니다:

비실현 가능 불확실성 (Non-viable uncertainty): teacher 가 여러 개의 다음 토큰에 확률을 할당하지만, 오직 하나 (또는 전혀 없음) 만 올바른 해법으로 이어지고 나머지는 "죽은 길 (dead ends)"입니다.
유익한 다양성 (Benign diversity): teacher 가 여러 개의 실현 가능한 해법 경로를 가진 토큰에 확률을 할당하며, 이 모든 경로가 정확성을 유지합니다.

기존 방법들은 이를 혼동하여, 유익한 분기점에서는 감독을 과도하게 낮추거나, teacher 가 궤도 이탈하는 비실현 가능한 분기점에서는 감독을 낮추지 못하게 할 수 있습니다. 핵심 문제는 비실현 가능한 불확실성과 실현 가능한 다양성을 구별하여 토큰을 적절히 가중치할 수 있는 신호를 식별하는 것입니다.

방법론

1. 분기 실현 가능성 진단 (Branch-Viability Diagnostic)

teacher 신뢰도를 경험적으로 조사하기 위해, 저자들은 분기 실현 가능성 진단을 도입합니다.

설정: 단일 모델 (Qwen3-4B) 을 사용하여 "teacher" 프롬프트 (정답 포함) 와 "학생" 프롬프트 (일반 문제) 를 정의합니다.
과정:
1. 학생이 올바른 "척추 (spine)" 롤아웃을 생성합니다.
2. teacher 는 이 척추를 따라 높은 모호성 위치 (높은 절단 엔트로피) 를 식별합니다.
3. 이러한 위치에서 teacher 는 대안적인 다음 토큰을 제안합니다.
4. 모델은privileged 정보 없이 학생 프롬프트 + 척추 접두어 + 강제된 대안에서 계속되도록 강제됩니다.
5. 레이블링: 강제된 연속이 올바른 답을 회복하지 못하면, 해당 분기는 **실제 불확실 (real-uncertain, 낮은 신뢰도)**로 레이블링됩니다. 답을 회복하면 **다양성 (diversity, 유익함)**으로 레이블링됩니다.
발견: 진단은 강력한 **위치 구조 (positional structure)**를 드러냅니다. 실제 불확실 분기는 시퀀스 초기에 집중되는 반면, 실현 가능한 다양성 분기는 후반부에 더 흔하게 나타납니다.
예측력: 방향성 있는 시퀀스 내 위치 점수 (초기 위치 = 낮은 신뢰도) 는 실제 불확실 후보와 다양성 후보를 구분하는 데 AUROC 0.83을 달성합니다. 반면, 표준 국소 불확실성 지표 (예측 엔트로피, MC-dropout 상호 정보량, 디리클레 정밀도) 는 우연 수준 (AUROC $\le$ 0.57) 에 근접합니다.

2. 위치 가중 온-정책 자기 증류 (PW-OPSD)

초기 고모호성 위치가 구조적으로 덜 신뢰할 수 있다는 발견에 영감을 받아, 저자들은 PW-OPSD를 제안합니다.

핵심 메커니즘: PW-OPSD 는 표준 OPSD 구성 (학생 롤아웃,privileged teacher 통과, 어휘별 클리핑된 순방향 KL) 을 유지하지만, 외부 토큰 집계 방식을 수정합니다.
가중치 스케줄: 균일한 가중치 대신, 시퀀스 내 정규화된 위치에 기반한 증가하는 시그모이드 함수를 적용합니다:
$w_{i,t} = w_{min} + (1 - w_{min})\sigma\left(\frac{r_{i,t} - \tau}{s}\right)$
여기서 $r_{i,t}$ 는 정규화된 위치입니다.
설계 근거: 이 스케줄은 진단이 높은 비신뢰도를 나타내는 초기 토큰에는 낮은 가중치를, 신뢰도가 높은 후기 토큰에는 높은 가중치를 할당합니다. 바닥값 ( $w_{min}$ ) 은 초기 토큰이 완전히 무시되지 않도록 보장합니다.
축소 (Reduction): 이 방법은 배치 간 평균화 전에 각 시퀀스 내 토큰 손실을 평균화하여, 더 긴 롤아웃이 단순히 토큰 수 때문에 기울기를 지배하지 않도록 합니다.
효율성: 추가적인 teacher 통과, 보조 검증기, 복잡한 샘플링이 필요하지 않습니다. 가중치는 위치의 결정론적 함수입니다.

주요 기여

신뢰도 진단: 이 논문은 추론에서 teacher-토큰 신뢰도가 순수한 국소적 (엔트로피 의존) 이 아니라 **궤도 구조적 (위치 의존적)**임을 경험적으로 입증하는 분기 실현 가능성 진단을 도입합니다.
엔트로피 대비 위치의 우위: 저자들은 국소 불확실성 점수 (AUROC $\le$ 0.57) 보다 위치가 teacher 신뢰도의 훨씬 강력한 예측 변수 (AUROC 0.83) 임을 보여주며, 엔트로피 기반 적응형 가중치가 최적이라는 가정에 도전합니다.
PW-OPSD 알고리즘: 순방향 KL 손실에 위치 의존적 가중치를 적용하여 추가적인 계산 오버헤드 없이 추론 성능을 향상시키는 간단하고 플러그 앤 플레이 방식의 OPSD 수정안입니다.
일반화: 이 방법은 모델별 하이퍼파라미터 튜닝 없이 다양한 모델 계열 (Qwen, Llama, OLMo) 과 규모 (4B~8B 파라미터) 에서 검증되었습니다.

실험 결과

주요 평가 (Qwen3-4B)

최대 생성 길이 38,912 토큰으로 MATH-500, AIME 2024, AIME 2025, HMMT 2025 에서 평가:

AIME 2024: PW-OPSD (적당) 는 표준 OPSD 대비 Avg@12 를 +1.0 점 향상시켰습니다 (76.20 대 75.19).
AIME 2025: PW-OPSD 는 Avg@12 를 +1.1 점 향상시켰습니다 (67.78 대 66.67).
HMMT 2025: "공격적" 스케줄 변형은 OPSD 대비 Avg@12 를 +1.48 점 향상시켰습니다.
베이스라인: 엔트로피를 사용하는 적응형 베이스라인 (EOPD) 이나 정책 경사 보상 (REOPOLD) 은 일반적으로 AIME 벤치마크에서 OPSD 보다 낮은 성능을 보였으며, 위치 진단이 식별한 신뢰도 패턴을 포착하지 못했습니다.
MATH-500: 성능은 OPSD 와 비교 가능하게 유지되어 (거의 포화 상태), 이 방법이 더 쉬운 작업에서 성능을 저하시키지 않음을 나타냅니다.

교차 모델 전이

Qwen3-4B 진단에서 파생된 "적당" 스케줄을 재튜닝 없이 다음에 적용했습니다:

DeepSeek-R1-Distill-Llama-8B: Avg@12 를 +0.35 점 향상.
Olmo-3-7B-Think: Avg@12 를 +0.50 점 향상.
이는 위치 신뢰도 사전 지식이 모델 계열 및 아키텍처 전반에 걸쳐 견고함을 시사합니다.

제거 실험 (Ablation Studies)

스케줄 민감도: 다양한 스케줄 구성 (부드러움, 날카로움, 공격적) 이 모두 평균적으로 OPSD 를 능가하여 위치 가중치 개념의 견고성을 확인했습니다.
축소 대 위치: 개선은 위치 가중치와 시퀀스별 손실 축소의 공동 적용에 의존한다는 것이 발견되었습니다. 하나의 축만 적용하면 전체 성능 향상을 얻지 못했습니다.

중요성 및 주장

이 논문은 추론 증류에서 teacher-토큰 신뢰도가 궤도 구조적이라고 주장합니다. 주요 중요성은 적응형 증류를 위해 국소 모호성 (엔트로피) 에서 구조적 신뢰도 (위치) 로 초점을 이동하는 데 있습니다.

소극적 주장: 저자들은 PW-OPSD 가 성능을 향상시키지만, 그 이득은 "소극적"이라고 명시합니다 (예: AIME 에서 +1.0~+1.1 점). PW-OPSD 는 "플러그 앤 플레이" 근사치이며, 향후 위치 기반 순방향 및 역방향 KL 혼합과 같은 더 정교한 목적 함수를 설계할 여지가 있음을 인정합니다.
추가 비용 없음: 주요 주장은 이러한 개선이 추가적인 teacher 계산이나 보조 검증 단계 없이 달성된다는 점으로, 기존 증류 파이프라인에 매우 효율적인 수정임을 의미합니다.
이론적 통찰: 이 작업은 증류를 신뢰도 가중 위험으로 해석하는 이론적 해석을 제공합니다. 여기서 목표는 teacher 의 타겟이 올바른 해법을 유지하는 데 도움이 될 사후 확률에 따라 토큰을 가중치하는 것입니다. 진단은 위치가 이 사후 확률에 대한 강력한 대리 변수임을 경험적으로 입증합니다.

코드는 https://github.com/SaFo-Lab/PW-OPSD에서 이용 가능합니다.