Mismatch Matters: On-Policy Distillation Beyond Token Agreement

Official implementation of TIDE (Token-level Independent Deficit–Excess correction).

On-policy distillation optimizes token agreement with the teacher, and students learn to game it: they fall into repetitive loops that score near-perfect agreement while the response as a whole is wrong. We call this degenerate agreement, and it motivates looking at mismatch instead. Mismatch is not one phenomenon but two, and they need opposite treatments.

• Student-excess tokens are generated by the student but assigned near-zero probability by the teacher. Their log-ratio corrections grow unbounded and destabilize the update. TIDE replaces the log-ratio with the bounded Hellinger shaping h(a) = 2 (exp(a/2) - 1) > -2, so no single rejected token can dominate a step.
• Student-deficit tokens are preferred by the teacher but rarely sampled by the student, so a sampled-token estimator never sees them and the teacher's reasoning patterns never transfer. TIDE injects them analytically over the teacher's top-K support, with a gradient that does not vanish as the student probability goes to zero.
• Quantile gating restricts each branch to the most informative positions via batch-adaptive thresholds. Positions in the matched band receive exactly zero advantage.

Across math reasoning benchmarks and several Qwen3 teacher–student pairs, TIDE beats standard OPD and recent token-selection and reward-shaping baselines. The gap widens under strong teacher–student mismatch, where it lifts Avg@8 from 6.9% to 20.3% while cutting average response length by 3.6x.

Paper: arXiv:2608.09836 · Repository: https://github.com/yzc-666/TIDE

Repository layout

tide/
├── verl/                              # vendored verl fork (0.7.0.dev) with TIDE integrated
│   ├── verl/utils/tide_reward.py      # the objective: shaping, gates, tide_scores
│   ├── verl/workers/actor/dp_actor.py # per-candidate advantages (strategy="tide")
│   ├── verl/workers/fsdp_workers.py   # teacher forward: top-K log-probs
│   ├── verl/workers/config/rollout.py # tide_keep_neg / tide_keep_pos / tide_deficit_lambda
│   ├── verl/trainer/ppo/ray_trainer.py
│   └── tests/utils/test_tide_reward_on_cpu.py
├── on_policy_distillation.sh          # TIDE / on-policy distillation entry point
├── grpo.sh                            # GRPO entry point (RL baseline)
├── scripts/val/eval/                  # generation + grading for math benchmarks
├── scripts/infer/                     # vLLM rollout and data-prep helpers
├── datasets/test_data/                # benchmark metadata (parquet files not included)
└── LlamaFactory/                      # third-party, used only for the SFT baseline

TIDE is a new top_k_strategy inside the actor, so it reuses the fork's existing distillation plumbing: the teacher is loaded as verl's reward model and returns sampled-token log-probs plus top-K log-probs each step, and dp_actor turns them into detached per-candidate advantages consumed by the token_reward_direct estimator.

Installation

The fork tracks upstream verl 0.7.0.dev; its own requirements apply unchanged.

conda create -n tide python=3.10 -y
conda activate tide
cd verl
pip install -e .[vllm,math]

Expect CUDA 12.x, PyTorch 2.x, vLLM >= 0.8.5, and Ray >= 2.41. Verify the objective without a GPU:

cd verl
python -m pytest tests/utils/test_tide_reward_on_cpu.py -q

Data

Training and evaluation read verl-format parquet files. The benchmark JSON under datasets/test_data/ documents the expected schema, but the parquet files themselves are not shipped — build them with verl's preprocessing scripts (verl/examples/data_preprocess/) and place them as:

datasets/dapo-math-17k.parquet           # training set
datasets/test_data/AIME24/test.parquet   # and AIME25, AMC23, MATH-500, ...

Training

Point the script at a student and a teacher, then launch. Defaults reproduce the main setting: 8 GPUs on one node, DAPO-Math-17k, K = 16, 4 rollouts per prompt, 7168 response tokens.

export MODEL_ROOT=/path/to/hf_models
export ACTOR_MODEL_PATH=${MODEL_ROOT}/DeepSeek-R1-Distill-Qwen-1.5B   # student
export REWARD_MODEL_PATH=${MODEL_ROOT}/JustRL-DeepSeek-1.5B           # teacher
export PROJECT_PATH=/path/to/run_artifacts

bash on_policy_distillation.sh

Checkpoints, validation logs and Hydra output all land under $PROJECT_PATH. Runs resume automatically from the latest checkpoint of the same EXPERIMENT_NAME; use NEW_RUN=1 to force a fresh run.

Evaluation

scripts/val/eval/ runs generation and grading for the math benchmarks. Generation runs one vLLM instance per GPU (the GPU list is at the top of main(), eight by default) and writes jsonl rollouts under eval_outputs/<model basename>/; grading scores them into grading_results.json.

cd scripts/val/eval

# One model, or EVAL_MODEL_FOLDER=<dir of checkpoints> to sweep a whole run.
EVAL_MODEL_PATH=/path/to/checkpoint python gen_vllm.py --disable-thinking
EVAL_RUN_NAME=$(basename /path/to/checkpoint) python grade.py

The task list and sampling settings are the constants at the top of gen_vllm.py; reported numbers use temperature 0.7, top-p 0.95, and Avg@8 on AIME24 / AIME25 / AMC23. Grading is rule-based by default — python grade.py --enable_model_verifier additionally loads a verifier model from EVAL_VERIFIER_PATH.

Acknowledgements

Built on verl. The bundled fork keeps upstream's Apache 2.0 license, and TIDE's own files are released under the same terms.

Citation

@article{yu2026mismatch,
  title={Mismatch Matters: On-Policy Distillation Beyond Token Agreement},
  author={Yu, Zichao and Yu, Chengzhi and Xu, Shengze and Han, Yujin and Jiang, Bingqing and Wang, Xu and Zou, Difan},
  journal={arXiv preprint arXiv:2608.09836},
  year={2026}
}