Official code release for the paper

When Are Teacher Tokens Reliable? Position-Weighted On-Policy Self-Distillation for Reasoning. Xiaogeng Liu, Xinyan Wang, Yingzi Ma, Yechao Zhang, Chaowei Xiao.

This repository reproduces the four training methods (OPSD, EOPD, REOPOLD, PW-OPSD) and the long-context reasoning evaluation suite (MATH-500, AIME 2024, AIME 2025, HMMT 2025) on three base models: Qwen3-4B, DeepSeek-R1-Distill-Llama-8B, and Olmo-3-7B-Think.

Entries report mean ± across-seed sample standard deviation over three evaluation seeds. The Avg@12 column is the equal-weight mean of the four per-benchmark Avg@12 values.

Across both base models, PW-OPSD with the Moderate schedule (w_min, tau, s) = (0.25, 0.30, 0.10) (held fixed across models) delivers the highest Avg@12 among the four compared methods.

1. Create the conda env (Python 3.10 + PyTorch 2.8 + TRL 0.26 + vLLM 0.11):
   conda env create -f code/environment.yml conda activate pwopsd
2. Set the WandB API key, or run WandB in offline mode:
   export WANDB_API_KEY=<your wandb api key> # or, to skip the API-key check entirely: export WANDB_MODE=offline
   The launchers abort early with a clear error if neither WANDB_API_KEY nor WANDB_MODE=offline|disabled is set.
3. (Optional) Override the training dataset:
   export PWOPSD_TRAIN_DATASET="siyanzhao/Openthoughts_math_30k_opsd" # default
   The trainer reads only the problem and solution string columns from the dataset's train split, so any Hugging Face dataset that exposes those two columns can be plugged in via PWOPSD_TRAIN_DATASET.
4. (Optional) For HMMT February 2025 evaluation, point at a local parquet derived from the MathArena hmmt_feb_2025 release:
   export HMMT25_PARQUET_PATH=/path/to/hmmt25_clean/train.parquet

The launchers resolve REPO_ROOT automatically from the launcher script's own location (one directory above launchers/), so they should be run from the repository top level or via an absolute path:

Each launcher invokes ea_opsd_train.py with the shared hyperparameters (LoRA r=64, alpha=128, lr=5e-6, distillation temperature 1.1, forward-KL clip 0.05, fixed teacher) and dispatches to the corresponding method via the --method flag. The PW-OPSD launcher additionally passes the Moderate schedule --position_w_min 0.25 --position_tau 0.30 --position_s 0.10, and the REOPOLD launcher passes the published reward-floor and phase-mask hyperparameters.

All trainings use effective batch size 32 on 4 GPUs and run for 150 optimizer steps with a save every 25 steps. The launchers merge and evaluate checkpoint-100 — the paper protocol selects the step-100 snapshot for all reported numbers. The per-GPU batch / gradient-accumulation split is chosen per base-model size (Qwen3-4B uses bs=4 x ga=2, DSR1-L8B and Olmo-3-7B-Think use bs=2 x ga=4). The Moderate position schedule itself is identical across the three models.

Each launcher trains the LoRA adapter at outputs/<run_name>/ and then merges it into a single HuggingFace-loadable checkpoint at outputs/<run_name>_merged/, then writes a .ready_for_eval sentinel. <run_name> follows {model_tag}_{method} (e.g. qwen3_4b_base_pwopsd, dsr1_l8b_pwopsd, olmo3_7b_think_pwopsd).

DSR1-Llama-8B and Olmo-3-7B-Think are released as reasoning-distilled models whose chat templates unconditionally append an open <think> tag and ignore enable_thinking=False. To match OPSD's non-thinking-mode prompt format on these bases, the data collator detects an unclosed trailing <think> and appends </think>, restoring the intended enable_thinking=False semantics. The patch is a no-op for templates (e.g. Qwen3-4B base) that already honor the flag.

Accepted <method_name> values are any of the <run_name>s listed above; accepted <dataset> values are math500, aime24, aime25, hmmt25. The eval launcher writes JSONs of the form eval_results/<method_name>_38k.json (or eval_results/<method_name>_<dataset>_38k.json) and per-run logs under eval_results/logs/. Override HOSTLABEL to disambiguate logs across parallel workers.

evaluate_math.py reports Pass@12 / Avg@12 / Maj@12 with math-equivalence clustering (sympy + math_verify). The INVALID cluster participates in the plurality count and is scored incorrect when selected (formal definitions in the paper appendix).

recompute_majn.py re-aggregates Maj@12 over existing eval JSONs using the same fixed clustering rule.

• MATH-500 — HuggingFaceH4/MATH-500
• AIME 2024 — HuggingFaceH4/aime_2024
• AIME 2025 — yentinglin/aime_2025
• HMMT February 2025 — locally cleaned parquet derived from the MathArena hmmt_feb_2025 release (https://github.com/eth-sri/matharena). The evaluator expects a parquet with at least two string columns, problem and answer, holding the 30 HMMT February 2025 problems. Set HMMT25_PARQUET_PATH to point at it (default ./data/hmmt25_clean/train.parquet).

The launchers default HF_HUB_OFFLINE=0 and TRANSFORMERS_OFFLINE=0, so Hugging Face downloads happen on first use; export HF_HUB_OFFLINE=1 / TRANSFORMERS_OFFLINE=1 to force the offline cache once the assets are downloaded. The training launchers pull the base model identified by MODEL=... inside the script (Qwen/Qwen3-4B, deepseek-ai/DeepSeek-R1-Distill-Llama-8B, or allenai/Olmo-3-7B-Think) and the dataset selected by PWOPSD_TRAIN_DATASET. The MATH-500, AIME 2024, and AIME 2025 benchmarks are downloaded automatically by evaluate_math.py from their Hugging Face dataset IDs; only HMMT 2025 needs the local parquet described above.

All paper experiments use 4×H100 80GB.

This code is released under the MIT License; see LICENSE for the full text. Upstream model and dataset assets remain under their original licenses as declared on Hugging Face and GitHub.