Dynamic Important Example Mining for Reinforcement Finetuning (DIEM)

Reference implementation of DIEM, a fully automated framework that makes data utilisation adaptive throughout reinforcement finetuning (RFT). DIEM adds two components to every optimisation step:

1. Dynamic importance measuring — a gradient-alignment estimator that approximates each sample's marginal contribution to policy improvement (paper Eq. 4).
2. Dynamic reweighting — a constrained batch reweighting scheme that maximises aggregate utility while preserving the update's gradient magnitude, so the optimisation stays stable (paper Eq. 6–7).

Because the gradient a sample contributes in policy-gradient RFT is G_z = A_z · ∇log π_z, scaling that sample's advantage by w_z reweights its gradient contribution by exactly w_z. DIEM therefore needs no change to the loss, the optimiser, or the rollout engine: it is a drop-in advantage estimator, registered as grpo_diem.

Paper: Dynamic Important Example Mining for Reinforcement Finetuning https://arxiv.org/abs/2608.29252.

diem/
  advantage.py               DIEM advantage estimator (Eq. 4 + Eq. 6-7)
  reward.py                  rule-based math reward with robust answer extraction
  integrate.py               attaches DIEM to a stock veRL install, no source edits
  main_ppo.py                training entrypoint (veRL's main_ppo + DIEM)
scripts/
  train.sh                   generic launcher; every hyperparameter is an env var
  experiments/               one script per reported experiment
data/
  train_unified.parquet      training mixture, 14,973 math problems
  eval_bench5.parquet        evaluation suite, 985 questions
  build_train.py             regenerates the training mixture from source datasets
tests/
  test_diem.py               mathematical properties of the estimator

git clone <this repository> && cd diem
pip install -r requirements.txt

DIEM does not patch veRL's source. diem/integrate.py does two things at import time, and diem/main_ppo.py calls it for you:

• registers the estimator under the name grpo_diem, so algorithm.adv_estimator=grpo_diem resolves to it;
• wraps veRL's advantage stage so the response token ids and rollout log probabilities reach the estimator — no built-in estimator needs those fields, so veRL does not forward them by default.

The reward function is supplied through veRL's official reward.custom_reward_function hook, which scripts/train.sh sets.

Both parquet files are included, so nothing has to be downloaded before training.

Training — data/train_unified.parquet, 14,973 problems: GSM8K train (7,473) plus a 7,500-problem sample of DAPO-Math-17k. Both halves sit behind one instruction and are tagged data_source="math", so a single rule-based verifier grades the whole mixture. Regenerate with python3 data/build_train.py (seeded, byte-stable).

Evaluation — data/eval_bench5.parquet, 985 questions:

Evaluation runs inside the training loop every TESTFREQ steps, sampling VALN answers per question at temperature VALTEMP, and reports accuracy per benchmark.

Each experiment is one script. Point MODEL at a HuggingFace id or a local checkpoint directory, choose the GPUs, and run:

bash scripts/experiments/qwen3-1.7b-base.sh    
bash scripts/experiments/qwen2.5-3b.sh
bash scripts/experiments/qwen3-4b-base.sh  
bash scripts/experiments/qwen2.5-7b.sh

# override anything from the command line
SEEDS=1 GPUS=4,5 bash scripts/experiments/qwen2.5-3b.sh
MODEL=/path/to/local/checkpoint bash scripts/experiments/qwen2.5-7b.sh

Checkpoints, the trainer log and DIEM's per-step weight diagnostics land in outputs/<run name>/.

To train an GRPO baseline under the identical protocol, pass grpo instead of grpo_diem and unset the DIEM_* variables:

MODEL=Qwen/Qwen2.5-3B MICRO=8 GPUMEM=0.45 bash scripts/train.sh grpo_qwen2.5_3b grpo 0,1 2

@misc{tan2026dynamicimportantexamplemining,
      title={Dynamic Important Example Mining for Reinforcement Finetuning}, 
      author={Haoru Tan and Sitong Wu and Yanfeng Chen and Shizhen Zhao and Yang-Tian Sun and Tianjia Liu and Chirui Chang and Shaofeng Zhang and Samm Sun and Xiuzhe Wu and Ruobing Xie and Xiaojuan Qi},
      year={2026},
      eprint={2608.29252},
      archivePrefix={arXiv},
}