기술 요약: 온-폴리시 증류와 최선 N 번 롤아웃 교사 선택 (BRTS)

1. 문제 정의

온-폴리시 증류 (OPD) 는 대규모 언어 모델 (LLM) 을 위한 데이터 효율적인 사후 학습 방법으로 등장했으며, 여기서 학생 모델은 교사의 토큰별 로그 확률을 사용하여 자신이 샘플링한 궤적에 대해 감독을 받습니다. OPD 는 학생이 실제로 방문하는 상태와 감독을 정렬함으로써 노출 편향을 완화하지만, 표준 구현은 높은 분산의 감독 신호로 고통받습니다.

식별된 핵심 제한 사항은 다음과 같이 두 가지입니다:

1. 노이즈가 있는 학생 컨텍스트: 감독은 불완전한 학생이 생성한 접두어에 대해 계산되므로, 국소적인 교사 피드백이 덜 유익한 노이즈가 있는 추론 상태로 벗어날 수 있습니다.
2. 확률적 교사 롤아웃: 표준 OPD 는 일반적으로 프롬프트당 단일 교사 롤아웃에 의존합니다. 교사가 확률적이기 때문에 단일 샘플링된 궤적은 잘못되었거나, 정보가 부족하거나, 학생의 현재 추론 행동과 잘 정렬되지 않을 수 있습니다. 이로 인해 특히 교사의 샘플링된 출력물이 정확성과 스타일 측면에서 크게 변하는 어려운 프롬프트에서 신뢰할 수 없는 지도가 발생합니다.

결과적으로, 표준 OPD 는 오류가 있는 교사 신호를 전파하거나 학생을 완전하고 정확한 추론 경로에 노출시키지 못해 위험한 추론 벤치마크에서의 효과를 제한할 수 있습니다.

2. 방법론: BRTS 프레임워크

저자들은 OPD 루프 내에서 감독에 사용되기 전에 신뢰할 수 있는 교사 궤적을 선별하도록 설계된 **BRTS(Best-of-N Rollout Teacher Selection)**를 제안합니다. BRTS 는 선택된 고품질 교사 궤적으로 구성된 교사 컨텍스트 감독 분기를 통해 표준 학생 컨텍스트 증류를 보완합니다.

2.1 정확성 및 정렬 인식 궤적 선별

각 프롬프트에 대해 BRTS 는 $N$개의 무조건부 교사 롤아웃 $\{y_{T,i}\}_{i=1}^N$ 풀을 샘플링합니다. 선택 과정은 다음과 같은 우선순위 계층을 따릅니다:

1. 우선순위 1: 정확성. 시스템은 추출된 최종 답변이 정답 ($y^*$) 과 일치하는 롤아웃을 필터링합니다.
2. 우선순위 2: 학생 정렬. 여러 개의 올바른 롤아웃이 존재하는 경우, BRTS 는 학생의 상위-$K$ 후보 집합과 토큰 수준의 중첩이 가장 높은 것을 선택합니다. 이는 선택된 궤적이 학생의 현재 추론 분포와 호환되도록 보장합니다.
3. 정답 유도 복구 (Tier-2). 모든 무조건부 교사 샘플이 실패하는 경우 (즉, 올바른 것이 하나도 없는 경우), BRTS 는 복구 단계를 호출합니다. 이는 정답을 검증 신호로 조용히 주입하면서 교사가 자연스러운 유도를 생성하도록 지시하는 수정된 프롬프트를 구성합니다. 새로운 롤아웃을 샘플링하고, 만약 올바르면 이를 선택합니다. 이 또한 실패하면, 시스템은 학생과의 중첩이 가장 높은 무조건부 롤아웃으로 대체합니다.

2.2 교사 컨텍스트 감독

학습 목적 함수는 두 가지 손실 구성 요소를 결합합니다:

