학생 모델이 불확실해하는 지점에서 분기해 증류 감독을 넓히는 BranchOPD

무엇인가

온폴리시 증류(OPD)는 학생 모델이 스스로 생성한 궤적 위에서 교사의 토큰 확률을 감독 신호로 쓴다. 교사가 쓴 접두사만 보여주는 오프라인 지식 증류와 달리 추론 시점의 분포와 학습 분포를 맞출 수 있다는 점이 장점이다. 문제는 궤적 하나가 각 의사결정 지점에서 가능한 여러 연속 중 단 하나만 실현한다는 것이다. 교사가 접두사에서 전체 다음 토큰 분포를 알고 있어도, 선택되지 않은 토큰이 이끌었을 다단계 역사에는 감독이 전혀 들어가지 않는다. 수학 추론에서는 초반 한 스텝이 이후 전개를 제약하고, 에이전트 과제에서는 한 행동이 이후 모든 관측을 바꾸기 때문에 이 공백이 크다. 독립 궤적을 여러 개 더 뽑으면 커버리지는 늘지만, 추가 토큰 대부분이 동일하거나 거의 같은 접두사를 다시 생성하는 데 낭비된다.

어떻게 동작하나

BranchOPD는 이 낭비를 줄이면서 커버리지를 넓히려고 접두사 공유 분기 샘플링을 제안한다. 입력마다 학생이 먼저 주 궤적을 하나 생성하고, 그 위에서 최대 B개의 분기 지점을 고른다. 분기점 점수는 두 신호를 섞는다. 하나는 학생 엔트로피로 계산하는 불확실성 점수 U_t(상위 20개 토큰 확률 기준)이고, 다른 하나는 국소 128토큰 윈도에서 -Δlog p_t를 누적한 교사-학생 불일치 점수 D_t다. 두 점수를 궤적 내에서 순위 정규화한 뒤 v_t = α·R(U_t) + (1-α)·R(D_t)로 결합한다(기본 α=0.5). 수학 과제에서는 토큰 위치를 직접 고르고, 에이전트 과제에서는 먼저 턴 단위 점수 v_turn(턴 내 상위 η=0.2 비율 점수의 평균)으로 턴을 고른 뒤 그 턴 안의 상위 3개 후보 위치 중 하나에서 분기한다. 기본 설정은 B=3, 분기당 k=2개 연속 샘플링이라 입력당 최대 1+Bk=7개 궤적이 나온다. 에이전트 과제에서는 분기 전에 환경을 스냅샷이나 행동 재생으로 해당 상태로 되돌려, 재개된 연속이 중단 없는 롤아웃과 같은 분포를 갖도록 한다. 첫 토큰이 주 궤적과 같은 분기도 버리지 않고 유지한다. 이후 교사가 주 궤적과 분기의 모든 생성 토큰을 채점하고, 공유 접두사는 손실에 한 번만 들어간다. 논문은 이 분기 추정량이 조건부로 편향되지 않음을 명제로 제시한다(공유 접두사 조건 하에서 기대값이 원래 기대값과 같다).

무엇과 다른가

논문은 여기에 프로토콜 경계 복구(PBR)라는 보조 손실을 더한다. OPD에서 흔한 실패 모드는 부적절한 종료, 즉 최종 답을 낸 뒤에도 중복 내용이나 잘못된 포맷 토큰을 계속 뱉는 현상이다. PBR은 고정 탐지 규칙으로 각 궤적을 스캔해 (접두사 s, 토큰 집합 C, 방향 b) 제약을 뽑고, b=1이면 p_θ(C|s)를 높이고 b=0이면 낮추는 이진 교차엔트로피 형태의 손실을 건다. 예컨대 수학 답이 완결된 뒤 EOS 확률을 올리고, 에이전트 응답 앞에 붙는 불필요한 닫는 생각 태그를 억제한다. 탐지 자체는 미분하지 않으며, 이 보조항은 주 궤적과 분기에 똑같이 적용된다. 저자들은 비교 실험이 분기 효과만 분리하도록 베이스라인에도 PBR을 적용한 버전(+PBR)을 함께 보고한다.

어떻게 쓰나

