Per-example gradients: a new frontier for understanding and improving optimizers

Dit paper toont aan dat het berekenen van per-exemplaar gradiënten in deep learning frameworks efficiënt mogelijk is en nieuwe inzichten biedt voor het optimaliseren van algoritmen, zoals de cruciale plaatsing van de 'sign'-operatie in signSGD en de superioriteit van het gemiddelde boven de variantie in de Adam-voorconditionering.

De kern

Titel: Het Geheim van de Individuele Leerling: Waarom Optimizers Slimmer Kunnen Worden

Stel je voor dat je een enorme klas vol studenten hebt die allemaal een moeilijke wiskundetoets maken. De leraar (de computer) wil weten hoe de klas het doet, zodat hij de volgende les beter kan aanpakken.

In de traditionele manier van lesgeven (zoals de meeste AI-modellen nu werken), doet de leraar het volgende:

1. Hij kijkt naar de toetsen van een hele groep (een "mini-batch").
2. Hij telt alle fouten bij elkaar op en deelt ze door het aantal studenten.
3. Hij krijgt zo één gemiddeld resultaat.
4. Hij past zijn lesmethode aan op basis van dat ene gemiddelde.

Het probleem? Door te kijken naar alleen het gemiddelde, verliest de leraar de details. Hij ziet niet wie precies welke fout maakte. Misschien heeft de ene student een domme slordigheid gemaakt, terwijl de andere echt een fundamenteel misverstand heeft. Maar omdat ze in het gemiddelde "verdwijnen", kan de leraar niet precies weten hoe hij moet corrigeren.

Het Nieuwe Idee: Kijk naar Iedere Student Apart

De auteurs van dit paper zeggen: "Wacht even! Waarom kijken we niet naar de fouten van elke student individueel?"

Vroeger dachten wetenschappers dat dit te duur was. Het zou te veel geheugen kosten om de resultaten van 64 studenten tegelijk op te slaan in plaats van alleen het gemiddelde. Het was alsof je dacht dat je 64 aparte notitieboekjes nodig had, terwijl je er maar één nodig had.

Maar dit paper toont aan dat dit idee verkeerd was. Met slimme trucjes (zoals het "chirurgisch" aanpassen van de computercode) kunnen we nu de individuele resultaten bekijken zonder dat het de computer veel extra werk kost. Het is alsof we een magische bril opzetten die ons laat zien wat elke student apart deed, terwijl we toch maar één lesplannetje hoeven bij te houden.

Wat levert dit op? Twee grote ontdekkingen

De auteurs hebben dit nieuwe inzicht gebruikt om twee populaire leermethoden (optimizers) te verbeteren.

1. De "Ja/Nee" Leraar (SignSGD)
Stel je voor dat een leraar alleen kijkt of een antwoord "goed" of "fout" is, en niet hoeveel punten er afgehaald zijn. Dat is wat de methode SignSGD doet: hij kijkt alleen naar het teken van de fout (plus of min).

• De oude manier: De leraar kijkt eerst naar het gemiddelde van de hele klas, en zegt dan: "Oké, de klas gaat in de richting van 'fout', dus we gaan 'goed' doen."
• De nieuwe ontdekking: De auteurs ontdekten dat dit verkeerd is. Je moet eerst naar de individuele studenten kijken, hun richting bepalen, en dan pas het gemiddelde nemen.
• De analogie: Stel je voor dat je een groep mensen vraagt om te duwen.
  • Oude manier: Je vraagt eerst "Hoe hard duwt iedereen gemiddeld?" en zegt dan "Duw naar links". Maar als de ene naar links duwt en de andere naar rechts, is het gemiddelde nul. Je krijgt geen kracht.
  • Nieuwe manier: Je zegt tegen iedereen: "Duw naar links als je dat denkt, en naar rechts als jij dat denkt." Daarna tel je op. Zo krijg je een veel sterkere en duidelijkere richting.
  • Conclusie: Laat de "richting-bepaling" pas gebeuren nadat je alle individuele meningen hebt gehoord, niet voordat.

2. De "Gemiddelde" vs. "Verspreiding" Leraar (Adam)
De populaire methode Adam probeert te voorspellen hoe snel hij moet leren. Hij kijkt naar twee dingen:

1. Hoe groot is de gemiddelde fout? (Het "signaal").
2. Hoe verschillen de fouten van elkaar? (De "ruis" of verspreiding).

