Adaptive Test-Time Compute Allocation with Evolving In-Context Demonstrations

1. Problem Statement

Large Language Models (LLMs) have shown that increasing test-time compute (e.g., generating multiple candidate responses and selecting the best one) can significantly improve performance, often surpassing the gains from training-time scaling. However, existing test-time scaling methods suffer from two primary limitations:

1. Static Generation Distributions: Most methods (e.g., Best-of-N, Self-Consistency) allocate compute by sampling repeatedly from a fixed conditional distribution $p(\cdot|x)$. Once a question is difficult, simply resampling from the same distribution yields diminishing returns.
2. Uniform or Pre-defined Allocation: Many approaches either allocate the same number of samples to all questions (uniform) or use static predictors to allocate samples based on estimated difficulty. They fail to dynamically adjust how the model generates responses based on the evolving context of the test set itself.

The core problem is to design a framework that jointly decides where to spend compute (allocation) and how to use that compute to reshape the model's generation distribution (adaptation) during inference, without modifying model parameters.

2. Methodology

The authors propose a unified framework called Self-Evolving Test-Time Allocation, which operates in two distinct stages. The method leverages In-Context Learning (ICL) dynamically constructed from the test set itself to shift the sampling distribution.

A. Core Intuition

The key insight is that effective test-time scaling requires not just more samples, but better samples. By conditioning the generation of a difficult query on successful responses from semantically similar queries (which have already been solved), the model can shift its sampling distribution toward more promising reasoning trajectories.

B. Algorithm Overview

The process is divided into two stages:

1. Warm-up Phase (Difficulty Estimation & Pool Construction)

• Action: A small, fixed budget of compute is allocated to every query in the test set.
• Mechanism: The model generates $K$ responses for each question using a standard zero-shot prompt (fixed distribution).
• Outcome:
  • Solved Queries: If a correct response is found (verified by a reward oracle/ground truth), the query is marked as solved and removed from the active set.
  • Demonstration Pool ($D_{test}$): The correct responses from these solved queries are stored. This forms an initial pool of "demonstrations" derived directly from the test set.
  • Purpose: This resolves easy cases early (saving compute) and seeds the pool with high-quality examples for the next stage.

2. Adaptive Distribution Sampling Phase

• Action: Focuses compute resources on the remaining "unresolved" (hard) queries.
• Mechanism:
  • Semantic Retrieval: For each active query $x$, the system retrieves the top-$P$ most semantically similar queries ($z$) from the solved set $D_{test}$ using cosine similarity on embeddings.
  • Dynamic Prompting: An ICL prompt $\Pi(x)$ is constructed by pairing the target query $x$ with the successful responses $\tilde{y}(z)$ of its neighbors.
  • Distribution Shift: The model generates new responses for $x$ conditioned on this prompt: $y \sim p(\cdot | x, \Pi(x))$.
  • Iteration: As more queries are solved in this phase, $D_{test}$ grows, providing richer and more relevant demonstrations for subsequent queries. The process repeats until all queries are solved or the compute budget is exhausted.

Key Distinction: Unlike previous methods that sample from $p(\cdot|x)$, this method samples from $p(\cdot|x, \Pi(x))$, where $\Pi(x)$ evolves as the test set is processed.

3. Key Contributions

1. Unified Framework: The paper introduces a novel framework that unifies compute allocation (deciding which questions to work on) with distributional adaptation (changing how the model generates answers).
2. Self-Evolving ICL: It proposes a mechanism to use the test set itself as a dynamic source of in-context examples. As the system solves problems, it "learns" which examples are most helpful for similar problems, effectively performing test-time adaptation without parameter updates.
3. Efficiency: The method achieves higher performance with significantly fewer output tokens compared to baselines, as it avoids redundant sampling on hard queries by shifting the generation distribution.
4. Modularity: The framework is agnostic to the specific sampling intervention (e.g., it could theoretically use temperature tuning instead of ICL) and the allocation strategy.

4. Experimental Results

The method was evaluated across four diverse benchmarks: MATH-500, LiveCodeBench, MinervaMath, and GPQA-Diamond, using various model families (Gemini-2.5, GPT-4.1 Mini, GPT-5 Nano).

• Performance Gains:
  • The proposed method (Ours (Adpt)) consistently outperformed uniform sampling (Unif) and adaptive allocation without distribution shifts (Elim).
  • On the MATH-500 benchmark, the method achieved significantly higher coverage (e.g., ~96% vs. ~92% for baselines) with fewer total tokens.
  • On LiveCodeBench and MinervaMath, the gains were particularly pronounced for reasoning-heavy tasks, demonstrating the value of semantic similarity in guiding code/math generation.
• Token Efficiency:
  • The method required substantially fewer output tokens to reach the same coverage levels as baselines.
  • Analysis shows that by resolving "unsolvable" queries early via ICL shifts, the method prevents the model from wasting tokens on long, hallucinated reasoning chains that occur under static sampling.
• Ablation Studies:
  • Fixed vs. Evolving: Using a fixed set of demonstrations performed worse than the evolving pool, confirming the importance of updating the context as the test set is processed.
  • Random vs. Semantic: Randomly selecting demonstrations (Ours (Rand)) yielded inconsistent results, often underperforming baselines on reasoning tasks. This highlights that semantic similarity is critical for effective distribution shifting.
  • Robustness: The method remained robust even when using noisy reward models (simulating imperfect verification).

5. Significance and Implications

• Rethinking Test-Time Scaling: The paper challenges the notion that test-time scaling is merely about "generating more samples." It argues that the quality of the generation distribution is equally important and can be dynamically improved.
• Cost-Effective Inference: By achieving higher accuracy with fewer tokens, this approach offers a practical pathway to deploying powerful reasoning capabilities in LLMs without the prohibitive costs of massive parameter scaling or extensive fine-tuning.
• Self-Improving Systems: The concept of using the test set to bootstrap its own context (self-evolving ICL) suggests a new paradigm where inference systems become more capable as they process more data, even without external training.
• Generalizability: The framework is not limited to math or code; the principles of adaptive allocation and semantic conditioning could apply to any domain where LLMs face heterogeneous difficulty levels.

In summary, this work demonstrates that adaptive test-time compute allocation combined with evolving in-context demonstrations is a superior strategy for improving LLM reasoning, offering a balance of high performance and computational efficiency that static methods cannot match.