Rethinking Reward Supervision: Rubric-Conditioned Self-Distillation

기술 요약: 보상 감독의 재고: 루브릭 조건부 자기 증류 (RCSD)

문제 정의

추론 언어 모델(LLM)의 사후 학습(Post-training)은 현재 두 가지 패러다임에 크게 의존하고 있으며, 두 방식 모두 미세한 피드백을 제공하는 데 있어 상당한 한계를 가집니다.

1. 검증 가능한 보상을 활용한 강화 학습 (RLVR): GRPO(Group Relative Policy Optimization)와 같은 방법들은 일반적으로 평가 피드백을 희소한 스칼라 신호(예: 정답 여부 또는 단일 점수)로 압축합니다. 이는 검증 가능한 작업에는 효과적이지만, 모델이 왜 특정 궤적이 성공했거나 실패했는지에 대한 이유를 파악하는 것을 방해하여, 모델이 올바른 중간 추론 단계와 잘못된 단계를 구분하지 못하게 만드는 "신용 할당 병목 현상(credit-assignment bottleneck)"을 초래합니다.
2. 온-폴리시 자기 증류 (On-Policy Self-Distillation, OPSD): 이러한 방법들은 희소한 보상을 교사 모델(teacher model)로부터 오는 조밀한 토큰 수준의 가이드로 대체합니다. 그러나 기존 방식들은 교사가 단 하나의 특권적인 참조 궤적(특정 골드 솔루션)에 조건화되어 있습니다. 이는 감독을 특정 추론 경로에 종속시켜, 답을 도출하는 방식이 지나치게 구체화되도록 만들며 다양한 대안적 도출 가능성을 무시할 수 있습니다. 또한, 학생 모델이 이 단일 경로에서 약간만 벗어나더라도, 교사는 국소적인 수정보다는 전역적인 재구성을 권장하게 될 수 있습니다.

또한, 루브릭(구조화된 평가 기준)은 스칼라 점수보다 더 풍부하고 해석 가능한 품질 표현을 제공함에도 불구하고, 이전 연구들은 루브릭의 판단을 단일 스칼라 보상으로 통합함으로써 이러한 구조적 이점을 최적화 과정에서 대부분 버려왔습니다.

방법론: 루브릭 조건부 자기 증류 (RCSD)

저자들은 **루브릭 조건부 자기 증류 (RCSD)**를 제안합니다. 이 프레임워크는 루브릭을 온-폴리시 자기 증류를 위한 특권적이고 구조화된 교사 측 감독(teacher-side supervision)으로 취급합니다. RCSD는 루브릭 피드백을 스칼라로 붕괴시키는 대신, 교사 모델이 루브릭의 기준별 정보를 바탕으로 학생의 자체 샘플링된 궤적에 대해 조밀한 토큰 수준의 가이드를 제공하도록 합니다.

이 프레임워크는 다음과 같은 2단계 파이프라인을 통해 작동합니다.

1단계: 루브릭 생성기 학습

인스턴스별 루브릭을 얻는 데 드는 높은 비용을 해결하기 위해, 저자들은 전용 루브릭 생성기를 학습시킵니다.

• 교사 정책 (Teacher Policy): 적절한 평가 기준을 추론하기 위해 질문($x$)과 참조 답변($y^\star$)을 특권적 입력으로 받습니다.
• 학생 정책 (Student Policy): 질문($x$)만을 관찰하며, 교사의 기준을 근사하는 루브릭($r$)을 생성하는 법을 배웁니다.
• 목표: 학생은 자신의 샘플링된 루브릭 궤적을 따라 교사의 토큰 분포를 일치시키도록 온-폴리시 증류를 통해 학습됩니다. 이를 통해 인스턴스별 기준 생성 과정을 재사용 가능한 모델로 아모티제이션(amortize)합니다.

2단계: 루브릭 조건부 추론

학습된 루브릭 생성기는 추론 모델을 가이드하는 데 사용됩니다.

• 교사 정책 (Teacher Policy): 질문($x$)과 생성된 루브릭($r$)을 특권적 컨텍스트로 받습니다.
• 학생 정책 (Student Policy): 질문($x$)만을 관찰하며 추론 궤적($\hat{y}$)을 생성합니다.
• 목표: 학생은 자신의 분포와 루브릭에 조건화된 교사의 분포 사이의 Kullback-Leibler (KL) 발산을 최소화합니다.
  $$\mathcal{L}_{reason} = \mathbb{E}_{\hat{y} \sim p_S} \left[ \sum_{t=1}^{|\hat{y}|} D_{KL} \left( p_T(\cdot \mid x, r, \hat{y}_{<t}) \parallel p_S(\cdot \mid x, \hat{y}_{<t}) \right) \right]$$
