Test-Time Augmentation for LLMs: When Input Diversity Beats Output Diversity at Matched Compute

Analysis

Why This Paper Matters

As large language models (LLMs) become more capable, the cost of inference at scale has emerged as a critical bottleneck. A growing body of work, exemplified by Snell et al., has shown that scaling test-time compute—spending more computation during inference—can improve accuracy without retraining. However, most existing methods focus on output-level diversity: sampling multiple outputs from the same prompt and aggregating them (e.g., self-consistency, best-of-n). This paper challenges that default by asking a simple but profound question: at a fixed compute budget, is it better to diversify the inputs or the outputs?

The answer, according to the authors, is that input diversity often wins. By systematically comparing input-level test-time augmentation (paraphrasing, reformatting, etc.) against output-level sampling under matched compute, the paper provides a new lens for optimizing inference. This matters because it suggests that many current test-time scaling strategies may be leaving accuracy on the table. For practitioners deploying LLMs in production, the implication is that prompt augmentation could be a more efficient lever than output sampling.

Moreover, the paper arrives at a time when the community is actively searching for compute-optimal inference recipes. If input diversity is indeed more effective, it could reshape how we design inference pipelines, allocate compute, and even train models to be robust to input variations.

Technical Contributions

• Matched-compute comparison: The authors design experiments that hold total inference compute constant while varying the allocation between input augmentation and output sampling. This controlled setup isolates the effect of diversity source.
• Input augmentation strategies: They explore a range of input-level transformations, such as paraphrasing, formatting changes, and possibly translation or style shifts, to generate diverse prompts.
• Output sampling baselines: They compare against standard output-level methods like temperature sampling, top-k sampling, and self-consistency with majority voting.
• Empirical evaluation: The study spans multiple LLMs and tasks, providing evidence that input diversity can yield higher accuracy than output diversity at the same compute cost.

Results

The abstract does not report specific metrics, but the central finding is that input-level augmentation outperforms output-level sampling at matched compute. This suggests that the marginal benefit of generating diverse prompts exceeds that of generating diverse outputs. The result likely holds across reasoning and knowledge-intensive tasks, though exact numbers and statistical tests are not provided in the abstract. The paper also implies that the optimal mix may depend on the task and model, but the general trend favors input diversity.

Significance

This work has the potential to shift the focus of test-time compute research from output sampling to input augmentation. It provides a practical guideline for practitioners: when you have a fixed compute budget, consider investing in prompt diversity before output diversity. It also raises important questions about why input diversity is more effective—perhaps because it explores different regions of the input space, leading to more varied reasoning paths. Ultimately, this could lead to more efficient and accurate LLM deployments, and inspire new methods for automatic prompt augmentation and compute allocation.