From Weeks to Hours: Fast and Principled SFT Data Generation for Reasoning LLMs

AI summary

Overview

Research area: Machine learning — post-training and data curation for large language models, specifically supervised fine-tuning (SFT) of reasoning models.

Technical level: Intermediate. The paper combines a formal geometric argument (loss-basin curvature, gradient-difference bounds) with a practical data pipeline; the theory can be skimmed while the pipeline remains understandable.

Scope: The paper proposes SFTProbe, a method that predicts the difficulty and diversity of reasoning data from a short LoRA adaptation and loss on a small prefix of each reasoning trace, replacing the expensive generate-then-filter paradigm for building SFT datasets for reasoning LLMs.

What This Paper Is About

Building good reasoning data for supervised fine-tuning currently means generating huge pools of long chain-of-thought traces from strong teacher models and then filtering them with LLM-based heuristics for difficulty and diversity — a process the authors say often takes weeks of computation. The paper asks whether the quality of a reasoning trace can be predicted directly, from the geometry of the pretrained model's loss landscape, instead of being discovered through exhaustive generation, filtering, and trial training. The goal is to produce small, high-quality SFT datasets in hours rather than weeks while matching or beating heavily filtered large datasets.

Key Contributions

1. SFTProbe, a direct quality-prediction pipeline. Instead of generate-then-filter, the method performs a brief LoRA-based adaptation on a small set of fully generated examples, then evaluates loss on only 5% of each reasoning trace to build a compact "loss signature" used to estimate both difficulty and diversity.

2. A geometric justification via a gradient-difference bound. Theorem 3.1 upper-bounds the pairwise full-rank gradient difference of two training examples using their loss differences at checkpoints along LoRA directions plus curvature terms; Corollary 3.2 shows the bound is tightest when one checkpoint is the pretrained model and the other is the curvature inflection point, with the optimal perturbation scale ε*₂ = 2√(δ/c_{H,i}). This implies directions with smaller curvature should be perturbed more.

3. A difficulty score and a diversity criterion computable from prefixes. Difficulty is the loss on the first 100 tokens at the inflection point along a random direction; diversity comes from clustering loss signatures computed over the first 1k tokens across 2r perturbed checkpoints.

4. A simple teacher-selection criterion and a new physics dataset. The best teacher is the one minimizing average early-token ("rephrasing") loss under the pretrained student. The authors also release a curated physics reasoning dataset built from olympiad problems (1967–2018, 3,084 problem–solution pairs) mixed with PhysicsEval problems to form 40k problems.

Main Findings

• Small generated sets match or beat a large filtered set. On mathematical reasoning, 5k generated examples match the highly filtered 89k-example OpenThoughts-Math dataset, and 7.5k generated examples outperform it by 2.5% in accuracy.
• Large speedups. The paper reports 53x speedup for the 5k setting and 41x for the 7.5k setting relative to OpenThinker-7B (Figure 1). The abstract frames the gain as weeks to hours; the conclusion states curation time drops "from weeks to minutes."
• Medical data selection gain. Selecting 1k examples from the pre-generated M23k dataset with SFTProbe gives a 1.1% increase over the best baseline with Llama3.1-8B-Instruct.
• Medical data curation gain. Starting from 37,698 decontaminated questions, SFTProbe selected QwQ-32B as teacher and generated 1k reasoning examples; training on them matches the full curated M23k dataset within 0.6%, closes 76% of the gap between random sampling and full-dataset training, and requires 7% of the compute cost.
• Math data curation gain. 3.4% increase in math reasoning performance over the best baseline; SFTProbe remains faster than data-selection baselines even when teacher generation time is included.
• Physics results. With Qwen2.5-7B-Instruct as student: base model 3.58; Google/Gemma-4-26B-A4B-it 9.65; PRIME-RL/P1-30B-A3B 9.05; SFTProbe with 1k data 9.14 (± 0.29), 2k data 9.32 (± 0.11), 5k data 9.76 (± 0.74) — beating a model described as capable of winning silver at IPhO 2025, from a much weaker base model.
• Inflection point is where difficulty separates best. Loss evaluated at the pretrained parameters correlates poorly with difficulty, while loss at perturbed parameters correlates well; the best separation occurs at the curvature inflection point. Within the basin, loss of hard problems grows faster than that of easy ones, magnifying the gap by up to 5x.
• Ablations. Replacing inflection points with checkpoints along the training trajectory costs 0.7%; additionally using checkpoints from the brief LoRA trajectory instead of the full training trajectory costs a further 0.8%. Reported averages (ID / OOD / Avg over 10 sets): prior method [12] 69.3% / 52.9% / 57.8%; SFTProbe 68.7% / 53.1% / 57.8%; without inflection search 67.9% / 52.5% / 57.1%; without LoRA 68.1% / 51.3% / 56.3%.
• Hyperparameter insensitivity. LoRA rank 4 with 200 iterations gives 68.7% / 53.1% / 57.8%; rank 2 gives 68.5% / 52.6% / 57.3%; 100 iterations gives 68.9% / 51.9% / 57.0%; 400 iterations gives 69.0% / 52.0% / 57.1%. Only 100–400 iterations are needed to identify dominant directions.
• Teacher selection signal. The higher a teacher's rephrasing loss under the pretrained student, the worse the student performs after fine-tuning on that teacher's data (Figure 4, NuminaMath, 1000 questions).

Methodology in Plain English

