JunliC/cs260-opsd

Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models

---

Overview

On-Policy Self-Distillation (OPSD) trains a single model to act as both student and teacher by conditioning on different contexts — the student sees only the problem, while the teacher additionally sees the ground-truth solution — and performs token-level distribution matching along the student's own on-policy trajectories.

Updates

• Mar 18, 2026: Released updated code.

  (1) Fixed chat template and zero2 bugs (see template issue), we re-ran experiments with updated results (detailed results & ablations updated on arxiv/blog). The fixes yield improved OPSD performance, most notably on Qwen3-1.7B.

  (2) Added a new training stabilization strategy 🚀: per-token point-wise KL clipping. We find style tokens (such as 'wait', 'think') can exhibit 6–15× higher KL divergence than math-related tokens, and dominates the training signal. Clipping stablizes training and improves performance.

• Mar 3, 2026: Initial code release.

Installation

conda env create -f environment.yml
conda activate opsd

pip install flash-attn==2.8.3 --no-build-isolation

If you encounter difficulties installing flash-attn, you can check the version matching your CUDA and PyTorch versions from the flash-attention releases page.

The code uses trl's experimental GOLD trainer as a base.

Repository Structure

├── opsd_trainer.py          # OPSDTrainer: core self-distillation trainer
├── data_collator.py         # Data collator for self-distillation
├── opsd_train.py            # OPSD training entry point
├── sft_train.py             # SFT baseline training entry point
├── grpo_train.py            # GRPO baseline training entry point
├── accelerate.yaml          # Accelerate config (multi-GPU)
├── scripts/
│   ├── run_opsd.sh          # Example launch script for OPSD
│   ├── run_sft.sh           # Example launch script for SFT
│   └── run_grpo.sh          # Example launch script for GRPO
└── eval/
    ├── evaluate_math.py     # Evaluation script (vLLM)
    └── run_eval.sh          # Example evaluation script

Quick Start

Reproduce results on Qwen3-1.7B (🚀 training only takes ~15 minutes on 4×H100 and peaks within 100 steps):

bash scripts/run_opsd_1b.sh

Evaluation: (evaluation takes ~ 30-50 minutes on 4xh100 for each checkpoint)

cd eval
bash run_eval.sh

Evaluation Results across Tasks on Qwen3-1.7B

AIME24	AIME25	HMMT25
Step Avg@12 Base 51.5% 25 51.4% 50 52.8% 75 54.4% 100 57.2%	Step Avg@12 Base 36.7% 25 42.5% 50 43.9% 75 40.6% 100 41.1%	Step Avg@12 Base 23.1% 25 24.7% 50 27.8% 75 26.9% 100 29.2%

Evaluation settings: temperature=1.0, thinking mode enabled, max new tokens=38912, top-p=none, top-k disabled, min-p=0, presence penalty=0, num samples=12

Non-Thinking Mode

OPSD can also run in non-thinking setting where both the Qwen student and teacher are enabled_thinking=False during training (--student_thinking False --teacher_thinking False) and evaluated with non-thinking inference (--no_thinking), with faster evaluation time than thinking mode.

Training:

bash scripts/run_opsd_4b_nonthink.sh
bash scripts/run_opsd_8b_nonthink.sh

Evaluation:

cd eval
bash run_eval_nonthink.sh

Evaluation Results with Non-Thinking Mode across Models

Qwen3-8B (--jsd_token_clip 1e-7)

AIME24	AIME25	HMMT25
Step Avg@12 Base 26.4% 50 49.7% 75 45.3% 100 38.3%	Step Avg@12 Base 19.7% 50 35.0% 75 26.9% 100 27.5%	Step Avg@12 Base 10.8% 50 18.3% 75 17.5% 100 15.3%

Qwen3-4B (--jsd_token_clip 1e-6)

AIME24	AIME25	HMMT25
Step Avg@12 Base 23.1% 50 20.3% 75 27.5% 100 31.1% 150 32.8%	Step Avg@12 Base 21.4% 50 21.4% 75 20.8% 100 21.1% 150 21.9%	Step Avg@12 Base 10.8% 50 11.1% 75 13.1% 100 16.4% 150 14.4%

Qwen3-1.7B (--jsd_token_clip 1e-6)

AIME24	AIME25	HMMT25
Step Avg@12 Base 11.9% 50 15.0% 75 13.9% 100 12.5%	Step Avg@12 Base 9.2% 50 6.2% 75 8.3% 100 8.1%	Step Avg@12 Base 5.0% 25 7.2% 50 5.8% 75 5.0%

Evaluation settings: temperature=1.0, non-thinking mode, num samples=12.

Key OPSD arguments

Argument	Default	Description
--fixed_teacher	False	Fix the teacher to the initial policy (step 0). Requires --use_peft. Note ❗ If you disable PEFT, the teacher will keep updating at every training step, which may make training unstable. Our main results use the fixed teacher, which is currently implemented with LoRA adapter weights.
--use_tinker_loss	False	Use sampled-token policy-gradient objective instead of full-vocabulary JSD. More memory efficient. Currently no clipped implemented for this variant, could be unstable.
--max_completion_length	—	Student generation length for distillation. We use 1024 in our main experiments.
--beta	—	Interpolation weight for the JSD mixture distribution. Beta=0 means forward KL and 1 means reverse KL.
--jsd_token_clip	0.05	Clip the JSD loss for each token to a maximum value. This can improve stability by preventing stylistic tokens from dominating the training signal. Note when clipping is applied, the loss can be negative due to positive KL summand being capped.
--reason_first	False	Prepend an explicit rationalization to the teacher context before distillation.
--run_config	None	Custom name suffix for the output directory and WandB run.

SFT Baseline

See scripts/run_sft.sh.

GRPO Baseline

See scripts/run_grpo.sh.

Acknowledgements

Our implementation builds on TRL GOLD Trainer. We sincerely thank @simran135 and @beanie00 for identifying the prompt template bugs and the zero-2 issue, respectively!

Citation

If you find this useful, please consider citing:

@article{zhao2026self,
  title={Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models},
  author={Zhao, Siyan and Xie, Zhihui and Liu, Mengchen and Huang, Jing and Pang, Guan and Chen, Feiyu and Grover, Aditya},
  journal={arXiv preprint arXiv:2601.18734},
  year={2026}
}