기술 요약: 온도 스케일링 온-폴리시 셀프 디스틸레이션 (TS-OPSD)

1. 문제 정의

검증 가능한 보상을 활용한 강화 학습(RLVR)은 대규모 언어 모델(LLM), 특히 수학 및 논리 영역에서의 추론 능력을 크게 향고시켰습니다. 그러나 이러한 진보는 **엔트로피 붕괴(entropy collapse)**라고 불리는 결정적인 병목 현상에 직면해 있습니다. 정책 최적화가 진행됨에 따라 모델의 출력 분포는 점점 더 집중(sharp)되어 다음과 같은 현상이 발생합니다:

• 롤아웃 다양성 감소: 모델이 반복적인 궤적을 생성하여 솔루션 공간의 탐색을 제한합니다.
• 학습 신호 약화: 효과적인 그래디언트 정보가 결정적인 추론 분기점에서의 소수 고엔트로피 토큰에 집중됩니다. 정책이 붕괴되면 이러한 신호는 사라져 추가적인 개선을 저해합니다.

기존의 해결책들은 다음과 같은 방식으로 이를 완화하려고 시도합니다:

1. RL 목적 함수 수정: 엔트로피 정규화나 적응형 클리핑을 사용합니다. 이러한 방법은 보상 최대화와 상충할 수 있는 외부 제약을 부과합니다.
2. 샘플링 온도 조정: 롤아웃 수집 중에 온도를 높입니다. 이는 행동 정책(behavior policy)을 변화시키지만, 모델의 근본적인 파라미터화된 분포를 변경하지는 않으므로, 행동 정책과 최적화된 정책 사이에 미스매치를 생성합니다.

이러한 접근 방식은 모델 파라미터에 대해 외부적인 상태로 남아 있습니다. 본 논문은 온도의 탐색적 효과가 지속적인 RL을 위한 더 나은 초기값을 제공할 수 있도록 모델 가중치에 직접 **내재화(internalized)**되어야 한다고 주장합니다.

2. 방법론: TS-OPSD

저자들은 외부 교사, 특권 데이터, 또는 추가적인 추론 비용 없이 엔트로피를 복구하도록 설계된 경량화된 "정책 재가열(policy reheating)" 방법인 **온도 스케일링 온-폴리시 셀프 디스틸레이션(TS-OPSD)**을 제안합니다.

핵심 메커니즘

TS-PODS는 RL 단계 사이의 개입으로서 작동합니다. 엔트로피가 붕괴된 체크포인트($\theta_{RL-collapse}$)에서 시작하여, 이 방법은 **셀프-티처(self-teacher)**를 구축하고 더 부드러운 분포를 학생 모델($\theta_{reheat}$)로 증류합니다.

1. 셀프-티처 구축: 학생의 현재 로짓 $z_\theta$가 주어졌을 때, 고온 소프트맥스 변환($T > 1$)을 로짓에 적용하여 셀프-티처 분포 $p_T$를 생성합니다. 결정적으로, 티처 로짓은 스톱 그래디언트(stop-gradient) 처리되며 외부 모델은 사용되지 않습니다.
   $$p_T(v | x_{<t}) = \frac{\exp(z_{\theta,v}/T)}{\sum_u \exp(z_{\theta,u}/T)}$$
2. 온-폴리시 디스틸레이션: 학생은 $q_\theta$ 자체에서 샘플링된 온-폴리시 롤아웃을 사용하여, 자신의 현재 분포 $q_\theta$와 고온 셀프-티처 $p_T$ 사이의 쿨백-라이블러(KL) 발산을 최소화하도록 훈련됩니다.
   $$\mathcal{L} = D_{KL}(p_T \parallel q_\theta) \quad \text{(Forward KL)}$$
   $$\mathcal{L} = D_{KL}(q_\theta \parallel p_T) \quad \text{(Reverse KL)}$$

구조적 특성

• 순위 보존(Rank Preservation): $p_T$는 $q_\theta$의 단조 증가 멱함수 변환($p_T(v) \propto q_\theta(v)^{1/T}$)이므로, 토큰 순위가 보존됩니다. 티처와 학생은 동일한 로컬 서포트(local support)와 의미적 선호 순서를 공유하며, 오직 분포의 날카로움(sharpness)만 다릅니다.
• 내재화(Internalization): 표준 디코딩 온도 조정과 달리, TS-OPSD는 모델 가중치를 업데이트하여 더 부드럽고 높은 엔트로피를 가진 행동을 영구적으로 인코딩합니다. 이는 훈련-추론 간의 미스매치를 제거합니다.
• 그래디언트 역학:
  • Forward KL (FKL): 그래디언트는 $q_\theta(v) - p_T(v)$ 형태를 취합니다. 이는 모델의 자체적인 사전 판단을 바탕으로 억제된 대안들의 가중치를 재조정하는 "구조적 평활화(structured smoothing)"를 수행합니다.
  • Reverse KL (RKL): 셀프-온도 결합 하에서, RKL을 최소화하는 것은 국소적으로 학생의 엔트로피를 최대화하는 것과 동등하며, 이때 온도는 그래디언트 크기에 대한 스칼라 승수로만 작용합니다.

