Search papers, labs, and topics across Lattice.
This paper conducts a theoretical analysis of in-context search in large language models (LLMs), modeling it as approximate inference over reasoning traces. The authors demonstrate that when self-reflection effectively identifies early mistakes, in-context search can exponentially enhance performance, allowing models to solve complex problems with significantly fewer attempts. Conversely, when reflections fail to localize errors, the benefits diminish, revealing critical insights into the conditions under which in-context search is most effective.
In-context search can exponentially boost LLM performance by leveraging self-reflection to correct early mistakes, but only under specific conditions.
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.
