Technische Samenvatting: TrajFusion voor het Schalen van Wiskundig Redeneren

Probleemstelling

Large Language Models (LLMs) hebben aanzienlijke vooruitgang geboekt in wiskundig redeneren, grotendeels door middel van Supervised Fine-Tuning (SFT) op hoogwaardige oplossingspaden (trajectories). Het dominante paradigma voor het construeren van deze data is Rejection Sampling Fine-Tuning (RFT). Bij RFT genereert een docentmodel meerdere kandidaat-redeneerpaden voor een gegeven probleem, waarbij alleen de paden met de juiste uiteindelijke antwoorden worden behouden voor training; alle onjuiste trajecten worden weggegooid.

Hoewel effectief, behandelt deze aanpak supervisie als een binaire filter. Het sluit systematisch door de docent gegenereerde fouten uit, wat een kloof creëert in hoe redeneerfouten tijdens de training worden gemodelleerd. De auteurs stellen dat onjuiste trajecten waardevolle diagnostische informatie bevatten — zoals veelvoorkomende computationele slips, ontbrekende aannames en logisch plausibele maar ongeldige argumentpaden — die kunnen helpen om modellen te laten herkennen en specifieke foutmodi te vermijden, met name bij langere of ambigue redeneertaken. Huidige methoden behandelen deze fouten als ruis die moet worden weggegooid in plaats van als gestructureerde supervisie die kan worden benut.

Methodologie: TrajFusion

Het paper stelt TrajFusion voor, een fine-tuning strategie die RFT herformuleert van een binair filtermechanisme naar een proces voor de constructie van gestructureerde supervisie. TrajFusion integreert geselecteerde onjuiste redeneertrajecten met correcte trajecten in één enkel gefuseerd trainingsmonster, waarbij expliciet een trial-and-error redeneerproces wordt gemodelleerd.

Kerncomponenten

1. Sampling en Analyse: Voor elk probleem $x$ samplet het docentmodel $K$ trajecten. Deze worden verdeeld in correcte ($Y^+$) en onjuiste ($Y^-$) sets met behulp van een automatische verifieerder.
2. Adaptieve Selectie: De methode bepaalt adaptief hoeveel onjuiste trajecten worden opgenomen op basis van twee statistieken:
   • Foutpercentage ($r(x)$): De frequentie van onjuiste pogingen.
   • Foutdiversiteit ($u(x)$): Het aantal unieke eindantwoorden onder de onjuiste pogingen (gemeten via Shannon-entropie of het aantal unieke antwoorden).
     Het aantal behouden onjuiste trajecten $k(x)$ wordt berekend als:
     $$k(x) = \min(k_{max}, \lfloor \alpha \cdot r(x) \cdot u(x) \rfloor)$$
     Dit zorgt ervoor dat problemen met diverse en frequente fouten rijkere supervisie ontvangen, terwijl problemen met schaarse of oninformatieve fouten van nature degraderen naar vanilla RFT.
3. Constructie van Gefuseerde Monsters: De geselecteerde onjuiste trajecten worden geïntercaleerd met het correcte traject met behulp van vaste reflectie-prompts (bijv. "Wacht... dat lijkt niet te kloppen. Laten we het opnieuw proberen."). De uiteindelijke structuur is:
   $$T(x) = [y_{e1}, \rho_1, \dots, y_{ek}, \rho_k, y^*]$$
   waarbij $y_{ei}$ onjuiste paden zijn, $\rho_i$ reflectie-prompts zijn en $y^*$ de correcte oplossing is.
4. Training: De gefuseerde sequenties worden getraind met standaard next-token prediction (SFT) zonder wijzigingen aan de modelarchitectuur of het trainingsdoel.

Belangrijkste Bijdragen

