Nexusformer: Nonlinear Attention Expansion for Stable and Inheritable Transformer Scaling

Probleemstelling

De huidige schaalwetten voor Large Language Models (LLMs) vereisen doorgaans dat grotere modellen vanaf nul worden getraind, wat leidt tot enorme rekenkosten en het verlies van kennis die in kleinere modellen is opgedaan. De auteurs identificeren de lineaire projecties in de Multi-Head Attention (MHA) mechanismen (Query, Key, Value) als de primaire bottleneck.

• Lineaire Expressiviteitsbeperking: Standaard Transformers beperken feature-extractie tot vaste lineaire subspaces. Dit dwingt diepere lagen om complexe semantische ontkoppeling te doen, wat de expressiviteit beperkt.
• Stijve Capaciteit: Het vergroten van de dimensies van deze lineaire projectiematrices veroorzaakt significante verdelingsverschuivingen (distribution shifts). Hierdoor moet het model vaak opnieuw worden getraind, wat de "inheritability" (erfbaarheid) van kennis verhindert.
• Huidige oplossingen: Bestaande methoden zoals LoRA zijn fundamenteel lineair, en groeigerichte benaderingen zoals Tokenformer behouden de onderliggende lineaire beperkingen. Er ontbreekt een mechanisme dat zowel expressiviteit verhoogt als niet-destructieve groei ondersteunt.

Methodologie: Nexusformer

De auteurs introduceren Nexusformer, een architectuur die de standaard lineaire Q/K/V-projecties vervangt door een Nexus-Rank laag. Dit is een compacte, niet-lineaire expansielaag die feature-projecties transformeert naar niet-lineaire manifolds in hogere dimensies.

1. Nexus-Rank Architectuur (Drie-staps mapping):
In plaats van een directe lineaire projectie $D \to D$, gebruikt Nexusformer een hiërarchische structuur met de dimensies $D < M < A$:

• Stap 1 (Uplift): Projectie van de invoer $D$ naar een tussenruimte $M$ (waarbij $M > D$), gevolgd door een niet-lineaire activatie (GeLU). Dit fungeert als een semantische ontkoppelingslaag.
• Stap 2 (Expansie): Projectie van $M$ naar een "over-capaciteit" ruimte $A$ (waarbij $A > M$), gevolgd door een tweede niet-lineaire activatie. Dit is de primaire reservoir voor het modelleren van hoge-orde interacties.
• Stap 3 (Reductie): Projectie van $A$ terug naar de originele dimensie $D$.
  De uiteindelijke output $\tilde{Q}$ bevat diepe features afgeleid van deze dubbele niet-lineaire transformaties, wat zorgt voor een veel sterkere expressieve kracht dan lineaire projecties.

2. Dubbele-as Groeimechanisme met Zero-Initialisatie:
Om stabiele groei mogelijk te maken zonder bestaande kennis te verliezen, introduceert Nexusformer een strategie voor structurele expansie:

• Dual-Axis Growth: Het model groeit gelijktijdig langs twee assen: de tussen-capaciteitsdimensie $M$ en de expansiedimensie $A$.
• Zero-Initialisatie: Nieuwe parameterblokken ($W_{new}$) worden geinitialiseerd op nul. Omdat deze blokken bij het moment van expansie geen bijdrage leveren aan de output, blijft het gedrag van het model identiek aan het voorgetrainde model.
• Seamless Inheritance: Dit stelt het model in staat om direct door te gaan met training (continued training) om de nieuwe capaciteit te leren, zonder dat er sprake is van een "catastrophic forgetting" of de noodzaak tot volledige hertraining.

Belangrijkste Bijdragen

1. Nexusformer Architectuur: Een structurele vervanging van lineaire attention-projecties door diepe, niet-lineaire Nexus-Rank lagen om de representatieve capaciteit te verhogen.
2. Niet-destructieve Groeistrategie: Een mechanisme voor "lossless" parameterexpansie via zero-initialisatie en dubbele-as groei, waardoor kennis efficiënt kan worden geërfd.
3. Statistische Analyse van Groeidynamiek: Een niet-parametrische analyse die een stabiele convergentietrajectie onthult. De auteurs leiden een geometrische schaalwet af die prestaties accuraat voorspelt op basis van de expansieschaal.

Resultaten

Experimenten zijn uitgevoerd op taalmodellerings- en redeneerbenchmarks (zoals PIQA, HellaSwag, ARC) over schalen van 170M tot 640M parameters.

• Efficiëntie: Bij progressieve schaling (van 240M naar 440M) behaalt Nexusformer dezelfde perplexiteit als de concurrent Tokenformer, maar met 41,5% minder trainingsrekenkracht.
• Prestaties: Nexusformer presteert consistent beter dan baselines (Qwen3, Tokenformer, Transformer++) op alle schalen. Op de 640M schaal behaalt het bijvoorbeeld 15,8% betere resultaten op ARC-Easy dan Tokenformer.
• Stabiliteit: De analyse toont aan dat zero-initialisatie leidt tot een stabiel "divergeer-then-converge" traject. Dit wordt gekwantificeerd met een nieuwe metriek, de Radiale Energie ($R$), die de afstand tot een ideaal gealigneerde toestand meet.
• Schaalwet: Er is een empirische relatie gevonden: $\ln(PPL) = -0.0991 \cdot |\ln(R)| + 2.4804$. Dit betekent dat een kleinere radiale energie (betere alignement) correleert met een lagere perplexiteit (betere prestaties).
• FLOPs: Ondanks de complexere drie-staps mapping, is het aantal FLOPs per token voor Nexusformer lager dan of gelijk aan dat van standaard Transformers, voornamelijk door de efficiëntere verdeling van de rekenlast.

Significantie en Impact

Dit paper biedt een fundamentele doorbraak in hoe Transformer-modellen kunnen worden geschaald:

• Van "Train from Scratch" naar "Inheritable Scaling": Het bewijst dat het mogelijk is om modellen structureel te vergroten en bestaande kennis te behouden, wat de totale trainingskosten drastisch verlaagt.
• Overcoming the Linear Bottleneck: Het toont aan dat het doorbreken van de lineaire beperkingen in attention-projecties via niet-lineaire manifolds leidt tot superieure expressiviteit en redeneervermogen.
• Geometrische Controle: De introductie van de radiale energie $R$ als controlevariabele biedt een nieuw kader voor het voorspellen en sturen van groeitrajecten, wat verder gaat dan traditionele schaalwetten die alleen kijken naar parameters, data en rekenkracht.

Kortom, Nexusformer stelt een nieuwe standaard voor voor efficiënte, groeibare foundation modellen die niet afhankelijk zijn van het volledig opnieuw trainen bij elke schaalvergroting.