Expanding the Role of Diffusion Models for Robust Classifier Training

Dit onderzoek toont aan dat het integreren van interne representaties van diffusiemodellen als aanvullend leersignaal in adversarial training, naast het gebruik van gegenereerde synthetische data, leidt tot robuustere beeldclassificatoren met beter ontkoppelde kenmerken.

De kern

Stel je voor dat je een leerling traint om foto's te herkennen, bijvoorbeeld om een kat van een hond te onderscheiden. In de wereld van kunstmatige intelligentie (AI) is dit een lastige taak, want er zijn "truuks" (zogenaamde adversarial examples) die zo klein en onzichtbaar zijn voor het menselijk oog, dat ze de AI volledig in de war kunnen brengen. Een kat kan er voor de computer uitzien als een auto als je er een paar onzichtbare pixels aan toevoegt.

Om dit op te lossen, gebruiken wetenschappers een techniek genaamd Adversarial Training. Dit is als een vechtschool voor AI: je laat de computer oefenen met deze "truuks" zodat hij er sterker tegen wordt.

Tot nu toe was de beste manier om deze AI sterker te maken, het gebruik van Diffusiemodellen (de technologie achter AI zoals DALL-E of Midjourney) om nieuwe, nep-foto's te maken. Het idee was simpel: "Laten we miljoenen nep-foto's maken en die gebruiken om de AI te trainen."

Maar in dit nieuwe onderzoek van Pin-Han Huang en zijn team (van de Nationale Universiteit van Taiwan) zeggen ze: "Wacht even, we gebruiken die diffusiemodellen verkeerd! We kijken alleen naar de foto's die ze maken, maar we negeren wat er in hun hoofd gebeurt terwijl ze die foto's maken."

Hier is de uitleg in simpele taal, met een paar creatieve vergelijkingen:

1. De "Geheime Kracht" van de Diffusie

Stel je een diffusiemodel voor als een kunstenaar die een schilderij maakt door eerst een pot vol modder te hebben en er langzaam water bij te gieten totdat het beeld helder wordt.

• De oude manier: De onderzoekers keken alleen naar het eindresultaat (het mooie schilderij) en gebruikten dat om hun AI te trainen.
• De nieuwe manier: De onderzoekers kijken naar de tussenstappen terwijl de kunstenaar werkt. Ze zeggen: "Kijk eens, in die modderige tussenstappen zit al een heel duidelijk idee van hoe een kat eruitziet, zelfs als het nog wazig is!"

Die "wazige tussenstappen" bevatten een soort geheime kennis over de wereld. Ze zijn minder gevoelig voor die kleine truuks die de AI in de war brengen. Het is alsof je niet alleen naar het eindantwoord van een wiskundig probleem kijkt, maar ook naar de stap-voor-stap uitleg van een slimme leraar die je helpt de logica te begrijpen.

2. De Nieuwe Methode: "Spiegelen" in plaats van "Kopiëren"

De onderzoekers hebben een nieuwe techniek bedacht, die ze Diffusion Representation Alignment noemen.

