SAKI가 최대 결합으로 교사 감독 위치를 스스로 고른다

무엇인가

온폴리시 증류(OPD)는 학생이 자기 궤적에서 생성한 상태에 교사 감독을 걸어 학습-추론 상태 불일치를 줄인다. 문제는 학생이 약할 때다. 초반 실수가 누적되면 교사가 자기 정책 아래에서는 거의 생성하지 않을 접두사로 롤아웃이 흘러가고, 교사 분포는 그 지점에서도 정의되긴 하지만 교사를 강하게 만든 추론 행동을 대표하지 못하는 조건부 신호가 된다. 선행 연구 TRB는 학생 중심 KL 신뢰영역 안에서 학생 분포와 교사 분포를 기하 보간한 행동 분포를 만들어 "어디서 감독을 받을지"를 개선했지만, 방문한 접두사마다 적용하는 목적 함수는 기존 역KL 그대로였다. 역KL은 학생 분포로 가중되므로 학생 지지가 낮은 교사 선호 토큰은 그래디언트가 약하게 들어간다.

어떻게 동작하나

SAKI(Supervision Allocation with KL-constrained Interpolation)는 이 빈틈을 롤아웃 과정에서 나오는 사건으로 메운다. 먼저 접두사마다 q_β,t(v) = p_t(v)^(1-β_t) T_t(v)^β_t / Z_t(β_t) 형태의 기하 보간 행동 분포를 만들고, D_KL(q_β,t ∥ p_t) ≤ ε을 만족하는 가장 큰 β_t를 고른다. ε=0이면 학생 분포 그대로, 신뢰영역이 충분히 크면 교사 분포가 된다. 이 q_t를 최대 결합(maximal coupling)으로 실현한다. 학생이 z_t ~ p_t를 제안하면 min(1, q_t(z_t)/p_t(z_t)) 확률로 수용하고, 거부되면 잔차 분포 r_t(v) ∝ [q_t(v) − p_t(v)]_+ 에서 교정 토큰을 뽑는다. 논문의 Proposition 1은 이때 교정 확률이 정확히 TV(p_t, q_t)임을, Proposition 2는 최대 결합이 주변분포 p와 q를 갖는 모든 결합 가운데 개입 확률을 최소화함을 보인다. Corollary 1은 신뢰영역 제약 아래 Pr(C_t=1) ≤ √(ε/2)를 준다. 즉 신뢰영역 반경 ε 하나가 롤아웃 편차와 개입 빈도, 나아가 특수 감독의 빈도를 동시에 통제한다.

무엇과 다른가

감독 라우팅은 이 실현된 수용/교정 이벤트를 그대로 재사용한다. 수용 위치에서는 샘플된 토큰에 대한 역KL 손실 L_RKL,t(v) = −A_t(v) log π_θ(v|h_t), A_t(v) = −sg[log p_t(v) − log T_t(v)]를 유지한다. 교정 위치에서는 같은 접두사에서 교사의 최고확률 토큰 v*_t = argmax_v T_t(v)에 대한 음의 로그가능도 L_TM,t = −log π_θ(v*_t|h_t)로 갈아탄다. 전체 손실은 L = (1/N_valid)[Σ_t M_t(1−C_t)L_RKL,t + Σ_t M_t C_t L_TM,t]로, 유효 응답 토큰 마스크 M_t와 교정 지시자 C_t로 두 항을 나눈다. 논문은 수용 위치의 역KL 감독이 정확히 학생-가이드 겹침 질량 min{p, q} 위에서 작동함을 보인다(식 17). 교정 토큰은 다음 오토리그레시브 접두사를 결정하고 교사 최빈 토큰은 파라미터 업데이트를 결정하므로, 하나의 결합 과정이 궤적 구성과 감독 라우팅을 동시에 지배한다. 확률적으로 교사에서 샘플링하는 대신 결정적 최빈 토큰을 쓰는 선택은 1.7B 학생에서 Mean@8 +0.62, Pass@8 +1.56으로 더 나았다(부록 F, 표 8). ε은 처음 50스텝 동안 0.02에서 0으로 선형 어닐링되어, 최종적으로는 정확히 학생 롤아웃 역KL 학습으로 되돌아간다.

어떻게 쓰나

토큰마다 학생과 교사 분포가 모두 필요하므로 순차 실행은 비싸다. 저자들은 엔진 내부에 상주하는 speculative 블록 검증기를 구현했다. 매 웨이브마다 동결된 롤아웃 학생이 K=8개 토큰을 드래프트하고, 교사가 해당 제안 접두사들을 한 번에 배치 검증한다. 엔진 안에서 블록 위치 전체의 q_t를 만들고 신뢰영역 계수 β_t를 배치로 풀며, 최대 결합 수용 여부를 왼쪽에서 오른쪽으로 평가한다. 연속 수용된 최장 접두사를 커밋하고, 첫 거부가 j번째에서 나면 c_j ∝ [q_j − p_j]_+ 를 샘플해 커밋한 뒤 이후 speculative 토큰과 KV 상태를 버리고 교정된 접두사에서 재개한다. 제안 접두사가 첫 거부까지는 실제 접두사와 일치하고 교정 이후 상태를 재사용하지 않기 때문에 순차적 exact-q 샘플링과 같은 자기회귀 법칙·결합 사건을 갖는다(부록 I). 같은 64프롬프트 × 8응답, 최대 길이 7,168, K=8 워크로드에서 초당 커밋 토큰이 3,276개로 외부 루프 구현의 776개 대비 4.22배였고, 학생 전용 비가이드 생성의 이론 상한 대비 42% 이상을 유지했다.

