Rewriting Pre-Training Data Boosts LLM Performance in Math and Code

This paper introduces SwallowCode and SwallowMath, two openly licensed datasets created through a systematic four-stage rewriting pipeline that transforms public code and math data into high-quality, self-contained examples, significantly boosting LLM performance in program synthesis and mathematical reasoning within a fixed training budget.

The Gist

Imagine you are trying to teach a brilliant but inexperienced apprentice (a Large Language Model, or LLM) how to be a master chef. The apprentice is smart, but their current skills are limited because the recipe books they've been reading are messy, full of typos, missing ingredients, and sometimes even dangerous instructions.

This paper, "Rewriting Pre-Training Data Boosts LLM Performance in Math and Code," is about a team of researchers who decided: "Instead of just throwing away the bad recipe books, let's hire a team of expert editors to fix them, clean them up, and rewrite them so they are perfect."

Here is the story of how they did it, broken down into simple concepts.

1. The Problem: The "Noisy" Library

Currently, AI models learn by reading massive amounts of data from the internet. Think of this data as a giant, chaotic library.

• The Code Library: It has millions of Python scripts, but many are broken, have bad formatting, or are missing crucial context (like a recipe that says "add salt" but doesn't say how much or when).
• The Math Library: It has math problems, but they are buried under website ads, timestamps, and confusing formatting.

Previous methods tried to fix this by filtering: they would just throw away the bad books and keep only the "good" ones. But this is like throwing away 90% of the library because a few pages are torn. You lose a lot of potential knowledge.

2. The Solution: The "Transform-and-Retain" Method

The researchers, from Tokyo Tech and AIST, introduced a new philosophy: Don't just filter; fix.

They created two new datasets: SwallowCode (for programming) and SwallowMath (for math). Instead of deleting the messy data, they used a super-smart AI (Llama-3.3) to act as a "Chief Editor" to rewrite the data into something perfect.

🐍 SwallowCode: The Code Refinery

Imagine a pile of messy, handwritten code snippets. The researchers built a four-stage assembly line to clean them up:

1. The Syntax Check (The Safety Inspector): First, they check if the code actually runs. If it has a typo that crashes the program, it gets tossed.
2. The Style Check (The Fashion Police): They use a tool called pylint to check if the code looks professional. Is the indentation right? Are the variable names clear? If it looks sloppy, it gets a "red pen" correction.
3. The Style Rewrite (The Editor): Here, the AI steps in. It rewrites the code to follow strict style guides (like Google's). It renames confusing variables, adds clear comments, and organizes the code so it's easy to read.
   • Analogy: It's like taking a rough draft of a novel and turning it into a polished bestseller with perfect grammar and flow.
4. The Optimization Rewrite (The Engineer): This is the magic step. The AI looks at the logic.
   • Does the code rely on a library that doesn't exist? Fix it.
   • Is the algorithm slow and inefficient? Make it faster.
   • Is the example too simple to be useful? Make it a real-world problem.
   • Result: The code is now "self-contained" (it works on its own) and "optimized" (it's the best way to solve the problem).

The Result: They turned a messy pile of 16 billion code tokens into a pristine, high-quality dataset. When they taught an AI using this new dataset, it became 17% better at solving coding puzzles than models trained on the old, filtered data.

➗ SwallowMath: The Math Tutor

For math, the process was similar but focused on clarity.

• The Problem: Math problems found online often have website headers, "Posted on: [Date]" footers, and missing steps in the solution.
• The Fix: The AI acts like a strict but helpful tutor. It strips away the website junk, fills in the missing context, and rewrites the solution to be a clear, step-by-step explanation.
• The Result: The AI learned to solve math problems with 12% higher accuracy on standard tests.

3. Why This Matters: The "Data Quality Gap"

The paper points out a big issue in the AI world: The biggest, most powerful AI models (like those from big tech companies) are trained on secret, high-quality data that we can't see. The open-source community is stuck with "noisy" public data.

This creates a "Data Quality Gap." It's like one student has a private tutor and a library of perfect textbooks, while another student has to study from a messy, torn-up library.

By releasing SwallowCode and SwallowMath for free, the researchers are handing the open-source community a set of "perfect textbooks." They proved that you don't need a bigger computer or a smarter model architecture; you just need better data.

4. The Takeaway: Quality Over Quantity

The most important lesson from this paper is that more data isn't always better; better data is.

• Old Way: "Let's read 100 books, even if 50 of them are garbage."
• New Way: "Let's take those 100 books, hire an editor to fix the bad ones, and teach the student with the 100 perfect books."

The researchers showed that by spending time and computing power to rewrite the data, they got much better results than just filtering out the bad stuff. They turned a "good enough" dataset into a "world-class" dataset, proving that if you feed an AI high-quality, clean, and logical information, it will learn to think and create much better.

In short: They didn't just clean the house; they renovated the whole building, and now the AI lives in a much better home.

Technical Summary

1. Problem Statement

The performance of Large Language Models (LLMs) in specialized domains like program synthesis and mathematical reasoning is fundamentally constrained by the quality of their pre-training corpora.

• Limitations of Current Methods: Existing public datasets (e.g., The-Stack-v2 for code, Finemath-4+ for math) rely primarily on rule-based extraction or exclusionary filtering (discarding low-quality samples). This approach retains noisy, redundant, or stylistically inconsistent data and fails to address structural defects like missing context or inefficient algorithms.
• The "Data Quality Gap": Leading frontier models (e.g., Qwen3, DeepSeek-V3) do not release their pre-training corpora, leaving the open-source community without access to high-quality, curated data comparable to state-of-the-art models.
• Need for Transformation: Simply filtering data is insufficient; the paper argues for a "transform-and-retain" approach that actively rewrites low-quality data into high-quality, self-contained examples rather than discarding it.

2. Methodology

The authors introduce two openly licensed datasets, SwallowCode and SwallowMath, constructed via a systematic, multi-stage pipeline that leverages Large Language Models (LLMs) for rewriting.

A. SwallowCode (Code Domain)

Built from The-Stack-v2 (Python subset), SwallowCode (≈16.1B tokens) is constructed using a four-stage pipeline:

1. Syntax Validation: Compiles code using Python's compile() to remove invalid syntax (reducing samples by ~9.7%).
2. Linter-Based Filtering: Uses pylint to enforce coding standards (score ≥ 7.0) and applies a heuristic penalty for overly verbose comments. This reduces the dataset by ~34.3%.
3. Style-Guided Code Rewriting (SGCR): An LLM (Llama-3.3-70B-Instruct) rewrites code to enforce the Google Python Style Guide. It adds docstrings, type hints, modularizes functions, and standardizes naming conventions.
4. Self-Contained Optimization Rewriting (SCOR): A second LLM stage ensures semantic correctness and efficiency. It:
   • Resolves missing dependencies (making code self-contained).
   • Replaces inefficient algorithms (e.g., naive recursion) with optimized ones.
   • Transforms trivial snippets into instructive examples.

Note: The pipeline uses a two-stage approach (SGCR then SCOR) rather than a single pass to prevent "instruction drift" where the model prioritizes one objective over another.

B. SwallowMath (Math Domain)

Built from Finemath-4+ (≈2.3B tokens), this dataset applies a tailored rewriting pipeline to:

1. Clean: Remove web headers, footers, privacy notices, and irrelevant metadata (timestamps).
2. Restore Context: Fill in missing information in incomplete questions or answers.
3. Reformat: Rewrite derivation steps to be concise yet comprehensive, providing clear, step-by-step solutions.

C. Experimental Setup

• Base Model: Continual pre-training on Llama-3.1-8B (and Qwen2-7B for generalizability checks).
• Budget: Fixed at 50 billion tokens (with 16% code or ~4.8% math, rest multilingual text) to ensure fair comparison.
• Ablation Studies: Systematic testing of each pipeline stage to isolate performance gains.
• Decontamination: Rigorous checks against benchmarks (HumanEval, GSM8K, MATH) to ensure no data leakage or self-contamination from the rewriting LLM.

3. Key Contributions

1. Transform-and-Retain Paradigm: Demonstrates that rewriting low-quality data yields significantly higher performance gains than exclusionary filtering.
2. Open Datasets: Releases SwallowCode and SwallowMath under the Llama 3.3 Community License, along with the full pipeline code, prompts, and checkpoints.
3. Novel Pipeline Design: Introduces a specific two-stage LLM rewriting process (Style-Guided + Self-Contained Optimization) that addresses both syntactic/style issues and semantic/algorithmic inefficiencies.
4. Empirical Validation: Proves that data quality improvements directly translate to downstream task performance without relying on undisclosed training data or complex model architectures.

4. Results

Under a fixed 50B token budget, continual pre-training with the rewritten datasets significantly outperformed baselines (Stack-Edu for code, Finemath-4+ for math):

Code Performance (Llama-3.1-8B):

• HumanEval: +17.0 points improvement over Stack-Edu.
• HumanEval+: +16.1 points improvement over Stack-Edu.
• Generalizability: On Qwen2-7B, SwallowCode improved HumanEval by +10.3 points over Stack-Edu.

Math Performance (Llama-3.1-8B):

• GSM8K: +12.4 points improvement over Finemath-4+.
• MATH: +7.6 points improvement over Finemath-4+.

Ablation Insights:

• Filtering: Syntax and linter filtering provided modest gains (~1-2 points).
• Rewriting: The LLM rewriting stages (SGCR and SCOR) were the primary drivers of performance, contributing a total of ~14 points (7-9 from SGCR, 5-6 from SCOR).
• MBPP Note: The authors excluded the MBPP benchmark from final results because SGCR's enforcement of standard naming conventions (snake_case) caused mismatches with MBPP's non-standard test cases, artificially lowering scores.

5. Significance and Impact

• Bridging the Data Gap: Provides the open-source community with high-quality, pre-processed data that rivals the proprietary datasets used by frontier models.
• Data-Centric AI: Shifts the focus from "bigger models" to "better data," showing that systematic data refinement is a cost-effective way to boost capabilities.
• Reproducibility: By releasing the full pipeline, prompts, and checkpoints, the authors enable the community to replicate results and adapt the "transform-and-retain" methodology to other domains or base models.
• Efficiency: While the rewriting process requires significant upfront compute (approx. 43k H100 GPU hours for the full pipeline), the resulting datasets offer higher data efficiency, allowing models to reach higher performance levels with fewer training tokens compared to noisy baselines.

In conclusion, this paper establishes that systematic rewriting of pre-training data is a superior strategy to simple filtering, offering a scalable and reproducible path to enhancing LLM capabilities in code and mathematics.