실험은 수학 추론과 에이전트 3종 벤치마크에서 이뤄졌다. 수학은 DAPO-Math-17K로 Qwen3-8B를 Qwen3-4B-Base에 증류하고 MATH-500, AIME 2024, AIME 2025에서 문제당 16개 샘플 평균 정확도(avg@16)를 측정했다. 에이전트는 ALFWorld(학습 140개 ID, 미학습 134개 OOD 성공률), WebShop(held-out 128개), Search-QA(exact-match)를 쓴다. BranchOPD는 Vanilla OPD 대비 수학 평균 18.55점, ALFWorld 평균 성공률 13.83점을 개선했다. 모든 방법에 PBR을 붙인 조건에서도 수학 평균 44.37점으로 Vanilla OPD와 Multi-Rollout OPD(n=7)를 각각 5.43점, 5.28점 앞선다. 수학에서는 KD, GRPO, EOPD, SPOT을 11.27, 8.81, 6.40, 5.94점 차로 앞서는데, SPOT 대비 격차가 MATH-500에서 1.36점인 반면 AIME24와 AIME25에서는 8.54점, 7.91점으로 훨씬 크다. GRPO와 SPOT이 정답 검증을 쓰는 것과 달리 BranchOPD는 교사 감독만 쓴다는 점을 저자들은 강조한다. 에이전트 3종 평균은 46.39 SR로 KD, GRPO, TCOD을 12.87, 2.62, 2.60점 앞선다. ALFWorld는 Vanilla OPD(+PBR) 49.27%에서 56.54%로 올라 교사 모델(51.79%)까지 넘어섰고, WebShop은 30.47%에서 33.59%로 올라 역시 교사(31.25%)를 앞선다. Search-QA 개선폭은 0.47점이다.

전제와 한계

분기 위치와 예산에 대한 절제 실험도 있다. 엔트로피와 불일치를 섞는 방식이 가장 좋아서, 무작위 분기 대비 ALFWorld ID/OOD 성공률을 7.14/8.96점 올리고 수학 평균을 40.11에서 44.37로 끌어올렸다. 두 신호를 각각 단독으로 쓰면 이 이득의 대부분을 회복하지만 전부는 아니다. 턴 시작 지점에서 분기하는 방식은 기본 설정보다 ALFWorld에서 2.14/1.49점 뒤처져, 턴 안 어디에서 갈라지느냐도 중요하다. 예산 면에서는 k=1에서 k=2로 늘리면 세 지표 모두 개선되지만 k=3은 추가 이득이 없고, k=2에서 분기를 2개에서 3개로 늘리면 수학 0.43점, ALFWorld 0.73점이 더 붙는다. 즉 여러 의사결정을 적은 수의 연속으로 덮는 편이 낫다. 행동 커버리지 측정에서 선택된 분기 지점의 대안 확률 질량(1-π(a0|s))은 ALFWorld 0.63, WebShop 0.69였고, 분기의 74.2%/47.7%가 주 궤적과 다른 행동 순서를, 55.5%/37.5%가 주 궤적에 없던 행동을 포함했다. 주 궤적이 실패했을 때 분기 중 최소 하나가 성공하는 비율은 WebShop 5.0%에서 수학 42.7%까지 나타난다. 토큰 비용은 ALFWorld에서 궤적당 평균 2,813토큰이라 독립 7개면 19,691토큰이 필요한데 분기 설정은 15,924토큰으로 19.1% 절약하고, WebShop은 7,014 대 4,824토큰으로 31.2% 절약한다. PBR 단독 효과도 커서 수학에서 Vanilla OPD를 25.82에서 38.94 평균으로 올렸고, ALFWorld ID/OOD를 7.14/5.97점 개선했다. 반복 페널티(28.70~30.40), 과장 샘플 필터링(33.20), 조기 종료(22.05) 같은 일반적 처방보다 PBR(38.94)이 일관되게 낫다.

실무 관점에서 이 논문이 주는 시사점은 명확하다. 증류 파이프라인에서 롤아웃 수를 늘리는 대신 어디에 예산을 쓸지 고르는 것만으로 같은 토큰 예산에서 더 넓은 행동 분포를 감독할 수 있다는 것이다. 특히 에이전트 과제처럼 환경 상호작용 비용이 큰 경우, 접두사 공유로 학생 생성 토큰을 20~30% 줄이면서 성능을 올린다는 점은 학습 비용 계산에 직접 반영할 만하다. 다만 분기 재개를 위해 환경 상태를 스냅샷이나 행동 재생으로 복원할 수 있어야 한다는 전제가 붙고, PBR은 탐지 규칙을 과제 프로토콜에 맞춰 직접 정의해야 한다. 또한 논문이 보고한 토큰 절감 수치는 교사 채점, 역전파, 환경 재생 비용을 제외한 학생 생성 토큰만 센 값이므로, 실제 학습 비용을 추정할 때는 이 항목들을 따로 계산해야 한다. 평가는 과제당 궤적 하나로, 분기 없이 이뤄진다는 점도 배포 시 고려할 부분이다.

저자들이 명시한 한계와 전제는 다음과 같다. 토큰 절감 수치는 교사 채점, 역전파, 환경 재생을 제외한 값이다. 에이전트 과제에서 분기 재개가 중단 없는 롤아웃과 같은 분포를 가지려면 환경을 정확히 복원해야 하며, 이는 스냅샷이나 행동 재생으로 구현한다고 밝히고 있다. PBR의 탐지 규칙은 고정되어 있고 미분되지 않으며, 제약 집합 C는 어휘의 진부분집합이어야 한다. 또한 분기점 선택은 주 궤적에 의존할 수 있지만 그 선택 자체가 조건부 집합에 포함되므로 추정 편향은 없다는 것이 논문의 주장이며, 이는 환경 복원이 정확하다는 전제 위에서 성립한다.