전제와 한계

실험은 Qwen3-0.6B-Base와 Qwen3-1.7B-Base 학생을 Qwen3-4B-Base-GRPO 교사로부터 DAPO-Math-17K에서 증류한다. 학생 규모별로 같은 체크포인트, 프롬프트, 200스텝 예산, 64프롬프트 × 8응답 롤아웃 배치를 쓰고, MATH-500, HMMT-Feb26, AIME 2026, AIME 2025, AMC 2023, Minerva Math, OlympiadBench 일곱 개 벤치마크에서 온도 1.0, 최대 길이 16,384로 8개 응답을 샘플링해 Mean@8과 Pass@8을 잰다. 1.7B 학생은 TRB의 27.9/44.6에서 29.0/47.5로(+1.1/+2.9), 0.6B 학생은 17.2/33.6에서 18.4/35.6으로(+1.2/+2.0) 올랐다. SKD 대비로는 1.7B에서 +4.5/+6.3, 0.6B에서 +6.0/+6.0이다. Mean@8은 14개 학생-벤치마크 조합 중 13개에서 TRB를 앞섰다. 배치 대조군에서는 같은 교사모드 예산을 쓰되 위치를 무작위로 고른 Random-TM이 28.30/45.50, 국소 TV(p_t, q_t)로 가중해 고른 TV-Weighted-TM이 28.24/46.77이었고, SAKI는 29.00/47.50으로 TV-Weighted-TM을 Mean@8 +0.76, Pass@8 +0.73 앞섰다. 213개 프롬프트에서 뽑은 2,048개 위치 고정 프리픽스 프로브에서는 교사 최빈 토큰 확률 C1이 50스텝에서 TRB 대비 4.63%p, Top-16 집합 확률 C16이 4.54%p 높았고, 교사모드 감독이 꺼진 뒤 149스텝의 역KL만 돈 200스텝 시점에도 각각 4.26%p, 2.94%p 격차가 남았다. Random-TM 대비로는 200스텝에서 C1 +0.648, C16 +0.816이었고, C16 이점은 최저 갈등 사분위의 +0.25%p에서 최고 갈등 사분위의 +1.32%p로 커졌다(Q4−Q1 +1.07, 95% CI [0.87, 1.29]). 이 증가분은 주로 교사가 지지하는 비-argmax 토큰(+1.28%p)이 끌어올렸다.

실무 관점에서 이 논문이 주는 시사점은 명확하다. 교사 모델을 온라인으로 띄울 수 있는 환경이라면, 롤아웃 분포를 교사 쪽으로 당기는 것과 손실 함수를 토큰별로 바꾸는 것을 별개 하이퍼파라미터로 다룰 필요가 없다는 것이다. 신뢰영역 반경 ε이 롤아웃 편차와 개입 빈도를 함께 묶고, 최대 결합의 수용/교정 이벤트가 추가 임계값 없이 감독 위치를 정한다. 다만 이 방식은 매 생성 위치에서 교사 전체 어휘 로짓이 필요하므로, 엔진 내부 speculative 검증 같은 시스템 최적화 없이 그대로 붙이면 처리량이 크게 떨어진다. 또한 교정 위치에 걸리는 교사모드 감독은 ε 어닐링과 함께 사라지도록 설계되어 있어, 학습 후반부에는 사실상 순수 학생 롤아웃 역KL과 동일한 목적함수가 된다는 점을 스케줄 설계 시 염두에 둬야 한다.

논문에는 별도의 한계(Limitations) 절이 없다. 본문에서 확인되는 전제와 제약은 다음과 같다. 첫째, exact-q 롤아웃은 모든 생성 위치에서 학생과 교사 분포를 동시에 요구하므로 교사 추론 비용을 전제로 한다. 둘째, 교정 감독은 일시적이며 ε을 0으로 어닐링하면 학생 롤아웃 역KL로 정확히 복귀한다. 셋째, 실험은 Qwen3 계열 0.6B·1.7B 학생과 단일 Qwen3-4B-GRPO 교사, DAPO-Math-17K, 200스텝 예산, 수학 추론 벤치마크 일곱 개로 한정되며 다른 도메인이나 교사-학생 조합으로의 일반화는 검증되지 않았다. 넷째, 개입 빈도 상한 √(ε/2)는 Pinsker 부등식에서 나온 값으로, 부록 J.2에 따르면 실제 관측 교정 확률은 이 상한보다 낮게 유지되고 ε이 0이 되면 사라진다.