Technische Samenvatting: Waarom Feedback-Augmented Self-Distillation faalt bij Retrieval-Interleaved Search Agents

Probleemstelling

Retrieval-interleaved zoekagents, die afwisselen tussen redeneren en extern instrumentgebruik (zoeken), kampen met uitdagingen tijdens de training door de schaarste van resultaatgestuurde beloningen (outcome-level rewards). Hoewel Reinforcement Learning (RL) en dichte supervisiesignalen (bijv. process rewards) zijn voorgesteld, biedt On-Policy Self-Distillation (OPSD) een veelbelovend alternatief door één Large Language Model (LLM) te gebruiken als zowel student als docent. In OPSD wordt de docent geconditioneerd op "geprivilegieerde informatie" (zoals succesvolle rollouts) om dichte token-niveau supervisie te bieden aan de ongeconditioneerde student.

Het blijft echter onduidelijk of dit paradigma effectief kan worden uitgebreid naar retrieval-interleaved agents. In tegenstelling tot standaard taalmodellering genereren zoekagents trajecten waarbij de acties van het model de daaropvolgende opgehaalde observaties (omgevingsoutputs) bepalen, wat een mismatch creëert tussen standaard token-niveau distillatie en de voor effectieve retrieval vereiste lange-horizon zoekbeslissingen. Dit artikel onderzoekt of Feedback-Augmented Self-Distillation (FA-SD) — waarbij succesvolle rollouts worden gebruikt als geprivilegieerde demonstraties om de zelf-docent te conditioneren — betrouwbare supervisie kan bieden voor deze agents.

Methodologie

1. Feedback-Augmented Self-Distillation (FA-SD)

De auteurs instantiëren FA-SD voor zoekagents. Het proces omvat:

• Rollout Generatie: Het samplen van een groep trajecten voor een gegeven vraag.
• Geprivilegieerde Feedback: Als er binnen de groep een succesvol traject bestaat, wordt dit gebruikt als feedback ($f$) om de prompt van de "zelf-docent" te verrijken.
• Distillatie Doelstelling: De ongeconditioneerde studentpolicy ($\pi_\theta$) wordt getraind om de tokenverdeling van de feedback-verrijkte zelf-docent ($\pi_\theta(\cdot | x, f)$) te matchen. De verliesfunctie minimaliseert de divergentie (bijv. KL-divergentie) tussen de student en de stop-gradient zelf-docent, waarbij opgehaalde observaties worden gemaskeerd en de focus ligt op assistent-gegenereerde tokens.

2. Diagnostiek van Falende Modi

Om te begrijpen waarom FA-SD faalt, ontleden de auteurs de instabiliteit van het supervisiesignaal in twee componenten:

• Model Inconsistentie: De zelf-docent en de student delen dezelfde continu bijgewerkte parameters, waardoor de docent-target tijdens de training verschuift (drift).
• Prompt Inconsistentie: De docent is geconditioneerd op feedback-demonstraties, terwijl de student dat niet is. Dit creëert een distributionele mismatch die mogelijk niet zuiver kan worden overgedragen.

3. Voorgestelde Mitigaties

Om Model Inconsistentie aan te pakken, evalueren de auteurs twee regularisatiestrategieën voor de zelf-docent:

• Fixed Reference Teacher: De docent wordt bevroren op de initiële parameters van de student aan het begin van de training.
• Exponential Moving Average (EMA) Teacher: De parameters van de docent worden bijgewerkt als een voortschrijdend gemiddelde van de studentparameters ($\theta_{EMA} = \alpha \theta_{EMA} + (1-\alpha)\theta_t$), waardoor de drift van de docent-target wordt vertraagd.

4. Externe Docent Extensie

De auteurs testen ook of de feedback-augmentatie truc generaliseert naar Multi-Rollout On-Policy Distillation (MOPD) met een externe docent (FA-MOPD), waarbij de externe docent geconditioneerd is op succesvolle rollouts terwijl de student ongeconditioneerd blijft.

Belangrijkste Bevindingen en Resultaten

1. Falen van Standaard FA-SD

Over zeven open-domain QA-benchmarks (NQ, HotpotQA, etc.) heen, falen de standaard FA-SD varianten (unclipped en PPO-stijl clipped) om verbetering te handhaven.

