Technische Samenvatting: Verwijdering van Gesteunde Normalisatie en Ankeren van Schaal in Pre-Norm Transformers

Probleemstelling

Normalisatielagen (bijv. LayerNorm, RMSNorm) zijn standaardcomponenten in transformer-architecturen; ze stabiliseren activaties en maken het trainen van diepe netwerken mogelijk. Deze lagen introduceren echter steekproefafhankelijke, niet-lineaire berekeningen die aanzienlijke overhead veroorzaken tijdens inferentie. Het blijft onduidelijk of deze per-token statistieken strikt noodzakelijk zijn gedurende de volledige inferentiefase, met name in pre-norm transformers. Hoewel eerdere werken hebben geprobeerd normalisatie te verwijderen via initialisatietricks of sequentiële fine-tuning, ontbreekt er inzicht in de distincte functionele rollen van interne versus finale normalisatielagen, specifiek wat betreft hoe ze de schaal van de pre-logit-representatie controleren.

Methodologie

De auteurs stellen TaperNorm voor, een benadering voor gesteunde normalisatieverwijdering die is ontworpen voor pre-norm transformers. Deze methode stelt het model in staat om tijdens training over te gaan van standaard normalisatie naar geleerde, steekproefonafhankelijke lineaire of affiene kaarten, waardoor normalisatiebewerkingen tijdens inferentie kunnen worden verwijderd.

Kernmechanisme: TaperNorm

TaperNorm interpolatie dynamisch tussen een genormaliseerde tak en een vaste schalingstak met behulp van een globale poort $g \in [0, 1]$. Voor een token-vector $h$:
$$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$$
Waarbij:

• $r(h)$ de RMSNorm-schaalfactor is.
• $D\gamma$ en $D\tilde{\gamma}$ leerbare diagonale schalingsmatrices zijn.
• $c$ een per-laag scalair is.
• $g$ een globale poort is die begint bij 1 (standaard normalisatie) en tijdens training cosinus-afgezwakt wordt naar 0.

Trainingsprocedure:

1. Warmup: De poort $g$ wordt op 1 gehouden. Het model traint als een standaard pre-norm transformer.
2. Kalibratie: Aan het einde van de warmup wordt het scalair $c$ berekend om de kleinste-kwadratenverschil tussen de genormaliseerde tak en de schalingstak te minimaliseren, zodat een soepele overgang zonder distributieverandering wordt gegarandeerd.
3. Afwakking: De poort $g$ wordt afgezwakt naar 0. De parameters van de schalingstak ($\tilde{\gamma}$) blijven trainbaar.
4. Inferentie-invoeging: Zodra $g=0$, wordt de TaperNorm-laag een eenvoudige lineaire/affiene kaart ($c \cdot h D\tilde{\gamma}$). Deze bewerking kan wiskundig worden ingevoegd in aangrenzende lineaire projecties (bijv. QKV-projecties of MLP-weights), waardoor de expliciete normalisatiebewerking uit het inferentiegrafiek wordt verwijderd.

Analyse van Schaalankering

Het artikel identificeert een kritiek onderscheid tussen interne en finale normalisatielagen:

• Finale Normalisatie als Anker: De finale normalisatielaag fungeert als een impliciet schaalanker. Wiskundig is deze 0-homogeen, wat betekent dat het de radiale gradiëntcomponent bij de output verwijdert. Zonder dit anker kan het minimaliseren van cross-entropy het model ertoe aanzetten de logit-magnitudes onbeperkt te verhogen ("logit chasing") om de loss te verlagen, wat leidt tot instabiliteit.
• Expliciet Alternatief: Wanneer de finale normalisatie wordt verwijderd (in "All-Taper"-instellingen), introduceren de auteurs een fixed-target scale loss ($L_{aux}$). Deze hulploss straft afwijkingen van tokenschalen van een vaste doelwaarde $s_{tgt}$ af, en biedt een expliciete radiale herstellende kracht die de neiging tot logit-chasing tegenwerkt.

Twee experimentele configuraties worden getest:

1. Internal-Taper: Interne bloknormalisators worden afgezwakt naar lineaire kaarten, maar de finale normalisatie blijft actief (impliciet anker).
2. All-Taper: Alle normalisators, inclusief de finale, worden afgezwakt. Dit vereist de fixed-target scale loss om stabiliteit te behouden (expliciet anker).