• De oude wijsheid: Mensen dachten dat de verspreiding (de ruis) het belangrijkst was. Alsof je dacht dat het belangrijkst is om te weten hoe chaotisch de klas is.
• De nieuwe ontdekking: De auteurs ontdekten dat het juist het gemiddelde is dat het belangrijkst is. De "ruis" (het verschil tussen studenten) maakt de lesmethode juist onstabiel en traag.
• De analogie: Stel je voor dat je een boot bestuurt in een storm.
  • Als je alleen kijkt naar hoe hard de golven op en neer gaan (de verspreiding), ga je paniekzaaien en de boot onnodig heen en weer sturen.
  • Als je kijkt naar de gemiddelde richting van de stroming (het signaal), kun je de boot rustig en snel laten varen.
  • Conclusie: De beste AI-optimizers moeten zich focussen op het echte signaal (de gemiddelde fout), en niet te veel last hebben van de ruis (de verschillen tussen individuele voorbeelden).

Waarom is dit belangrijk?

Vroeger dachten we dat het te duur was om naar individuele voorbeelden te kijken. Dit paper zegt: "Nee, dat is het niet!" Door slimme programmeertrucs te gebruiken, kunnen we nu de "ruis" en het "signaal" van elke individuele data-punt zien.

Dit opent de deur voor:

• Snellere AI: Modellen leren sneller omdat ze niet worden afgeleid door ruis.
• Stabielere AI: Lessen worden minder vaak onderbroken door vreemde pieken in de resultaten.
• Nieuwe ideeën: We kunnen nu experimenten doen die vroeger onmogelijk leken, omdat we de "microscopische" details van het leerproces kunnen zien.

Samenvattend:
Dit paper is als het vinden van een nieuwe manier om naar een klas te kijken. In plaats van alleen naar het gemiddelde cijfer te kijken, kijken we naar elke leerling. Hierdoor ontdekken we dat we onze lesmethode (de optimizer) moeten aanpassen: we moeten de richting van de individuele leerlingen eerst bepalen voordat we samenvatten, en we moeten focussen op het echte signaal in plaats van de chaos. Dit maakt het leren van AI niet alleen sneller, maar ook veel slimmer.

Technische samenvatting

Probleemstelling

In de huidige diepe leerpraktijk worden gradienten doorgaans berekend op basis van een mini-batch, waarbij de individuele gradienten van elke steekproef in de batch worden gemiddeld voordat ze naar de optimizer worden gestuurd. Dit betekent dat optimizers alleen toegang hebben tot een schatting van de verwachting van de gradienten ($\mathbb{E}[\nabla f]$), maar niet tot de volledige verdeling van de per-example gradienten.

Hoewel statistieken over deze verdeling (zoals variantie, covariantie of niet-lineaire transformaties zoals het teken van de gradient) waardevol zouden kunnen zijn voor het begrijpen van het trainingsgedrag en het ontwerpen van betere algoritmen, wordt het berekenen ervan traditioneel gezien als:

1. Computatie-intensief: Het vereist het opslaan van $B$ individuele gradienten, wat leidt tot een lineaire toename van het geheugengebruik ($O(B \cdot P)$, waarbij $P$ het aantal parameters is).
2. Complex: Het vereist vaak aangepaste implementaties die de standaard reverse-mode automatic differentiation (AD) omzeilen.

Dit beperkt onderzoekers in het verkennen van een groot deel van de ontwerpruimte voor trainingsalgoritmen, vooral bij grote schaalmodellen.

Methodologie

De auteurs presenteren een technische aanpak om per-example gradientstatistieken efficiënt te berekenen, met name voor sequentie-georiënteerde architecturen zoals Transformers.

1. Geheugenbottlenecks en Architecturale Kansen:

   • De paper analyseert het geheugengebruik tijdens de forward en backward pass. Ze tonen aan dat bij Transformers (en andere sequentie-modellen) de kosten voor het opslaan van activeringen (activations) tijdens de forward pass vaak groter zijn dan de kosten voor het opslaan van individuele gradienten.
   • Door slim gebruik te maken van bestaande checkpoints van input-activaties, kunnen individuele gradienten tijdelijk worden opgeslagen zonder de piekgeheugenniveaus te verhogen.

2. Computational Graph Surgery (Chirurgie van het berekeningengraf):

   • In plaats van na te denken over het berekenen van gradienten en deze vervolgens te middelen, manipuleren de auteurs het computationele graf direct.
   • Ze gebruiken JAX en de vmap-functionaliteit om gradienten per voorbeeld te vectoriseren.
   • Cruciaal is hun "chirurgie"-benadering: ze analyseren de graaf van de backpropagation en injecteren de gewenste niet-lineaire operatie $\phi$ (bijv. kwadrateren of tekenfunctie) voordat de som-reductie over de batch plaatsvindt.
   • Voor dichtbevolkte lagen (dense layers) in MLP's en Transformers kan dit worden gedaan met verwaarloosbare overhead door eigenschappen van tensorcontracties (bijv. het kwadrateren van input en co-tangenten voorafgaand aan de matrixvermenigvuldiging).

