Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization

Probleemstelling

Deep learning-modellen, vooral overgeparameteriseerde neurale netwerken, hebben de neiging om tijdens het trainen met (Stochastic) Gradient Descent ((S)GD) te vallen in een simpliciteitsbias (simplicity bias). Dit betekent dat het model eerst "snelle" of eenvoudige kenmerken leert (die vaak dominant en goed zichtbaar zijn in de data) en pas later "trage" of complexere kenmerken. Hoewel dit vaak leidt tot goede prestaties, kan het ertoe leiden dat het model ondervertegenwoordigde maar cruciale kenmerken over het hoofd ziet of overfittet op de snelle kenmerken, wat de generalisatie op de originele data-distributie (in-distribution) beperkt.

Bestaande methoden zoals Sharpness-Aware Minimization (SAM) hebben aangetoond dat ze beter generaliseren dan (S)GD door vlatere minima te vinden en kenmerken uniformer te leren. Echter, SAM is computationally duur (het verdubbelt de traintijd). De centrale vraag van dit paper is: Kunnen we de verdeling van de traindata aanpassen om (S)GD (en zelfs SAM) te sturen naar oplossingen met een betere generalisatie, zonder de complexiteit van SAM te hoeven kopiëren?

Methodologie: USEFUL

De auteurs stellen USEFUL (UpSample Early For Uniform Learning) voor, een methode die de traindata-distributie tijdelijk aanpast om de simpliciteitsbias te verminderen. De methode bestaat uit drie hoofdstappen:

1. Identificatie van snelle voorbeelden:

   • Het model wordt kort getraind op de originele data (in de vroege fasen van het trainen).
   • Op basis van de theorie dat voorbeelden met "snelle" kenmerken eerder goed gescheiden worden in de outputruimte, wordt er k-means clustering toegepast op de outputvectoren van het model voor elke klasse.
   • Dit splitst de data in twee clusters: een cluster met voorbeelden die snelle kenmerken bevatten (laag verlies, goed gescheiden) en een cluster met voorbeelden die trage kenmerken bevatten (hoger verlies, moeilijker te leren).

2. Eenmalig Upsampling (One-shot Upsampling):

   • In plaats van dynamisch te herschalen of herhaaldelijk te herschikken, worden de voorbeelden uit de cluster met de trage kenmerken (de rest van de data) één keer oversampled (bijvoorbeeld met een factor $k=2$).
   • Dit verhoogt de frequentie van de moeilijke voorbeelden in de trainset, waardoor het model gedwongen wordt deze kenmerken eerder en uniformer te leren.

3. Herstarten van het trainen:

   • Het model wordt opnieuw getraind (vanaf het begin of vanaf het checkpoint) op deze aangepaste, vaste data-distributie.

Theoretische Bijdrage

De auteurs bieden een rigoureuze theoretische analyse van een twee-laags CNN:

• Vergelijking GD vs. SAM: Ze bewijzen dat SAM kenmerken leert met een meer uniforme snelheid dan GD. GD neigt ertoe om snelle kenmerken extreem vroeg te leren en trage kenmerken te verwaarlozen tot veel later in het trainingsproces.
• Simpliciteitsbias: Ze tonen aan dat de simpliciteitsbias van GD leidt tot een suboptimale generalisatie omdat het model overfittet op de snelle kenmerken voordat het de trage kenmerken kan integreren.
• Effect van Upsampling: Ze bewijzen dat het vergroten van de sterkte van trage kenmerken (via upsampling) de gradiëntrichting van GD verandert zodat deze meer lijkt op die van SAM, waardoor de trage kenmerken sneller en effectiever worden geleerd.

Resultaten

USEFUL werd uitgebreid getest op diverse datasets (CIFAR-10/100, STL-10, CINIC-10, Tiny-ImageNet) en architecturen (ResNet, VGG, DenseNet, ViT, MLP):

• Prestatieverbetering: USEFUL verbetert de generalisatieprestaties van zowel SGD als SAM aanzienlijk.
  • Op CIFAR-10 met ResNet18 bereikt SAM + USEFUL + TrivialAugment state-of-the-art resultaten (laagste testfout).
  • Interessant genoeg presteert SGD + USEFUL vaak beter dan SAM alleen, wat aantoont dat de data-aanpassing de simpliciteitsbias effectief wegneemt.
• Efficiëntie: In tegenstelling tot SAM, dat de traintijd verdubbelt, kost USEFUL slechts een kleine toename in traintijd (ongeveer 1.3x tot 1.6x t.o.v. standaard SGD), omdat het slechts één keer hersample en daarna standaard traint.
• Robuustheid: De methode werkt goed over verschillende batchgroottes, leersnelheden en modelarchitecturen.
• Vergelijking met Random Upsampling: USEFUL presteert significant beter dan willekeurig oversampled data, wat bevestigt dat het selecteren van de juiste (trage) voorbeelden cruciaal is.
• Aanvullende voordelen: De methode vermindert ook "forgetting scores" (het aantal keren dat een voorbeeld na correcte classificatie weer fout wordt geclassificeerd) en leidt tot vlatere minima (gemeten via Hessian-eigenwaarden).

Significantie en Impact

Dit paper is significant omdat het voor het eerst aantoont dat het modificeren van de traindata-distributie een krachtige tool is om de in-distribution generalisatie te verbeteren door de inherente simpliciteitsbias van optimizers te mitigeren.

• Nieuwe inzichten: Het biedt een complementair perspectief op waarom SAM werkt (door het verminderen van simpliciteitsbias) en biedt een praktische, goedkope manier om dit effect na te bootsen.
• Praktische toepasbaarheid: USEFUL is lichtgewicht, vereist geen complexe hyperparameter-tuning (behalve het kiezen van het juiste epoch om te splitsen) en kan worden gecombineerd met bestaande augmentatiestrategieën en optimizers.
• Toekomstperspectief: Hoewel de focus ligt op in-distribution generalisatie, tonen de auteurs aan dat de methode ook voordelen biedt voor out-of-distribution (OOD) scenario's, zoals bij spurious correlations en langstaart-distributies.

Kortom, USEFUL biedt een elegante oplossing om de leerkracht van neurale netwerken te optimaliseren door de data-distributie slim aan te passen, waardoor modellen sneller en beter leren van complexe patronen zonder de rekenkosten van geavanceerde optimizers als SAM volledig te hoeven dragen.