S2O: Enhancing Adversarial Training with Second-Order Statistics of Weights

Dit artikel introduceert S²O, een nieuwe methode die de robuustheid van neurale netwerken verbetert door de onafhankelijkheidsaanname van gewichten in PAC-Bayesiaanse frameworks te versoepelen en de tweede-orde statistieken van deze gewichten te optimaliseren.

De kern

Stel je voor dat je een zeer slimme, maar nogal nerveuze student hebt die een examen moet doen. Deze student is een Neuraal Netwerk (een type kunstmatige intelligentie). Hij is goed in het herkennen van katten en honden op foto's, maar hij is erg kwetsbaar voor "trucs".

Als iemand een foto van een kat een heel klein beetje verandert (bijvoorbeeld door een paar pixels te verschuiven die het menselijk oog niet ziet), kan de student ineens denken: "Dit is een hond!" en dat met 100% zekerheid zeggen. Dit noemen we een adversariaal voorbeeld (een tegenstander die de student probeert te misleiden).

Om deze student sterker te maken, gebruiken we Adversarial Training. Dit is als een trainingskamp waar de student wordt geoefend met deze trucs, zodat hij leert ze te doorzien. Maar de onderzoekers van dit papier merkten op dat deze training nog niet perfect was. Ze hadden een nieuwe manier bedacht om de student nog slimmer en robuuster te maken: S2O.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude idee: Alles is losjes

In de traditionele training behandelden we de "hersenen" van de student (de gewichten of weights) als losse, onafhankelijke individuen. Het was alsof je dacht: "Elke neuron in zijn hoofd denkt alleen maar aan zijn eigen ding en heeft niets te maken met zijn buren."

De onderzoekers zeggen: "Nee, dat klopt niet!" In een echt brein (en in een goed werkend neuraal netwerk) werken neuronen samen. Ze hebben een relatie met elkaar. Als neuron A een beetje verandert, heeft dat invloed op neuron B. Ze zijn gecorrreleerd.

2. De nieuwe aanpak: S2O (Tweede-orde statistiek)

De onderzoekers introduceren een methode genaamd S2O (Second-Order Statistics Optimization).

• De Analogie van het Orkest:
  Stel je een symfonieorkest voor.
  • Oude methode: De dirigent luistert alleen of elke violist op zichzelf goed speelt.
  • Nieuwe methode (S2O): De dirigent luistert naar de harmonie tussen de violisten. Hij kijkt niet alleen naar de individuele noten, maar naar hoe ze samen klinken. Als de violisten te veel in elkaars vaart zitten (te veel correlatie), klinkt het orkest rommelig en is het makkelijker om de muziek te verstoren met een klein geluidje.

S2O zorgt ervoor dat de "relaties" tussen de neuronen in balans blijven. Het zorgt voor een betere harmonie in de hersenen van de AI.

3. Hoe doen ze dit? (De wiskundige magie)

De onderzoekers hebben een wiskundige formule bedacht (gebaseerd op een theorie genaamd PAC-Bayes) die laat zien dat als je de "correlatie" tussen de neuronen optimaliseert, de AI veel moeilijker te misleiden is.

Ze gebruiken een slimme truc (de Laplace-benadering) om deze relaties te meten zonder dat de computer uren moet rekenen. Het is alsof je in plaats van elke speler in het orkest één voor één te meten, gewoon kijkt naar het totale geluidsniveau en de resonantie van de zaal om te weten of het orkest goed samenwerkt.

4. Het resultaat: Een onwrikbare AI

Wat gebeurt er als je deze S2O-methode toepast?

• Sterker tegen trucs: De AI wordt veel moeilijker te misleiden. Zelfs als iemand een foto manipuleert, blijft de AI kalm en ziet hij nog steeds een kat.
• Beter in het algemeen: Interessant genoeg wordt de AI niet alleen sterker tegen trucs, maar ook beter in het herkennen van normale foto's. Het is alsof de student door het leren van de trucs, de stof zelf beter begrijpt.
• Werkt overal: Het werkt niet alleen op simpele netwerken, maar ook op de allermodernste, complexe systemen (zoals ViT en DeiT, die worden gebruikt voor beeldherkenning).

Samenvatting in één zin

De onderzoekers hebben ontdekt dat AI-modellen sterker worden als we niet alleen kijken naar hun individuele "hersencellen", maar vooral naar hoe die cellen samenwerken; door die samenwerking (de correlatie) te optimaliseren, maken we de AI onkwetsbaarder voor trucs en slimmer in het algemeen.

Het is de verschuiving van "elk voor zich" naar "samen sterk".

Technische samenvatting

Probleemstelling

