AEGIS: Anchor-Enforced Gradient Isolation for Knowledge-Preserving Vision-Language-Action Fine-Tuning

Titel

AEGIS: Anker-afgedwongen Gradiëntisolatie voor Kennisbehoud bij Vision-Language-Action (VLA) Fine-tuning

1. Het Probleem: Cross-Modale Gradiënt Asymmetrie

Het aanpassen van voorgeöefende Vision-Language Models (VLM's) voor robotbesturing (Vision-Language-Action of VLA) leidt vaak tot catastrofaal vergeten van de oorspronkelijke visuele en semantische vaardigheden (zoals het beantwoorden van vragen, VQA).

• De Oorzaak: Er bestaat een fundamentele geometrische mismatch tussen de trainingssignalen:
  • VLM Pre-training: Gebaseerd op Cross-Entropy (CE) voor tokenvoorspelling. Dit creëert een complexe, hoog-dimensionale semantische manifold met gradiënten die energie verspreiden over vele singulariteitsrichtingen.
  • Robot Fine-tuning: Gebaseerd op Flow-Matching met een Mean Squared Error (MSE) verlies voor continue motorische acties. Dit genereert gradiënten met een laag-rang (low-rank) profiel, geconcentreerd in een smalle fysische subspace.
• Het Effect: Wanneer deze hoog-magnitude, laag-rang MSE-gradiënten terugpropageren door het VLM, overschrijven ze de delicate semantische gewichten. Standaard optimalisatoren behandelen deze gradiënten indiscriminatie, wat leidt tot een snelle erosie van het VQA-vermogen binnen enkele duizenden stappen.
• Bestaande Oplossingen en Hun Tekortkomingen:
  • Stop-Gradient: Onderbreekt de gradiëntstroom volledig. Dit behoudt VQA, maar gooit de rijke continue supervisie weg en dwingt het gebruik van discrete proxy-tokens af (wat de representatie beperkt).
  • LoRA (Low-Rank Adaptation): Beperkt de update-capaciteit tot een laag-rang subspace. Echter, als de destructieve gradiënt in deze subspace projecteert, wordt deze toch uitgevoerd. LoRA vertraagt het vergeten, maar voorkomt het niet.
  • Co-training: Het mengen van VQA-data met robotdata verdubbelt de rekentijd en vereist toegang tot de originele dataset.

2. Methodologie: AEGIS

AEGIS (Anchor-Enforced Gradient Isolation System) lost dit probleem op door de richting van de gradiënt te manipuleren in plaats van de grootte of de data te veranderen. Het is een buffer-vrij, laag-voor-laag orthogonaal gradiëntprojectie-ramek.

De methode bestaat uit drie hoofdbestandsdelen:

1. Statisch Wasserstein-Anker (Static Wasserstein Anchor):

   • Voordat fine-tuning begint, berekent AEGIS per laag statistieken (gemiddelde $\mu$ en variantie $\sigma^2$) van de geactiveerde neuronen van het voorgeöefende model op een set VQA-data.
   • Dit vormt een referentiepunt op de Wasserstein-2 (W2) manifold van activatieverdelingen. Het anker gebruikt volledige attention-masks (beelden + tekst) om de volledige semantische manifold te vangen.

2. Wasserstein-2 Transport Penalty:

   • Tijdens elke trainingsstap berekent AEGIS de huidige activatiestatistieken van het model.
   • Het berekent de kwadratische W2-afstand (via de Bures-metric) tussen de huidige statistieken en het statische anker.
   • Dit genereert een referentiegradiënt ($g_{ot}$) die de richting aangeeft waarin het model moet bewegen om de oorspronkelijke activatieverdeling te herstellen. Belangrijk: Dit wordt niet als verliesfunctie toegevoegd, maar alleen gebruikt om een gradiëntrichting te genereren.

3. Laag-voor-laag Orthogonale Gradiënt Projectie (Layer-wise OGP):

   • AEGIS voert een sequentiële dubbele backpropagatie uit: eerst voor de taakgradiënt ($g_{task}$, van de robot-taak) en daarna voor de anker-gradiënt ($g_{ot}$).
   • Voor elke transformer-laag wordt de Gram-Schmidt-orthogonalisatie toegepast:
     • Als de taakgradiënt een destructieve component heeft (d.w.z. een negatieve projectie op de anker-gradiënt), wordt deze component verwijderd.
     • Formule: $g_{final} = g_{task} - \alpha \cdot g_{ot}$, waarbij $\alpha$ de projectiecoëfficiënt is.
   • Dit "buigt" de taakgradiënt weg van de destructieve richting, zodat deze orthogonaal staat op de anker-manifold, terwijl de constructieve energie behouden blijft.

3. Belangrijkste Bijdragen

• Identificatie van de Oorzaak: Het paper formaliseert "cross-modal gradient asymmetry" als de geometrische oorzaak van vergeten bij VLA-finetuning, waarbij het toont hoe laag-rang MSE-gradiënten hoog-dimensionale semantische manifolds vernietigen.
• AEGIS Framework: Een nieuwe, buffer-vrije methode die een statisch W2-anker gebruikt en per-laag orthogonale projectie toepast. Het vereist geen replay buffer, geen co-training data en geen discrete tokenisatie.
• Geometrische Inzichtelijkheid: Het demonstreert dat de destructieve component van gradiënten klein is (<1% van de energie), maar cumulatief desastreus is. AEGIS verwijdert alleen deze specifieke "slab" van energie.

4. Resultaten

De experimenten zijn uitgevoerd op PaliGemma2-3B met data van LIBERO (robotmanipulatie).

• VQA Behoud:
  • Naive Fine-tuning: VQA-verlies stijgt drastisch (van ~0.39 naar ~0.78), wat betekent dat het model zijn visuele redeneervermogen volledig verliest.
  • LoRA: Vertraagt het verlies, maar het verlies neemt toch gestaag toe.
  • Stop-Gradient: Behoudt VQA perfect, maar het model leert geen continue acties (alleen discrete tokens).
  • AEGIS: Behoudt de VQA-prestaties bijna perfect (verlies blijft rond de basislijn, ~0.374), terwijl het model wel effectief leert om continue robotacties te voorspellen.
• Actie Learning: AEGIS convergeert even snel als naive fine-tuning op de robot-taak. Er is geen meetbaar verlies aan leersnelheid voor de robottaak.
• OOD Evaluatie (OK-VQA): Na training scoort AEGIS 60.23% op een out-of-distribution VQA-test, wat statistisch gelijk is aan het voorgeöefende model (60.15%). Naive fine-tuning zakt naar 57.36%.
• Efficiëntie: AEGIS verwijdert gemiddeld slechts 0.62% van de gradiëntenergie. De "throttle rate" (hoe vaak er ingegrepen wordt) is ongeveer 51% van de lagen, maar de impact is zeer gericht.

5. Betekenis en Conclusie

AEGIS biedt een fundamentele doorbraak in het fine-tunen van grote multimodale modellen voor robotica.

• Paradigmaverschuiving: In plaats van te vertrouwen op data-mixing (co-training) of het beperken van het model (LoRA/Stop-Gradient), lost AEGIS het probleem op op het niveau van de gradiëntgeometrie.
• Efficiëntie: Het elimineert de noodzaak voor dure replay buffers of het mengen van datasets, wat de rekentijd en complexiteit verlaagt.
• Toekomstperspectief: De methode is schaalbaar naar grotere modellen en kan worden gecombineerd met andere technieken (zoals LoRA voor capaciteitsbeperking, terwijl AEGIS de richting regelt).

Kortom, AEGIS bewijst dat het mogelijk is om een VLM te leren om complexe robotacties uit te voeren zonder de rijke semantische kennis die het model al bezit, door simpelweg de destructieve componenten van de gradiënten chirurgisch te verwijderen.