[논문 리뷰] On-Policy Context Distillation for Language Models

Tianzhu Ye, Li Dong외 3명|arXiv (Cornell University)|2026. 02. 12.

Topic Modeling인용 수 0

한 줄 요약

OPCD는 컨텍스트-조건부 교사를 온-정책 샘플링과 역 KL로 모방하도록 학생을 학습시키며, 수학, 게임 및 도메인 작업 전반에 걸친 체험 지식과 시스템 프롬프트의 내재화를 가능하게 하고, 오프-정책 컨텍스트 증류를 능가합니다.

ABSTRACT

Context distillation enables language models to internalize in-context knowledge into their parameters. In our work, we propose On-Policy Context Distillation (OPCD), a framework that bridges on-policy distillation with context distillation by training a student model on its own generated trajectories while minimizing reverse Kullback-Leibler divergence against a context-conditioned teacher. We demonstrate the effectiveness of OPCD on two important applications: experiential knowledge distillation, where models extract and consolidate transferable knowledge from their historical solution traces, and system prompt distillation, where models internalize beneficial behaviors encoded in optimized prompts. Across mathematical reasoning, text-based games, and domain-specific tasks, OPCD consistently outperforms baseline methods, achieving higher task accuracy while better preserving out-of-distribution capabilities. We further show that OPCD enables effective cross-size distillation, where smaller student models can internalize experiential knowledge from larger teachers.

연구 동기 및 목표

• 오프-정책 컨텍스트 증류의 한계(노출 편향 및 모드 문제)를 해결한다.
• 컨텍스트-조건부 교사에 대항해 학생의 자체 궤적에서 학습하도록 On-Policy Context Distillation(OPCD)을 제안한다.
• 수학, 게임 및 도메인 과제에 걸쳐 체험 지식 증류와 시스템 프롬프트 증류에 대해 OPCD를 입증한다.
• OPCD가 작은 모델이 더 큰 교사로부터 학습하고 망각을 줄이는 크로스-사이즈 증류를 지원함을 보인다.

제안 방법

• 온-정책 샘플링을 사용하여 학생과 맥락 인식 교사 간의 역 KL 발산을 최소화한다.
• 탑-k 토큰 근사를 통해 토큰 수준 D_KL을 계산하고 모드 추구 행동을 촉진한다.
• 맥락 없이 학생이 응답을 생성하도록 학습시키고, 이후 맥락에 조건화된 교사의 분포에 맞춘다.
• 유연한 교사 구성 허용(동결된 교사를 가진 교사-학생 또는 가중치를 공유하는 자기 증류).
• 수학 문제, 텍스트 기반 게임, 의학 및 안전 프롬프트를 포함한 체험 지식 및 시스템 프롬프트 증류 과제에서 오프-정책 컨텍스트 증류 기준선과 비교하여 평가한다.

실험 결과

연구 질문

• RQ1 온-정책 컨텍스트 증류가 일시적 맥락 지식을 모델 매개변수에 내재화할 수 있는가?
• RQ2 OPCD가 도메인 전반에 걸친 체험 지식 응집 및 시스템 프롬프트 증류를 개선하는가?
• RQ3 더 크고 어쩌면 동결된 교사로부터 OPCD를 통해 작은 학생 모델이 혜택을 얻을 수 있는가?
• RQ4 오프-정책 방법과 비교하여 OPCD가 분포 밖 작업에서의 망각을 완화하는가?

주요 결과

• OPCD가 수학 문제와 텍스트 기반 게임에서 테스트 정확도 면에서 오프-정책 컨텍스트 증류를 능가한다.
• OPCD는 분포 밖(OOD) 성능을 더 잘 내고 분포 내 정확도는 유지한다.
• OPCD를 활용한 시스템 프롬프트 증류가 오프-정책 기준선보다 의학 및 안전 작업 정확도가 더 높다.
• OPCD는 효과적인 크로스-사이즈 증류를 가능하게 하며, 작은 모델이 큰 동결된 교사로부터 혜택을 받는다.
• 온-정책 학습은 더 안정적인 개선을 제공하고 OOD 데이터에서의 망각을 오프-정책 방법과 비교하여 완화한다.