Diepe neurale netwerken (DNN's) zijn kwetsbaar voor adversarial examples: kleine, voor mensen onzichtbare verstoringen in de invoer die leiden tot foutieve voorspellingen met hoge zekerheid. Adversarial training (AT) is momenteel een van de meest effectieve methoden om de robuustheid van deze netwerken te verbeteren. Dit wordt doorgaast geformuleerd als een min-max optimalisatieprobleem waarbij de modelgewichten worden geoptimaliseerd met behulp van methoden zoals Stochastic Gradient Descent (SGD).

Een fundamentele beperking in bestaande theoretische kaders, met name binnen de PAC-Bayesian analyse, is de aanname dat de gewichten van het model statistisch onafhankelijk zijn (vaak gemodelleerd als een sferische Gaussische verdeling). De auteurs betogen dat deze aanname onrealistisch is voor complexe DNN's, waar gewichten vaak onderling gecorreleerd zijn. Het negeren van deze correlaties leidt tot minder nauwkeurige generalisatiegrenzen en beperkt de mogelijkheid om de robuustheid verder te optimaliseren.

Methodologie

De paper introduceert een nieuw paradigma genaamd Second-Order Statistics Optimization (S2O). De kern van de methode is het behandelen van modelgewichten als stochastische variabelen waarbij de tweede-orde statistieken (correlaties) expliciet worden gemodelleerd en geoptimaliseerd.

De aanpak bestaat uit drie hoofdcomponenten:

1. Theoretische Uitbreiding (PAC-Bayes):

   • De auteurs lossen de aanname van onafhankelijke gewichten op door een niet-sferische Gaussische verdeling voor de perturbaties te introduceren, gekenmerkt door een correlatiematrix $R$.
   • Ze leiden een verbeterde PAC-Bayesian robuuste generalisatiegrens af. Deze grens bevat termen die afhankelijk zijn van de determinant en de spectrale norm van de correlatiematrix van de gewichten.
   • De theorie toont aan dat het minimaliseren van de spectrale norm en het maximaliseren van de determinant van de correlatiematrix leidt tot een strakkere (betere) generalisatiegrens, wat impliceert dat het optimaliseren van deze statistieken de robuustheid en generalisatie van het model verbetert.

2. Schatting via Laplace-benadering:

   • Het direct berekenen van de correlatiematrix is computatierijk. Om dit op te lossen, gebruiken de auteurs de Laplace-benadering.
   • Hierbij wordt de posterior-verdeling van de gewichten benaderd met een Gaussische verdeling rond het Maximum A Posteriori (MAP) schatting. De covariantiematrix wordt geschat als het omgekeerde van de Hessiaan van de verliesfunctie.
   • Om de rekentijd beheersbaar te houden, wordt gebruikgemaakt van een Kronecker-gefactoriseerde benadering (geïnspireerd door Botev et al. en Ritter et al.). Dit maakt het mogelijk om de inverse Hessiaan efficiënt te schatten zonder de volledige matrix op te slaan.

3. S2O Regularisator:

   • Gebaseerd op de bovenstaande theorie, wordt een nieuwe regularisator toegevoegd aan de adversarial training loss-functie.
   • In plaats van de complexe termen direct te optimaliseren, wordt de Frobenius-norm van de genormaliseerde post-activatie covariantiematrix ($\|A\|_F^2$) geminimaliseerd.
   • De nieuwe objectief functie wordt: $\tilde{J}_{adv} = J_{adv} + \alpha (\|A_x\|_F^2 + \|A_{x'}\|_F^2)$, waarbij $A_x$ en $A_{x'}$ de geschatte correlaties zijn voor respectievelijk schone en adversarial data, en $\alpha$ een hyperparameter is.

Belangrijkste Bijdragen

• Theoretisch: De eerste afleiding van een robuuste generalisatiegrens binnen het PAC-Bayes-raamwerk die expliciet rekening houdt met de correlatie tussen gewichten (tweede-orde statistieken). Dit levert een strakkere grens op dan eerdere methoden die onafhankelijkheid aannamen.
• Algoritmisch: De ontwikkeling van de S2O-methode, een efficiënte regularisatietechniek die de tweede-orde statistieken van gewichten optimaliseert tijdens het trainen, zonder de complexiteit van volledige tweede-orde optimalisatie.
• Empirisch: Uitgebreide experimenten die aantonen dat S2O niet alleen op zichzelf werkt, maar ook naadloos integreert met state-of-the-art adversarial training technieken (zoals TRADES, AWP, en methoden met DDPM-gegenereerde data).

Resultaten

De auteurs hebben hun methode getest op diverse datasets (CIFAR-10/100, SVHN, Tiny-ImageNet, Imagenette) en architecturen (PreAct ResNet, WideResNet, ViT-B, DeiT-S).

• Verbeterde Robuustheid: S2O verhoogt consistent de weerstand tegen verschillende adversarial attacks (FGSM, PGD, CW, Auto Attack) in vergelijking met standaard adversarial training. Bijvoorbeeld, op CIFAR-10 met een $\ell_\infty$-aanval steeg de robuuste nauwkeurigheid van 52,77% (AT) naar 55,11% (AT+S2O).
• Generalisatie: De methode verbetert ook de "clean accuracy" (nauwkeurigheid op schone data), wat de vaak waargenomen trade-off tussen schone en robuuste nauwkeurigheid vermindert.
• Synergie: S2O werkt synergistisch met andere geavanceerde methoden. De combinatie van S2O met TRADES en AWP leverde de beste resultaten op, en verbeterde ook methoden die gebruikmaken van gegenereerde data (DDPM).
• Veelzijdigheid: De methode bleek effectief voor zowel CNN's als Vision Transformers (ViT), wat aantoont dat het principe van het optimaliseren van gewichtsstatistieken architectonisch onafhankelijk is.
• Zwarte Doos en BPDA: S2O toonde ook verbeterde weerstand tegen black-box (transfer) attacks en BPDA-attacks (waarbij de verdediging niet differentieerbaar is), wat aangeeft dat de verbeterde robuustheid echt is en niet het gevolg van "gradient masking".

Betekenis

Deze paper is significant omdat het een fundamentele aanname in de theoretische analyse van DNN-robustheid (onafhankelijkheid van gewichten) daagt en corrigeert. Het biedt een nieuwe theoretische onderbouwing voor het belang van gewichtscorrelaties. Praktisch gezien introduceert S2O een efficiënte en plug-and-play regularisatietechniek die de prestaties van bestaande defensiemethoden direct verbetert zonder ingrijpende wijzigingen in de trainingspipeline. Het stelt de gemeenschap in staat om DNN's te trainen die niet alleen robuuster zijn tegen aanvallen, maar ook beter generaliseren, door het bewust optimaliseren van de statistische eigenschappen van de interne representaties van het netwerk.