CoDE-Stop

This is the repository of our paper: Early Stopping for Large Reasoning Models via Confidence Dynamics

User Guide

Environment Variables

The user should set at least the following variables in env_vars.py:

• CACHE_DIR: Huggingface cache directory where models will be loaded from
• NLTK_CACHE_DIR: cache for NLTK (location to download the sentence tokenizer)
• OPENAI_API_KEY: for using gpt-4o-mini as an LLM judge for certain benchmarks
• RESULTS_DIR: location for method results to be logged and cached

Quick Start

The inference.py script is the main program entrypoint. Use the --help option to learn the available models, datasets, methods, and other parameters. The METHOD_HYPERPARAMS dictionary lists the required hyperparameters for every method, and the METHOD_DEPENDENCIES dictionary lists the dependencies between early-stopping methods (reuse of cached rollouts to reduce total compute).

srun $PYTHON_EXEC inference.py \
    --model qwen-4b \
    --dataset gpqad \
    --method think_or_not \
    --num-chunks 20 \
    --chunk-idx $SLURM_ARRAY_TASK_ID \
    -B 32768 \
    -N 5 \
    --conf-threshold 0.95 \
    --stop-word Alternatively \
    --ewt True \
    --cp-cache False \
    --cond True \
    --degen-threshold 9.5 \
    --ramp-steps 1 \
    --ramp-min 0.9

Repository Structure

• Infrastructure
  • models.py: Loading HuggingFace models and configuring their parameters. get_model_and_tokenizer(model_id) loads a model from the MODEL_REGISTRY that has already been downloaded to the CACHE_DIR specified in env_vars.py
  • dataset_loader.py: Loads a math reasoning dataset. get_dataset(dataset_name) loads a dataset from the DATASET_REGISTRY
  • llm_eval.py: answer parsing and evaluation logic, including using an LLM judge. Evaluation method determined by dataset.
• Methods
  • method_prompts.py: Simple prompting-based approaches: step-by-step (the vanilla prompt), chain-of-draft, Nothinking, and budget-prompt.
  • method_think_or_not.py: Think-or-Not
  • method_wheels: Stop Spinning Wheels
  • method_ansconv_consistency: the Answer Consistency method from Answer Convergence
  • method_eat.py: Entropy-After-Think
  • method_deer.py: Dynamic Early Exit in Reasoning Models
  • method_codestop.py: Our method
• Scripts
  • basic_results.py: Prints accuracy and cost metrics for all models and rollout hyperparameters for a given dataset.
  • compact.py: Useful for compacting the cache files, if multiple GPUs were used for generation.
  • inference.py: The main program entrypoint, taking a model, dataset, and method hyperparameters as inputs

Citation

If you use CoDE-Stop in your research, please cite our paper:

@misc{hosseini2026earlystoppinglargereasoning,
    title={Early Stopping for Large Reasoning Models via Confidence Dynamics}, 
    author={Parsa Hosseini and Sumit Nawathe and Mahdi Salmani and Meisam Razaviyayn and Soheil Feizi},
	journal={arXiv preprint arXiv:2604.04930},
    year={2026},
}