3. 주요 기여

1. 정책 재가열 프레임워크: 엔트로피 붕괴 이후의 개입으로서, 온도의 탐색적 효과를 모델 파라미터에 내재화하는 TS-OPSD를 도입했습니다.
2. 이론적 통찰: 셀프-온도 결합 하에서 티처-학생 간의 미스매치가 최소화(순위 보존)되며, RKL이 직접적인 엔트로피 최대화 업데이트로 감소함을 입증했습니다.
3. 실험적 검증: 엔트로피가 붕괴된 체크포인트를 재가열하는 것이 표준적인 지속 훈련이나 롤아웃 레벨의 온도 조정보다 지속적인 RL을 위한 더 우수한 초기값을 제공함을 보여주었습니다.
4. 역학 분석: 엔트로피 회복이 두 단계(느린 평활화 후 급격한 불연속적 점프)로 일어나는 "엔트로피 점프" 현상을 규명하고, 이러한 파라미터 업데이트의 가역성을 분석했습니다.

4. 실험 결과

실험은 Qwen3-4B-Base 및 Qwen3-8B-Base를 대상으로 DAPO-Math-17K 데이터셋을 사용하여 수행되었으며, MATH-500, AIME24/25/26, OlymMATH에서 평가되었습니다.

• 성능 향상: TS-OPSD(특히 PR-FKL)는 표준 DAPO 및 롤아웃 재가열(DAPO w/ RR) 모두를 일관되게 능가했습니다.
  • Qwen3-8B-Base에서 PR-FKL은 표준 DAPO의 51.60, 롤아웃 재가열의 52.66과 비교하여 Avg pass@8 54.74를 달성했습니다.
  • 성능 향상은 어려운 벤치마크(AIME, OlymMATH)에서 가장 두드러졌습니다.
• 엔트로피 회복: 재가열 후, 모델은 더 높은 엔트로피 초기값으로부터 RL을 재개하여 억제되어 있던 표준 DAPO보다 더 높은 최종 보상을 달성했습니다.
• 롤아웃 재가열과의 비교: 롤아웃 중 샘플링 온도를 높이는 것(RR)이 유의미한 이득을 주었으나, 정책 재가열(PR)이 더 우수한 결과를 냈습니다. 이는 이점이 단순히 노이즈가 섞인 샘플링에서 오는 것이 아니라, 더 부드러운 정책 분포 자체가 최적화를 위한 더 나은 기반으로 작용하기 때문임을 시사합니다.
• 능력 보존: 분석 결과 TS-OPSD는 **서포트 보존 확률 재분배(support-preserving probability redistribution)**를 수행함을 확인했습니다:
  • 중간 표현(Intermediate Representations): 베이스 모델과 재가열된 모델 간의 선형 CKA 점수가 1.0에 가깝게 유지되어, 내부 표현이 거의 변하지 않았음을 나타냈습니다.
  • 추론 능력: MATH-500 및 GPQA에 대한 태스크 정확도가 디스틸레이션 과정 동안 안정적으로 유지되어, 엔트로피 회복이 추론 기술을 저하시키지 않았음을 확인했습니다.
  • 후보 집합(Candidate Sets): Top-10 토큰 중첩 및 확률 질량 보유율이 1.0에 가까워, 모델의 의미적 선호도는 유지하면서 마진(margin)만 완화되었음을 보여주었습니다.

5. 의의 및 주장

본 논문은 엔트로피 복원이 RL 목적 함수로부터 분리될 수 있으며, 정책이 너무 결정론적이 될 때 삽입할 수 있는 모듈형 중간 단계 연산임을 주장합니다.

• 모듈형 개입: TS-OPSD는 핵심 RL 알고리즘을 변경하지 않고도 엔트로피 붕파 병목 현상을 넘어 추론 지향적 RL을 확장할 수 있는 간단하고 경량화된 메커니즘을 제공합니다.
• 데이터 희소 환경에서의 안정성: 데이터가 줄어든 실험(2048 프롬프트)에서 표준 RL 훈련은 심각한 엔트로피 붕괴로 인해 불안정(그래디언트 스파이크, 보상 저하)해졌습니다. TS-OPSD는 안정화 메커니즘으로 작용하여 더 길고 안정적인 훈련을 가능하게 했습니다.
• 계산 효율성: 재가열 개입은 전체 RL 훈련 시간을 약 15~17% 증가시킵니다. 배치 크기로 정규화했을 때, 샘플당 비용은 표준 RL 단계와 비슷하거나 더 낮습니다.
• 한계점: 저자들은 실험이 특정 모델 규모(4B/8B)에 국한되어 있으며, 재가열 중단 기준(엔트로피 임계값)이 원칙적이기보다는 경험적이라는 점을 언급했습니다. 이들은 향에 연구에서 OPSD와 RL이 교대로 나타나는 계층적 스케줄을 탐구할 것을 제안합니다.

요약하자면, TS-OPSD는 온도를 "내재화"하는 것이 붕괴된 정책의 탐색 능력을 효과적으로 복구하여, 획득된 추론 구조를 보존하면서도 RL이 더 높은 성능 한계에 도달할 수 있도록 함을 입증합니다.