OPMD는 교사의 중간 계산을 학생이 쓸 수 있게 증류한다

무엇인가

이 논문은 온폴리시 증류가 학생이 스스로 생성한 접두사에 대해 교사의 출력 분포만 감독한다는 점을 문제로 잡는다. 다음 토큰 분포가 비슷해도 모델 내부에서 증거를 표현하고 변환하는 방식은 다를 수 있고, 출력 수준의 일치는 학생이 어떤 중간 계산을 배웠는지 유일하게 결정하지 못한다. 복잡한 추론에서는 예측을 만드는 데 쓰인 잠재 상태 정보가 전달되지 않을 수 있다. 저자들은 AMC 2023에서 RL 최적화 모델의 잠재 상태를 베이스 모델에 패칭하면 정확도가 오르는 것을 보여 중간 상태에 실행 가능한 정보가 있음을 확인한다. 다만 MSE로 학습한 사영기는 교사 상태를 학생 공간에 옮겨도 생성을 크게 망가뜨렸고, 과제로 학습한 사영기는 패칭 깊이 전반에서 이득을 냈다. 따라서 기하학적 정렬만으로는 부족하고, 학생의 후속 계산이 실제로 쓸 수 있는 형태여야 한다.

어떻게 동작하나

OPMD는 두 단계로 이 문제를 푼다. 1단계는 학생 친화적 교사 구성이다. 교사와 학생 백본은 고정하고, 선택한 층 쌍마다 학습된 사영기가 교사 은닉 상태를 학생 은닉 공간의 잔차 업데이트로 주입한다. 이때 사영기는 교사와 학생의 은닉 상태를 직접 MSE로 맞추지 않고, 고정된 학생이 만들어내는 출력 행동을 기준으로 최적화한다. 토크나이저가 같으면 학생 친화적 교사와 원 교사의 전체 어휘 역KL을 최소화하고, 토크나이저가 다르면 학생 토크나이저로 토큰화한 교사 응답에 대한 SFT로 사영기를 학습한다. 이렇게 하면 교사 정보가 학생의 기존 계산 인터페이스 안에서 사용 가능한 중간 상태로 표현된다.

무엇과 다른가

2단계는 온폴리시 내면화다. 현재 학생이 자기 정책으로 응답을 샘플링하고, 같은 학생 생성 접두사에서 학생 친화적 교사와 네이티브 학생의 지정 층 잠재 상태를 비교한다. 네이티브 학생의 잠재 상태를 교사 친화적 교사의 잠재 상태에 stop-gradient를 둔 MSE로 맞추며, LoRA 파라미터를 업데이트하고 백본은 고정한다. 온폴리시 실험의 2단계에서는 별도 과제 손실을 쓰지 않고 잠재 상태 목적식만 사용한다. SFT 변형은 교사 응답 접두사에 대한 메커니즘 감독과 토큰 예측 NLL을 결합하되 로짓 수준 KL은 쓰지 않는다. 추론 시에는 교사와 사영기를 모두 제거하고 네이티브 학생만 남긴다.

어떻게 쓰나

주요 수학 실험은 DAPO-Math-17K로 학습하고 AMC, AIME 2024, AIME 2025로 평가했다. 교사는 Qwen3-4B와 RL-Math(Qwen3-4B-Non-Thinking-RL-Math), 학생은 Qwen3-1.7B-Base와 Qwen2.5-3B-Instruct를 썼다. 비교 대상은 출력 수준 OPD, 교사 생성 응답 SFT, OPRD-Bridge다. OPMD는 평가한 세 교사-학생 설정 모두에서 증류 베이스라인 중 가장 높은 평균 정확도를 기록했다. AMC에서는 OPD 대비 각각 3.76, 7.04, 2.66 퍼센트포인트 올랐다. RL-Math에서 Qwen3-1.7B-Base로 증류할 때 AIME 2024와 AIME 2025도 각각 4.37, 3.95 포인트 개선됐다. 교사를 Qwen3-4B에서 RL-Math로 바꾸면 OPD의 AMC 정확도는 35.31%에서 33.12%로 떨어졌지만 OPMD는 39.07%에서 40.16%로 올랐다. 1,024 업데이트까지 늘리면 AMC와 AIME 2025는 최고 정확도가 1,024에서 나왔고 AIME 2024는 768에서 정점을 찍었다. OPMD는 평가한 모든 업데이트 수에서 OPD보다 짧은 응답으로 더 높은 정확도를 냈고, 256 업데이트에서도 OPD의 512·768 업데이트보다 적은 토큰으로 더 높은 정확도를 보였다. 사영기 사전학습을 한 투영 MSE 방식은 계속 학습하면 정확도가 떨어져 초기 정렬만으로 안정적 전이가 보장되지 않음을 보였다.

