Towards Theoretical Understanding of Transformer Test-Time Computing: Investigation on In-Context Linear Regression

Scaling test-time computation during language model inference, such as generating intermediate thoughts or sampling multiple candidate answers, has proven effective in improving model performance. While these techniques inherently rely on the stochastic nature of inference to explore diverse reasoning paths, prior theoretical works typically build on a deterministic decoding framework, overlooking the stochastic nature of practical language model inference. This work takes an initial step to bridge this gap by establishing a new theoretical framework, incorporating randomness and sampling directly into the decoding analysis. To demonstrate the framework's effectiveness, we apply it to the canonical in-context linear regression task with continuous and binary coefficients, simulating decoding via noise injection and sampling to analyze widely adopted inference techniques. We validate our theoretical findings through numerical simulations, with additional experiments on real-world tasks substantiating the framework's potential for practical applications.

Large language models (LLMs) perform better when given more time to "think" — for example, by generating multiple answers and picking the most common one. However, we lack mathematical understanding of why these techniques work, since existing theory assumes AI models behave deterministically, while real models are inherently random. We built a new theoretical framework that incorporates this randomness, and used it to prove when and why sampling-based methods outperform deterministic ones: majority voting escapes "dead ends" that trap greedy strategies, while averaging multiple answers prevents overfitting to noisy data. We also developed a practical tool to predict high-compute performance from cheap experiments, helping practitioners choose efficient inference settings.