CompoDistill: Attention Distillation for Compositional Reasoning in Multimodal LLMs

Probleemstelling

Multimodale Grootte Taalmodellen (MLLMs) presteren uitstekend op diverse taken, maar hun enorme rekenkosten maken ze onpraktisch voor veel real-world toepassingen. Kennisdistillatie (Knowledge Distillation - KD) is een veelbelovende aanpak om grote 'leraar'-modellen te comprimeren naar kleinere 'student'-modellen.

Echter, de auteurs stellen vast dat bestaande KD-methoden weliswaar goed presteren op Visuele Vraagbeantwoording (VQA) (gericht op objectherkenning), maar falen bij het distilleren van Visuele Perceptie (gericht op complexe relaties tussen objecten, attributen en context). Dit wordt zichtbaar in Compositional Reasoning (CR) taken, waar studenten vaak niet beter presteren dan modellen die alleen met Supervised Fine-Tuning (SFT) zijn getraind.

De kernoorzaak van dit falen is geïdentificeerd als visuele attentie-misalignement: de student-modellen slagen er niet in om de rijke visuele aandachtspatronen van de leraar te kopiëren. Waar de leraar zich richt op de relevante beeldgebieden voor een specifieke tekstvraag, focust de student vaak op irrelevante gebieden, wat leidt tot een gebrek aan vermogen om objectrelaties te begrijpen.

Methodologie: CompoDistill

Om dit probleem op te lossen, stellen de auteurs CompoDistill voor, een nieuw KD-framework dat specifiek gericht is op het verbeteren van visuele perceptie door attentie-misalignement aan te pakken. Het framework bestaat uit twee kernmodules en een drie-fasen trainingsstrategie.

1. Kernmodules

• Visual ATtention Alignment (VAT) Module:
  • Doel: De visuele aandacht van de student expliciet aligneren met die van de leraar.
  • Mechanisme: De module berekent een verliesfunctie (cosine distance) tussen de attentiematrices van de leraar en de student, specifiek gericht op de visuele tokens binnen de "visuele begrijpende lagen" (intermediate layers) van het LLM.
  • Group Layer Matching: Omdat leraren vaak meer lagen hebben dan studenten, wordt een "one-to-many" strategie gebruikt. Elke laag van de student wordt gekoppeld aan een groep opeenvolgende lagen van de leraar. Dit zorgt voor een robuustere overdracht van kennis dan een simpele 1-op-1 mapping.
• Teacher Adapter Fetch (TAF) Module:
  • Doel: De incompatibiliteit tussen de visuele feature-ruimte van de leraar en die van de student overbruggen.
  • Mechanisme: De TAF-module gebruikt de bevroren (frozen) adapter van de leraar om visuele features naar de taalruimte te projecteren. Een lichte, trainbare MLP (Multi-Layer Perceptron) wordt toegevoegd om de dimensies aan te passen aan de student. Dit zorgt ervoor dat de student de visuele input verwerkt via dezelfde "lens" als de leraar, waardoor de overdrage van de attentie-mechanismen effectief wordt.

2. Drie-fasen Trainingsstrategie

1. Distilled Pre-Training (DPT): De visuele feature-ruimte wordt afgestemd op de taalruimte. De adapter van de student wordt geïnitieerd met de TAF-module (leraar-adapter) en getraind met taalmodelleringsverlies en KL-divergentie, terwijl de LLM en visuele encoder bevroren blijven.
2. Distilled Fine-Tuning (DFT): De kernfase voor visuele perceptie. De student LLM en adapter worden fine-getuned met het toevoegen van het VAT-verlies (attentie-distillatie) om de attentiepatronen van de leraar te leren.
3. Supervised Fine-Tuning (SFT): Een laatste fase om de instructie-opvolging te versterken en de overgedragen kennis te consolideren, waarbij alleen het autoregressieve taalverlies wordt gebruikt.

Belangrijkste Bijdragen

1. Analyse van Visuele Perceptie: Voor het eerst wordt aangetoond dat bestaande KD-methoden in MLLMs falen in het distilleren van visuele perceptie (relaties en context) ondanks succes in herkenning. De auteurs identificeren visuele attentie-misalignement als de primaire oorzaak.
2. CompoDistill Framework: Een nieuw framework dat de VAT en TAF modules combineert om zowel visuele herkenning als perceptie effectief over te dragen.
3. Empirische Validatie: Uitgebreide experimenten tonen aan dat het aanpakken van attentie-misalignement direct leidt tot betere prestaties in complexe redeneertaken.

Resultaten

De auteurs evalueren CompoDistill op diverse benchmarks voor VQA (herkenning) en CR (perceptie):

• Compositional Reasoning (CR): CompoDistill (2B parameters) bereikt prestaties die concurreren met 4B-modellen en significantly outperforms bestaande KD-methoden (zoals LLaVA-KD en LLaVA-MoD) en SFT-baselines. Het verbetert de prestaties op datasets zoals SugarCrepe, SADE, BiVLC en Winoground aanzienlijk.
• Visual Question Answering (VQA): Het model behoudt sterke prestaties op VQA-taken (zoals GQA, VQAv2), wat aantoont dat het verbeteren van perceptie niet ten koste gaat van herkenning.
• Relational Hallucinations: CompoDistill vermindert relationele hallucinaties (fouten in het begrijpen van relaties tussen objecten) aanzienlijk, wat resulteert in prestaties die bijna gelijk zijn aan die van de leraar op benchmarks zoals R-Bench en Reefknot.
• Efficiëntie: Het model bereikt deze resultaten met slechts 1,2 miljoen trainingsvoorbeelden, wat veel efficiënter is dan andere methoden die grotere datasets vereisen.
• Generalisatie: De methode werkt ook effectief met verschillende backbone-architecturen (bijv. MobileLLaMA) en grotere leraren (tot 7B parameters).

Significantie

Dit werk biedt een nieuw perspectief op de beperkingen van kennisdistillatie in multimodale modellen. Het benadrukt dat het simpelweg kopiëren van output (logits) niet voldoende is voor complexe visuele redeneertaken; de interne attentie-mechanismen moeten worden afgestemd.

CompoDistill stelt een nieuwe standaard voor het ontwikkelen van efficiënte, kleine MLLMs die niet alleen objecten kunnen herkennen, maar ook de complexe relaties en context in beelden begrijpen. Dit is cruciaal voor het toepassen van geavanceerde AI in real-world scenario's waar rekenkracht beperkt is, maar diepgaand visueel inzicht vereist is.