PaperlayerView exact revision on arXiv

Source-linked AI summary

Thinking Augmented Pre-training

Liang Wang, Nan Yang, Shaohan Huang, Li Dong, Furu Wei

arXiv:2509.20186v4cs.CLcs.LG

TL;DR

High-quality training data is limited, and some valuable tokens are difficult to learn from a single prediction step. TPT augments text with automatically generated thinking trajectories; experiments report broader performance gains and a 3x reduction in tokens needed for the same base-model performance.

  • Problem

    Human-authored web data is finite, while some valuable tokens encode complex, multi-step reasoning that is difficult to learn in one next-token prediction step.

    Source: §1 · p. 2

  • Method

    TPT generates thinking trajectories with an open-source model, appends them to existing text, and trains on the augmented samples with standard next-token prediction.

    Source: §1 · p. 2, §2 · p. 3

  • Results

    Across model sizes and math, code, and general-reasoning benchmarks, TPT improves performance and reduces the tokens needed for the same base-model performance by a factor of 3.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Trajectories tend to be longer for difficult, reasoning-intensive text, allocating more training compute to those samples during training.

    Source: §1 · p. 2

    The current approach uses a single prompt template; automatic prompt optimization and more powerful thinking-generation models remain areas for further investigation.

    Source: §7 · p. 10

Abstract

from arXiv · show

