Half the Nonlinearity Is Wasted: Measuring and Reallocating the Transformer's MLP Budget

Dit onderzoek toont aan dat een aanzienlijk deel van de niet-lineariteit in transformer-MLP's overbodig is en dat het dynamisch vervangen ervan door lineaire substituten, gebaseerd op context, de rekenkosten kan verlagen of zelfs de prestaties kan verbeteren.

De kern

De helft van de "niet-lineaire" kracht wordt verspild: Een simpele uitleg van het onderzoek

Stel je voor dat je een enorm, superintelligent kantoor hebt (een AI-model) waar duizenden werknemers (de neurale lagen) elke dag hard werken om teksten te begrijpen en te voorspellen. De directeur van dit kantoor heeft altijd gedacht: "Elke werknemer moet elke dag zijn volledige, complexe takenpakket uitvoeren. Als we één stap overslaan, crasht het hele systeem."

Dit onderzoek van Peter Balogh komt met een schokkend nieuwsbericht voor die directeur: "Nee, dat klopt niet. De helft van de tijd doen ze eigenlijk niets anders dan een simpele, lineaire berekening. We verspillen dus enorm veel energie aan werk dat niet nodig is."

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Kantoor en de "Grote Berekening"

In een AI-model zoals GPT-2 of Pythia is er een specifieke afdeling, de MLP (Multilayer Perceptron). Dit is de afdeling die de echte "denkwerk" doet.

• De oude aanname: Elke werknemer op deze afdeling moet elke dag een ingewikkelde, niet-lineaire puzzel oplossen (met veel krullen en bochten in de logica).
• De ontdekking: Het onderzoek toont aan dat voor de meeste woorden, in de meeste situaties, deze "ingewikkelde puzzel" eigenlijk gewoon een rechte lijn is. Het is alsof je een dure, krachtige tractor gebruikt om een postzegel van de muur te halen. Je kunt dat ook met een simpele handbeweging doen.

2. De "Poortwachter" (De Gate)

De onderzoekers hebben een slimme poortwachter bedacht. Dit is een heel klein, simpel algoritme (zoals een simpele ja/nee-vraag) dat kijkt naar wat er op dat moment gebeurt.

• Hoe werkt het? De poortwachter kijkt niet naar welk woord er is (bijvoorbeeld: "de" of "hond"). Hij kijkt naar de context: "Wat is er net gebeurd? Wat zeggen de andere woorden?"
• Het resultaat: Als de context simpel is, zegt de poortwachter: "Gebruik de simpele, snelle lijn." Als de context complex is, zegt hij: "Gebruik de dure, krachtige tractor."
• Het effect: In het GPT-2 model kon ongeveer 40% van de werkkracht worden vervangen door een simpele lijn, zonder dat de kwaliteit van de tekst verslechterde. Op sommige plekken werd het model zelfs beter omdat het niet meer probeerde om ingewikkelde dingen te doen die het niet nodig had.

3. De Grootste Misvatting: "Woorden zijn niet de sleutel"

Eerst dachten de onderzoekers: "Ah, misschien zijn het de 'kleine' woorden (zoals 'de', 'en', 'is') die simpel zijn, en de 'inhoudswoorden' (zoals 'kip', 'vliegen', 'liefde') die complex zijn."

• De test: Ze maakten een lijstje van woorden die altijd complex leken.
• De teleurstelling: Toen ze die lijstje gebruikten op een ander stuk tekst (bijvoorbeeld van Wikipedia naar een fictieboek), werkte het helemaal niet meer.
• De les: Het woord "de" kan soms heel complex zijn (als het in een rare zin staat) en soms heel simpel. Het hangt niet af van het woord zelf, maar van waar het staat. Je kunt niet op voorhand zeggen: "Woord X is altijd moeilijk." Je moet kijken naar de situatie.

4. Architectuur is Koning: Het Verschil tussen GPT-2 en Pythia

Het onderzoek toont een interessant verschil tussen twee soorten kantoren (architecturen):

• GPT-2 (Het oude kantoor): Hier werken de mensen in een rij. Eén doet iets, dan de ander. Hier bleek dat de meeste werknemers in het midden van de rij eigenlijk maar simpele taken hadden. Je kon ze vervangen door simpele lijnen en het kantoor liep soepeler.
• Pythia (Het moderne kantoor): Hier werken mensen parallel. Hier is het werk soms echt complexer en moeilijker om te vereenvoudigen. Maar zelfs hier bleek dat op de "midden-afdelingen" veel werk verspild werd.

5. Het Experiment: "De Chirurgische Operatie"

Om te bewijzen dat dit echt werkt, hebben ze een experiment gedaan:

