Post-Training Augmentation Invariance

Titel: Post-Training Augmentation Invariance

Auteurs: Keenan Eikenberry, Lizuo Liu, Yoonsang Lee (Dartmouth College)

---

1. Het Probleem

Groot-schalige voorgeprogrammeerde modellen (zoals Vision Transformers, CLIP, DINOv2) vormen de basis van moderne machine learning pipelines. Deze modellen zijn getraind op enorme datasets en leveren universele features die goed overdraagbaar zijn naar downstream-taken. Echter, deze modellen zijn niet per se invariant voor specifieke augmentaties die cruciaal kunnen zijn voor bepaalde toepassingen (bijvoorbeeld rotatie-invariantie voor satellietbeelden of robustheid tegen ruis).

De centrale vraag van dit werk is: Kunnen we een voorgeprogrammeerd model invariant maken voor specifieke augmentaties na het trainen, zonder de bestaande capaciteiten van het model te corrumperen?

Bestaande oplossingen zoals het volledig hertrainen van het model of standaard fine-tuning zijn vaak computatief prohibitief of leiden tot een degradatie van de geleerde representaties. Fine-tuning kan de geometrie van de latente ruimte verstoren, waardoor de oorspronkelijke features onbruikbaar worden voor andere taken.

2. Methodologie

De auteurs stellen een raamwerk voor dat bestaat uit het toevoegen van lichte "adapter" netwerken aan de latente ruimte van een bevroren (frozen) voorgeprogrammeerd netwerk $F$. Het doel is om een adapter $E_\theta$ te leren die augmentaties compenseert, terwijl de structuur van de oorspronkelijke latente verdeling behouden blijft.

Kernconcepten:

• Augmented Encoders: Probabilistische encoders die het proces formaliseren van het willekeurig augmenteren van een invoer en vervolgens coderen.
• Post-Training Augmentation Invariance: Een adapter $E_\theta$ is invariant als de samengestelde netwerk $E_\theta \circ F$ invariant is voor een augmentatie $t$, maar tegelijkertijd de metriek van de oorspronkelijke latente verdeling $F_\sharp \mu_X$ behoudt (tot op een toelaatbare afbeelding $V$, zoals een isometrie).

Twee Nieuwe Verliesfuncties:

De paper introduceert twee verliesfuncties gebaseerd op Optimal Transport (OT) om deze doelen te bereiken:

1. Markov-Wasserstein (MaWa) Minimalisatie:

   • Dit is een verankerde Mean Squared Error (MSE) loss.
   • Het forceert dat de encoder op niet-augmenteerde invoer dicht bij de identiteit blijft (verankering) en dat augmented views naar hun corresponderende schone features worden getrokken.
   • Formeel wordt dit gezien als het minimaliseren van de afstand in de Markov-Wasserstein metriek tussen de augmented encoder en de identiteit.
   • Dit voorkomt "collapse" (waarbij alles naar één punt wordt gemapt), een veelvoorkomend probleem bij simpele MSE-losses voor invariantie.

2. Wasserstein Correlatie (WaCo) Maximalisatie:

   • In plaats van een afstand tussen netwerken te minimaliseren, maximaliseert deze loss de correlatie van de gezamenlijke verdeling die wordt gegenereerd door de augmented encoder.
   • Het maximaliseert de "Wasserstein correlation" (een genormaliseerde versie van Wasserstein afhankelijkheid).
   • Dit zorgt ervoor dat de encoder invariant wordt voor augmentaties zonder de metriek van de invoerverdeling te verstoren.
   • Een uniek voordeel is dat WaCo ook kan worden gebruikt voor dimensionaliteitsreductie terwijl de lokale isometrie behouden blijft.

Vergelijking met Bestaande Methoden:

De auteurs vergelijken hun methode met:

• SimCLR (Contrastive Learning): Dwingt augmented views dicht bij elkaar, maar verstoort de geometrie van de oorspronkelijke latente ruimte drastisch, waardoor de features voor andere taken onbruikbaar worden.
• HSIC (Hilbert-Schmidt Independence Criterion): Leidt tot een nog ernstigere corruptie van de latente ruimte en faalt in het behouden van invariantie zonder structuurverlies.

3. Belangrijkste Bijdragen

1. Formalisatie: Een strikte wiskundige definitie van post-training augmentation invariance, inclusief het concept van "collision rate" (hoe vaak augmentaties features over klassengrenzen verplaatsen).
2. Nieuwe Objecten: Introduceert "augmented encoders" als fundamenteel object van studie.
3. Nieuwe Losses: Proposeert MaWa-minimalisatie en WaCo-maximalisatie als effectieve trainingstargets.
4. Empirisch Bewijs: Demonstreert dat eenvoudige, één-laagse MLP-adapters, getraind met deze losses, prestaties verbeteren zonder fine-tuning van het basisnetwerk.

4. Resultaten

De experimenten zijn uitgevoerd op diverse voorgeprogrammeerde modellen (DINOv2, SwAV, CLIP, ResNet50) en datasets (STL10, TinyImageNet) met augmentaties zoals rotatie, affiene transformaties en ruis.

• Prestatieverbetering:

  • Op STL10 met DINOv2 features: Een netwerk zonder adapter ($C \circ F$) daalt in nauwkeurigheid van 98% naar 71% bij willekeurige rotaties. Met de MaWa-adapter ($C \circ E_\theta \circ F$) blijft de nauwkeurigheid op 94%.
  • Voor ruis-invariantie: De nauwkeurigheid stijgt van 58% naar 86%.
  • Deze resultaten worden bereikt zonder fine-tuning van de weights van $F$.

• Structuurbehoud:

  • MaWa en WaCo behouden de metriek van de oorspronkelijke latente ruimte bijna perfect (hoge $R^2$ waarden, lage Normalized RMSD).
  • In tegenstelling hiermee leiden SimCLR en HSIC tot een totale corruptie van de latente ruimte (zeer lage $R^2$, hoge afwijkingen), wat betekent dat de oorspronkelijke features niet meer bruikbaar zijn.

• Collision Rate Analyse:

  • De methode werkt het beste bij zelf-supervised modellen (zoals DINOv2) waar augmentaties coherent zijn in de latente ruimte (lage collision rate).
  • Bij gesuperviseerde modellen met hoge collision rates (zoals ResNet50) kan MaWa falen, maar WaCo blijft robuuster.

• Dimensionaliteitsreductie:

  • WaCo kan succesvol worden gebruikt om de dimensie te verlagen (bijv. van 384 naar 96) terwijl de invariantie en classificatieprestaties behouden blijven.

5. Betekenis en Conclusie

Dit werk biedt een praktische en efficiënte oplossing voor het aanpassen van grote, voorgeprogrammeerde modellen aan specifieke domeinvereisten (zoals rotatie- of ruisinvariantie) zonder de dure en riskante stap van volledige hertraining of fine-tuning.

• Efficiëntie: De methode vereist slechts het trainen van een zeer kleine adapter (één verborgen laag), wat rekenkundig goedkoop is.
• Veiligheid: Het garandeert dat de oorspronkelijke kennis van het model behouden blijft, wat essentieel is voor multi-task learning.
• Toekomstperspectief: De resultaten suggereren dat deze aanpak veelbelovend is voor adversarial robustness, door adversarial attacks te behandelen als augmentaties.

Kortom, de paper bewijst dat het mogelijk is om "na-trainen" (post-training) augmentatie-invariantie toe te voegen aan state-of-the-art modellen door middel van lichte adapters en optimal transport-gebaseerde verliesfuncties, terwijl de integriteit van de oorspronkelijke representatie volledig behouden blijft.