Technische Samenvatting: IAPO: Input Attribution-Aware Policy Optimization voor het gebruik van tools door kleine multimodale agenten

Probleemstelling

Het artikel behandelt de uitdaging van het trainen van Small Language Model (SLM) agenten om effectief tools te gebruiken in multimodale scenario's, specif wrong specifiek voor visuele vraagbeantwoording (VQA) met betrekking tot grafieken en tabellen. Hoewel Reinforcement Learning (RL) methoden zoals Proximal Policy Optimization (PPO) en Group Relative Policy Optimization (GRPO) de prestaties van agenten hebben verbeterd, kampen ze met aanzienlijke beperkingen bij de toepassing op multimodale SLM's:

1. Inadequate beloningsontwerp voor multimodale taken: Bestaande methoden vertrouwen vaak op exacte overeenstemming met grondwaarheid-tooltrajecten of vooraf gedefinieerde formaten. Echter, multimodale taken laten vaak meerdere geldige tool-gebruikspaden toe (bijv. het highlighten, maskeren of tekenen van bounding boxes kan allemaal geldig zijn om te focussen op relevante regio's in een grafiek). Rigide formaatcontroles slagen er niet in deze flexibiliteit te vatten.
2. Spaarzame en breekbare beloningen: Huidige benaderingen maken meestal gebruik van binaire, op resultaten gebaseerde beloningen (correct/incorrect definitief antwoord). Deze bieden spaarzame feedback en falen in het sturen van het tussenliggende besluitvormingsproces van het model. Als gevolg hiervan kunnen SLM's leren om accidentele shortcuts of oppervlakkige correlaties te exploiteren om een correct eindantwoord te bereiken, zonder werkelijk robuuste tool-selectiebeleid te leren die geworteld is in relevante visuele bewijslast.
3. Suboptimale prestaties van Vanilla GRPO: Zelfs met vanilla GRPO blijven multimodale SLM's aanzienlijke aantallen foutief gebruikte tool-aanroepen produceren, wat aangeeft dat standaard RL-algoritmen onvoldoende zijn voor het leren van de genuanceerde besluitvorming die vereist is voor multimodale tool-gebruik.

Methodologie: Input Attribution-Aware Policy Optimization (IAPO)

Om deze problemen aan te pakken, stellen de auteurs IAPO voor, een RL-algoritme dat de input-attributie van het studentmodel afstemt op die van een sterker teacher-model. De kerngedachte is dat de bijdrage van specifieke inputcomponenten aan een gegenereerde tool-aanroep sterk correleert met de correctheid van de tool-aanroep.

Kerncomponenten:

1. Integrated Gradient (IG) Scores:

   • De input prompt wordt gedecomposeerd in semantische blokken: pure tekstcontext ($P$), gebruikersquery ($R$), afbeelding placeholder ($I$) en tool-definities ($T$).
   • De methode berekent Integrated Gradient scores om te kwantificeren hoeveel elk semantisch blok bijdraagt aan de log-waarschijnlijkheid van de gegenereerde tool-aanroepsequentie.
   • Specifiek voor een tool-aanroepsequentie $y'$, is de score $s(E)$ de log-waarschijnlijkheid onder het huidige beleid. De gradiënt van deze score met betrekking tot de input-embeddings wordt berekend, en de token-niveau attributie wordt gedefinieerd als de norm van het elementgewijze product van de gradiënt en de embedding ($o_i = \|g_i \odot e_i\|_2$).
   • Deze token-niveau attributies worden geaggregeerd om blok-niveau scores ($O_k$) voor elk semantisch type te produceren.

2. Attributie-gebaseerde Straf (Penalty):

   • Een sterker teacher-model (bijv. een grotere SLM of een met GRPO getraind model) genereert referentie IG scores ($\hat{b}$) voor dezelfde input.
   • De IG scores van het studentmodel ($\hat{a}$) worden vergeleken met die van de teacher met behulp van de Kullback-Leibler divergentie om een attributie-straf te vormen: $P_A = D_{KL}(\hat{b} \parallel \hat{a})$.
   • Deze straf moedigt de student aan om zich te concentreren op dezelfde inputcomponenten (bijv. specifieke tool-definities of query-delen) die de teacher als relevant beschouwt.

3. Gemodificeerde Beloningsfunctie:

   • De uiteindelijke beloning voor een traject $i$ combineert de standaard op resultaten gebaseerde beloning ($R_{out}$) met de attributie-straf:
     $$R_{IAPO}^i = R_{out}^i - \lambda P_A^i$$
   • Hierbij controleert $\lambda$ de sterkte van de attributie-afstemming. Dit dichte signaal stuurt het beleid aan om niet alleen de taak nauwkeurigheid te maximaliseren, maar ook om de interne besluitvormingsprocedure af te stemmen op de attributiepatronen van de teacher.

4. Trainingsframework:

   • De methode gebruikt een twee-beurten rollout: het model genereert tool-aanroepen om een afbeelding te bewerken, en redeneert vervolgens over de bewerkte afbeelding om een definitief antwoord te produceren.
   • Training maakt gebruik van de GRPO-doelstelling, maar gebruikt de gemodificeerde $R_{IAPO}$ om groep-relatieve voordelen te berekenen.

Belangrijkste Bijdragen

1. Nieuwe Integrated Gradient Scores: Het artikel introduceert een methode om de bijdrage van verschillende input-semantische blokken (tekst, afbeelding, tool-definities) aan de generatie van tool-aanroepen te kwantificeren, wat een dicht supervisiesignaal biedt.
2. IAPO Algoritme: De auteurs stellen het eerste RL-algoritme voor dat specifiek is ontworate voor SLM tool-gebruik en gebruikmaakt van attributie-gebaseerde straffen. Door de IG-scores van student en teacher via KL-divergentie te vergelijken, draagt IAPO de tool-aanroepbekwaamheid en de besluitvormingslogica van de teacher over aan de student.
3. Empirische Validatie: Het artikel toont aan dat vanilla GRPO leidt tot aanhoudend foutief tool-gebruik in multimodale SLM's. Het toont verder aan dat het aanpassen van beloningsontwerpen van tekst-alleen modellen onvoldoende is, terwijl IAPO de prestaties aanzienlijk verbetert door input-attributie af te stemmen.

Experimentele Resultaten

De methode werd geëvalueerd op visuele vraagbeantwoordingstaken waarbij het Qwen2.5-VL-3B model als policy werd gebruikt en een Qwen2.5-VL-7B model (getraind met GRPO) als teacher.

• Prestatiewinst: IAPO verbeterde de nauwkeurigheid van visuele vraagbeantwoording met gemiddeld 3% over zes testsets (inclusief ChartQA, VWTQ en CharXiv) vergeleken met bestaand visueel tool-gebruik werk (VTool-R1, TRM, ToRL).
• Tool Selectie Nauwkeurigheid: Casestudies toonden aan dat zelfs na 100 stappen van vanilla GRPO training, de SLM significante fouten maakte in de tool-selectie (bijv. het toepassen van rij-gebaseerde tools op kolom-deelbare grafieken). IAPO verminderde deze fouten aanzienlijk, waarbij het succespercentage van de tool-aanroep van het model de 100% naderde in latere trainingsfasen.
• Attributie Afstemming: De studie toonde aan dat IG-scores voor correcte inputcomponenten toenamen tijdens de training en uiteindelijk de scores van foutief gebruikte componenten overtroffen, wat correleert met de reductie in foutief tool-gebruik. Het 7B teacher-model vertoonde een duidelijkere scheiding in IG-scores tussen correcte en incorrecte inputs, wat het gebruik van een sterkere teacher voor begeleiding valideert.
• Generalisatie: IAPO demonstreerde superieure generalisatie op out-of-distribution benchmarks (CharXiv), wat suggereert dat attributie-afstemming helpt om overfitting op oppervlakkige prompt-correlaties te voorkomen.

Betekenis en Claims

Het artikel beweert dat IAPO een effectiever framework biedt voor het trainen van multimodale SLM-agenten dan bestaande RL-methoden. De betekenis ligt in:

• Het gat dichten voor SLM's: Het pakt de specifieke moeilijkheid aan waarmee SLM's kampen bij het leren van spaarzame, binaire beloningen door een dicht, procesmatig supervisiesignaal af te leiden uit input-attributie.
• Verder gaan dan binaire beloningen: Het daagt de afhankelijkheid van exacte overeenstemming of binaire resultaatbeloningen in multimodale settings uit, door voor te stellen dat het afstemmen van hoe een model naar input kijkt (attributie) cruciaal is voor robuust tool-gebruik.
• Teacher-Student Distillatie in RL: Het slaagt erin een sterker teacher-model te gebruiken, niet alleen voor output-imitatie, maar ook voor het begeleiden van het interne besluitvormingsproces (attributiepatronen) van een kleiner student-model tijdens online RL-training.

De auteurs concluderen dat hoewel grotere modellen van nature betere tool-selectiebekwaamheden en attributiepatronen vertonen, IAPO deze capaciteiten effectief overdraagt aan kleinere modellen, waardoor ze kunnen letten op de meest relevante visuele bewijslast en een state-of-the-art prestatie kunnen leveren in multimodale tool-gebruikstaken.