Optimal Stopping of Self-Refining Foundation Models

The paper Optimal Stopping of Self-Refining Foundation Models by Kim Hammar, Tansu Alpcan, and Emil C. Lupu addresses the problem of deciding when to stop the self-refinement process of foundation models. The abstract states: "Foundation models can improve their outputs through a self-refinement process driven by external feedback. In this process, the model is embedded in an iterative loop where it generates outputs, receives feedback from verifiers, and refines its responses through in-context learning. Following a novel approach, we formalize this process as an optimal stopping problem where the number of refinement iterations is decided based on expected improvement relative to cost. We derive optimal stopping policies and show that they can be efficiently computed through stochastic approximation. To evaluate our approach experimentally, we apply it to a coding benchmark for foundation models. The empirical results show that our stopping policies are significantly more cost-efficient than stopping policies proposed in prior work."

The paper introduces a decision-theoretic model of the self-refinement process. The authors state: we formulate self-refinement as an optimal stopping problem in which the number of refinement iterations is decided based on the expected improvement relative to the cost of invoking the foundation model. The contributions are summarized as: We present a novel formulation of self-refinement in foundation models as an optimal stopping problem and We derive optimal stopping policies and validate them on a coding benchmark. The results show that they are more cost-efficient than stopping policies proposed in prior work.

The self-refinement use case is described as follows: "We consider a scenario where a foundation model is used to solve a task specified by an instruction in natural language... To evaluate the solution produced by the model, we associate it with a score x ∈ [0, 1], with x = 1 being the optimal score... We assume that the model is embedded in an iterative loop where it refines its output based on feedback... Each iteration invokes the model with the current score x ∈ [0, 1] as feedback, which leads to a new output that receives an updated score x′ and incurs a cost c > 0."

The optimal stopping formulation is formalized as a discrete-time dynamical system: "We formalize the self-refinement use case described above as a discrete-time dynamical system where the score evolves as a Markov process (xk)N k=0. At each stage k ∈ 0, 1,..., N − 1, two controls are available: (S)top and (C)ontinue. Control u = S in state x yields a payoff g(x) ≥ 0 that quantifies the quality of the output and terminates the process. Conversely, control u = C incurs a cost c > 0 and transitions the system to the next state according to xk+1 = f (xk, wk). The objective is: µ⋆ ∈ arg max E g(xτµ) − cτµ subject to (1). The optimal value function satisfies the Bellman equation: Vk⋆ (x) = max g(x), −c + E Vk+1 ⋆ (f (x, wk))."

The paper identifies the stopping problem components through system identification. The continuation cost c is defined as the average monetary cost of performing a self-refinement iteration with values: HAIKU 4.5 has c = 0.01, GEMINI FLASH-LITE 3.1 has c = 0.0025, and GPT CODEX MINI 5.1 has c = 0.005. The payoff function is defined as g(x) = βx, where β > 0 is a weighting factor. The system function f is estimated using a Gaussian process model: xk+1 = f (xk, wk) = min 1, max xk, q(xk) + wk, where q is an unknown score function estimated via a Gaussian process prior with mean m(x) = x and a Matérn covariance function.

From the identified models, the authors extract two structural observations: Observation 1 (Monotonicity). The identified model q̃ [cf. (6)] is nondecreasing on [0, 1] and Observation 2 (Diminishing returns). The difference q̃(x) − x [cf. (6)] is nonincreasing on [0, 1].

Based on these observations, the paper derives structural properties of the optimal stopping policy. Proposition 2 states: "There exist an optimal policy µ⋆ and thresholds α0, α1,..., αN −1 ∈ [0, 1] such that µ⋆k (x) = S if x ≥ αk, C if x < αk. Proposition 3 states: The optimal thresholds [cf. (8)] are nonincreasing in the stage k, i.e., α0 ≥ α1 ≥ · · · ≥ αN −1. Proposition 4 states: The optimal stopping sets satisfy S0⋆ = S1⋆ = · · · = SN⋆ −1. Equivalently, the optimal thresholds are stage-independent, i.e., α0 = α1 = · · · = αN −1 = α⋆."

The computation of the optimal policy is described: "By Props. 2–4, this computation reduces to optimizing the stopping threshold α, which we use to parameterize a threshold-based stopping policy µα; cf. (8). We implement this optimization using three methods: simultaneous perturbation stochastic approximation (SPSA) [37], the cross-entropy method [38], and differential evolution [39]. The results show: All methods converge to similar expected values, which suggests that they identify optimal or near-optimal thresholds. In all cases, convergence is achieved within seconds."

