Thinking Augmented Pre-training (2509.20186v1)

Abstract: This paper introduces a simple and scalable approach to improve the data efficiency of LLM training by augmenting existing text data with thinking trajectories. The compute for pre-training LLMs has been growing at an unprecedented rate, while the availability of high-quality data remains limited. Consequently, maximizing the utility of available data constitutes a significant research challenge. A primary impediment is that certain high-quality tokens are difficult to learn given a fixed model capacity, as the underlying rationale for a single token can be exceptionally complex and deep. To address this issue, we propose Thinking augmented Pre-Training (TPT), a universal methodology that augments text with automatically generated thinking trajectories. Such augmentation effectively increases the volume of the training data and makes high-quality tokens more learnable through step-by-step reasoning and decomposition. We apply TPT across diverse training configurations up to $100$B tokens, encompassing pre-training with both constrained and abundant data, as well as mid-training from strong open-source checkpoints. Experimental results indicate that our method substantially improves the performance of LLMs across various model sizes and families. Notably, TPT enhances the data efficiency of LLM pre-training by a factor of $3$. For a $3$B parameter model, it improves the post-training performance by over $10\%$ on several challenging reasoning benchmarks.

• The paper presents TPT, an automated method that augments text with expert-like reasoning trajectories to improve LLM learning.
• It demonstrates a 150× data efficiency improvement on benchmarks like GSM8k and MATH by dynamically allocating compute to complex tokens.
• Empirical evaluations across various model sizes and regimes highlight TPT’s scalability and its robust performance on reasoning tasks.

Thinking Augmented Pre-training: Methodology, Empirical Results, and Implications

Motivation and Problem Statement

The exponential growth in LLM pre-training compute has not been matched by a commensurate increase in high-quality, human-authored data. As the web’s data pool is increasingly exhausted, maximizing the utility of available data becomes a central challenge. A key bottleneck is that certain high-value tokens—often the result of complex, multi-step reasoning—are difficult for models to learn via standard next-token prediction, especially under limited model capacity or data constraints. Existing data engineering pipelines focus on curation, filtering, and synthetic data generation, but less attention has been paid to augmenting the learnability of existing data.

Thinking Augmented Pre-Training (TPT): Method

TPT is a universal, scalable data augmentation methodology that appends automatically generated “thinking trajectories” to existing text data. For each document $d$ in the pre-training corpus, a trajectory $t$ is generated by prompting an open-source LLM to simulate an expert’s in-depth reasoning about $d$. The augmented sample $x = [d; t]$ is then used for standard next-token prediction training. The process is fully automated, requires no human annotation, and is agnostic to document structure.

Key properties of TPT:

• Scalability: Operates at the document level, with no need for online rollouts or token-level RL, in contrast to methods like RPT.
• Dynamic Compute Allocation: Complex tokens are decomposed into stepwise reasoning, effectively allocating more training compute to challenging samples.
• LLM-Friendly Data Format: Transforms noisy, web-crawled data into a format more conducive to LLM learning.

The prompt for trajectory generation is designed to elicit deep, non-trivial reasoning, leveraging the Feynman technique to ensure explanations are comprehensive and accessible.

Empirical Evaluation

Pre-training Under Abundant Data

Two 8B-parameter models (LLaMA-3-8B architecture) were trained from scratch on 100B tokens: one on vanilla data, the other on TPT-augmented data. Despite the vanilla model seeing 3× more unique documents, the TPT model achieved:

• GSM8k: 50.1% (vs. 19.2% vanilla)
• MATH: 21.8% (vs. 9.1% vanilla)
• Average across 5 tasks: 43.9% (vs. 26.2% vanilla)

The TPT-8B model, trained on 100B tokens, matched or exceeded the performance of LLaMA-3.1-8B, which was trained on 15T tokens—a 150× data efficiency improvement for mathematical reasoning.

Pre-training Under Constrained Data

With only 10B unique tokens (4 epochs for vanilla, 1 for TPT), TPT again outperformed vanilla pre-training, especially on reasoning tasks. The vanilla model plateaued as unique tokens were exhausted, while TPT continued to improve, demonstrating superior data efficiency in low-resource regimes.

Mid-training and SFT

TPT was applied as mid-training to open-source models (Qwen2.5 and LLaMA-3, 1.5B–7B parameters), followed by SFT on the Mixture-of-Thoughts dataset. Across 10 challenging benchmarks (math, code, general reasoning), TPT consistently outperformed both vanilla and OpenR1 baselines. Notably, LLaMA-3B’s AIME24 score increased from 5.8% (OpenR1) to 18.6% (TPT), a 3× improvement.

Ablation Studies

• Trajectory Generation Strategies: Custom back-thinking models and random focus point prompts yielded only marginal gains over the default method, at the cost of increased complexity.
• Model Size for Trajectory Generation: Surprisingly, using a smaller model (1.5B) for thinking generation sometimes outperformed larger models, echoing findings from OpenThoughts.
• Scaling Mid-training: Increasing the mid-training token budget led to continual performance gains, with no observed saturation at 100B tokens.
• SFT Data Size: TPT-augmented checkpoints provided a stronger foundation for SFT, with improvements persisting across multiple epochs.

Analysis of Thinking Trajectories

Analysis of trajectory lengths across domains, reasoning intensity, and target audience revealed:

• Mathematics and Physics: Produced the longest trajectories, aligning with their high reasoning demands.
• Reasoning Intensity: Positive correlation with trajectory length; “Advanced Reasoning” samples received ~50% more tokens than “No Reasoning.”
• Target Audience: “Expert” documents had shorter trajectories than “Undergraduate,” possibly due to denser, less stepwise exposition.

This dynamic allocation of training compute acts as a natural up-sampling mechanism for high-value data, without manual heuristics or explicit filtering.

Theoretical and Practical Implications

TPT demonstrates that augmenting data with explicit reasoning trajectories can dramatically improve the learnability of complex tokens, especially in reasoning-intensive domains. The method is orthogonal to existing data curation and rewriting techniques, and can be integrated into any stage of LLM training (pre-training, mid-training, SFT).

Strong empirical results—such as achieving LLaMA-3.1-8B-level math performance with 1/150th the data—underscore the potential for TPT to reduce the compute and data requirements for high-performance LLMs. The finding that smaller models can generate more effective thinking trajectories for downstream learning challenges the assumption that only the strongest models should be used for synthetic data generation.

Future Directions

• Scaling: Applying TPT to larger models and corpora.
• Prompt Optimization: Exploring automatic prompt engineering to further improve trajectory quality.
• Trajectory Generation Models: Systematic paper of the relationship between the generator and the target model.
• Integration with RL and Other Data Engineering: Combining TPT with RL-based methods and advanced data selection strategies.

Conclusion

Thinking Augmented Pre-Training (TPT) provides a simple, scalable, and effective approach to enhancing LLM data efficiency by augmenting text data with automatically generated reasoning trajectories. The method yields substantial improvements in reasoning benchmarks, particularly in mathematics and code, and is robust across model sizes and training regimes. TPT’s dynamic allocation of training compute to high-value tokens offers a principled mechanism for maximizing the utility of limited data, with significant implications for the future of efficient LLM scaling and data engineering.