Vocabulary Dropout for Curriculum Diversity in LLM Co-Evolution

Official code for the COLM 2026 paper. Vocabulary dropout (VD) applies a hard, non-stationary Bernoulli mask to the proposer's output logits during co-evolutionary self-play (R-Zero). By removing a random subset of the vocabulary each batch, it prevents the proposer from collapsing onto fixed question templates, sustaining curriculum diversity and improving the downstream solver.

• Paper (arXiv): https://arxiv.org/abs/2604.03472
• alphaXiv: https://www.alphaxiv.org/abs/2604.03472

Citation

@misc{dineen2026vocabularydropoutcurriculumdiversity,
      title={Vocabulary Dropout for Curriculum Diversity in LLM Co-Evolution},
      author={Jacob Dineen and Aswin RRV and Zhikun Xu and Ben Zhou},
      year={2026},
      eprint={2604.03472},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2604.03472},
}

What is here

This repository builds on R-Zero and its vendored verl trainer. Vocabulary dropout is implemented as a drop-in vocabulary constraint on the proposer's sampling. Core additions:

• verl/utils/vocab_constraint.py — the VocabDropout mask (fresh per-batch Bernoulli mask, retention probability alpha, optional linear annealing) plus the other constraint methods (topk, random, adaptive, rotating).
• verl/utils/diversity_controller.py — the compression-ratio adaptive controller used for the annealing analysis.
• verl/workers/rollout/vllm_rollout_spmd.py, verl/workers/rollout/config.py — wiring the constraint into vLLM sampling via allowed_token_ids.
• verl/trainer/core_algos.py, verl/workers/actor/config.py — the KL-Cov and Clip-Cov entropy-control losses (Cui et al. 2025) used as baselines.
• question_generate/, question_evaluate/, evaluation/, scripts/ — the co-evolution training loop, curriculum generation, and evaluation.
• scripts/compute_vendi.py, scripts/compute_epiplexity.py, scripts/compute_diversity_metrics.py — the diversity metrics reported in the paper.

Setup

Python 3.10+ with CUDA is required (flash_attn and vllm are pinned to CUDA builds; the pins in requirements.txt are the exact freeze used for the paper's runs).

pip install -r requirements.txt
# Secrets are read from the environment (never commit tokens):
export HF_TOKEN=...                 # HuggingFace access token (required)
export HUGGINGFACENAME=your-hf-user # HF username/org (required for training)
export STORAGE_PATH=/path/to/storage # defaults to ./rzero_artifacts
export WANDB_API_KEY=...            # optional, for logging

HF_TOKEN and HUGGINGFACENAME are required for training runs. After each proposer phase, the filtered curriculum is pushed to the Hub as a public dataset under HUGGINGFACENAME/<run_name> and the solver phase reads it back from there (edit private=False in question_evaluate/upload.py if you want your generated curricula private).

Paper experiments use 2 H200 GPUs per run via scripts/main.sh, with one GPU for GRPO training and one for the vLLM inference server (optional third and fourth arguments, defaulting to 0 and 1).

Running

Every run goes through scripts/main.sh <base_model> <run_name>, which alternates proposer (questioner) and solver GRPO across TOTAL_ITERATIONS (default 5) co-evolution iterations and then evaluates. Behaviour is set with environment variables.

Vocabulary dropout (our method)

export STORAGE_PATH=/path/to/storage
VOCAB_METHOD=vocab_dropout \
QGEN_VOCAB_METHOD=vocab_dropout \
VOCAB_RATIO=0.75 \
bash scripts/main.sh Qwen/Qwen3-8B-Base vd75_8b

• VOCAB_METHOD / QGEN_VOCAB_METHOD — proposer constraint during GRPO training and during curriculum generation. Set both to vocab_dropout for VD.
• VOCAB_RATIO — the retention probability alpha. Each token survives with probability alpha under a fresh Bernoulli mask per batch. The paper uses 0.75 at 8B and 0.85 at 4B. VOCAB_RATIO=1.0 retains every token, i.e. no dropout — identical to the baseline.
• VOCAB_ALPHA_END — target alpha for linear annealing across iterations (appendix); -1.0 (default) keeps alpha static.
• SOLVER_VOCAB_METHOD / SOLVER_VOCAB_RATIO — optionally apply dropout during solver training too (default none / 1.0).

Baseline (R-Zero, no dropout)

VOCAB_METHOD=none QGEN_VOCAB_METHOD=none SOLVER_VOCAB_METHOD=none \
bash scripts/main.sh Qwen/Qwen3-8B-Base baseline_8b

This is equivalent to VOCAB_METHOD=vocab_dropout VOCAB_RATIO=1.0 (alpha = 1 means no masking).

Entropy-control baselines (Cui et al. 2025)

Swap the proposer's GRPO policy loss with LOSS_MODE, keeping vocabulary dropout off.

KL-Cov:

VOCAB_METHOD=none QGEN_VOCAB_METHOD=none SOLVER_VOCAB_METHOD=none \
LOSS_MODE=kl_cov \
bash scripts/main.sh Qwen/Qwen3-8B-Base klcov_8b

Clip-Cov (hyperparameters match Cui et al.'s 7B headline):

VOCAB_METHOD=none QGEN_VOCAB_METHOD=none SOLVER_VOCAB_METHOD=none \
LOSS_MODE=clip_cov CLIP_COV_RATIO=2e-4 CLIP_COV_LB=1.0 CLIP_COV_UB=5.0 \
bash scripts/main.sh Qwen/Qwen3-8B-Base clipcov_8b

The loss implementations are compute_policy_loss_kl_cov and compute_policy_loss_clip_cov in verl/trainer/core_algos.py; the knobs (loss_mode, clip_cov_ratio, clip_cov_lb, clip_cov_ub) are in verl/workers/actor/config.py.

Diversity metrics

Compute the paper's diversity metrics on generated curricula. Both scripts read iteration outputs from $STORAGE_PATH/generated_question/ (written during training) and take the run name via --experiment:

python scripts/compute_vendi.py --experiment <run_name> --all_iterations   # semantic (Vendi score)
python scripts/compute_epiplexity.py --experiment <run_name>               # functional (prequential MDL)

License

Apache License 2.0 (see LICENSE). The verl/ directory contains modified code from verl (Bytedance, Apache 2.0), with the original license headers retained.

Acknowledgements

Built on R-Zero (Huang et al., 2025) and verl. Please cite those works as well when using this codebase.