Source-linked AI summary
ShortOPD: Recovering Pruned LLMs with Short-to-Long On-Policy Distillation
Qingyu Zhang, Qianhao Yuan, Hongyu Lin, Yaojie Lu, Xianpei Han, Le Sun, Ming Xu, Jiarui Li
TL;DR
Structured pruning은 주로 recognition benchmark로 검증되지만, compressed model은 deployment에 필요한 free-form generation에서 자주 실패한다. ShortOPD는 repetition-gated, truncation-aware on-policy distillation을 사용해 rollout horizon을 조정하며, 더 적은 rollout token으로 고정된 short 및 long horizon과 2점 이내의 차이를 보인다.
문제
Structured pruning은 deployment에 필요한 free-form generation 품질에 대한 평가와 recovery 근거가 충분하지 않다.
방법
ShortOPD는 repetition-gated, truncation-aware on-policy distillation을 사용해 usable prefix가 향상됨에 따라 rollout horizon을 short에서 long으로 조정한다.
결과
ShortOPD는 더 적은 rollout token을 사용하면서 고정된 short 및 long horizon과 2점 이내의 차이를 보이며, 평가한 pruning setting 전반에서 standard recovery baseline을 능가한다.
시사점 및 한계
Recovery 품질은 compressed model이 방문한 visited states에서 학습하고, 현재 사용할 수 있는 generation length에 rollout budget을 맞추는 데 달려 있다.
시사점 및 한계
Qwen3 model, 연구한 pruning setting, compression ratio, scale을 넘어선 일반성은 아직 확립되지 않았다.
Abstract
from arXiv · showhide
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.
1 서론
ShortOPD는 student 자신의 states에서 pre-compression model을 distilling하면서 rollout horizon을 조정해 반복 suffixes를 피함으로써, 구조적으로 pruning된 LLM을 복구한다. 광범위한 generation quality를 효율적으로 회복하지만, recovery 성능은 capability 전반을 포괄하는 corpus coverage에 좌우된다.
- 동기: 모든 benchmark에서 추가 sampling을 사용하면 Pass@k가 상승하며, GSM8K에서는 k=64일 때 91%에 도달한다. 이는 unpruned teacher의 greedy 성능 88%를 웃돌며, 올바른 trajectory가 여전히 sampling 가능함을 보여준다.Compressed model에서 사라진 능력은 소거된 것이 아니라 선택 확률이 낮아진 것이므로, 자체 sampling distribution에서 recovery를 수행할 근거가 된다.
- 동기: OPD는 label, verifier, external teacher 없이 student가 sampling한 rollout을 대상으로, frozen pre-compression self의 dense next-token supervision을 사용해 학습한다.이는 on-policy states와 token-level teacher distributions를 직접 결합한다.
- 동기: 초기 fixed-H=2048 OPD에서는 rollout의 55–75%가 suffix loop로 끝나는 반면 distillation loss는 약 step 80까지 낮은 수준에 머물러, rollout 및 teacher compute를 낭비한다.이러한 반복 suffixes는 긴 fixed horizon에서 teacher–student correction을 거의 추가로 제공하지 않는다.
- 방법: ShortOPD는 repetition-gated, truncation-aware horizon control을 사용하며, smoothed effective-length target과 점진적인 budget update를 통해 repetition이 발생하면 budget을 줄이고 truncation이 높으면 늘린다.이 controller는 exponential moving average를 사용해 horizon이 급격히 변하지 않도록 한다.
- 실험: ShortOPD는 여덟 개 과제군에 걸쳐 pruning하지 않은 teacher의 generation score를 약 삼분의 이까지 회복하며, 동일한 budget에서 SFT w/o KD, KD, SeqKD, sparse-reward RLVR을 능가한다.고정 rollout horizon과 비교해 최대 8192토큰까지 두 포인트 이내의 성능을 보이면서 rollout token은 최대 71% 적게 생성한다.
- 실험: Math 또는 code prompt를 제거하면 해당 capability에서 recovery 성능이 크게 저하되므로, recovery-corpus coverage는 핵심적인 design choice가 된다.On-policy distillation은 student 자신의 search가 방문한 states만 복구한다.
2 관련 연구
선행 연구는 depth 및 width 기반 compression을 통해 structured pruning을 연구해 왔으며, distillation 연구는 고정된 teacher output과 student inference state 사이의 불일치를 점점 더 다루고 있다. 이러한 연구 흐름은 pruning 이후 recovery를 위한 배경을 제공한다.
- Structured pruning과 recognition-generation gap: Structured pruning은 influence scoring, layer merging, learned mask와 같은 기준을 사용해 transformer block 전체 또는 layer 내부의 결합된 구조를 제거한다.Depth pruning에는 ShortGPT [5]와 LaCo [33]가 포함되며, width pruning에는 Sheared LLaMA [6]와 Minitron [7] [37]이 포함된다.
- Distillation과 on-policy recovery signal: Distillation method는 고정된 teacher-generated response에서 student-state supervision으로 발전했으며, imitation-learning KD [39], MiniLLM [40], divergence generalization [41] [42], generalized on-policy KD [23] 등이 이에 해당한다.Sequence-level KD [26]는 고정된 teacher-generated response로 학습하므로 off-policy인 반면, on-policy approach는 student state에서 distillation을 수행한다.
3 방법
ShortOPD는 frozen pre-compression teacher의 dense on-policy distillation과 repetition-gated short-to-long rollout schedule을 결합해 구조적으로 압축된 LLM을 복구한다. 압축된 model 자체의 state를 사용하면서 usable prefix에 맞춰 horizon을 조정하며, ground-truth response나 reward verification을 요구하지 않는다.
- Structured pruning: Pruning 실험은 낮은 Block Influence layer를 사용해 parameter를 약 25% 제거하지만, ShortOPD는 구조적으로 압축된 student와 그 pre-compression self만을 가정한다.Original model은 frozen teacher로 유지된다.
- On-policy self-distillation: OPD는 student가 생성한 trajectory에서 frozen pre-compression teacher를 실행하고, 모든 response position에서 next-token distribution을 일치시킨다.Loss는 on-policy이고 dense하며 reward-free다. clipped importance weighting과 top-100 logit 및 aggregated tail mass를 사용한다.
- Short-to-long schedule: ShortOPD는 recovery 전에 긴 horizon이 repetitive suffix로 진입할 수 있고, 반대로 영구적으로 짧은 horizon은 정당한 long generation을 잘라낼 수 있으므로 각 rollout budget을 조정한다.Global response_length는 padding 및 context ceiling으로 유지되고, step별 budget H_t는 H_max 내에서 변한다.
- Generation-side feedback: Controller는 terminal token에서 심한 repetition을 감지하고, OPD divergence와 teacher NLL로 그 시작점을 정교화하며, 살아남은 prefix를 effective length로 취급한다.이 통계는 기존 rollout token과 teacher score를 재사용하며, 현재 batch에는 생성된 모든 token에 대해 수정되지 않은 dense OPD loss가 계속 적용된다.
- Repetition-gated horizon control: High repetition은 다음 horizon을 관측된 usable length 쪽으로 줄이는 반면, clean truncation은 현재 horizon이 binding일 때만 증가를 허용한다.Hysteresis band와 EMA smoothing이 oscillation을 방지하며, repetition과 truncation이 겹치면 high repetition이 우선한다.
- Design space: ShortOPD는 on-policy student rollout, dense teacher distribution, repetition-gated horizon control을 결합하며, 이 차원 중 하나를 제외하는 recovery method와 다르다.Ground-truth response와 reward verifier를 모두 요구하지 않는다.
4 실험
8개 자유 형식 생성 태스크군에서 ShortOPD는 구조적 pruning을 적용한 모델의 성능을 크게 회복하며, 다양한 규모와 pruning 구조에서 offline recovery baseline을 일관되게 능가한다. Short-to-long controller는 training cost를 크게 낮추면서도 긴 고정 rollout과 비슷한 품질을 달성하며, 추가적인 on-policy 노출과 domain-matched data는 recovery를 더욱 향상한다.
- 주요 결과: 48.46 Avg는 pruned model의 5.71에서 거의 9배까지 회복하며, 가장 강력한 off-policy baseline보다 18 points 높다. 또한 ShortOPD는 code와 4개 open-ended task 모두에서 성능을 향상한다.8B에서 ShortOPD는 no-thinking model과 Thinking model에 대해 각각 42.84와 53.71 Avg를 달성하며, Thinking evaluation의 모든 generation task에서 선두를 차지한다.
- Closed-loop recovery trajectory: ShortOPD는 초기 repetitive warm-up을 줄여 post-warm-up regime에 약 40–50 steps 만에 도달하며, H=2048인 Vanilla fixed는 대략 80 steps가 필요하다.repetition이 높은 구간에서 controller는 budget을 2048에서 1024로 낮춘 뒤, clean truncation이 우세해지면 다시 2048로 되돌린다.
- 효율성: 8.5 hours와 250M rollout tokens만으로 fixed 8192-token rollout과 비슷한 Avg를 달성하며, 후자의 품질은 schedule에 따라 48.5에서 50.2 범위에 그친다.fixed 2048과 비교하면 ShortOPD는 mean generation time을 29% 줄이며, fixed 8192와 비교하면 76% 적은 time을 사용한다.
- Exposure에 따른 scaling: ShortOPD는 exposure에 따라 확장되며, 1 epoch 후 48.46이던 Avg를 3 epochs 후 55.41까지 높인다. 이는 teacher의 73.7%에 해당하며, math, code, open-ended task 전반에서 폭넓은 향상을 보인다.향상은 GSM8K, MATH-500, HumanEval, MBPP, Alpaca, QA, MT-Bench에 걸쳐 나타난다.
- Recovery signal: ShortOPD에서는 on-policy states가 중요하다. SeqKD는 19.9 Avg points를 잃고, KD는 동일한 frozen teacher를 사용함에도 ShortOPD보다 17.9 points 낮다.모든 method는 동일한 compressed student에서 시작하고 동일한 step budget과 prompt distribution을 사용하며, OPD와 ShortOPD는 student가 생성한 states에서 teacher를 모방한다.
- Ablations와 일반성: Domain-matched recovery data와 pruning structure는 모두 중요하다. math를 제거하면 mathematical recovery가 거의 사라지고, code를 제거하면 execution score가 붕괴하며, ShortOPD는 width-only와 hybrid configuration 전반에서 KD보다 6.44–9.01 Avg points 높다.실험은 math, code, open-ended instruction data를 포괄하며, 평가는 8개 task family의 free-form generation으로 제한된다.
5 결론
Structured pruning은 잠재된 능력을 보존하면서도 생성 과정에서 이를 발현하지 못할 수 있어, recovery를 분포 복구 문제로 만든다. Pruning 설정과 model mode 전반에서 on-policy self-distillation은 초기 rollout의 품질이 낮음에도 standard recovery baseline을 크게 능가한다.
- 5 결론: On-policy self-distillation은 Qwen3 model, mode, 그리고 depth·width·hybrid pruning 전반에서 SFT w/o KD, SeqKD, KD를 크게 능가한다.조건을 맞춘 Math+Code 대조 실험에서도 sparse-reward RLVR을 능가한다.
- 5 결론: Structured pruning은 model의 search space에 잠재된 능력을 남겨둘 수 있지만 생성 과정에서 이를 발현하지 못하게 하며, recovery를 분포 복구로 규정한다.
- 5 결론: On-policy self-distillation의 주요 비용은 초기 rollout의 낮은 품질이다.
6 한계 · 부록
이 연구의 범위는 4B 및 8B Qwen3 모델에서의 BI-depth, width-only, hybrid pruning으로 제한되어 있어, 더 폭넓은 설정에서의 일반성은 아직 밝혀지지 않았다. 압축이 경량 복구에 지나치게 공격적이 되는 조건을 규명하려면 추가 평가가 필요하다.
- 6 한계: 실험은 4B 및 8B Qwen3 모델에서의 BI-depth, width-only, hybrid pruning을 다룬다.
- 6 한계: 일반성을 확립하고 압축이 경량 복구에 언제 지나치게 공격적이 되는지 파악하려면 더 다양한 모델군, 압축 비율, 모델 규모가 필요하다.
A 학습 세부사항 · B 종단 주기 루프 검출기
이 논문은 고정된 학습 및 baseline 프로토콜에 따라 dense on-policy distillation으로 구조적으로 pruning된 student를 복구하고, ShortOPD는 심각한 종단 반복을 검출해 teacher가 확인한 loss 신호로 각 rollout의 effective length를 정교화한다.
- A 학습 세부사항: Pruning은 held-out corpus의 BI calibration을 사용해 36개 transformer block 중 9개를 제거하며, 동결된 원본을 teacher로, compressed model을 actor의 초기값으로 사용한다.Qwen3-4B-Thinking과 Qwen3-8B 모두 동일한 36 →27 depth reduction을 사용하며, calibration은 backbone별로 수행한다.
- A 학습 세부사항: 주요 recovery는 rollout group 8, batch 64, 4096-token prompt 및 response limit, temperature 0.8, learning rate 2 × 10−6, 710 training steps를 사용한다.목적 함수는 α = 0.5인 top-100+tail generalized JSD이며, importance weight는 2.0에서 clipping하고 policy-gradient loss는 사용하지 않는다.
- A 학습 세부사항: Offline SFT w/o KD, SeqKD, KD baseline은 동일한 prompt distribution, 8192-token sequence, batch 64, learning rate 2 × 10−6, 710 step을 사용한다.KD는 temperature 1.0의 forward KL을 사용하며 hard-label cross-entropy 항은 포함하지 않는다.
- A 학습 세부사항: Efficiency 실험은 주요 corpus, objective, batch size, rollout count, learning rate를 재사용하면서 ShortOPD를 2048 및 8192-token의 고정 response ceiling과 비교한다.Controller는 재시작 시 현재 budget과 모든 EMA state를 checkpoint에 저장한다.
- A 학습 세부사항: Dense OPD는 항상 generated token으로 학습하며, 독립적인 repetition-rate threshold만 shrink gate를 열 수 있고 low loss만으로는 shortening을 트리거할 수 없다.Probe는 loss를 mask하지 않는다.
- B 종단 주기 루프 검출기: Detector는 마지막 W=512개의 valid response token에서 period 1–10을 검사하고 shifted token comparison을 사용해 truncation이 cycle 중간을 잘라도 검출할 수 있게 한다.Raw response end에서 높은 agreement가 나타나면 반복된 material 뒤에 clean ending이 이어지는 경우를 terminal repetition으로 분류하지 않는다.
- B 종단 주기 루프 검출기: 심각한 loop에는 terminal agreement η=0.9, 최소 3개 cycle, 64-token minimum tail, 그리고 128-token tail 또는 response의 30% 중 하나가 필요하다.Detector는 terminal tail을 가장 잘 설명하는 가장 긴 구간을 선택한 뒤, 더 높은 agreement와 더 짧은 period를 우선해 onset을 structural effective-length candidate로 반환한다.
- B 종단 주기 루프 검출기: ShortOPD는 OPD loss와 teacher NLL에 대한 absolute 또는 relative threshold를 처음으로 만족하는 32-token window를 찾아 detector의 onset을 국소적으로 정교화한다.이 refinement는 boundary를 바꾸지만 repetition의 존재 여부를 결정하지는 않는다.
C 교사--학생 일치에서의 Conditional OPD gradient
고정된 sampled prefix를 조건으로 하면 generalized-JSD OPD gradient는 teacher–student agreement에서 사라지고, 그 근처에서는 선형으로 감소한다. Top-100-plus-tail aggregation과 clipped importance weighting은 이 stationary point를 유지하며, 구간화한 agreement만으로도 signal이 0이 되기에 충분하다.
- C 교사--학생 일치에서의 Conditional OPD gradient: 유도는 backpropagation 동안 고정된 하나의 샘플 prefix를 조건으로 하므로, 구현된 on-policy update와 일치하며 forward KL이 아니라 generalized JSD를 사용한다.두 KL term의 비로그 성분은 미분 과정에서 상쇄된다.
- C 교사--학생 일치에서의 Conditional OPD gradient: 정확한 교사–학생 일치에서는 conditional logit 및 parameter gradients가 0이며, retained-support positivity와 bounded local Jacobian 아래 그 근처에서 선형적으로 사라진다.α=0.5에서는 local slope가 forward-KL gradient의 사분의 일이다.
- C 교사--학생 일치에서의 Conditional OPD gradient: Top-100 teacher bins와 집계된 tail을 사용하면 분석을 binned distributions로 확장할 수 있으므로, coarse-grained agreement만으로도 update signal이 제거된다.tail bin 내부의 차이는 loss에 드러나지 않으므로, 이는 full-vocabulary agreement보다 약한 조건이다.
- C 교사--학생 일치에서의 Conditional OPD gradient: Clipped importance weighting은 그 bounded multiplier가 divergence와 gradient를 함께 스케일링하고, 두 값 모두 정확한 일치에서 0이므로 stationary point를 유지한다.이는 구현의 loop states에 해당하는 조건이다.
D 평가 세부 사항
평가는 수학, 코드, open-ended, multiple-choice 벤치마크 전반에서 task-specific decoding, response cap, answer normalization을 사용한다. Main 및 matched 비교는 통일된 no-thinking protocol로 다시 실행하며, 남은 matched Math+Code 결과는 Table 11에 보고한다.
- Main evaluation protocol: GSM8K는 1,319개의 모든 test example에 exact/normalized matching을 적용하고, MATH-500은 sampled 16,384-token protocol에서 boxed-answer extraction과 normalization을 적용한다.4B Instruct main 결과는 512-token GSM8K cap을 적용한 no-thinking decoding을 사용한다. MATH-500은 temperature 0.7, top-p 0.8, top-k 20, minimum-p 0을 사용하며, repetition stopping은 detected loop만 truncate한다.
- Comparison protocols: Width/hybrid 및 matched Math+Code 비교는 통일된 no-thinking rerun을 사용하며, GSM8K와 7개의 non-MATH task에는 greedy decoding을, MATH-500에는 sampled decoding을 적용한다.Response cap은 GSM8K에서 512 token, code 및 open-ended generation에서 768 token이다. Width/hybrid multiple-choice 평가는 thinking을 비활성화하고, zero-shot MMLU, 4096-token maximum length, mean log-probability를 사용한다.
- Matched Math+Code results: Table 11은 matched Math+Code 비교의 남은 full-evaluation 결과를 보고하며, 여기에는 1–10 scale의 judge score와 multiple-choice column 및 MC Avg의 accuracy percentage가 포함된다.이 결과는 앞서 설명한 matched comparison protocol을 사용한다.
E 반복과 pruning depth: 전체 sweep
전체 36-layer pruning sweep에서 repetition은 모델이 coherent한 동안에만(k≤12) 지배적인 degeneration mode이며, 더 깊은 pruning에서는 incoherent text가 생성된다. 이 전환은 정성적으로도 확인되며, 25% compression에서는 알아볼 수 있는 prefix와 반복적인 suffix가 나타나고, 더 깊은 compression에서는 mixed-token collapse가 발생한다.
- E 반복과 pruning depth: 전체 sweep: 모델이 coherent한 동안에는 n-gram loop를 유지할 만큼 충분한 구조를 보존하므로, looping이 지배적인 degeneration mode가 된다.이는 coherent regime 안에서만 repetition이 제거된 depth의 증가를 따라가는 이유를 설명한다.
- E 반복과 pruning depth: 전체 sweep: k≤12까지 repetition은 증가하고 distinct-2 diversity는 단조롭게 감소하는 반면, k≈13을 넘어 pruning하면 output이 더 이상 loop를 형성하지 못하는 incoherent text로 붕괴한다.전체 sweep은 192-prompt probe에서 Qwen3-4B-Instruct의 36개 layer를 모두 다룬다.
- E 반복과 pruning depth: 전체 sweep: 25% compression에서는 generation이 알아볼 수 있는 prefix를 유지하지만 repetitive suffix가 발달하며, 더 깊은 compression에서는 대신 incoherent mixed-token regime에 진입한다.Table 12의 발췌문은 정상적인 generation에서 repetition을 거쳐 incoherent collapse로 이어지는 전환을 보여준다.
F 객관식 인식 결과 · F.1 폭 및 하이브리드 프루닝
객관식 인식은 논문의 주요 자유 형식 생성 평가를 보완하는 진단으로, held-out ARC-Challenge, HellaSwag, MMLU, WinoGrande를 다룬다. 폭 및 하이브리드 프루닝은 생성보다 인식을 더 잘 보존하지만, ShortOPD의 인식 성능 향상이 더 작다는 점은 고정 선택지 정확도가 충분한 대리 지표가 아님을 보여준다.
- F 객관식 인식 결과: ARC-Challenge, HellaSwag, MMLU, WinoGrande에 대한 held-out 객관식 복원은 논문의 주요 생성 목표가 아니라 보완적 진단으로 기능한다.이 벤치마크들은 복원 코퍼스에 포함되지 않으며, 부록에서 객관식 표를 별도로 정리한다.
- F.1 폭 및 하이브리드 프루닝: 평가에서는 thinking을 비활성화하고, 고정된 assistant ANSWER: prefix 뒤의 후보 문자들을 mean token log-probability로 채점하며, zero-shot MMLU는 4096 tokens로 제한한다.이는 인식 결과에 대한 평가 프로토콜이다.
- F.1 폭 및 하이브리드 프루닝: 프루닝된 4B-width 및 8B-hybrid 모델에서는 생성 Avg 점수가 각각 17.23과 16.17에 불과함에도 51.14 및 69.41 MC Avg가 유지된다.폭 및 하이브리드 프루닝에서는 생성보다 인식 성능이 덜 저하된다.
- F.1 폭 및 하이브리드 프루닝: ShortOPD는 KD보다 MC 점수에서 0.87–2.08만큼 앞서지만, 생성 성능 향상은 6.44–9.01에 달하며, 이는 고정 선택지 인식이 생성의 충분한 대리 지표가 아님을 보여준다.인식 비교는 프루닝된 폭 및 하이브리드 모델에 대해 보고된다.
F.2 깊이 가지치기 sanity check
깊이 가지치기 sanity check는 특히 SFT에서 초기화하고 3 epoch 동안 학습했을 때 ShortOPD가 multiple-choice recognition을 크게 회복함을 보여준다. SFT와 ShortOPD를 결합하면 teacher의 recognition score에 도달한다.
- F.2 깊이 가지치기 sanity check: SFT-init ShortOPD는 Qwen3-4B-Instruct에서 Teacher의 76.7과 일치하는 76.8 MC Avg를 달성한다.SFT-init ShortOPD는 가지치기된 student가 아니라 1 epoch SFT w/o KD checkpoint에서 시작한다.
- F.2 깊이 가지치기 sanity check: 3 epoch one-epoch ShortOPD는 1 epoch 후 67.6에서 71.6 MC Avg로 상승하며, Pruned보다 22 points above에 해당한다.Table 14는 held-out multiple-choice recognition task에서 candidate log-likelihood scoring을 사용한다.
- F.2 깊이 가지치기 sanity check: recognition task에서는 generation performance가 Table 3에서 더 약함에도 SFT와 KD가 각각 74.9/74.4 MC Avg로 유리하다.Score는 thinking을 비활성화한 candidate log-likelihood scoring에서의 accuracy다.