ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation

Structured pruning is a hardware-friendly way to compress LLMs, but it is mostly validated on multiple-choice recognition tasks, while the same compressed checkpoints can collapse on the free-form generation that deployment actually requires. Two observations trace this gap. First, greedy \textsc{pass}@$1$ nearly vanishes after compression, yet \textsc{pass}@$k$ recovers substantially under repeated sampling: useful generations are demoted, not erased. Second, the recoverable regime fails mainly through suffix repetition. Recovery should therefore train on the compressed model's own on-policy states with dense token-level supervision, which On-Policy Distillation (OPD) provides by reusing the pre-compression model as a frozen teacher. However, long on-policy rollouts spend early recovery budget on low-information repetitive suffixes, delaying loss descent. To mitigate this waste, we propose \textbf{\shortopd}, a short-to-long OPD schedule that detects teacher-confirmed repetitive suffixes, treats the surviving prefix as each rollout's effective length, and allocates future rollout budgets to the effective lengths the policy can currently use. Across math, code, and open-ended generation, \shortopd\ raises the compressed model's score to about $9\times$ its unrecovered value and $1.6$--$4.4\times$ standard recovery recipes (SFT w/o KD, KD, and SeqKD), and it matches a fixed $8192$-token rollout horizon within two points using a quarter of the training time ($8.5$ vs.\ $35.9$ hours) and $71\%$ fewer rollout tokens. We hope this recipe helps move structured pruning beyond marginal gains on perplexity and multiple-choice benchmarks, a step closer to deployment-ready generation quality.

한 줄 요약

ShortOPD는 압축된 LLM의 생성 성능을 회복하기 위해 반복 후缀을 피하고 짧은 생성부터 점진적으로 길이를 늘리는 on-policy distillation 기법이다.

핵심 기여도

ShortOPD는 압축된 Qwen3-4B-Instruct 모델의 생성 점수를 약 9배 향상시킴.
기존 회복 방법(SFT w/o KD, KD, SeqKD) 대비 1.6–4.4배 성능 개선.
8192 토큰의 고정 롤아웃 헤ORIZON을 2점 이내로 달성하면서 71% 적은 토큰 수로 학습.
8.5시간에 학습 완료 (기존 35.9시간 대비 4분의 1 소요).

핵심 아이디어

기존 압축된 LLM은 생성 과정에서 반복적인 후缀을 생성하며, 유용한 정보는 상대적으로 뒤로 밀려난다. 이는 단순히 모델 내부에 잠재된 능력이 아닌, 생성 과정에서의 분포 문제이다. 따라서, 압축된 모델이 생성한 on-policy 상태에 대해 밀집된 토큰 수준의 지도가 필요하다. 이에 따라, ShortOPD는 압축 전의 동일 모델을 고정된 teacher로 사용하는 on-policy distillation(OPD)을 기반으로, 반복 후缀을 탐지하고 짧은 생성부터 점진적으로 길이를 늘리는 방식으로 학습 자원을 효율적으로 할당한다. 핵심 아이디어는 반복된 후缀은 정보가 적기 때문에 초기 학습 단계에서 이를 피하고, 유의미한 prefix에 집중하는 것이다.

기술적 접근법

**OPD**: 압축된 모델이 자신의 분포에서 샘플링한 rollout을 teacher(압축 전 모델)와 비교하며, 토큰 단위로 분포를 맞춘다.
**ShortOPD**: 반복 후缀을 탐지하는 컨트롤러를 도입. 반복이 높으면 rollout 길이를 줄이고, 반복이 줄어들면 길이를 늘린다.
**EMA 기반 조절**: 반복과 truncation을 기반으로 EMA(지수 이동 평균)를 사용해 rollout budget을 조정.
**데이터셋**: 45,447개의 math/code/instruction prompt로 학습.
**롤아웃 헤ORIZON**: 최대 8192 토큰, ShortOPD는 이와 유사한 성능을 71% 적은 토큰 수로 달성.

주요 결과

**GSM8K**: pass@k=64에서 91% (teacher greedy: 88%)
**전체 8개 태스크**: ShortOPD는 압축되지 않은 teacher의 2/3 수준의 생성 점수를 회복.
**기존 방법 대비 성능**: SFT w/o KD, KD, SeqKD 대비 1.6–4.4배 향상.
**학습 효율성**: 8192 토큰 롤아웃 헤ORIZON과 동일한 성능을 71% 적은 토큰 수로 달성.
**학습 시간**: 8.5시간 (기존 35.9시간 대비 4분의 1).

의의 및 한계

**의의**: 압축된 LLM의 생성 성능 회복을 위한 on-policy distillation 기법을 제안하며, 기존 압축 평가가 recognition 중심이었다는 한계를 보완.
**실용적 가치**: 압축된 모델을 실제 배포 환경에 적용 가능하게 만드는 기술적 기반 제공.
**한계**: 반복 후缀이 아닌 다른 유형의 생성 오류(예: 문법 오류, 논리적 불일치)에 대한 회복은 명시되지 않음. 또한, 특정 압축 기법(Block-Influence depth pruning)에 기반한 연구이므로 다른 압축 방법과의 호환성은 추가 연구 필요.

실용적 활용

**LLM 압축 후 배포**: ShortOPD는 압축된 LLM을 실제 생성 작업에 안정적으로 사용할 수 있도록 도와준다.
**하드웨어 제약 환경**: 압축된 모델을 유지하면서도 생성 성능을 회복할 수 있어, 자원 제한된 장치에서 유용.
**대규모 모델 서빙**: 서빙 비용을 줄이면서도 사용자 질의에 대한 생성 품질을 유지할 수 있는 기술적 기반 제공.