[Paper] 언어 모델을 위한 On-Policy Context Distillation

Source: arXiv - 2602.12275v1

개요

이 논문은 On‑Policy Context Distillation (OPCD) 를 소개한다. 이는 언어 모델이 일반적으로 프롬프트에서만 보는 유용한 지식을 “기억”하도록 가르치는 새로운 방법이다. 모델이 자체 생성 출력으로부터 학습하도록 하면서, 더 풍부한 컨텍스트에 접근할 수 있는 교사의 지도를 받게 함으로써, OPCD는 규모가 작거나 훈련이 덜 된 모델도 사실적 지식과 절차적 노하우를 내부화하게 하면서도 새로운 입력을 처리하는 능력을 희생하지 않는다.

Key Contributions

• On‑Policy Distillation for LMs – 고전적인 온‑폴리시 강화학습 증류 아이디어를 컨텍스트 기반 교육과 결합하여, 정적 데이터셋이 아니라 학생이 자신의 궤적에서 학습하도록 합니다.
• Reverse KL Objective – 역 Kullback‑Leibler 손실을 사용해 학생의 분포를 컨텍스트 조건화된 교사와 정렬시켜, 학생이 교사의 “사고 과정”을 채택하도록 장려합니다.
• Experiential Knowledge Distillation – 모델이 자체 과거 해결 과정(예: 이전 수학 단계, 게임 움직임)에서 재사용 가능한 지식을 추출하고 통합할 수 있음을 보여줍니다.
• System Prompt Distillation – 최적화된 프롬프트(종종 수작업으로 만들거나 프롬프트 엔지니어링을 통해 발견)를 모델 가중치에 내장할 수 있음을 입증하여, 추론 시 외부 프롬프트가 필요 없게 합니다.
• Cross‑Size Distillation – 작은 학생 모델이 훨씬 큰 교사 모델로부터 경험적 지식을 물려받을 수 있음을 검증하여, 효율적인 모델 배포를 가능하게 합니다.
• Broad Empirical Coverage – 수학적 추론, 텍스트 기반 게임, 도메인 특화 작업 전반에 걸친 벤치마크에서 일관되게 강력한 베이스라인을 능가하면서도 OOD(분포 외) 성능을 유지합니다.

Methodology

1. Teacher & Student Setup – 교사 모델은 전체 컨텍스트(예: 프롬프트와 외부 지식)를 받아 다음 토큰에 대한 확률 분포를 생성합니다. 학생 모델은 프롬프트만 보고(추가 컨텍스트 없이) 작동합니다.
2. On‑Policy Trajectory Generation – 학생은 학습 데이터에서 자신의 출력 시퀀스(“정책”)를 샘플링합니다. 이 자체 생성된 궤적이 학습 예제가 됩니다.
3. Reverse KL Distillation – 학생이 생성한 각 토큰에 대해 손실은 역 KL 발산 KL(teacher || student)입니다. 이는 교사가 가능하다고 판단한 토큰에 대한 확률 질량을 학생이 늘리도록 하여, 누락된 컨텍스트 상황에서 교사의 추론을 모방하도록 가르칩니다.
4. Iterative Refinement – 이 과정이 반복됩니다: 학생이 개선되고 더 나은 궤적을 생성하면, 교사(고정되거나 천천히 업데이트됨)는 계속해서 컨텍스트 지침을 제공합니다.
5. Applications –
   • Experiential Knowledge: 교사는 자신의 과거 해결 과정(trace)에 접근할 수 있는 모델 버전이며, 학생은 그 트레이스를 파라미터에 내재하도록 학습합니다.
   • System Prompt: 교사는 바람직한 행동을 이끌어내는 설계된 프롬프트로 구동되며, 학생은 프롬프트 없이도 그 행동을 재현하도록 학습합니다.

전체 파이프라인은 가볍습니다: 교사와 학생에 대한 순전파만 필요하고 외부 보상 모델이 필요 없으며, 일반 GPU 클러스터에서 실행할 수 있습니다.