• De vergelijking: Stel je voor dat je een student (de AI die we trainen) hebt die een examen doet.
  • Vroeger: Je gaf de student duizenden extra oefenexamens (de nep-foto's).
  • Nu: Je geeft de student ook een spiegel. In die spiegel ziet de student hoe een "slimme, ervaren leraar" (het diffusiemodel) naar hetzelfde probleem kijkt.
  • De student moet niet alleen het antwoord leren, maar ook nabootsen hoe de leraar denkt. De leraar denkt namelijk op een manier die veel moeilijker te verwarren is dan de student.

Door de student te dwingen om op die "slimme manier" te denken, wordt de student veel sterker tegen de truuks van de tegenstander.

3. Waarom werkt dit zo goed? (De Twee Voordelen)

Het onderzoek toont aan dat deze methode twee dingen doet die samenwerken:

1. De Nep-Foto's (De Oefening): Ze zorgen voor variatie. Het is alsof je een sporter traint met verschillende soorten ballen en windrichtingen. Dit helpt de AI om een strakke, efficiënte strategie te ontwikkelen (in technisch jargon: "low-rank representations").
2. De Spiegel (De Denkstijl): De "geheime kennis" uit het diffusiemodel zorgt ervoor dat de AI diverse en robuuste ideeën ontwikkelt. Het helpt de AI om de belangrijkste kenmerken (bijv. "oren" en "snor") los te koppelen van ruis (bijv. "achtergrondkleur"). Dit noemen ze "ontwarren" (disentangling).

De Metafoor:
Stel je voor dat je een detective traint.

• De nep-foto's zijn duizenden nieuwe moordzaken om op te lossen.
• De spiegel is een mentor die de detective leert om niet te kijken naar de afleidingen (de ruis), maar alleen naar de essentiële aanwijzingen.
• Als je alleen duizenden moordzaken geeft, wordt de detective misschien snel, maar hij blijft soms steken in details.
• Als je alleen een mentor hebt, begrijpt hij de theorie, maar heeft hij weinig praktijkervaring.
• Samen krijg je een detective die zowel snel als onverwoestbaar is.

4. Wat hebben ze bewezen?

Ze hebben dit getest op bekende datasets (CIFAR-10, CIFAR-100 en ImageNet). De resultaten zijn duidelijk:

• De AI wordt beter in het herkennen van echte foto's.
• De AI wordt veel sterker tegen die "truuks" die hem proberen te misleiden.
• Het werkt zelfs als je de AI traint op heel grote datasets (zoals ImageNet).

Conclusie

Kortom: Dit papier zegt dat we diffusiemodellen (de AI's die foto's maken) niet alleen moeten zien als fabrieken voor nep-foto's. We moeten ze ook zien als meesters in het begrijpen van de wereld.

Door de "denktrant" van deze meesters te laten "spiegelen" in onze eigen AI, maken we onze AI niet alleen slimmer, maar ook veel weerbaarder tegen aanvallen. Het is een nieuwe recept voor het bouwen van onverslaanbare kunstmatige intelligentie.

Technische samenvatting

Probleemstelling

Machine learning-modellen zijn kwetsbaar voor adversarial examples: invoer die met semantisch onwaarneembare ruis is verstoord, maar die de voorspellingen van het model drastisch verandert. Adversarial Training (AT) is momenteel de meest effectieve verdediging om robuustheid te bereiken, maar deze methode lijdt onder het fenomeen van robuust overfitting (waarbij de robuustheid op de testset verslechtert ondanks een stabiele nauwkeurigheid op schone data).

Bestaande verbeteringen, zoals DM-AT (Diffusion Model Adversarial Training), gebruiken diffusiemodellen voornamelijk als generators voor synthetische data om het trainingsproces te verrijken. Echter, dit paper stelt de vraag of de interne representaties van diffusiemodellen (die betekenisvolle kenmerken van de data coderen) ook direct kunnen worden benut om de training van robuuste classifiers te verbeteren, los van de generatie van synthetische data.

Methodologie

De auteurs onderzoeken het idee dat de denuisings-doelstelling van diffusiemodellen hen in staat stelt om robuuste semantische kenmerken te vangen, zelfs uit gedeeltelijk beschadigde beelden. Ze stellen een nieuwe methode voor genaamd Diffusion Representation Alignment (DRA).

1. Observatie van Diffusierepresentaties:

   • Analyse toont aan dat representaties uit diffusiemodellen (geëxtraheerd uit de tussenliggende activeringen tijdens het denoisen) diverser zijn en minder gevoelig voor irrelevante hoogfrequente ruis dan representaties uit standaard supervisie of pixel-reconstructie-modellen (zoals MAE).
   • Deze representaties vertonen een intrinsieke, zij het beperkte, robuustheid.

2. Diffusion Representation Alignment (DRA):

   • In plaats van alleen synthetische data te genereren, aligneren de auteurs de representaties van de classifier met de representaties van een bevroren diffusiemodel.
   • Architectuur: Een extra projectiehoofd (een MLP) wordt toegevoegd om de classifier-representaties ($h_{CLS}$) af te stemmen op de diffusierepresentaties ($h_{DR}$) die worden geëxtraheerd bij een optimale timesteps.
   • Trainingsdoel: De totale loss functie combineert de standaard adversarial loss ($L_{AT}$) met een regularisatieterm die de gelijkenis tussen de classifier- en diffusierepresentaties maximaliseert:
     $$L_{AT-DRA} = L_{AT} + \lambda L_{DRA}$$
     Waarbij $L_{DRA} = -\text{sim}(g_{proj}(h_{CLS}), h_{DR})$ en $\text{sim}$ de cosine-afstand is.

3. Analyse van Representaties:

   • De auteurs gebruiken Sparse AutoEncoders (SAEs) om te analyseren of de representaties makkelijker te ontrafelen (disentangle) zijn.
   • Ze gebruiken Classification Dimensions (via PCA) om te onderzoeken hoe de modeldimensies worden benut voor robuustheid versus schone nauwkeurigheid.

Belangrijkste Bijdragen

1. Nieuwe Rol voor Diffusiemodellen: Het paper toont aan dat diffusiemodellen niet alleen nuttig zijn als data-generators, maar dat hun interne representaties waardevolle, robuuste en diverse feature priors bieden die direct kunnen worden gebruikt als een aanvullend leersignaal.
2. Complementaire Mechanismen: De analyse onthult dat synthetische data en representatie-alignatie verschillende rollen spelen:
   • Synthetische data bevordert het leren van laag-rang (low-rank) representaties met sterke generalisatie-eigenschappen.
   • Representatie-alignatie moedigt het model aan om de beschikbare representatiedimensies effectiever te benutten om robuuste kenmerken te coderen (die niet noodzakelijk laag-rang zijn).
   • Samen leiden ze tot een ontkoppelde (disentangled) representatie die makkelijker te interpreteren is en robuuster.
3. Verbeterde Robuustheid: De combinatie van DRA met bestaande DM-AT-methoden resulteert in consistente verbeteringen in zowel schone nauwkeurigheid als adversarial robustness.

Resultaten

Experimenten zijn uitgevoerd op CIFAR-10, CIFAR-100 en ImageNet met verschillende architecturen (WRN, ViT, ConvNeXt).

• Prestatieverbetering: De methode DM-AT + DRA overtreft de state-of-the-art DM-AT methode op alle datasets.
  • Op CIFAR-10 (met 20M synthetische afbeeldingen) steeg de schone nauwkeurigheid van 92,44% naar 93,14% en de AutoAttack-robustheid van 67,31% naar 67,83%.
  • Op ImageNet (met ConvNeXt-B) steeg de schone nauwkeurigheid van 74,49% naar 76,03% en de robustheid van 54,44% naar 56,07%.
• Representatiekwaliteit: Metingen van uniformiteit en alignatie tonen aan dat DRA de "alignment-uniformity frontier" verbetert, wat betekent dat het model zowel robuust als divers in zijn kenmerken wordt.
• Ontkoppeling: SAE-analyse toont aan dat modellen getraind met DRA een lagere reconstructiefout hebben, wat aangeeft dat de geleerde representaties makkelijker te ontrafelen zijn in sparse features.

Significantie

Dit werk is significant omdat het een paradigmaverschuiving voorstelt in het gebruik van diffusiemodellen voor veiligheidskritieke AI. In plaats van diffusiemodellen alleen te zien als "data-verrijkingstools", positioneert het paper ze als bronnen van kwalitatief hoogwaardige, robuuste semantische kennis.

De bevinding dat synthetische data en representatie-alignatie complementair werken, biedt een nieuwe "recept" (recipe) voor het trainen van robuuste classifiers. Het suggereert dat toekomstige verdedigingen tegen adversarial attacks niet alleen moeten focussen op het genereren van betere beelden, maar ook op het benutten van de interne structuur en representaties van generatieve modellen om de leerprocessen van discriminatieve modellen te sturen. Dit resulteert in modellen die niet alleen robuuster zijn, maar ook beter georganiseerde (ontkoppelde) interne representaties hebben, wat de interpretatie en betrouwbaarheid van AI-systemen ten goede komt.