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

2026-08-10

AI Summary

This work addresses the challenge of improving the accuracy of large language models under fixed inference compute budgets, rather than merely increasing output diversity. It proposes a test-time input augmentation (TTA) strategy that systematically explores three classes of input-side enhancements—semantic paraphrasing, lexical perturbation, and visual transformation—combined with chain-of-thought prompting and prediction aggregation. Empirical results demonstrate, for the first time, that this approach significantly outperforms conventional self-consistency on five out of six benchmark tasks, achieving statistically significant accuracy gains at comparable computational cost. The method yields an average cost-effectiveness improvement of approximately 1.8×, exhibits particular efficacy for medium-scale models, and achieves Pareto dominance in the trade-off between cost and performance.

Abstract

Test-time scaling improves LLM accuracy but multiplies inference cost, making the accuracy gained per unit of compute the metric that matters in deployment. Self-consistency is one of the established approaches, which spends this budget entirely on the output side by sampling repeated reasoning paths. We study Test-Time Augmentation (TTA), which extends self-consistency by also perturbing the input, aggregating predictions across transformed versions of the input, and ask whether input-side diversity converts compute into accuracy more efficiently than output-side diversity. We perform a systematic, matched-compute comparison: we evaluate three simple input-side strategies (semantic rephrasing, lexical perturbations, and visual transformations) across six datasets covering general and multilingual knowledge, mathematical reasoning, multi-modal question answering, and sentiment classification, against chain-of-thought prompting and self-consistency. Semantic rephrasing delivers consistent and statistically significant accuracy gains while Pareto-dominating self-consistency on cost-effectiveness, delivering roughly 1.8X more accuracy per dollar and outperforming it on five of six tasks. We further analyze the number of augmentations, multi-modal strategies, and base model scaling, finding that TTA is most cost-effective for mid-tier models where a stronger model is unavailable or too expensive. Our findings indicate that for current mid-tier LLMs, varying the input converts inference compute into accuracy more efficiently than varying the reasoning path alone. The TTA implementation is available at https://github.com/aws-samples/sample-genai-reflection-for-bedrock.

Problem

• Test-Time Augmentation
• Large Language Models
• Inference Efficiency
• Input Diversity
• Accuracy per Compute

Innovation

• Test-Time Augmentation
• Input Diversity
• Self-Consistency
• Cost-Effectiveness
• Large Language Models