The starting assumption is that fine-tuning a pretrained LLM stays inside a small, bowl-shaped ("near-parabolic") region of the loss landscape, and that this region is not symmetric — it curves differently in different directions.

1. Find the important directions cheaply. Train a low-rank LoRA adapter for a short time on a small subset of examples — not to actually improve the model, but just to discover which directions in parameter space matter most. A QR decomposition of the adapter yields an orthogonal basis of these directions.
2. Walk each direction to its inflection point. Perturb the model further and further along each direction, fitting a cubic spline to the resulting loss curve and solving for where the second derivative changes sign. That point marks the boundary of the bowl in that direction. Directions with flatter curvature get pushed further out.
3. Score diversity with loss signatures. For each example, record the loss on only its first 1k tokens at 2r perturbed checkpoints (two per direction). Standardize the vectors and cluster them with FAISS; examples in the same cluster will take similar gradient steps during real training, so they are redundant. Sample evenly across clusters.
4. Score difficulty with one perturbed evaluation. Record the loss on just the first 100 tokens at the inflection point along a random direction — no LoRA needed here. Cluster scores within each data source into two groups and discard the easier group; then allocate the total budget across sources in proportion to their average remaining difficulty via a softmax.
5. Pick the teacher with the lowest rephrasing loss. Ask each candidate teacher to rephrase a sample of problems (capped at 100 tokens) and measure the loss of those rephrasings under the pretrained student. The teacher the student already understands best lets it spend capacity on solving rather than parsing.
6. Generate fully only at the end. Full reasoning traces are produced only for the small surviving subset. The paper uses 100 tokens as the difficulty horizon and 1k tokens as the diversity horizon.

Why This Matters

The work shifts data curation from search to prediction: instead of paying for enormous generation and LLM-judge filtering, it derives quality signals from a cheap probe of the pretrained model's local loss geometry. If it holds up, it makes high-quality reasoning post-training accessible to groups without the compute budget for large-scale generate-then-filter pipelines, and it offers a principled account of why early-token loss is informative rather than a purely empirical heuristic.

Real-world applications:

• Clinical and medical reasoning assistants trained on curated domain data (the paper evaluates on the M23k medical dataset and the 37,698-question decontaminated medical pool).
• STEM education and tutoring systems, using the released olympiad-level physics dataset and the math pipeline over NuminaMath.
• Enterprise domain adaptation, where teams must build small, high-quality SFT sets for a specific vertical under tight compute budgets.
• Competitive / olympiad problem-solving model development, given the physics results against an IPhO 2025 silver-medal-capable model.

Industry relevance: The authors include researchers from Optum AI, and the pipeline's selling point — cutting curation from weeks to hours and from full-dataset compute to 7% of it — targets exactly the cost bottleneck that limits how often companies can refresh domain-specific reasoning models.

Future Directions

• Beyond the near-parabolic assumption. The theory relies on a smooth loss with bounded curvature and a low-rank dominant subspace; how far the guarantees extend outside the basin, or for models whose fine-tuning leaves it, is open.
• Choosing and mixing multiple teachers. Teacher selection is currently a single scalar criterion (rephrasing loss); whether combining teachers, or selecting per-problem rather than globally, improves data quality is untested here.
• Extending to reinforcement learning and other post-training regimes. The paper evaluates only SFT; whether the same geometric probes predict usefulness for RL-based post-training is not reported.
• Robustness of the hyperparameters across scales and domains. The ablations cover LoRA rank (2 vs. 4) and training iterations (100–400) on a limited set of tasks; behavior at larger student models, higher ranks, and unseen domains such as code or law is not reported.

Target Audience

Machine learning researchers and practitioners working on LLM post-training and data curation; engineers building domain-specific reasoning models who need faster dataset construction; and readers interested in loss-landscape geometry and its practical consequences. Familiarity with LoRA, SFT, and basic loss-landscape concepts makes the paper considerably easier to follow.

Authors’ abstract

Reasoning-oriented post-training enables large language models (LLMs) to solve complex tasks via multi-step inference. Supervised fine-tuning (SFT) on high-quality reasoning traces is a particularly efficient approach, but its effectiveness depends critically on careful data curation. Existing pipelines rely on a costly generate-then-filter paradigm: large pools of long reasoning traces are produced by multiple teacher models and subsequently filtered for difficulty and diversity using additional LLMs, often requiring weeks of computation. In this work, we propose SFTProbe, a fundamentally different approach that bypasses this pipeline by directly predicting the quality of reasoning data. We show that a brief LoRA-based adaptation, combined with evaluating loss on only 5% of each reasoning trace, suffices to estimate both difficulty and diversity. SFTProbe is grounded in a geometric view of post-training, where pretrained LLMs lie in a low-rank, anisotropic loss basin. By probing this basin via low-rank perturbations, we construct compact loss signatures whose clustering captures diversity, while loss at curvature inflection points provides a robust proxy for difficulty. We also provide a simple criterion for teacher selection. Empirically, our approach reduces data curation time from weeks to hours while maintaining or improving performance. Remarkably, on mathematical reasoning, 5k generated examples match the highly filtered 89k-example in OpenThoughts-Math dataset, and 7.5k generated examples outperforms it by 2.5% in accuracy. We further introduce a high-quality physics reasoning dataset, and use it to train a model that outperforms an IPhO 2025 silver-medal-winning model. The code and the dataset are released at https://bigml-cs-ucla.github.io/SFTProbe-page/.