• Formalisering van een Supervisiekloof: De auteurs identificeren en formaliseren de systematische uitsluiting van door de docent gegenereerde onjuiste trajecten in RFT, en stellen dat deze fouten diagnostische waarde bieden voor het modelleren van redeneerfouten.
• TrajFusion Framework: Een gestructureerd supervisiekader dat adaptief onjuiste en correcte trajecten fuseert. Het behandelt foutdiversiteit als een expliciet controlesignaal, waardoor de methode robuust kan schalen over problemen van variërende moeilijkheid zonder hulpmodellen, kritiekgeneratoren of externe moeilijkheidslabels te vereisen.
• Empirische Validatie: De methode vertoont consistente verbeteringen ten opzichte van vanilla RFT over meerdere basismodellen (LLaMA3-8B, DeepSeekMath-7B), dataregimes (15K tot 100K voorbeelden) en contextlengtes (tot 32K).

Experimentele Resultaten

De auteurs evalueerden TrajFusion op zes wiskundige redeneerbenchmarks (GSM8K, MATH, CollegeMath, DeepMind-Mathematics, OlympiadBench, TheoremQA) en competitie-niveau datasets (AIME24, AIME25, MATH-500).

• Prestatiewinst: TrajFusion presteerde consequent beter dan vanilla RFT. Op de MATH dataset met LLaMA3-8B (100K data) behaalde TrajFusion 51,4% nauwkeurigheid vergeleken met 48,5% voor RFT. Op DeepSeekMath-7B bereikte het 52,3% versus 50,2%.
• Impact op Complexiteit: De verbeteringen waren het meest uitgesproken op uitdagende, langvormige redeneerbenchmarks (MATH, OlympiadBench, TheoremQA), waar redeneerpaden langer zijn en foutmodi diverser zijn. Winsten op eenvoudigere datasets zoals GSM8K waren kleiner maar nog steeds positief.
• Data-efficiëntie: TrajFusion behaalde competitieve resultaten met slechts 15K trainingsvoorbeelden, waarmee het baselines versloeg die aanzienlijk grotere datasets gebruikten (bijv. MMIQC met 2,3M voorbeelden).
• Schaalbaarheid naar Lange Context: Wanneer toegepast op een 32K-contextversie van Qwen2.5-Math-7B met langvormige data gedestilleerd van DeepSeek-R1, verbeterde TrajFusion de Pass@1 op AIME24 van 40,0% (RFT) naar 43,3%.
• Ablatie-studies:
  • Het verhogen van het aantal gefuseerde fouten ($k_{max}$) verbeterde de prestaties tot op een bepaald punt (afnemende meeropbrengst waargenomen bij $k_{max}=6$).
  • Het selecteren van fouten op basis van diversiteit (HighDiv) leverde betere resultaten op dan selectie op basis van enkel frequentie (LowDiv), wat suggereert dat structureel onderscheidende foutmodi informatiever zijn.
  • Modellen getraind met TrajFusion genereerden aanvankelijk langere sequenties, maar werden geleidelijk concizer terwijl ze de frequentie van reflectietokens verhoogden, wat duidt op een efficiënter gebruik van de structurele patronen.

Betekenis en Claims

Het paper claimt dat TrajFusion een eenvoudige maar effectieve manier biedt om wiskundig redeneren te schalen door gebruik te maken van supervisiesignalen die normaal gesproken worden weggegooid. Door rejection sampling te herformuleren als een proces voor de constructie van gestructureerde trial-and-error trajecten, biedt de methode rijkere supervisie zonder architecturale wijzigingen of complexe reinforcement learning-pipelines.

De auteurs benadrukken dat TrajFusion bijzonder waardevol is voor uitdagende en langvormige redeneerproblemen, waarbij het onderscheiden van geldige van gebrekkige redeneringen essentieel is. Zij stellen dat het adaptieve karakter van TrajFusion de robuustheid waarborgt: het verrijkt de supervisie wanneer fouten informatief zijn en keert veilig terug naar standaard RFT-gedrag wanneer dat niet het geval is. Deze aanpak overbrugt de kloof tussen binaire correctheid-filtering en de genuanceerde realiteit van redeneerfouten, waardoor modellen in staat worden gesteld om uit hun eigen fouten te leren in een gesuperviseerde setting.