Sample Complexity and Representation Ability of Test-time Scaling Paradigms

TL;DR

This paper provides a theoretical analysis of test-time scaling methods for large language models, revealing their sample efficiency differences and demonstrating how self-correction enhances multi-task capabilities.

Contribution

It establishes sample complexity bounds for self-consistency and best-of-$n$ strategies, and shows how self-correction enables Transformers to perform multi-task learning at test time.

Findings

Self-consistency requires $ heta(1/ riangle^2)$ samples, while best-of-$n$ needs $ heta(1/ riangle)$.

Self-correction with verifier feedback allows Transformers to simulate online learning.

Empirical validation confirms the effectiveness of self-correction methods.

Abstract

Test-time scaling paradigms have significantly advanced the capabilities of large language models (LLMs) on complex tasks. Despite their empirical success, theoretical understanding of the sample efficiency of various test-time strategies -- such as self-consistency, best-of-$n$, and self-correction -- remains limited. In this work, we first establish a separation result between two repeated sampling strategies: self-consistency requires $\Theta(1/\Delta^2)$ samples to produce the correct answer, while best-of-$n$ only needs $\Theta(1/\Delta)$, where $\Delta < 1$ denotes the probability gap between the correct and second most likely answers. Next, we present an expressiveness result for the self-correction approach with verifier feedback: it enables Transformers to simulate online learning over a pool of experts at test time. Therefore, a single Transformer architecture can provably solve multiple tasks without prior knowledge of the specific task associated with a user query, extending the representation theory of Transformers from single-task to multi-task settings. Finally, we empirically validate our theoretical results, demonstrating the practical effectiveness of self-correction methods.