1. Ze namen een getraind model.
2. Ze vervangen de "dure" werknemers in het midden van het kantoor door "simpele lijnen" (vaste matrices).
3. Ze gaven het model een klein beetje training om zich aan te passen.
   Het resultaat: Het model werd beter. Het was alsof je een auto hebt die te zwaar is, en je de zware motor uit het midden haalt en vervangt door een lichte veer. De auto rijdt sneller en zuiniger.

Samenvatting in één zin

Deze paper zegt dat we in AI-modellen veel te veel rekenkracht verspillen aan het doen van complexe berekeningen voor simpele situaties; als we slim kijken naar de context en alleen dan de "dure" kracht gebruiken, kunnen we modellen sneller, goedkoper en soms zelfs slimmer maken.

De metafoor:
Het is alsof je een restaurant hebt waar elke kok, voor elk gerecht, altijd de hele keuken voldraait met de duurste apparatuur. Het onderzoek zegt: "Nee, voor een simpele salade hoef je geen oven te gebruiken. Gebruik een simpele snijplank. En voor de complexe steak? Gebruik dan pas de oven." Door dit slim te regelen, bespaar je energie en wordt het eten soms zelfs lekkerder omdat de chef-koks niet oververmoeid raken door onnodig werk.

Technische samenvatting

Probleemstelling