This paper introduces a simple and scalable approach to improve the data efficiency of large language model (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.

1 INTRODUCTION

TPT augments training text with automatically generated thinking trajectories to make valuable but difficult-to-learn tokens more learnable. The paper reports improved data efficiency and model performance across training setups and benchmarks.

1 INTRODUCTION

  • Finite human-authored web data and valuable tokens that require complex multi-step reasoning make maximizing existing data utility an important challenge.With limited model capacity, such tokens may be learned by memorization rather than generalizable reasoning.Source: §1 · p. 2
  • TPT automatically adds generated thinking trajectories to pre-training data, without human annotation or constraints on document structure.Because explanation is often easier than generation from scratch, models trained on augmented data can surpass the thinking-generation model itself.Source: §1 · p. 2
  • Longer trajectories occur for more difficult, reasoning-intensive text, allocating more training compute to those samples during training rather than inference.The paper compares this training-time allocation with test-time scaling.Source: §1 · p. 2
  • Against RPT (Dong et al., 2025), Reasoning CPT (Ishibashi et al., 2025), and BoLT (Ruan et al., 2025), TPT uses a simpler approach and evaluates up to 100B tokens across broader training setups.The compared prior work is described as covering 150M–8B tokens and domain-specific evaluations; TPT requires neither meticulous prompt engineering nor EM-like iterative training.Source: §1 · pp. 2–3
  • TPT improves data efficiency and final performance across model sizes and math, code, and general-reasoning benchmarks, reducing required training tokens by a factor of 3.On a public dataset with 350k SFT samples, thinking-augmented mid-training provides consistent boosts, particularly for models not heavily pre-trained on reasoning-intensive data.Source: §1 · p. 3

2 THINKING AUGMENTED PRE-TRAINING

TPT appends automatically generated thinking trajectories to documents and trains on the augmented samples with standard next-token prediction. The approach is document-level, supports multiple training stages, and allocates more training compute to complex material.

2 THINKING AUGMENTED PRE-TRAINING

  • An off-the-shelf model generates a trajectory focused on complex, informative aspects; the original document and trajectory are concatenated and used with standard next-token prediction.The prompt asks the model to simulate an expert’s in-depth analysis and skip trivial details.Source: §2 · p. 3
  • TPT applies to pre-training from scratch and mid-training from existing checkpoints.The objective is the standard next-token prediction loss on the augmented sample.Source: §2 · p. 3
  • Document-level augmentation requires no online rollouts and can be applied to any text data, complementing other data-engineering methods.TPT transforms raw text into a format intended to facilitate more efficient learning.Source: §2 · pp. 3–4
  • Thinking augmentation decomposes difficult tokens into more explainable steps, allocating them more training compute; longer trajectories naturally up-sample high-value domains and documents.The paper describes this as analogous to test-time scaling, but applied during training.Source: §2 · p. 3

3 EXPERIMENTS

Across abundant-data, constrained-data, and mid-training settings, the experiments test whether thinking augmentation improves learning and downstream performance. TPT shows gains across these settings, with especially strong reported results on mathematical reasoning.

3.1 PRE-TRAINING UNDER ABUNDANT DATA

  • The abundant-data comparison trains two 8B models from scratch on 100B tokens, with each deduplicated sample used at most once; vanilla processes approximately 3× more raw documents.The models follow the LLaMA-3-8B architecture and train for 25k steps.Source: §3.1 · p. 4
  • In abundant-data pre-training, TPT initially performs similarly to vanilla, surpasses it after 20B tokens, and reaches a score comparable to LLaMA-3.1-8B at 100B tokens.Figure 2 plots pre-training loss and aggregated scores across five tasks; loss values are not directly comparable because data distributions differ.Source: §3.1 · pp. 4–5, Figure 2 · p. 4
  • GSM8k performance rises from 19.2% to 50.1%, while MATH scores rise from 9.1% to 21.8% for TPT-8B versus the vanilla baseline.The table reports base-model performance across five datasets and their average.Source: §3.1 · p. 5, Table 1 · p. 5
  • After SFT, TPT substantially improves performance across all evaluated challenging reasoning benchmarks and outperforms LLaMA-3.1-8B-Instruct on every benchmark.Table 2 reports performance after SFT on the 2B-token Mixture-of-Thoughts dataset.Source: §3.1 · p. 5, Table 2 · p. 5

3.2 PRE-TRAINING UNDER CONSTRAINED DATA

  • The constrained-data setup caps raw-document tokens at 10B and trains two 8B models for 40B tokens; vanilla sees four epochs while TPT sees the augmented data once.The 10B-token subset is selected by random sampling.Source: Figure 3 · p. 5, §3.2 · p. 6
  • With raw-document tokens constrained to 10B, vanilla performance plateaus or improves slowly as unique tokens run out, while TPT continues improving, especially on GSM8k and MATH.Figure 3 plots task scores against total training tokens.Source: Figure 3 · p. 5, §3.2 · p. 6

3.3 THINKING AUGMENTED MID-TRAINING

  • Mid-training covers Qwen2.5 and LLaMA-3 models from 1.5B to 7B parameters, followed by SFT on 350k Mixture-of-Thoughts samples and evaluation across math, code, and general-reasoning benchmarks.The comparison primarily uses OpenR1 models fine-tuned on the same dataset; inference allows up to 32k generated tokens.Source: §3.3 · pp. 6–7
  • The 3B LLaMA model’s AIME24 score rises from 5.8% to 18.6%, and TPT models improve across sizes from 1.5B to 7B parameters.The paper reports that gains are particularly pronounced for the LLaMA-3 series.Source: Table 3 · p. 6, §3.3 · p. 7

4 ANALYSIS OF THINKING PATTERNS

The analysis examines how generated thinking length varies with document characteristics. Longer trajectories are associated with mathematics and greater reasoning intensity, while expert-targeted documents do not necessarily receive longer trajectories.

4 ANALYSIS OF THINKING PATTERNS

  • The analysis uses metadata-stratified samples of 20k documents and groups trajectory lengths by domain, reasoning intensity, and target audience.Thinking trajectories were generated with DeepSeek-R1-Distill-Qwen-7B; Figure 4 lists the top 10 domains with the longest trajectories.Source: §4 · p. 7, Figure 4 · p. 7
  • Advanced Reasoning documents have approximately 50% more thinking tokens than No Reasoning documents, while Expert-targeted documents have shorter trajectories than Undergraduate-targeted ones.The authors suggest specialized concepts may not require more reasoning steps for comprehension.Source: §4 · p. 7
  • Mathematics and Physics documents exhibit notably longer thinking trajectories, and the analysis finds that high-value data naturally yields longer trajectories.Longer trajectories allocate more training compute to that content, effectively up-sampling it without manual heuristics or explicit quality filtering.Source: §4 · p. 7

5 ABLATION STUDIES

The ablations examine thinking-generation strategies, mid-training token budgets, and SFT epochs. They report continuing gains with more mid-training and a roughly 15-point AIME24 increase for LLaMA-3B after 100B tokens.

Thinking Trajectory Generation

  • Alternative generation strategies yield only marginal gains across most benchmarks, so the study retains its default strategy for simplicity and reproducibility.The alternatives are a custom back-thinking model and a prompt directing random focus; Table 4 evaluates models after 40B mid-training tokens and SFT.Source: §5 · p. 8, Table 4 · p. 8

Scaling Thinking Generation Model

  • The 1.5B DeepSeek-R1-Distill-Qwen generator outperforms the default 7B generator in the reported ablation.The authors say the relationship between the fine-tuned model and the thinking-generation model warrants further investigation.Source: §5 · p. 8

Impact of Mid-training Token Budget

  • After 100B thinking-augmented mid-training tokens, LLaMA-3B gains approximately 15 AIME24 points; scores also rise continually across model sizes and evaluated tasks as the token budget grows from 0 to 100B.Figure 5 defines the 0B point as direct SFT without thinking-augmented mid-training.Source: §5 · pp. 8–9, Figure 5 · p. 8

Impact of SFT Data Size

  • Increasing SFT epochs generally improves most benchmarks without serious overfitting at five epochs, while mid-trained checkpoints start SFT ahead of vanilla counterparts and retain that advantage.LiveCodeBench had not fully converged, and the study uses five epochs.Source: §5 · p. 9, Figure 6 · p. 9

6 RELATED WORK

Related work addresses scaling and data curation, synthetic and rewritten training text, and methods for eliciting reasoning. TPT uses open-source models to generate thinking trajectories for training across multiple stages.

Data Engineering for Large-Scale Pre-training

  • Large-scale pre-training relies on curated data, while modern pipelines transform raw web text through multiple stages such as extraction, deduplication, and quality filtering.The passage describes FineWeb and DCLM as examples of data curation and evaluation.Source: §6 · p. 9

Synthetic Data and Rewriting

Thought-Augmented Training

  • Reasoning CPT and BoLT mine hidden thoughts, but their reported evaluations are narrower in scale or scope than TPT’s 100B-token pre-training and mid-training experiments.Reasoning CPT scales to approximately 150M tokens in continual training and evaluates base models without post-training; BoLT reports math-task results with a 1B model and 8B tokens.Source: §6 · p. 10

Chain-of-Thought Reasoning

  • Chain-of-Thought prompting elicits intermediate reasoning steps, with later work exploring tree structures for search, backtracking, and self-correction.The passage notes that prompting-based CoT comes at the cost of increased inference time.Source: §6 · p. 10

Reinforcement Learning and Thinking

  • Reinforcement-learning methods such as OpenAI o1 and DeepSeek-R1 train models to generate long thinking trajectories; TPT instead uses open-source LLMs to generate trajectories for training across stages.The passage reports a positive correlation between generated-token length and task performance in the cited reinforcement-learning work.Source: §6 · p. 10

7 CONCLUSION

TPT augments existing text with thinking trajectories to improve pre-training data efficiency. Experiments up to 100B tokens report gains across model sizes and training configurations, especially on reasoning-intensive tasks.

Conclusion

  • Experiments with up to 100B training tokens show TPT improves data efficiency and final performance across model sizes and training configurations, with notable gains on reasoning-intensive tasks.The conclusion describes TPT as a simple, scalable approach that augments existing text data with thinking trajectories.Source: §7 · p. 10

Future Work

  • Scaling to larger training corpora and model sizes is identified as future work.The current approach also uses a single prompt template, and the authors suggest automatic prompt optimization and more powerful thinking-generation models for further study.Source: §7 · p. 10

Conclusion

  • The authors hope the findings encourage scalable data engineering that makes more efficient use of data in foundation models.This is framed as a research direction rather than an established outcome.Source: §7 · p. 10

A.1 TRAINING DETAILS

The training details specify the data mixtures, thinking-generation setup, and model and hardware configurations used for pre-training, mid-training, and SFT. They also describe dataset packing and response-only loss for SFT.

Pre-training and Mid-training Datasets

  • Pre-training and mid-training combine FineWeb-Edu with MegaMath-Web-Pro-Max, using a 0.125 sample weight for FineWeb-Edu and document-level exact-match deduplication.MegaMath-Web-Pro-Max is described as a filtered subset with approximately 70B tokens of math-intensive web content.Source: §A.1 · p. 13

Post-training Datasets

  • SFT uses the public 350k-example Mixture-of-Thoughts dataset; examples are packed without cross-example attention, and loss is computed only on assistant response tokens.The passage reports that the 1.5B reproduction remains below DeepSeek-R1-Distill-Qwen-1.5B and suggests the smaller dataset size as a likely explanation.Source: §A.1 · pp. 13–14

Thinking Trajectory Generation

  • Thinking generation truncates inputs at 2k tokens, allows up to 8k generated thinking tokens, and uses temperature 0.6 and top-p 0.9.DeepSeek-R1-Distill-Qwen-7B generates trajectories for mid-training; Qwen3-8B does so for pre-training from scratch, with the pipeline taking roughly 20k A100 GPU hours for 100B tokens.Source: §A.1 · p. 14

Training Hyperparameters

  • Training runs use MI300 GPUs; pre-training an 8B model for 100B tokens and mid-training a 7B model each take about one week.SFT of a 7B model for five epochs takes about 3.5 days; Table 5 summarizes the training hyperparameters.Source: §A.1 · p. 14, Table 5 · p. 14

Model Initialization

  • Pre-training uses randomly initialized LLaMA-3 architecture weights, while 1.5B, 3B, and 7B mid-training models start from Qwen2.5-Math-1.5B, Llama-3.2-3B, and Qwen2.5-Math-7B, respectively.For post-training context extension from 8k to 32k, the RoPE base frequency is multiplied by 16.Source: §A.1 · p. 14

A.2 EVALUATION DETAILS

The appendix specifies separate evaluation suites for base and instruction-tuned models, along with prompting, scoring, and implementation procedures.

A.2 EVALUATION DETAILS

  • Base models are evaluated on MMLU, MMLUPro, BoolQ, GSM8K, and MATH, with average performance reported across the five datasets.The reported metric is answer accuracy.Source: §A.2 · p. 14
  • Base-model answers are extracted with regular expressions, and outputs without a valid answer count as incorrect.The maximum generation length is 2k tokens; invalid options can depress multiple-choice scores, especially for less-trained models.Source: §A.2 · p. 14
  • Instruction-tuned models are tested on math, competition, science, general-reasoning, and coding benchmarks, including MATH-500, AIME24/25, HMMT, GPQA-Diamond, MMLU-Pro, HumanEval, LiveCodeBench, and JEEBench.Pass@1 is the main metric; coding outputs are automatically verified with test cases.Source: §A.2 · p. 15
  • Evaluations mainly use lighteval, with evalchemy scripts for HumanEval, HMMT, and JEEBench; generative tasks use temperature 0.6 and top-p 0.95.The 8B evaluation takes about one day on four A100 GPUs.Source: §A.2 · p. 15

A.3 PROMPT TEMPLATES

The prompt template selects a random focus point in each document and asks for an expert-style analysis of the surrounding context.

A.3 PROMPT TEMPLATES

  • The template places a reading marker between a context prefix and the remaining context, then closes with an end-of-context marker.Source: §2 · p. 3, §A.3 · p. 15
  • The prompt asks an expert to analyze complex, informative aspects in depth and skip trivial details, using the Feynman technique when possible.Source: §A.3 · p. 15
  • For each document, a position is randomly selected as the focus point, and the document is partitioned into two parts.Source: §A.3 · p. 15

A.4 ADDITIONAL RESULTS

The additional results report limited effects from vanilla mid-training, provide example thinking trajectories, and state the paper's use of language models.

A.4 ADDITIONAL RESULTS

  • Vanilla mid-training on MegaMath-Web-Pro-Max and FineWeb-Edu for 40B tokens leaves most benchmarks largely stable but shows a coding-performance decline.Source: Table 6 · p. 15, §A.4 · p. 15
  • The distillation-model comparisons are not direct because those models used a larger private dataset than Mixture-of-Thoughts.Source: §A.4 · p. 16

A.5 EXAMPLES

  • Tables 8–10 provide generated trajectory examples, including samples from the Figure 1 example, MegaMath-Web-Pro-Max, and FineWeb-Edu.The examples identify DeepSeek-R1-Distill-Qwen-7B as the thinking-generation model.Source: §A.5 · p. 16, Table 8 · p. 17, Table 9 · p. 18

A.6 USE OF LARGE LANGUAGE MODELS

  • The paper states that LLMs assisted with writing only, not with generating the research ideas or analyzing the results.Source: §A.6 · p. 16

Related papers

Loading 2509.20186v4…