Technical Summary: TrajFusion for Scaling Mathematical Reasoning

Problem Statement

Large Language Models (LLMs) have achieved significant progress in mathematical reasoning, largely through Supervised Fine-Tuning (SFT) on high-quality solution trajectories. The dominant paradigm for constructing this data is Rejection Sampling Fine-Tuning (RFT). In RFT, a teacher model generates multiple candidate reasoning paths for a given problem, and only those with correct final answers are retained for training; all incorrect trajectories are discarded.

While effective, this approach treats supervision as a binary filter. It systematically excludes teacher-generated errors, creating a gap in how reasoning failures are modeled during training. The authors argue that incorrect trajectories contain valuable diagnostic information—such as common computational slips, missing assumptions, and logically plausible but invalid argument paths—that could help models recognize and avoid specific failure modes, particularly in long-form or ambiguous reasoning tasks. Current methods treat these errors as noise to be discarded rather than structured supervision to be leveraged.

Methodology: TrajFusion

The paper proposes TrajFusion, a fine-tuning strategy that reframes RFT from a binary filtering mechanism into a structured supervision construction process. TrajFusion integrates selected incorrect reasoning trajectories with correct ones into a single fused training sample, explicitly modeling a trial-and-error reasoning process.

Core Components

1. Sampling and Analysis: For each problem $x$, the teacher model samples $K$ trajectories. These are partitioned into correct ($Y^+$) and incorrect ($Y^-$) sets using an automatic verifier.
2. Adaptive Selection: The method adaptively determines how many incorrect trajectories to include based on two statistics:
   • Error Rate ($r(x)$): The frequency of incorrect attempts.
   • Error Diversity ($u(x)$): The number of unique final answers among incorrect attempts (measured by Shannon entropy or unique answer count).
     The number of retained incorrect trajectories $k(x)$ is calculated as:
     $$k(x) = \min(k_{max}, \lfloor \alpha \cdot r(x) \cdot u(x) \rfloor)$$
     This ensures that problems with diverse and frequent errors receive richer supervision, while problems with sparse or uninformative errors naturally degenerate to vanilla RFT.
3. Fused Sample Construction: The selected incorrect trajectories are interleaved with the correct trajectory using fixed reflection prompts (e.g., "Wait... that doesn't seem right. Let's try again."). The final structure is:
   $$T(x) = [y_{e1}, \rho_1, \dots, y_{ek}, \rho_k, y^*]$$
   where $y_{ei}$ are incorrect paths, $\rho_i$ are reflection prompts, and $y^*$ is the correct solution.
4. Training: The fused sequences are trained using standard next-token prediction (SFT) with no changes to the model architecture or training objective.

Key Contributions

• Formalization of a Supervision Gap: The authors identify and formalize the systematic exclusion of teacher-generated incorrect trajectories in RFT, arguing that these errors provide diagnostic value for modeling reasoning failures.
• TrajFusion Framework: A structured supervision framework that adaptively fuses incorrect and correct trajectories. It treats error diversity as an explicit control signal, allowing the method to scale robustly across problems of varying difficulty without requiring auxiliary models, critique generators, or external difficulty labels.
• Empirical Validation: The method demonstrates consistent improvements over vanilla RFT across multiple base models (LLaMA3-8B, DeepSeekMath-7B), data regimes (15K to 100K examples), and context lengths (up to 32K).

Experimental Results

The authors evaluated TrajFusion on six mathematical reasoning benchmarks (GSM8K, MATH, CollegeMath, DeepMind-Mathematics, OlympiadBench, TheoremQA) and competition-level datasets (AIME24, AIME25, MATH-500).

• Performance Gains: TrajFusion consistently outperformed vanilla RFT. On the MATH dataset with LLaMA3-8B (100K data), TrajFusion achieved 51.4% accuracy compared to 48.5% for RFT. On DeepSeekMath-7B, it reached 52.3% vs. 50.2%.
• Impact on Complexity: The improvements were most pronounced on challenging, long-form reasoning benchmarks (MATH, OlympiadBench, TheoremQA), where reasoning paths are longer and failure modes are more diverse. Gains on simpler datasets like GSM8K were smaller but still positive.
• Data Efficiency: TrajFusion achieved competitive results with only 15K training examples, outperforming baselines that used significantly larger datasets (e.g., MMIQC with 2.3M examples).
• Scalability to Long Context: When applied to a 32K-context version of Qwen2.5-Math-7B using long-form data distilled from DeepSeek-R1, TrajFusion improved Pass@1 on AIME24 from 40.0% (RFT) to 43.3%.
• Ablation Studies:
  • Increasing the number of fused errors ($k_{max}$) improved performance up to a point (diminishing returns observed at $k_{max}=6$).
  • Selecting errors based on diversity (HighDiv) yielded better results than selecting based solely on frequency (LowDiv), suggesting that structurally distinct failure modes are more informative.
  • Models trained with TrajFusion initially generated longer sequences but gradually became more concise while increasing the frequency of reflection tokens, indicating more efficient use of the structural patterns.

Significance and Claims

The paper claims that TrajFusion offers a simple yet effective way to scale mathematical reasoning by leveraging supervision signals that are typically discarded. By reframing rejection sampling as a process of constructing structured, trial-and-error trajectories, the method provides richer supervision without architectural modifications or complex reinforcement learning pipelines.

The authors emphasize that TrajFusion is particularly valuable for challenging and long-form reasoning problems, where distinguishing valid from flawed reasoning is essential. They assert that the method's adaptive nature ensures robustness: it enriches supervision when errors are informative and safely reverts to standard RFT behavior when they are not. This approach bridges the gap between binary correctness filtering and the nuanced reality of reasoning failures, enabling models to learn from their own mistakes in a supervised setting.