Technische Samenvatting: LARK (Learnability-Grounded Trajectory Selection)

Probleemstelling

Het artikel behandelt de uitdaging van reasoning distillation (redeneer-destillatie), een paradigma waarbij een kleiner student Large Language Model (LLM) wordt gefinetuned om redeneertrajecten te imiteren die zijn gegenereerd door grotere teacher-modellen. Hoewel bestaande methoden voor dataselectie vaak opereren op het niveau van de vraag of vertrouwen op heuristieken zoals trajectkwaliteit, correctheid of modelvertrouwen, negeren ze vaak een cruciale factor: leerbaarheid (learnability).

Specifiek kunnen er voor één enkele vraag meerdere kandidaat-redeneertrajecten bestaan (gegenereerd door verschillende teachers, sampling-runs of stijlen). Zelfs als deze trajecten tot hetzelfde correcte eindantwoord leiden, leveren ze zeer verschillende trainingssignalen. Een traject dat "hoogwaardig" of "natuurlijk" lijkt, is niet noodzakelijkerwijs het traject waarvan de student het meest efficiënt kan leren. De centrale vraag is: Kunnen we een principieel criterium voor trajectselectie ontwerpen dat ervoor zorgt dat gedestilleerde redeneertrajecten leerbaar zijn voor het studentmodel?

Methodologie: LARK

De auteurs stellen LARK voor (Learnability-grounded Anchor-time Ranking), een methode die trajecten selecteert op basis van een optimalisatietheoretische definitie van leerbaarheid in plaats van statische kwaliteitsmetrieken.

1. Leerbaarheid als een optimalisatiedoel

LARK formuleert trajectselectie als een policy-optimalisatieprobleem. Het definieert een leerbaarheidsfactor $\rho(\theta; q)$, die de snelheid karakteriseert waarmee de trainingsverlies van de student afneemt. Deze factor wordt gedefinieerd als de ratio van de gekwadrateerde gradiëntnorm tot de loss:
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
waarbij $q$ de gewichtsvector is over $K$ kandidaat-trajecten. Onder de Polyak–Łojasiewicz (PL) conditie impliceert een hogere $\rho$ een snellere convergentie.

2. Anchor-Time Approximatie

Het direct optimaliseren van de tijd-geïntegreerde leerbaarheidsdoelstelling over het volledige fine-tuning traject is onhandelbaar. LARK introduceert de anchor-relatieve conditie, die stelt dat voor overgeparameteriseerde modellen in het "lazy training" regime, de leerbaarheidsgraad proportioneel blijft aan de waarde bij de initiële anchor-parameters $\theta_{ref}$. Dit maakt het mogelijk om de doelstelling te benaderen door $\rho(\theta_{ref}; q)$ te maximaliseren.

3. Lokale Linearisatie en Forward-Pass Proxy

Het maximaliseren van $\rho(\theta_{ref}; q)$ is problematisch omdat het een quasi-convexe functie is die inklapt naar een gedegenereerde "one-hot" oplossing (het selecteren van slechts één enkel traject). Om dit te voorkomen, gebruikt LARK:

• Lokale Linearisatie: Breidt $\rho$ uit rond de uniforme distributie $p$ (standaard SFT zonder selectie) met behulp van een eerste-orde Taylor-expansie.
• Forward-Pass Surrogate ($\hat{g}_k$): In plaats van exacte gradiënten te berekenen (die backward passes vereisen voor elke kandidaat), leidt LARK een berekenbare proxy $\hat{g}_k$ af die uitsluitend gebruikmaakt van forward passes. Deze proxy schat de lokale gradiënt van de leerbaarheidsgraad op basis van de token-niveau voorspellingsresiduen en cross-entropy loss van de student.
  • Specifiek bevoordeelt $\hat{g}_k$ trajecten waar de student "zelfverzekerd fout" is op een gestructureerde manier: de loss $\ell_k$ is hoog (wat duidt op slechte voorspelling), maar de residuele massa is geconcentreerd op een paar "bijna-fouten" tokens (wat een duidelijk correctiesignaal biedt).

4. $\chi^2$-Regularisatie Selectiebeleid