In transformer-modellen (zoals GPT-2 en Pythia) voeren de Multilayer Perceptrons (MLP's) op elke laag en voor elk token complexe niet-lineaire transformaties uit. De algemene aanname is dat deze niet-lineariteit essentieel is voor het leren van complexe taalfuncties; zonder deze zou het netwerk reduceren tot een lineaire kaart die taal niet kan modelleren.

Het paper onderzoekt echter de vraag of deze niet-lineariteit werkelijk overal nodig is. De auteurs stellen dat een aanzienlijk deel van deze berekeningen "verspild" is en vervangen kan worden door een vooraf berekende lineaire matrix zonder merkbare prestatieverlies. Het centrale probleem is het identificeren van wanneer niet-lineariteit nodig is en of dit voorspeld kan worden op basis van het token zelf (bijv. woordsoort) of de context.

Methodologie

De auteurs hebben een systematisch onderzoek uitgevoerd over zes modellen (162M tot 2,8B parameters), twee architecturen (GPT-2 en Pythia/GPT-NeoX) en drie corpora (WikiText-103, LAMBADA).

1. Lineaire Benadering van MLP's:
   Voor elke laag werd een lineaire surrogaat ($\hat{f}(x) = Wx + b$) gefit op de MLP-activaties via ridge-regressie (Tikhonov-regularisatie). De auteurs maten de "All-Linear" perplexity-kost wanneer een volledige laag werd vervangen door deze lineaire matrix.

2. Adaptieve Gating (Routing):
   In plaats van een laag volledig te vervangen, trainden ze een "gate" (een classifier) die per token beslist of de volledige MLP of de lineaire surrogaat moet worden gebruikt.

   • De meest effectieve gate was een eenvoudige logistische regressie met $d+1$ parameters (een enkel hypervlak in de activatieruimte).
   • De gate werd getraind als een binaire classificatieopdracht: "Is de lineaire benadering goed genoeg?" (gebaseerd op het verschil in cross-entropy verlies, $\delta$).

3. Analyse van Routing-Logica:
   Om te begrijpen wat de gate leert, werd de input van de MLP ontbonden in twee componenten:

   • Token-identiteit: De embedding van het token zelf.
   • Contextuele bijdrage: Alles wat door de attention-mechanismen en vorige lagen is toegevoegd.
     Er werden ook "No-Fly lijsten" (tokens die consistent niet-lineariteit nodig hebben) getest over verschillende corpora om generaliseerbaarheid te testen.

4. Progressieve Linearisatie:
   Als bewijs van concept werden meerdere lagen in het midden van het netwerk sequentieel vervangen door bevroren lineaire matrices, gevolgd door minimal fine-tuning. Dit werd vergeleken met een twee-fasen benadering waarbij de gate tijdens training meeleert.

Belangrijkste Bijdragen

1. Kwantificering van Lineariteit: Eerste systematische meting dat een groot deel van de MLP-berekeningen in transformer-modellen bijna lineair is.
2. Minimale Adaptieve Gating: Een simpele lineaire classifier ($d+1$ parameters) is voldoende om een groot deel van de activaties naar een lineaire weg te routeren met verwaarloosbare kosten.
3. Sterk Negatief Resultaat over Token-gebaseerde Routing: Het is onmogelijk om de behoefte aan niet-lineariteit te voorspellen op basis van het token zelf (woordsoort, frequentie, etc.). Cross-corpus correlaties zijn nihil ($r < 0.05$).
4. Context is Bepalend: De beslissing voor routing wordt uitsluitend bepaald door de contextuele bijdrage in de residual stream, niet door het token zelf.
5. Regularisatie door Linearisatie: In sommige lagen (vooral in GPT-2) verbetert het vervangen van de MLP door een lineaire benadering de prestaties, wat suggereert dat de originele MLP overfitting veroorzaakte.
6. Architectuur-afhankelijkheid: De mate van lineariseerbaarheid verschilt sterk tussen architecturen (GPT-2 vs. Pythia).

Resultaten

• Linearisatiekosten:

  • Bij GPT-2 Medium kunnen 70% van de lagen volledig lineair worden gemaakt met een perplexity-kost van minder dan 3%. Bij 4 van de 23 lagen verbetert de prestatie zelfs door het verwijderen van niet-lineariteit.
  • Bij GPT-2 Large (774M) zijn 11 van de 36 lagen beter dan het baseline-model na lineaire vervanging. Geen enkele laag kost meer dan 3,7% perplexity.
  • Pythia-modellen tonen hogere kosten, maar bij de 2,8B variant is er één laag (L3) die het baseline-model net verslaat (-0,13%).
  • Uitzondering: De eerste laag (Layer 0) in Pythia-2.8B is catastrofaal om te lineair maken (+513% kosten), wat wijst op een scherpe drempel in architectuur-afhankelijke complexiteit.

• Gating Prestaties:

  • De gate kan 25–56% van de MLP-berekeningen routeren naar de lineaire weg met minder dan 1% perplexity-kost.
  • In GPT-2 Medium leidt gating op laag 6 tot een verbetering van 0,06% in perplexity.
  • De gate werkt omdat de verdeling van "niet-lineariteits-behoefte" sterk scheef is: de meeste tokens zijn al bijna lineair; de gate hoeft alleen de zeldzame uitschieters te vangen.

• Token vs. Context:

  • De correlatie tussen token-identiteit en de noodzaak van niet-lineariteit is verwaarloosbaar ($r < 0.05$). Een "No-Fly lijst" van tokens die op het ene corpus niet-lineair nodig zijn, werkt niet op een ander corpus (soms zelfs negatief).
  • De contextuele component alleen levert een AUC van 0,605 op (tegenover 0,609 voor de volledige input), wat aantoont dat de context de enige relevante signalen bevat.

• Progressieve Linearisatie:

  • Door 5 lagen in het midden van GPT-2 Medium te vervangen door lineaire matrices en te fine-tunen, wordt de perplexity met 10,2% verbeterd ten opzichte van het originele model.
  • Een twee-fasen benadering (linearisatie + leren van de gate) leidt tot een verbetering van 17,3% (PPL 19,00), wat beter is dan een standaard fine-tuning control. Dit bewijst dat de niet-lineaire MLP's op deze lagen actief schadelijk waren (overfitting).

Significantie en Implicaties

1. Nieuw Inzicht in Transformer-architectuur: De paper weerlegt het idee dat niet-lineariteit overal en altijd nodig is. In plaats daarvan is niet-lineariteit een "selectief hulpbron" dat voornamelijk nodig is in de randlagen (ingangs- en uitgangslagen) en context-afhankelijk is in de middenlagen.
2. Architectuur is Cruciaal: Het verschil tussen GPT-2 (sequentieel: attention -> MLP) en Pythia (parallel: attention en MLP tegelijk) is groter dan het verschil tussen verschillende routing-strategieën. GPT-2-achtige architecturen lijken inherent meer geschikt voor linearisatie.
3. Efficiëntie en Toekomstige Ontwerpen:
   • Token-based routing is een doodlopende weg: Systemen die proberen te routeren op basis van het woord zelf (bijv. lookup tables) zullen niet generaliseren.
   • Niet-uniforme Capaciteit: Toekomstige modellen zouden kunnen profiteren van variabele MLP-groottes: grote, niet-lineaire MLP's aan de randen en kleine, lineaire of hybride MLP's in het midden.
   • Learned Routing tijdens Pretraining: In plaats van post-hoc gating, zou het model tijdens het trainen zelf kunnen leren wanneer niet-lineariteit nodig is, waardoor capaciteit wordt bespaard voor moeilijke gevallen.

Conclusie: De auteurs concluderen dat "de helft van de niet-lineariteit wordt verspild". Door de niet-lineaire capaciteit te herverdelen naar de lagen waar deze echt nodig is (de randen) en het midden lineair te houden, kunnen modellen efficiënter worden zonder prestatieverlies, en in sommige gevallen zelfs met verbeterde generalisatie.