Prompting Test-Time Scaling Is A Strong LLM Reasoning Data Augmentation

AI summary

Overview

• Research area: Natural Language Processing / Large Language Model reasoning, specifically inference-time data augmentation and supervised fine-tuning data efficiency.
• Technical level: Intermediate. The underlying recipe (wrap prompts, query a teacher, fine-tune a student) is simple, but the paper also leans on bias-variance reasoning, information-theoretic framing, and embedding/trigram diversity metrics.
• Scope in one sentence: The paper introduces Prompting Test-Time Scaling (P-TTS), an inference-time data augmentation strategy that turns only 90 seed AIME problems into up to 900 reasoning training examples via principled instructional wrappers, and shows that the resulting Qwen2.5 fine-tunes (7B and 32B) beat S1/S1.1 1K-shot baselines on AIME24, AIME25, MATH500, and GPQA-Diamond.

What This Paper Is About

Building strong reasoning LLMs usually requires thousands to millions of curated chain-of-thought examples, which is expensive and brittle: it demands prompt engineering, human verification of multi-step solutions, and constant refreshing as data distributions shift. P-TTS attacks this by treating the prompt itself as a scaling axis rather than the dataset. Starting from just 90 manually selected AIME 2022–2024 problems, it systematically wraps each question in fixed instructional templates ("principles") and reward-framing paraphrases, queries a teacher model (DeepSeek-R1) to elicit diverse reasoning traces, then fine-tunes Qwen2.5-Instruct students on the augmented corpus.

Key Contributions

1. The P-TTS framework. A simple inference-time reasoning data-augmentation method that expands a compact seed pool through principle-guided instructional wrapping, producing prompt variants that preserve the original problem's tokens while changing its instructional framing.
2. A 90-seed corpus that beats 1K-shot static prompts. The authors show that 90 seeds, expanded through P-TTS, outperform 1K-shot baselines, reducing curation cost by an order of magnitude.
3. Empirical evidence that prompt-space exploration is an underutilized scaling dimension. They demonstrate accuracy scaling from 90 → 360 → 450 → 900 examples (multipliers m ∈ {1, 4, 5, 10}) and report diminishing returns beyond roughly 6 prompt-augmentations per base question.
4. Released code and the augmented exemplar pool at https://github.com/VILA-Lab/PTTS to support reproducibility and transfer to new domains.

Main Findings

• Headline gains over S1 and S1.1 (1K-shot). P-TTS-7B and P-TTS-32B outperform the prior competitive baselines S1 and S1.1, with absolute accuracy gains of +26.66% and +30.00% on AIME'24 (7B), and +13.34% and +6.67% on AIME'25 (7B). P-TTS-32B yields gains of +23.33% and +16.63% on AIME'24, and +26.63% and +3.33% on AIME'25 (versus S1 and S1.1, respectively), with comparable or better performance on MATH500 and GPQA-Diamond.

• Full 900-example 7B model reaches 49.03% average. In the data-volume ablation (Table 5), the model trained on Core + Seed + six Reward-variant prompts (900 examples) scores 43.33 on AIME24, 26.67 on AIME25, 84.20 on MATH500, and 41.92 on GPQA-D, averaging 49.03% — above the 1K-example S1.1 baseline (38.99%) and the S1-7B baseline (37.15%). The base Qwen2.5-7B-Instruct scores 13.33 / 6.67 / 76.40 / 36.36, averaging 33.19%.

• Scaling is consistent across corpus sizes. Core (360 examples) averages 39.07%, Core+Seed (450) averages 42.80%, and the full 900-example setting averages 49.03%. The paper reports the largest average gain between 450 and 900 examples (+6.23%), with improvements on all benchmarks, and a notable AIME24 jump from 360 to 450 (+13.33%).

• Reward framing is the strongest single principle. Under 90 examples each, P-TTS_Reward yields the highest overall gain (+6.67%) over the seed baseline (31.17% vs. 24.50%), improving accuracy across all benchmarks. P-TTS_Penalty is also strong, especially on MATH500 (63.20), though with a drop on AIME25 (0.00). P-TTS_Correctness (25.02 avg.) and P-TTS_Think (27.25 avg.) offer more modest improvements; the Think variant underperforms on MATH500 (58.20).