Om te voorkomen dat de selector het leerbaarheidscriterium "hackt" door te ver af te wijken van de oorspronkelijke datadistributie (wat de generalisatie schaadt), legt LARK een $\chi^2$-regularisatie term op. Dit leidt tot een closed-form, budget-geparameteriseerde selectieregel:

• Trajecten worden gerangschikt op basis van hun score $\hat{g}_k$.
• De top $B$ trajecten worden geselecteerd.
• In tegen afzetting van harde top-$B$ truncatie, wijst LARK zachte gewichten toe die proportioneel zijn aan de score-marge boven de $(B+1)$-de traject. Deze wegingregel is afgeleid om de ondergrens van de leerbaarheidsverbetering te maximaliseren terwijl de $\chi^2$-divergentie van de uniforme prior onder controle wordt gehouden.

Belangrijkste Bijdragen

1. Formulering: De auteurs formuleren de selectie van redeneertrajecten als een leerbaarheids-gegrond policy-optimalisatieprobleem, waarbij ze de anchor-time leerbaarheidsgraad $\rho$ introduceren als een principieel doel.
2. Algoritme: Ze leiden een closed-form, budget-geparameteriseerde selectieregel af die een forward-pass surrogate $\hat{g}_k$ en $\chi^2$-regularisatie gebruikt. Dit vermijdt dure per-traject backward passes terwijl het theoretisch garanties biedt op foutmarges en generalisatie.
3. Empirische Validatie: De paper demonstreert dat LARK consistent beter presteert dan bestaande baselines (inclusief heuristische, kwaliteitsgebaseerde en student-zijde scoring methoden) over meerdere base-modellen (Qwen-2.5-7B/1.5B, Llama-3.2-3B) en reasoning benchmarks (AIME, AMC, GPQA, MATH-500).

Experimentele Resultaten

• Prestaties: LARK behaalde de hoogste gemiddelde Acc@5 over alle geteste model-budget settings. Bijvoorbeeld, op Qwen-2.5-7B met een budget van $B=3$, verbeterde LARK de prestaties met 7,07 procentpunt ten opzichte van de sterkste baseline (RSR).
• Efficiëntie: LARK is een forward-pass-only methode, waardoor het qua computationele kosten overeenkomt met de goedkoopste student-zijde baselines, terwijl exacte gradiënt-gebaseerde methoden aanzienlijk meer rekenkracht zouden vereisen.
• Diagnostische Analyse:
  • Loss Decay: Trajecten geselecteerd door LARK induceerden de snelste reductie in de SFT loss tijdens de training.
  • Voorspellende Geldigheid: De LARK-score $\hat{g}_k$ vertoonde een sterke positieve correlatie (Pearson $r=0,69$) met downstream accuratesse wanneer geaggregeerd per teacher-model.
  • Token-niveau Gedrag: Geselecteerde trajecten vertoonden lagere correct-token waarschijnlijkheden maar een hogere "wrong-token concentratie", wat bevestigt dat LARK voorbeelden selecteert waar de student fout zit, maar de fout gestructureerd en corrigeerbaar is.

Betekenis en Claims

De paper claimt dat leerbaarheid een onderscheidend en effectief principe is voor efficiënte reasoning distillation, los van trajectkwaliteit of moeilijkheidsgraad. Door selectie te funderen in de optimalisatiedynamiek van het studentmodel, identificeert LARK trajecten die het meest efficiënte leersignaal bieden.

De auteurs benadrukken dat hun methode:

• Generalisatie Behoudt: Door middel van $\chi^2$-regularisatie voorkomt het dat het selectiebeleid overfit op een smalle set "makkelijke" of "gehackte" trajecten.
• Theoretisch Onderbouwd Is: De selectieregel is afgeleid van eerste principes van optimalisatietheorie (PL-condities, gradient flow) met expliciete foutmarges.
• Praktisch Bruikbaar Is: Het vereist geen hyperparameter tuning voor de selectieregel zelf (de regularisatiestärke wordt bepaald door het budget $B$) en vermijdt de computationele overhead van backward passes.

Het werk suggereert dat toekomstige extensies dit leerbaarheids-gegronde perspectief kunnen toepassen op ruiziger kandidaat-pools en bredere reasoning domeinen, maar de huidige scope is strikt beperkt tot reasoning distillation in wiskundige probleemoplossing.