전제와 한계

다중 교사 실험에서는 RL-Math와 RL-Code를 Qwen3-1.7B-Base에 증류했고, 도메인별 프롬프트 25,276개를 사용했다. 1단계 128 업데이트, 2단계 512 업데이트를 수행한 뒤 HumanEval+ 164문제와 MBPP+ 378문제를 pass@1로 평가했다. M-OPD 대비 OPMD는 AMC 4.31, HumanEval+ 0.61, MBPP+ 6.61 퍼센트포인트를 올렸다. 공유 학생이 두 도메인 모두에서 좋아져 수학 이득이 코드 정확도를 희생한 것이 아니며, HumanEval+보다 MBPP+에서 이득이 큰 것은 메커니즘 감독의 효과가 과제마다 균일하지 않음을 보여준다. 교차 계열 SFT에서는 Qwen3-4B에서 Llama3.1-8B-Instruct로 증류하며 SFT와 SFT+MD를 비교했다. SFT+MD는 응답 NLL에 잠재 상태 MSE 계수 0.5를 더했고 로짓 KL은 쓰지 않았다. AMC는 25.31% 대 22.50%, AIME 2025는 1.81% 대 1.46%로 올랐지만 AIME 2024는 3.88% 대 4.33%로 떨어졌다. 즉 교사 텍스트 응답 너머의 정보를 더할 수 있으나 모든 과제에서 일관된 개선은 아니다.

층별 분석에서는 1,024 업데이트 OPMD 학생의 블록 1~7 또는 8~14 가중치를 원본 Qwen3-1.7B-Base에 패칭했다. 정확도는 14.84%에서 각각 31.88%, 34.06%로 올랐고 평균 응답 길이는 줄었지만, 전체 OPMD 학생의 40.68%에는 못 미쳤다. q17과 q22 사례에서 Full은 Q1보다 각각 62.50, 31.25 포인트 높은 보관 정확도를 보였고, Q1은 직접 경로를, Q2와 Full은 매개 경로를 선호했다. Q2가 경로 선호는 Full에 가까우면서 정확도는 낮다는 점은 서로 다른 층군이 추론 메커니즘의 서로 다른 측면을 담고, 전략 선택만으로 성공적 실행이 완전히 회복되지는 않음을 시사한다. 비용 실험은 두 장의 H100에서 수행됐고, OPMD는 768 학생 업데이트에서 38.96%로 OPD의 1,024 업데이트 36.50%를 앞섰다. OPMD 비용에는 1단계가 포함되며, 최적화된 vLLM 구현에서 2단계 업데이트 처리량은 OPD와 거의 같았다. RL-Math에서 1단계를 32 업데이트만 해도 AMC 40.00%로 96 업데이트의 40.16%와 비슷했다.

개발자에게 OPMD는 출력 로짓 증류가 정체되거나 교사와 학생의 구조·토크나이저가 달라 중간 표현 전이가 불안정할 때 참고할 만한 설계다. 다만 사영기 뱅크, 층 쌍 선택, 1단계 교사 구성, 2단계 LoRA 내면화, 추론 시 교사 제거라는 파이프라인을 직접 구현해야 하고, 토크나이저 불일치 시 SFT 경로를 써야 한다. 저자들이 밝힌 한계는 평가한 교사-학생 쌍과 과제가 제한적이라는 점이다. 더 큰 모델, 더 넓은 아키텍처, 수학·코드 이외 과제로 일반화되는지는 검증되지 않았다. 4.3절의 두 사례는 OPMD 학생이 Patch 1~7보다 나은 경우를 골라 분석한 것이라 일반적 층 특화나 인과 메커니즘을 입증하지 않는다. 또한 증류는 교사의 오류와 편향을 전달할 수 있고, 벤치마크 개선이 배포 환경의 신뢰성이나 안전성을 보장하지 않는다.