3. Implementatie:

   • Ze implementeren dit via aangepaste afgeleide regels in JAX (jaxpr surgery), wat toelaat dat de compiler (XLA) de operaties optimaliseert voor hardware (zoals TPUs), waardoor de overhead minimaal blijft.

Belangrijkste Bijdragen

1. Efficiënte Berekening: Het bewijs dat per-example gradientstatistieken berekend kunnen worden met verwaarloosbare geheugen- en rekentijd-overhead voor moderne architecturen (zoals Transformers), in tegenstelling tot de heersende mening dat dit prohibitief duur is.
2. Analyse van SIGNSGD: Een diepgaande studie naar de plaatsing van de sign-operatie in SIGNSGD-varianten.
3. Herinterpretatie van ADAM: Een nieuwe familie van ADAM-varianten die werken op per-example statistieken (in plaats van gemiddelde statistieken), en het inzicht dat de voorwaarde (preconditioner) beter gebaseerd is op het kwadraat van het gemiddelde dan op de variantie.

Resultaten

1. SIGNSGD en de Plaatsing van het Teken
De auteurs vergeleken drie varianten van SIGNSGD met momentum (EMA):

• SIGNEMA: sign -> EMA -> avg (Teken wordt als laatste toegepast).

• SIGNSGD: EMA -> sign -> avg.

• MICROSIGNSGD: EMA -> avg -> sign (Teken wordt als eerste toegepast op per-example gradienten).

• Vondst: SIGNEMA presteerde het beste, terwijl MICROSIGNSGD het slechtste presteerde (onstabiel en traag).

• Redenering: Het toepassen van de sign-functie op individuele gradienten (met lage signaal-ruisverhouding of SNR) versterkt het ruisniveau. Door het teken zo laat mogelijk toe te passen (na maximale middeling), wordt de SNR van het object dat de sign-functie passeert gemaximaliseerd, wat leidt tot stabielere training.

2. ADAM en Per-Example Statistieken
De auteurs introduceerden MICROADAM, waarbij de preconditioner $\nu$ wordt berekend als het gemiddelde van de kwadraten van de per-example gradienten ($\frac{1}{B}\sum g_i^2$) in plaats van het kwadraat van het gemiddelde ($(\frac{1}{B}\sum g_i)^2$).

• Verwachting vs. Realiteit: Conventionele wijsheid en de "square root scaling rule" suggereren dat de variantie ($\sigma^2$) de dominantie heeft in de preconditioner. De auteurs toonden echter aan dat bij ADAM het kwadraat van het gemiddelde ($\mu^2$) dominant is, zelfs bij kleine batches.
• Nieuwe Variants:
  • MICROADAMVAR: Een variant die puur op variantie ($\sigma^2$) focust. Dit bleek instabiel en slechter presterend dan standaard ADAM.
  • MICROADAMMSQ: Een variant die puur op het kwadraat van het gemiddelde ($\mu^2$) focust. Deze variant trainde stabieler en bereikte uiteindelijk een iets lagere validatieverlies dan standaard ADAM (3.05 vs 3.06).
• Conclusie: De preconditioner in ADAM werkt het beste wanneer deze wordt gedomineerd door het kwadraat van het gemiddelde van de gradientverdeling, niet door de variantie. Dit staat in contrast met de traditionele interpretatie van de schalingsregels.

Betekenis en Impact

• Nieuw Onderzoeksgebied: Het paper opent een nieuw dimension voor het optimaliseren van trainingsalgoritmen door toegang te geven tot de volledige verdeling van gradienten, wat eerder onbereikbaar was vanwege computatiekosten.
• Praktische Toepasbaarheid: De methoden zijn niet alleen theoretisch maar ook praktisch implementeerbaar met moderne tools zoals JAX, zonder dat dit de trainingskosten significant verhoogt.
• Fundamenteel Inzicht: De resultaten daagden bestaande aannames uit over hoe optimizers zoals ADAM en SIGNSGD werken. Het toont aan dat het verminderen van ruis (door middelen voordat niet-lineaire transformaties worden toegepast) cruciaal is, en dat de "variance" in ADAM minder belangrijk is dan het "mean squared" component voor stabiliteit en convergentie.
• Toekomstperspectief: Deze technieken kunnen worden uitgebreid naar andere statistieken (zoals Hessiaan-vector producten) en kunnen leiden tot een nieuwe generatie van optimizers die beter zijn afgestemd op de statistische eigenschappen van de gradientverdeling.

Kortom, dit werk bewijst dat het "openen" van de mini-batch voor per-example analyse niet alleen mogelijk is, maar ook essentieel voor het doorgronden en verbeteren van de fundamentele mechanismen van deep learning training.