수학 문제 풀이 AI, 근처 파라미터로 학습 더 정확해진다
원제 Better Supervision Is Nearby: Neighborhood On-Policy Self-Distillation
추천 10댓글 2
Key Point
기존 자기증류 방식의 고정된 감독 전략을 파라미터 근처 탐색으로 개선해 AIME 수학 올림피아드 벤치마크에서 2~3포인트 일관된 성능 향상을 달성했으므로, 수학 문제 풀이 AI의 정확도 향상이 어떻게 이루어지는지 실제로 확인할 수 있다.
핵심 요약
- 수학 추론 모델을 학습시키는 온정책 자기증류(OPSD)는 기준 해답을 보는 특권화된 교사가 학생이 샘플한 중간 단계를 감독한다.
- 기존 OPSD는 매 상태마다 고정된 파라미터 설정 하나만 사용했으나, 근처의 파라미터 변동이 같은 기준 맥락에서 추가 감독을 제공한다.
- 지역 파라미터 섭동은 기준과 정렬된 상이한 정정(correction)을 드러내며, 다양한 전문가들이 기준 위치의 여러 지점에서 이를 제공한다.
- 저자들은 이 정정을 학생 방문 상태의 감독으로 변환하는 근처-OPSD(N-OPSD)를 제안했다.
- 오프라인에서 탐욕적 선택이 각 위치에서 현재 최고 성능 이상의 기준 토큰 이득으로 보상받는 고정 전문가 풀을 구축한다.
- 온라인 라우팅에서 최고봉(MaxPeak) 메커니즘이 앵커 토큰을 선택하고, 분위수 선택으로 그와 일치하는 전문가를 고르며, 학생은 클리핑된 정방향 KL 목적으로 해당 전문가의 완전한 다음 토큰 분포를 학습한다.
- AIME 2024, AIME 2025, HMMT 2월 2025 벤치마크에서 N-OPSD는 기존 OPSD 대비 Qwen3의 1.7B, 4B, 8B 모델에서 각각 2.75, 1.67, 1.94포인트 향상을 보였다.
- 학생 프리픽스 연속으로 기준 궤적 외에도 풀 활용이 가능하며, 필터링된 기준 토큰 이득이 선택 기준임을 검증했고, 풀 내 겹침 고려와 상태별 라우팅으로 추가 개선된다.
- 추론 시에는 증류된 학생 모델만 사용하므로 추가 계산 비용이 없다.