Conformal Thinking: Risk Control for Reasoning on a Compute Budget

Xi Wang, Anushri Suresh, Alvin Zhang, Rishi More, William Jurayj, Mehrdad Farajtabar, Daniel Khashabi, Eric Nalisnick

Abstract:

Reasoning Large Language Models (LLMs) enable test-time scaling, with dataset-level accuracy improving as the token budget increases,
motivating adaptive reasoning---spending tokens when they improve reliability and stopping early when additional computation is unlikely to help.
However, setting the token budget, as well as the threshold for adaptive reasoning, is a practical challenge that entails a fundamental risk-accuracy trade-off. We re-frame the budget setting problem as risk control, limiting the error rate while minimizing compute.
Our framework introduces an upper threshold that stops reasoning when the model is confident (risking incorrect output) and a novel parametric \emph{lower} threshold that preemptively stops unsolvable instances (risking premature stoppage).
Given a target risk and a validation set, we use distribution-free risk control to optimally specify these stopping mechanisms.
Empirical results across diverse reasoning tasks and models demonstrate the effectiveness of our risk control approach, demonstrating computational efficiency gains from the lower threshold and ensemble stopping mechanisms, all while adhering to the user-specified risk target.

Lay Summary:

Current works on reasoning model early stopping have two issues:
- Stopping criteria / threshold value is hard to set, making it hard to use in practice.
- Only adopt the "stop when confidence" principle, often use up all budgets on unsolvable problems, where the confidence is not decreasing.
Our work utilizes a risk control framework and a validation set to solve the first problem, and a lower threshold mechanism thresholding the progress rate to solve the second problem

Link To Code: https://github.com/xidulu/reasoning_risk_control/

Primary Area: Deep Learning->Large Language Models

Keywords: Risk Control, LLM Reasoning, Early-Exit Algorithm

Submission Number: 20587