DRP: Distilled Reasoning Pruning with Skill-aware Step Decomposition for Efficient Large Reasoning Models

1. Probleemstelling

Large Reasoning Models (LRM's), zoals OpenAI's o1 en DeepSeek-R1, blinken uit in complexe taken door lange Chain-of-Thought (CoT) redeneertrajecten te genereren. Ze lijden echter aan "overthinking", waarbij ze buitensporig uitgebreide outputs produceren met redundante stappen, onnodige terugkeer en repetitieve lussen. Dit leidt tot:

• Hoge Inferentiekosten: Significant tokenverbruik en latentie.
• Afname van Nauwkeurigheid: Redundant redeneren kan het model naar verkeerde conclusies leiden.
• Distillatie-mismatch: Bij het distilleren van kennis van een "Short-CoT" leraar (beknopt, gepolijst) naar een "Long-CoT" student (uitgebreid, reflecterend), ontstaat er een leerbaarheidskloof. De student worstelt om de beknopte stijl van de leraar direct na te bootsen, wat vaak resulteert in slechte prestaties of dalingen in nauwkeurigheid.

Bestaande oplossingen zoals inferentie-tijd pruning (het vroegtijdig stoppen van de generatie) riskeren het redeneren voortijdig te onderbreken, terwijl standaard distillatie vaak faalt om de structurele kloof tussen uitgebreide studenten en beknopte leraren te overbruggen.

2. Methodologie: Distilled Reasoning Pruning (DRP)

DRP is een hybride raamwerk dat pruning op inferentietijd verenigt met supervised fine-tuning (SFT) distillatie. In plaats van de student simpelweg te trainen op het definitieve antwoord van de leraar of een korte CoT, gebruikt DRP de leraar om de eigen uitgebreide redeneersporen van de student te pruningen en te verfijnen.

Het proces bestaat uit drie fasen (geïllustreerd in Figuur 2 van het paper):

A. Op Vaardigheden Gebaseerde Stap-decompositie

• Het studentenmodel genereert een initiële, uitgebreide CoT-spoor ($T$).
• Een Teacher LLM (bijv. GPT-4o) decomposeert dit spoor in niet-overlappende segmenten.
• Kerninnovatie: In tegenstelling tot naïeve zinsplitsing, kiest de leraar een functionele vaardigheidslabel toe aan elke stap (bijv. "Lezen van de gegeven hoeveelheid", "Algebraïsche representatie", "Logische inferentie"). Dit zorgt ervoor dat stappen semantisch coherent en functioneel uitgelijnd zijn.

B. Stap-niveau Pruning en Herschrijven

Het lerarenmodel evalueert elke met vaardigheden gelabelde stap en past een van de vier acties toe:

1. Behouden: De stap onveranderd behouden.
2. Verwijderen: Redundante, niet-informatieve of terugkerende stappen verwijderen.
3. Herschrijven: De stap samenvatten terwijl de logische intentie behouden blijft.
4. Samenvoegen: Aanliggende stappen combineren tot een enkele atomaire eenheid.

• De leraar reconstrueert vervolgens een vloeiende, beknopte redeneerspoor ($\hat{T}$) dat de oorspronkelijke toon van de student behoudt maar de uitgebreidheid verwijdert. Het actualiseert ook de samenvatting van het definitieve antwoord ($\hat{A}$) indien nodig.

C. Supervised Fine-Tuning (SFT)

• Het studentenmodel wordt gefine-tuned op paren van $(Input, \hat{R})$, waarbij $\hat{R}$ het gepreinde, door de leraar verfijnde antwoord is.
• Het doel is om de negatieve log-waarschijnlijkheid van het verfijnde spoor te minimaliseren, waardoor de student leert op vaardigheden uitgelijnde, token-efficiënte redeneerstrategieën te internaliseren zonder de structurele diepgang die nodig is voor complexe problemen te verliezen.

3. Belangrijkste Bijdragen

1. DRP-raamwerk: Een nieuwe methode die de "leerbaarheidskloof" overbrugt door het redeneerpad van de student te pruningen met behulp van een leraar, in plaats van de student te dwingen een volledig ander redeneerstijl van de leraar na te bootsen.
2. Vaardigheidsbewuste Decompositie: Introduceert een mechanisme om redeneren te segmenteren in semantisch coherente, functioneel gelabelde eenheden. Dit biedt een stabiele basis voor nauwkeurige pruning en presteert beter dan naïeve splitsing.
3. Inzicht in Structurele Consistentie: Demonstreert dat effectieve trainings-CoT's zowel informatief als structureel consistent moeten zijn met het redeneerproces van de student. Het behoud van de stap-voor-stap structuur (zelfs als deze verkort is) is effectiever dan naïeve compressie.

4. Experimentele Resultaten

De auteurs evalueerden DRP op DeepSeek-R1-Distill-Qwen-modellen (1.5B en 7B) over in-domein (GSM8K) en out-of-distribution (OOD) benchmarks (MATH500, AIME24, AMC23).

Belangrijkste Bevindingen:

• Token-efficiëntie: DRP verminderde het tokenverbruik aanzienlijk zonder afbreuk te doen aan de nauwkeurigheid.
  • GSM8K (7B): Tokens gereduceerd van 917 naar 328 (-64%), terwijl de nauwkeurigheid verbeterde van 91,7% naar 94,1%.
  • AIME24 (7B): Tokens gereduceerd met 43% zonder daling in nauwkeurigheid (15/30 correct).
  • AMC (1.5B): Tokens gereduceerd van 4574 naar 3657, met een stijging in nauwkeurigheid van 22/40 naar 27/40.
• Vergelijking met Baselines:
  • vs. Prompt-based (TALE): DRP presteerde beter dan TALE op complexe OOD-taken waar de prestaties van TALE aanzienlijk daalden.
  • vs. SFT-based (CoT-Valve): DRP behaalde betere nauwkeurigheid en compressie. CoT-Valve leed vaak aan nauwkeurigheidsafwegingen bij compressie.
  • vs. RL-based (ThinkPrune): Hoewel ThinkPrune goed comprimeerde, behaalde DRP een hogere nauwkeurigheid, waarschijnlijk door hoogwaardige leraar-gidsing.
• Ablatiestudies:
  • Vaardigheidsdecompositie: Het verwijderen van vaardigheidslabels (standaard splitsing) verminderde de prestaties, wat bevestigt dat semantische uitlijning cruciaal is.
  • Directe Distillatie: Het direct trainen op korte leraar-CoT's (zonder het spoor van de student te pruningen) leidde tot scherpe dalingen in nauwkeurigheid op OOD-taken, wat bewijst dat structureel behoud vitaal is.
  • Robuustheid Leraar: De methode werkte effectief met verschillende leraren (GPT-4o, Gemini 2.0, DeepSeek-V3), hoewel GPT-4o de beste resultaten opleverde.

5. Betekenis en Conclusie

• Afweging Efficiëntie vs. Nauwkeurigheid: DRP daagt het idee uit dat efficiëntie ten koste moet gaan van nauwkeurigheid. Door het redeneerproces te verfijnen in plaats van alleen de output te verkorten, bereikt het beide.
• Schaalbaarheid: De methode is bijzonder effectief voor kleinere modellen (1.5B), waarbij beperkte capaciteit wordt gecompenseerd door hoogwaardige, gestructureerde supervisie.
• Robuustheid: Het mitigeert "degenererende lussen" (waarbij modellen vast komen te zitten in het genereren van oneindige tekst), wat leidt tot een stabieler inferentie.
• Toekomstige Impact: Dit werk biedt een blauwdruk voor het trainen van efficiënte redeneermodellen die complexe, langdurige taken kunnen hanteren zonder de rekenkundige overhead van "overthinking", waardoor geavanceerd redeneren toegankelijker en kosteneffectiever wordt.

Kortom, DRP lost de inefficiëntie van Large Reasoning Models op door een leraar te gebruiken om chirurgisch redundantie te verwijderen uit het eigen redeneerpad van een student, geleid door functionele vaardigheden, wat resulteert in modellen die sneller, goedkoper en vaak nauwkeuriger zijn.