ROM: Real-time Overthinking Mitigation via Streaming Detection and Intervention

Xinyan Wang¹, Xiaogeng Liu², Ming Pei¹, Chaowei Xiao²

¹University of Wisconsin-Madison   ²Johns Hopkins University

Abstract

Large Reasoning Models (LRMs) often reach a correct solution before their long Chain-of-Thought trace ends, yet continue with redundant verification, repeated attempts, or unnecessary exploration that wastes computation and can even overturn the correct answer. We frame this behavior as a latent productive-to-redundant transition and show it is directly reflected in hidden states: around first-correct-solution (FCS) boundaries, late-layer representations separate efficient from overthinking tokens, while boundary-permutation and position controls collapse. We propose ROM, a streaming intervention framework that monitors a frozen LRM with a lightweight hidden-state detector (~0.1% of backbone parameters) and intervenes at well-formed reasoning boundaries; Counterfactual Self-Correction (CSC) balances supervision with wrong→correct trajectories, preserving useful pre-FCS self-correction. Unlike prior adaptive early-exit methods, ROM extracts no intermediate answers, launches no probe decoding, and updates no backbone weights. Across five backbones from three model families and five reasoning benchmarks, against ten recent baselines under a shared protocol, ROM_CSC attains the highest accuracy in 19 of 25 model–benchmark settings, cuts response length by 28–77% (mean 45%) versus vanilla decoding, and is the only method on the accuracy–length Pareto front in every setting. The same MATH500-trained supervision transfers zero-shot across scales, families, and task domains, and end-to-end wall-clock latency drops by 46.5% with ~5% per-token overhead.

Project Structure

ROM/
├── rom/                        # Core package
│   ├── models.py               # StreamingHead, Qwen3WithHead
│   ├── dataset.py              # Dataset loading & hidden-state cache
│   ├── train.py                # Training pipeline
│   ├── eval.py                 # Evaluation (vLLM)
│   ├── env.py                  # Environment setup
│   └── utils/
│       ├── math.py             # Answer extraction & correctness checking
│       └── eval_helpers.py     # Streaming probs, backtracing, metrics
├── configs/
│   ├── train.yaml              # Training defaults (all hyperparameters)
│   └── eval.yaml               # Evaluation defaults
├── data/                       # CSC training data (downloaded, not tracked)
├── assets/
├── requirements.txt
├── LICENSE
└── README.md

Quick Start

Installation

pip install -r requirements.txt

Requires Python 3.11+, PyTorch >= 2.9.0, and a CUDA-capable GPU. All results were produced on a single NVIDIA A100 (80 GB).

Data

The CSC training data (740 efficient + 793 overthinking = 1,533 samples) is hosted on HuggingFace: xinyan-wang/ROM.

Download and place under data/:

# Using huggingface-cli
huggingface-cli download xinyan-wang/ROM --repo-type dataset --local-dir data

Training

All hyperparameters live in configs/train.yaml (20 epochs, AdamW, lr 5e-5, weight decay 0.1, effective batch 32, cosine schedule with 0.1 warmup, bf16). Run with defaults:

python -m rom.train

Override via CLI:

python -m rom.train --lr 1e-4 --num_train_epochs 30

Hidden states are cached on the first run, so training itself completes in under an hour with the backbone frozen.

W&B logging is enabled by default. Disable with --no_wandb.

Evaluation

We evaluate on MATH500 (held-out 100 problems), GSM8K (full test set, 1,319), AIME25 (all 30), GPQA-Diamond (full set, 198), and MMLU-Pro (70-question validation split), served via vLLM on a single A100 (80 GB) at temperature 0.6, top-p 0.95, top-k 20, seed 46, with n=3 samples per problem (n=10 on AIME25) and a uniform 8,192-token output budget.

python -m rom.eval

Override as needed:

python -m rom.eval --ckpt_path checkpoints/my_model.pt --test_data data/test_data/math500.jsonl

ROM triggers at decision threshold 0.5 (--threshold) and backtraces to the nearest well-formed reasoning boundary before appending the final-answer cue. Pass --no_backtrack to cut directly at the trigger token (the ablation row in the paper).

Citation

If you find ROM useful, please cite our paper 📝 and give us a ⭐!

@misc{wang2026romrealtimeoverthinkingmitigation,
      title={ROM: Real-time Overthinking Mitigation via Streaming Detection and Intervention},
      author={Xinyan Wang and Xiaogeng Liu and Ming Pei and Chaowei Xiao},
      year={2026},
      eprint={2603.22016},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.22016},
}

License

This project is licensed under the MIT License.