When Does In-Context Search Help? A Sampling-Complexity Theory of Reflection-Driven Reasoning

2026-07-07

AI Summary

This work investigates how context search can effectively enhance the performance of large language models in iterative reasoning, with a focus on the sampling complexity inherent in the critique-and-revise process. The authors model this process as approximate Bayesian inference over reasoning trajectories, where the base model provides a prior and self-reflection yields posterior updates. Theoretical analysis demonstrates that, under a local correction condition, only a polynomial number of sequential attempts are required to overcome the exponential decay in zero-shot success rates; moreover, when reflection reliably identifies early errors, context search yields exponential performance gains. This mechanism remains robust under approximate updates and can be efficiently learned via cross-entropy training. Key theoretical predictions are empirically validated on real large language models.

Abstract

Training large language models (LLMs) with extended reasoning has enabled in-context search, in which models iteratively generate, critique, and revise solution attempts. We provide a theoretical analysis of in-context search by modeling it as approximate inference over reasoning traces, where the base model defines a prior and self-reflection provides feedback for posterior updates, and study the resulting inference-time sampling complexity - the number of sequential attempts needed to achieve high success probability. We show that when reflections reliably localize early mistakes, in-context search can yield exponential improvements over the base model, solving problems with exponentially small zero-shot pass rates using only a polynomial number of sequential attempts, whereas when this property fails, conditioning on past attempts offers no asymptotic benefit over parallel sampling. We further show that these gains are robust and learnable: approximate posterior updates suffice, and cross-entropy training on search rollouts recovers the required behavior with polynomial sample complexity. Finally, we show that under a stagewise abstraction of reinforcement learning with verifiable rewards, the optimal policy extension implements the same posterior reweighting rule. We validate key qualitative predictions of the theory on real large reasoning models.

Problem

• in-context search
• reasoning
• sampling complexity
• self-reflection
• large language models

Innovation

• in-context search
• sampling complexity
• reflection-driven reasoning
• approximate inference
• posterior reweighting