• Reward pairs best with Correctness. In pairwise ablations of 180 examples each, P-TTS_{Reward ∪ Correctness} achieves the highest overall accuracy (39.02%), above Reward ∪ Penalty (34.38%) and Reward ∪ Think (32.82%).

• Leave-one-out confirms Reward's central role. On the 360-example Core set (20.00 / 20.00 / 80.40 / 35.86, avg. 39.06), removing Reward causes the largest drop to 35.40%. Removing Correctness (37.24) or Penalty (38.53) causes moderate degradation, and removing Step-by-Step thinking has minimal or slightly positive effects (39.29%).

• Diversity correlates with accuracy. Single P-TTS variants with higher semantic diversity (Diversity Gain) — P-TTS_Reward and P-TTS_Penalty — show larger accuracy improvements. P-TTS_Reward attains the highest trigram diversity across both final responses and reasoning traces. The paper reports P-TTS_Reward improving performance by approximately 6.7% on AIME25 and 5.0% on GPQA-Diamond over the seed baseline.

• DeepSeek-R1 is the best teacher among three tested. Claude-3-Opus, DeepSeek-R1, and OpenAI Omni-4 were each used to collect 90 answer–reasoning pairs and fine-tune Qwen2.5-7B-Instruct; the DeepSeek-R1-trained model consistently outperformed the Claude and Omni-4 counterparts. Claude-3-Opus and DeepSeek-R1 natively return aligned answer–reasoning pairs, while Omni-4 required an augmented prompt.

• Prompt placement matters, slightly. Positioning the instructional prompt at the beginning (R1) leads the model to focus more on the instructions and achieves better accuracy than placing it at the end (R4), though when the question is relatively short the difference is small.

• Out-of-domain generalization improves zero-shot. P-TTS enhances zero-shot generalization on Gaokao, Kaoyan, OlympiadBench, AMC23, GradeSchoolMath, and Minerva.

• Ordering perturbations only slightly counter position bias, per the authors' summary of their diversity analysis.

Methodology in Plain English

1. Seed selection. 90 problems from AIME 2022–2024 were chosen for reasoning density across algebra, combinatorics, number theory, geometry, and probability; reliable formatting with definitive three-digit answers; and reduced contamination risk from restricting to recent editions.

2. Wrap, don't rewrite. Four core principles — Reward, Penalty, Correctness, and StepByStep — each have a fixed template that is simply concatenated in front of the unmodified question. Removing the template deterministically recovers the original problem, so semantics are preserved. Six Reward paraphrases (R1–R6) vary the incentive magnitude (e.g., $200 vs. $200,000 vs. $500,000), whether the tip is conditional, and placement (before vs. after the question).

3. Elicit traces from a teacher. A teacher model (DeepSeek-R1) is queried on the null (unwrapped) prompt and on every wrapped variant, returning a reasoning trace, a full response, and an answer for each.

4. Assemble four corpus sizes. Single-P-TTS (90 examples per principle), Core-P-TTS (4 × 90 = 360), Seed+Core (450), and Full P-TTS (10 × 90 = 900).

5. Fine-tune students. Qwen2.5-Instruct (7B/14B/32B) models are trained with an SFT recipe adapted from s1: the student predicts the full assistant output (reasoning + answer) with token-level cross-entropy computed on assistant tokens only, with user tokens masked. Each dataset configuration trains a separate model so the contribution of each prompting strategy can be isolated.

6. Evaluate deterministically. Using lm-evaluation-harness with temperature 0 (greedy decoding), reporting accuracy (equivalent to pass@1) on AIME24 (30 problems), AIME25 (15 problems), MATH500 (500 problems), and GPQA-Diamond (198 PhD-level science questions from Biology, Chemistry, and Physics, with reported expert performance of 69.7%).

7. Measure diversity two ways. Semantic diversity via Diversity Gain — the mean squared L2 distance from each Core example to its nearest seed neighbor using OpenAI's text-embedding-ada-002 — and surface-level diversity via trigram diversity, the ratio of non-overlapping distinct word-level trigrams between texts.

Why This Matters

Impact on research. The paper reframes test-time prompting from a one-shot design choice into a scalable, stochastic process. It supplies a complementary result to low-resource supervision work like S1 and LIMO (which claim roughly 1K high-quality prompts as a competitive alternative to massive corpora) by showing that prompt-space variation, not just data curation, can substitute for scale. It also positions instruction-level manipulation as distinct from data-level augmentation methods like MetaMath and ReasoningMix.

