ReNIO: LLM 온-폴리시 증류를 위한 부정적 궤적 중요도 재가중치

학생이 생성한 잘못된 추론 궤적(incorrect SGO)이 모델의 탐색적·수정적 행동을 보존해 학습에 더 유용한 경우가 존재한다. ReNIO는 최종 정답을 보지 않고도 접두사 수준의 학생·교사 확률비를 이용해 그러한 부정적 궤적에 자동으로 높은 가중치를 부여해 OPD/OPSD 성능을 개선한다.

출발점과 기존 한계: LLM 후처리(즉, reasoning 성능 향상)에서 On-policy Distillation(OPD)은 학생 정책으로 샘플한 출력(SGO)에 대해 토큰 단위로 교사 분포를 맞추어 학습한다. 이 방식은 전체 롤아웃의 최종 정답을 기다리는 강화학습보다 더 촘촘한 지도신호를 제공하고, 짧은 접두사만으로도 학습 가능해 비용이 낮다. 그러나 표준 OPD는 모든 SGO에 동일 가중치를 부여하므로, 정보량이 매우 다른 궤적들을 구별하지 못한다.

해결 원리: ReNIO는 토큰 수준의 학생·교사 확률비 r_t=π_S(y_t|x,y_<t)/π_T(y_t|x,y_<t)의 로그 ℓ_t=log r_t를 사용해 '학생이 강하게 선택하고 교사가 낮게 평가한' 국지적 분기 결정을 식별한다. 핵심 전제는 잘못된 궤적들이 몇몇 '핵심 토큰'에서 학생의 과신으로 인해 잘못된 분기로 접어들고, 이런 토큰들이 유효한 교정 신호를 담고 있다는 점이다.

핵심 작동 방식: ReNIO는 세 단계로 동작한다. (1) 각 SGO의 접두사에서 토큰별 로그비 ℓ_t를 계산한다(입력: π_S 및 π_T의 prefix-conditioned probabilities → 연산: 로그비 계산 → 결과: 토큰 수준 불일치 신호). (2) ℓ_t 분포에서 임계값 τ를 넘는 토큰들을 핵심 토큰 집합 K(x,y)으로 선택한다(선택 이유: 분포는 긴 꼬리 형태로 대부분의 토큰은 ℓ_t≈0이다). (3) 선택된 토큰들의 로그비를 평균(로그 평균)해 지수화한 값을 샘플 가중치 w(x,y)로 계산하고, 배치 평균으로 나누어 ĥw(x,y)=w/ȳ_B로 정규화한다(의미: 배치 내 평균 가중치를 1로 유지해 전체 그레이디언트 스케일 보존).

전체 구조: ReNIO는 기존 OPD/OPSD 파이프라인을 수정하지 않고 샘플 수준 가중치 모듈만 도입한다. 입력으로는 프롬프트 x와 학생이 생성한 접두사 y_<t에서의 학생·교사 prefix-conditioned distributions {p_S^t,p_T^t}가 들어온다. 이 분포들의 토큰 로그확률을 사용해 ℓ_t를 계산하고, 임계값 τ를 넘는 t만 골라 기하평균 기반의 w(x,y)를 얻는다. 배치 내 평균으로 나누어 ĥw(x,y)를 만들고, 최종 목적함수는 ĥw(x,y)·𝒟(π_S(y|x)‖π_T(y|x))의 기댓값이다.

논문에서 명시한 한계: 하드웨어 제약으로 대규모 모델(더 큰 파라미터 규모)에 대한 효과 검증이 수행되지 않았다.