Belangrijkste Bijdragen

1. Gesteunde Normalisatieverwijdering: Een methode om RMSNorm/LayerNorm-modules af te zwakken naar vaste lineaire/affiene herschalingen, waardoor per-token statistieken tijdens inferentie kunnen worden geëlimineerd via gewichtsinvoeging.
2. Theorie van Output Schaalankering: Een analyse die aantoont dat de primaire, unieke rol van de finale normalisatielaag het verwijderen van radiale gradiënten is. Het artikel laat zien dat zonder dit anker cross-entropy onbeperkte logit-groei induceert, en dat een fixed-target scale loss deze stabiliserende rol effectief kan vervangen.
3. Empirische Validatie: Experimenten op TinyStories-pretraining en GPT-2-finetuning tonen aan dat interne normalisatie kan worden verwijderd met minimale toename van de validatie-loss. De benadering bereikt tot 1,18× hogere doorvoer bij KV-gecacheerde autoregressieve decoding na het invoegen van interne schalingen.

Resultaten

• Pre-training (TinyStories):
  • Internal-Taper (+aux): Bereikte validatie-lossen binnen 0,7–1,5% van de baseline over modelgroottes (1M tot 30M parameters).
  • All-Taper (+aux): Bereikte lossen binnen 0,8–1,8% van de baseline, wat aantoont dat een volledig norm-vrij model haalbaar is met expliciete schaalankering.
  • Dynamiek: Modellen zonder schaalanker vertoonden snelle groei in logit-magnitudes ("logit chasing"). De fixed-target scale loss onderdrukte dit gedrag succesvol en bracht de gradiëntmagnitudes in lijn met de verankerde baseline.
• Inferentie-doorvoer:
  • Op een NVIDIA H100 bereikte de Internal-Taper (gefuseerde) configuratie 1,15–1,18× hogere tokens-per-seconde-doorvoer vergeleken met de standaard RMSNorm-baseline. De Internal-Taper (niet-gefuseerde) configuratie (expliciete schalingsbewerkingen) bereikte 1,10–1,14× snelheidswinst.
• Fine-tuning (GPT-2):
  • De voorgestelde Internal-TaperLN en All-TaperLN methoden bereikten cross-entropy-lossen die concurrerend waren met de "FakeLN"-benadering (Baroni et al., 2025) en standaard Vanilla-modellen over de OpenWebText- en The Pile-datasets.
  • In tegenstelling tot sequentiële verwijderingsmethodes gebruikt TaperNorm een enkele globale poort om alle geselecteerde lagen gelijktijdig te verwijderen.

Betekenis en Beweringen

Het artikel beweert dat interne en finale normalisatielagen distincte rollen spelen in pre-norm transformers. Waar interne normalisators kunnen worden vervangen door geleerde steekproefonafhankelijke lineaire kaarten met verwaarloosbare prestatiedegradatie, vervult de finale normalisatielaag een unieke functie als output schaalanker.

De betekenis van dit werk ligt in:

1. Efficiëntie: Het biedt een praktische weg om computatievriendelijke normalisatiebewerkingen uit het inferentiegrafiek te verwijderen, wat meetbare doorvoerwinsten oplevert (tot 1,18×) zonder opnieuw te moeten trainen vanaf nul.
2. Theoretisch Inzicht: Het verduidelijkt dat de stabiliteit die door de finale normalisatielaag wordt geboden, grotendeels te wijten is aan het verwijderen van radiale gradiënten. Dit inzicht maakt het mogelijk "volledig norm-vrije" modellen te ontwerpen door het impliciete anker te vervangen door een expliciete fixed-target scale loss.
3. Robuustheid: De benadering behoudt standaard pre-norm-dynamiek vroeg in de training en gaat pas later over naar lineaire kaarten, wat stabieler lijkt dan methoden die proberen te trainen zonder normalisatie vanaf het begin.

De auteurs blijven bescheiden wat betreft de reikwijdte, en merken op dat experimenten beperkt zijn tot modellen tot 30M parameters en specifieke hardware-stacks. Ze erkennen dat robuustheid tegen activatie-uitbijters in regimes met miljarden parameters en de toepasbaarheid op andere doelen dan cross-entropy open vragen blijven voor toekomstig werk.