결과 및 발견

• 크로스‑사이즈 증류: 13B 교사 모델에서 증류된 1.3B 학생 모델은 수학 벤치마크에서 교사 성능의 75%를 달성했으며, 일반 1.3B 모델은 62%에 머물렀다.
• 프롬프트‑프리 추론: 시스템 프롬프트 증류 후, 학생 모델은 런타임에 프롬프트 없이도 교사의 프롬프트 강화 성능과 일치했으며, 추론 지연 시간을 약 30% 감소시켰다.
• OOD 견고성: 공격적인 파인튜닝과 달리, OPCD는 모델이 관련 없는 질의에 답할 수 있는 능력을 유지했으며, 증류된 지식이 기존 능력을 덮어쓰지 않고 통합됨을 확인했다.

Practical Implications

• Smaller Deployments: 소형 배포: 기업은 여전히 더 크고 비용이 많이 드는 시스템의 “경험”을 담은 컴팩트한 모델을 제공할 수 있습니다—엣지 디바이스, 모바일 앱, 혹은 비용에 민감한 SaaS에 유용합니다.
• Prompt‑Engineering Savings: 프롬프트 엔지니어링 비용 절감: 고성능 프롬프트가 발견되면(대개 비용이 많이 드는 RLHF나 수동 튜닝을 통해), OPCD는 그 행동을 모델에 내재시켜 실행 시 프롬프트 처리를 없애고 지연 시간을 줄입니다.
• Continuous Learning Pipelines: 지속 학습 파이프라인: 팀은 프로덕션 모델이 자체 솔루션 트레이스(예: 버그 수정 제안, 코드 완성)를 기록하도록 하고, 주기적으로 OPCD를 실행해 성공적인 패턴을 내재화함으로써 외부 데이터 정제 없이도 자체 개선 루프를 만들 수 있습니다.
• Domain Adaptation: 도메인 적응: 규제 산업(금융, 헬스케어, 법률)에서는 OPCD가 기본 모델의 일반 언어 능력을 유지하면서도 독점 지식 베이스를 모델에 삽입하는 방법을 제공합니다.
• Simplified Inference Stack: 단순화된 추론 스택: 외부 컨텍스트(프롬프트, 검색 모듈)의 필요성을 없애면서 OPCD는 추론 아키텍처를 간소화하여 확장성과 모니터링을 용이하게 합니다.

제한 사항 및 향후 작업

• 교사 의존성: 증류된 지식의 품질은 교사의 컨텍스트 처리 능력에 좌우됩니다; 설계가 부실한 프롬프트나 잡음이 섞인 과거 추적은 오류를 전파할 수 있습니다.
• 계산 오버헤드: 대규모 데이터셋에 대해 온‑폴리시 트래젝터리를 생성하는 비용이 많이 들 수 있지만, 전체 RLHF 파이프라인보다 여전히 저렴합니다.
• 지식 전이 범위: OPCD는 절차적이거나 프롬프트‑구동 행동에 뛰어나지만, 외부 기반이 필요한 고도로 사실적인 백과사전식 지식에는 한계가 있을 수 있습니다.
• 향후 방향: 저자들은 다중 교사 앙상블, 보존 대 획득을 균형 맞추는 적응형 KL 가중치, 그리고 증류 가능한 지식 범위를 넓히기 위한 검색‑증강 생성 통합을 탐구할 것을 제안합니다.

핵심 요약: 온‑폴리시 컨텍스트 증류는 프롬프트의 유연성과 컴팩트하고 자체 포함된 모델의 효율성 사이에 실용적인 다리를 제공하여, 언어‑모델 서비스에 전문성을 직접 삽입하려는 개발자에게 매력적인 도구가 됩니다.

저자

• Tianzhu Ye
• Li Dong
• Xun Wu
• Shaohan Huang
• Furu Wei

논문 정보

• arXiv ID: 2602.12275v1
• Categories: cs.CL
• Published: 2026년 2월 12일
• PDF: PDF 다운로드