• 학생 컨텍스트 손실 ($L_{stu-ctx}$): 학생이 생성한 접두어 조건 하에서 학생과 교사 분포 간의 역 KL 발산을 최소화하는 표준 OPD 목적 함수입니다. 이는 학생이 실제로 방문하는 상태에 대해 학생을 수정합니다.
• 교사 컨텍스트 손실 ($L_{tea-ctx}$): 선택된 교사 궤적 ($y'$) 에 대해 계산된 보조 손실입니다. 여기서 학생과 교사 분포 모두 교사의 접두어 조건 하에 있습니다. 이는 학생을 현재 고확률 영역 바깥에 있을 수 있는 일관되고 올바른 추론 경로에 노출시킵니다.

총 손실은 다음과 같이 정의됩니다:
$$L_{total} = L_{stu-ctx} + \lambda L_{tea-ctx}$$
여기서 $\lambda$는 가중치 계수이며 (경험적으로 10 으로 설정됨) 입니다.

2.3 상위-K 후보 방향

이러한 손실을 효율적으로 구현하기 위해 BRTS 는 상위-$K$ 집계를 사용합니다:

• 학생 컨텍스트 분기의 경우, 후보 토큰 집합은 학생 접두어 하에서 학생의 상위-$K$ 토큰에서 파생됩니다.
• 교사 컨텍스트 분기의 경우, 후보 집합은 선택된 교사 접두어 하에서 교사의 상위-$K$ 토큰에서 파생됩니다. 이를 통해 보조 분기는 학생이 현재 고려하지 않을 수 있는 교사가 선호하는 토큰을 감독하면서도 컴팩트한 후보 집합을 유지할 수 있습니다.

3. 주요 기여

논문은 다음과 같이 기여 사항을 요약합니다:

1. 높은 분산의 감독 신호 식별: 저자들은 OPD 의 표준 학생 컨텍스트 감독이 노이즈가 있는 학생 접두어와 확률적 교사 롤아웃에 의해 제한받으며, 이로 인해 어려운 프롬프트에서 교사 신호가 신뢰할 수 없게 된다고 강조합니다.
2. BRTS 도입: 신뢰할 수 있는 교사 궤적을 선택하고 증류하기 위한 경량 프레임워크입니다. 이는 "정확성 우선, 정렬 차선" 선택 규칙을 적용하며, 실패한 무조건부 샘플을 위한 정답 유도 대체 수단을 포함합니다.
3. 경험적 개선: BRTS 는 어려운 추론 벤치마크에서 표준 OPD 를 크게 능가합니다. 논문은 교사 후보의 품질, 복구 범위, 백본 선택이 성능에 어떻게 영향을 미치는지 보여주는 진단 분석을 제공합니다.

4. 실험 결과

실험은 1.5B 학생 모델 (DeepSeek-1.5B) 과 다양한 교사 구성 (JustRL-1.5B 및 DeepSeek-R1-Distill-Qwen-7B) 을 사용하여 AIME 2024, AIME 2025, AMC 2023에서 수행되었습니다.

• 성능 향상: BRTS 는 학생 전용 베이스라인보다 일관되게 개선된 성과를 보입니다. 이러한 이득은 더 어려운 AIME 2025 데이터셋에서 가장 두드러집니다. 예를 들어, AIME 2024 에서 네 개의 교사 후보를 사용하면 평균 정확도가 0.3917 인 베이스라인 대비 0.4000 을 달성했습니다.
• Tier-2 복구의 영향: 정답 유도 복구 단계 (Tier-2) 는 무조건부 샘플링이 실패하는 어려운 프롬프트에서 상당한 개선을 제공합니다. AIME 2025 에서 Tier-2 를 추가하면 다수결 투표 정확도가 Tier-1 만 사용할 때의 0.2748 에서 0.3133 으로 증가했습니다.
• 초기 학습 피크: BRTS 는 베이스라인에 비해 초기 학습 단계에서 더 높은 성능 피크를 달성하여, 더 빠른 수렴과 더 나은 컴퓨팅 - 성능 트레이드오프를 시사합니다.
• 강건성: 이 방법은 교사 백본을 교체할 때 (예: 1.5B 학생에게 7B 교사 사용) 도 효과적이며, 약간의 프롬프트 교란에도 취약하지 않으며, 이는 실제로 교사 샘플의 상관관계를 제거하는 데 도움이 됩니다.
• 계산 비용: 오버헤드는 적당합니다. 8×B200 GPU 에서 교사 후보 풀을 1 개에서 2 개로 늘리면 학습 단계당 약 59 초가 추가되는데, 이는 성능 향상을 고려할 때 수용 가능한 수준으로 간주됩니다.

5. 중요성과 주장

논문은 BRTS 가 어려운 작업에서 확률적 교사 신호의 신뢰성 부재라는 온-폴리시 증류의 중요한 병목 현상을 해결한다고 주장합니다. 교사 감독을 정확성 인식(오류 필터링) 과 정렬 인식(학생 분포 매칭) 이 모두 가능하도록 구조화함으로써, BRTS 는 확률적 교사 롤아웃을 구조화된 고품질 감독으로 변환합니다.

저자들은 BRTS 를 OPD 를 대체하는 것이 아니라, 학생들이 표준 감독 미세 조정과 관련된 종종 치명적인 망각이나 표준 OPD 의 높은 분산을 치르지 않고도 완전하고 신뢰할 수 있는 추론 경로에서 학습할 수 있게 하는 필수적인 보완재로 위치시킵니다. 이 방법은 검증 가능한 정답이 있는 도메인에서 추론 능력을 향상시키는 확장 가능하고 경량화된 확장으로 제시됩니다.