The experimental evaluation on EFFIBENCH compares the optimized threshold policy against fixed-iteration policies and the UCB policy. The results show: We observe that the optimized threshold policy µα achieves the highest expected value in all configurations. For fixed-iteration policies, fewer iterations perform better when β is small, while more iterations are preferred when β is large. The UCB policy automatically adapts to β but consistently underperforms the optimized threshold policy.

The discussion summarizes the main takeaways: 1) Self-refinement exhibits diminishing returns... 2) Optimal threshold-based stopping policies exist... 3) Optimal stopping improves over the state-of-the-art.

The conclusion states: "Foundation models can improve their outputs through a self-refinement process where they iteratively critique and refine their own outputs. We show that this process can be formulated as an optimal stopping problem in which the decision to continue refining or to stop is made sequentially based on the expected improvement relative to the refinement cost. We establish conditions for optimal threshold-based stopping policies and validate them on a coding benchmark across three frontier models. Empirical results demonstrate that our stopping policy significantly improves cost-efficiency compared to the state-of-the-art."

Improvements to AI Systems:

1. Adaptive Cost-Aware Self-Refinement: AI systems can now dynamically decide whether to continue refining their outputs (e.g., code, text, or reasoning) based on a learned threshold that balances expected quality gain against computational/monetary cost. This eliminates wasteful over-refinement and premature stopping, enabling real-time resource allocation.

2. Threshold-Based Stopping with Monotonicity Guarantees: By exploiting the identified structural properties (monotonicity and diminishing returns of the score function), AI systems can implement a simple, stage-independent threshold policy (α*) that is provably optimal. This replaces complex, ad-hoc stopping heuristics with a mathematically grounded, computationally efficient rule (computable in seconds via stochastic approximation).

3. Model-Agnostic Cost Optimization: The system can be applied to any foundation model (e.g., HAIKU, GEMINI, GPT-CODEX) by estimating its specific cost parameter (c) and score-transition function (f) via Gaussian process regression. This allows AI systems to self-tune their refinement behavior per model and per task without retraining.

4. Inference-Time Decision Making: The AI system can now make a binary decision at each refinement step—stop and output the current result, or continue and pay a cost—using only the current score x. This is a lightweight, online decision rule that does not require lookahead or full trajectory simulation, making it suitable for latency-sensitive applications.

5. Improved Value Estimation under Uncertainty: Using the Bellman equation and optimal stopping framework, the AI system can compute the expected net payoff (quality minus cost) for any given state, enabling it to prioritize refinement only when the expected marginal gain exceeds the cost. This is particularly useful in multi-step reasoning, code generation, and iterative problem-solving tasks.

6. Benchmark-Driven Calibration: The system can be calibrated on domain-specific benchmarks (like EFFIBENCH) to identify the optimal threshold α* for a given model and task distribution. This yields a plug-and-play policy that outperforms fixed-iteration and UCB-based baselines in cost-efficiency, as demonstrated empirically.

What the Improved AI System Can Do:

• Self-Refine Efficiently: It will automatically stop refining once the expected improvement falls below the cost of another iteration, avoiding unnecessary API calls or compute cycles.

• Adapt to Task Difficulty: It will use fewer iterations for easy tasks (where scores saturate quickly) and more for hard tasks (where scores are still improving), based on the learned score dynamics.

• Operate in Real-Time: It can make stop/continue decisions in milliseconds, using only the current score, without needing to simulate future states.

• Generalize Across Models: It can be quickly recalibrated for new foundation models by estimating their cost and score-transition parameters, ensuring optimal behavior without manual tuning.

• Provide Cost-Performance Trade-offs: It can explicitly trade off output quality against computational budget by adjusting the weighting factor β in the payoff function, allowing users to specify their preference (e.g., high quality at higher cost, or lower cost with acceptable quality).

Foundation models can improve their outputs through a self-refinement process driven by external feedback. In this process, the model is embedded in an iterative loop where it generates outputs, receives feedback from verifiers, and refines its responses through in-context learning. Following a novel approach, we formalize this process as an optimal stopping problem where the number of refinement iterations is decided based on expected improvement relative to cost. We derive optimal stopping policies and show that they can be efficiently computed through stochastic approximation. To evaluate our approach experimentally, we apply it to a coding benchmark for foundation models. The empirical results show that our stopping policies are significantly more cost-efficient than stopping policies proposed in prior work.