• Trainingsdynamiek: De pass rates blijven gedurende het grootste deel van de training nabij nul, met slechts een tijdelijke opleving voordat ze weer terugvallen.
• Decoding Collapse: De auteurs identificeren een specifieke falende modus genaamd decoding collapse. Modellen beginnen te vertrouwen op terugkerende reasoning-and-search output templates. Deze trajecten lijken divers, maar zijn grotendeels ongevoelig voor de invoervraag (input-agnostic). Bijgevolg wordt het KL-gebaseerde distillatiesignaal oninformatief omdat het model niet leert de specifieke vraag op te lossen, maar eerder een generieke template nabootst. Standaard geaggregeerde metrieken missen dit vaak omdat ze de input-specificiteit niet verifiëren.

2. De Zelf-Docent is Capabel, maar Internalisatie Faalt

Experimenten tonen aan dat de feedback-verrijkte zelf-docent tijdens inferentie vaak sterker is dan de ongeconditioneerde student. Dit weerlegt de hypothese dat het falen komt door een zwakke docent. In plaats daarvan is het een internalisatie-falen: de student kan het token-niveau distillatiesignaal afgeleid van de geprivilegieerde feedback niet betrouwbaar absorberen in zijn ongeconditioneerde policy.

3. Impact van Regularisatie

• Stabilisatie: Zowel de Fixed Reference als de EMA docent stabiliseren de training door te voorkomen dat de docent-target te snel verschuift.
• Prestaties: De EMA-docent levert de beste resultaten op. Het stelt het model in staat om te herstellen van tijdelijke regressies en bereikt een hogere uiteindelijke prestatie (gemiddelde score van 0.206 over benchmarks) vergeleken met de baseline (0.114).
• Warm-up Fase: De EMA-aanpak vereist een warm-up fase waarin de prestaties tijdelijk kunnen dalen voordat het supervisiesignaal stabiliseert en verbetert.

4. Generalisatie naar Externe Docenten

Wanneer de feedback-augmentatie truc wordt toegepast op MOPD met een sterke externe docent (FA-MOPD), verslechteren de prestaties vergeleken met standaard MOPD. Dit suggereert dat geprivilegieerde feedback prompt-inconsistentie of onnodige conditionering introduceert, zelfs wanneer een sterke externe docent aanwezig is, en niet rechtstreeks generaliseert naar externe-docent OPD.

Belangrijkste Bijdragen

1. Instantiatie van FA-SD: Het artikel definieert en implementeert Feedback-Augmented Self-Distillation voor retrieval-interleaved zoekagents, waarbij succesvolle rollouts worden gebruikt als geprivilegieerde demonstraties.
2. Identificatie van Decoding Collapse: De auteurs karakteriseren een falende modus waarbij modellen input-agnostische trajecten produceren via terugkerende templates, waardoor het distillatiesignaal oninformatief wordt. Dit fenomeen is verschillend van eenvoudige prestatie-regressie en wordt vaak gemist door standaard metrieken.
3. Analyse van Supervisie Inconsistentie: De studie ontleedt de instabiliteit van FA-SD in model-inconsistentie (drifting teacher) en prompt-inconsistentie (conditioning mismatch). Het toont aan dat het stabiliseren van de docent via EMA of vaste referenties cruciaal is voor effectieve training.
4. Negatieve Resultaten op Generalisatie: Het paper laat zien dat hoewel sterke externe docenten stabiele vroege supervisie bieden in MOPD, het toevoegen van geprivilegieerde feedback aan hen (FA-MOPD) de prestaties verslechtert, wat aangeeft dat feedback-augmentatie specifieke inconsistenties introduceert in de context van zoekagents.

Betekenis

Dit artikel biedt een kritische diagnose van waarom zelf-distillatie, een veelbelovende techniek voor dichte supervisie, faalt in de complexe setting van retrieval-interleaved agents. Het benadrukt dat decoding collapse en supervisie-inconsistentie de primaire barrières zijn, en niet enkel het gebrek aan een sterke docent. Door aan te tonen dat EMA-regularisatie de zelf-docent kan stabiliseren en prestaties kan herstellen, biedt het werk een praktisch pad voor het toepassen van zelf-distillatie op agentische taken. Bovendien waarschuwt het tegen de naïeve toepassing van feedback-augmentatie op externe-docent frameworks, wat suggereert dat de voordelen van geprivilegieerde informatie zeer gevoelig zijn voor de specifieke architectuur en conditionering van de distillatie-opstelling.