• 핵심 설계: 단일 참조 답변이 아닌 루브릭에 조건화됨으로써, 교사는 기준 인식형 가이드를 제공합니다. 이를 통해 최적화 신호는 단일 점수나 단일 참조 경로로 압축되는 대신, 부분적 정답의 다양한 차원(예: 올바른 단계 이후의 국소적 오류)을 구분할 수 있습니다. 저자들은 Forward KL 발산을 사용하여 학생이 루브릭에 조건화된 교사 분포의 서포트(support)를 커버하도록 유도함으로써, 유효한 추론 경로의 다양성을 보존합니다.

주요 기여

1. 새로운 감독 인터페이스: RCSD는 루브릭을 단순한 채점 도구가 아니라, 최적화 과정 동안 토큰 수준의 학습을 형성하는 구조화된 감독 인터페이스로 재정의합니다.
2. 2단계 파이프라인: 저자들은 태스크별 루브릭을 자동으로 학습하는 확장 가능한 방법(1단계)을 도입하고, 이를 통해 추론을 가이드하는 방법(2단계)을 제시하여, 증류 단계에서의 수동 루브릭 주석이나 고비용의 LLM-as-a-Judge 보상에 대한 의존성을 제거했습니다.
3. 기준 인식형 온-폴리시 학습: 단일 궤적에 조건화되는 표준 OPSD와 달리, RCSD는 기준에 조건화됨으로써, 학생의 자체 샘플링된 궤적 위에서도 특정 국소적 오류를 처벌하면서 올바른 단계는 보존하는 정밀한 신용 할당을 가능하게 합니다.

실험 결과

이 방법은 과학 추론 벤치마크(검증 가능한 작업인 GPQA-DIAMOND, SciBench, PIQA 및 비검증적/개방형 작업인 ResearchQA, RubricHub 포함)를 통해 평가되었습니다.

• 성능: RCSD는 벤치마크 전반에서 가장 높은 평균 점수(70.6)를 달성했습니다. 이는 베이스 모델인 Qwen3-8B보다 4.7포인트 높았으며, GRPO보다 평균 1.0포인트, OPSD보다 0.9포인트 높았습니다.
• 도메인 특이성: 이득은 루브릭 기반 및 과학적 추론 작업(예: ResearchQA에서 +8.2포인트, RubricHub에서 +4.9포인트)에서 가장 두드러졌는데, 이는 스칼라 보상이 품질을 포착하는 데 실패하는 영역입니다.
• 일반화: 이 방법은 의료 벤치마크(MedMCQA, PubMedQA)와 같은 도메인 외 데이터에서도 파괴적 망각 없이 견고한 일반화 능력을 보여주었습니다. 이는 학습된 기준이 일반적인 추론 구조를 포착하고 있음을 시사합니다.
• 절제 연구 (Ablation Studies):
  • 손실 함수 유형: Forward KL이 가장 효과적인 목적 함수로 나타났으며, Reverse KL이나 JS Divergence에 비해 가장 높은 평균 성능과 높은 학생 엔트로피(다양성)를 유지했습니다.
  • 루브릭 출처: 1단계에서 학습된 생성기가 생성한 루브릭을 사용하는 것이 실제 정답 참조 루브릭을 사용하는 것과 경쟁력 있는 결과를 냈으며, 이를 통해 2단계 접근 방식의 확장성을 입증했습니다.
  • 강건성: 일반적이거나 노이즈가 섞인, 혹은 축소된 루브릭 입력이 주어져도 방법론은 효과적으로 작동했으나, 인스턴스별로 학습된 루브릭이 가장 강력한 감독 신호를 제공했습니다.

의의 및 주장

본 논문은 RCSD가 개방형 추론 모델의 사후 학습을 위한 검증기 불필요(verifier-free) 접근 방식을 제공한다고 주장합니다. 최적화 과정 중에 기준 수준의 구조를 보존하고 루브릭을 조밀한 토큰 수준의 가이드로 변환함으로써, 이 방법은 스칼라 보상 최적화의 신용 할당 병목 현상과 참조 조건부 증류의 경로 특정 편향 문제를 해결합니다.

저자들은 RCSD가 자동 검증(정확한 일치 답변 또는 유닛 테스트)이나 스칼라 결과 보상으로 품질을 완전히 포착하기 어려운 작업에 특히 유용하다고 단언합니다. 응답이 어떤 속성을 충족해야 하는지(루브릭을 통해) 학습하게 함으로써, 특정 솔루션 경로를 모방하는 방식(how-to)이 아닌, RCSD는 더 유연하고 견고한 추론 개선을 촉진합니다. 이 연구는 복잡한 추론 작업에서 구조화된 다차원 피드백이 희소한 보상이나 단일 궤적 증류보다 우수한 감독 신호임을 시사합니다.