Stanford CS329A: Test-Time Compute Scaling for AI Agents

Overview of Test-Time Compute Scaling

Test-time compute scaling allows AI models to improve their performance without changing model parameters or undergoing additional fine-tuning. While traditional LLM development focuses on pre-training (high compute, long duration) and fine-tuning (lower compute), test-time scaling shifts the compute budget to the inference stage to elicit better responses from existing models.

Parallel Sampling and the Power Law of Coverage

Repeated parallel sampling involves asking a model the same input problem multiple times and using a verifier to select the correct response. This approach can enable smaller, inferior models (e.g., Llama 3-8B) to outperform larger, proprietary models (e.g., GPT-4o) on specific tasks.

The Scaling Law for Inference

Research indicates that the relationship between coverage (the fraction of problems solved by at least one sample) and the number of samples (k) follows an exponential power law. This allows engineers to predict the resources required to achieve a specific coverage target across various model sizes and domains.

The Long Tail of Hard Problems

This power law behavior is driven by the distribution of problem difficulty. Most problems are solved at "pass@1," but there is a "long tail" of complex problems that the model solves only rarely. Increasing the number of samples increases the probability of hitting these rare correct solutions for the hardest problems.

The Generation-Verification Gap

While models can often generate a correct answer within a large sample set, the ability to identify that answer is a separate challenge known as the generation-verification gap.

Verifiable vs. Non-Verifiable Domains

  • Verifiable Domains: In coding (via unit tests), math (via formal proofs), or language translation (via equivalency checks), a perfect or near-perfect verifier exists. For example, in CUDA code generation, the output can be verified by comparing it to the source PyTorch code output.
  • Non-Verifiable Domains: In domains without automated verifiers, methods like majority voting (selecting the most frequent answer) often plateau quickly and fail to capture the rare correct answers for hard problems. Even LLM-based reward models often leave a significant gap between current verification capabilities and the model's true coverage.

Optimizing Test-Time Compute: Parallel vs. Sequential

Beyond simple repeated sampling, compute can be scaled through sequential revisions and guided search.

Sequential Revision

Instead of parallel attempts, the model generates an initial approach and iteratively revises it. This sequential approach allows the model to look at the problem from different angles and refine its logic.

Outcome vs. Process Reward Models

  • Outcome Reward Models (ORM): Score the final answer of a response.
  • Process Reward Models (PRM): Score each individual step of the generation process. PRMs can be used to guide a beam search, where the model expands only the most promising reasoning paths based on step-by-step scores.

Pre-training vs. Test-Time Scaling

For easy and medium-difficulty problems, increasing test-time compute is often more compute-efficient than increasing pre-training tokens. However, for the hardest problems, larger models with more extensive pre-training still maintain a performance advantage, even with significant test-time compute budgets.

Archon: Inference-Time Architecture Search

Archon is a framework that treats inference scaling as an architecture design problem, using a Bayesian optimizer to find the optimal combination of models and techniques for a given compute budget.

Inference-Time Operations

Archon utilizes several key operations:

  • Generation: Standard sampling from one or more LLMs.
  • Fusion: Asking an LLM to synthesize a single final answer based on multiple candidate responses.
  • Critic: Using a model to describe the strengths and weaknesses of a response.
  • Ranker: Prompting a model to rank candidates by quality.
  • Unit Test Generation: Generating tests to verify the correctness of a solution.

Key Findings from Archon

  • Deep Architectures: Stacking multiple layers of critiques, rankers, and fusers (creating a "deep" inference architecture) significantly improves accuracy.
  • Performance Gains: By optimizing these architectures, open-source models can match or exceed frontier closed-source models. Archon reported an average 14.1% improvement in pass@1 accuracy over GPT-4 and Claude 3.5 Sonnet on reasoning, math, and coding tasks.
  • Generalization: Architectures optimized for specific tasks can also be designed as general-purpose systems that perform well across diverse benchmarks.

Sources