Real-world applications:

• Resource-constrained domains where collecting thousands of verified reasoning examples is infeasible — the paper explicitly frames P-TTS as practical for "resource-constrained or rapidly evolving domains."
• Rapidly shifting domains where static large-shot prompts go brittle and fixed exemplars bias the model toward spurious solution templates under domain shift.
• Math and science tutoring / evaluation pipelines that need chain-of-thought quality on AIME-style, olympiad, and PhD-level science questions without fresh annotation.
• Domain transfer to new verticals, since the code and augmented exemplar pool are released to facilitate rapid transfer.

Industry relevance. Reducing curation cost by an order of magnitude directly lowers the labeling and human-verification budget for teams building reasoning models. The finding that both exemplar diversity and scaled sampling schedules are critical drivers of improvement gives practitioners two concrete, cheap levers — template family and corpus multiplier — rather than a mandate to collect more raw data.

Future Directions

• Testing whether the 90-seed recipe transfers beyond math. The released exemplar pool is explicitly intended for "rapid transfer to new domains," but the reported benchmarks are math-centric plus GPQA-Diamond and the out-of-domain reasoning suite; generalizing the principle set to code, law, or medicine is untested here.
• Determining why Reward framing dominates. The paper shows the effect but the mechanism — whether it is semantic diversity, surface phrasing variety, or something about incentive framing itself — is framed as an analysis rather than settled.
• Pushing past the ~6 augmentation per question point. The authors report diminishing returns beyond roughly 6 prompt-augmentations per base question and call it a practical deployment point; whether other principles or schedules shift that saturation point is open.
• Combining prompt-space scaling with other scaling axes. The paper positions P-TTS as orthogonal to decoding-based inference-time scaling and to parameter/data scaling, leaving open how it composes with those approaches and with reinforcement-learning reasoning pipelines.

Target Audience

Researchers and engineers working on LLM reasoning, chain-of-thought data curation, and supervised fine-tuning under tight annotation budgets. It is most useful for practitioners who want a concrete, low-cost recipe for building reasoning training sets, and for academics studying inference-time scaling, prompt engineering as a control knob, and data-efficiency trade-offs in reasoning models. Readers need only a working familiarity with SFT and few-shot prompting to follow the method, though the diversity-metric sections assume some comfort with embeddings and n-gram statistics.

Authors’ abstract

Large language models (LLMs) have demonstrated impressive reasoning capabilities when provided with chain-of-thought exemplars, but curating large reasoning datasets remains laborious and resource-intensive. In this work, we introduce Prompting Test-Time Scaling (P-TTS), a simple yet effective inference-time data augmentation strategy for enhancing LLM reasoning through finetuning. Rather than collecting thousands or even millions of examples, P-TTS leverages a small pool of only 90 manually selected reasoning instances and systematically varies exemplar augmentation through principled instruction prompting intensities at test time to synthesize diverse reasoning trajectory contexts. Then we finetune the various sizes of Qwen-2.5 models on P-TTS data. Across a suite of mathematical reasoning AIME2024 & 25, MATH500, and GPQA-Diamond, our P-TTS-7B and 32B models outperform the prior competitive baselines like S1 and S1.1 (1K-shot), achieving absolute accuracy gains of +26.66% and +30.00% on AIME'24 (7B), and +13.34% and +6.67% on AIME'25 (7B); P-TTS-32B yields gains of +23.33% and +16.63% on AIME'24, and +26.63% and +3.33% on AIME'25 (vs. S1 and S1.1, respectively), with comparable or better performance on MATH500 and GPQA-Diamond. We further show that P-TTS enhances zero-shot generalization accuracy on out-of-domain reasoning benchmarks of Gaokao, Kaoyan, OlympiadBench, AMC23, GradeSchoolMath, and Minerva. Our analysis suggests that test-time scaling effectively explores the latent space of reasoning patterns, amplifying LLM problem-solving with minimal annotation overhead, and further unlocking the reasoning potential and capabilities of LLMs. Prompting Test-Time Scaling offers a practical, low-cost way to elicit LLM reasoning in resource-constrained or rapidly evolving domains.