Riepilogo Tecnico: Rimozione della Normalizzazione a Cancellazione e Ancoraggio della Scala nei Transformer Pre-Norm

Enunciato del Problema

I livelli di normalizzazione (ad es. LayerNorm, RMSNorm) sono componenti standard nelle architetture transformer, che stabilizzano le attivazioni e abilitano l'addestramento di reti profonde. Tuttavia, questi livelli introducono calcoli non lineari dipendenti dal campione che comportano un sovraccarico significativo durante l'inferenza. Non è ancora chiaro se queste statistiche per token siano strettamente necessarie durante l'intera fase di inferenza, in particolare nei transformer pre-norm. Sebbene lavori precedenti abbiano tentato di rimuovere la normalizzazione tramite trucchi di inizializzazione o fine-tuning sequenziale, manca una comprensione dei ruoli funzionali distinti dei livelli di normalizzazione interni rispetto a quelli finali, specificamente per quanto riguarda il modo in cui controllano la scala della rappresentazione pre-logit.

Metodologia

Gli autori propongono TaperNorm, un approccio di rimozione della normalizzazione a cancellazione progettato per i transformer pre-norm. Questo metodo permette al modello di transitare dalla normalizzazione standard a mappe lineari o affini apprese e indipendenti dal campione durante l'addestramento, consentendo la rimozione delle operazioni di normalizzazione in inferenza.

Meccanismo Principale: TaperNorm

TaperNorm interpola dinamicamente tra un ramo normalizzato e un ramo di scaling fisso utilizzando un gate globale $g \in [0, 1]$. Per un vettore di token $h$:
$$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$$
Dove:

• $r(h)$ è il fattore di scaling RMSNorm.
• $D\gamma$ e $D\tilde{\gamma}$ sono matrici di scaling diagonali apprendibili.
• $c$ è uno scalare per livello.
• $g$ è un gate globale che inizia a 1 (normalizzazione standard) e decade cosinusoidalmente a 0 durante l'addestramento.

Procedura di Addestramento:

1. Warmup: Il gate $g$ è mantenuto a 1. Il modello si allena come un transformer pre-norm standard.
2. Calibrazione: Alla fine del warmup, lo scalare $c$ viene calcolato per minimizzare la differenza ai minimi quadrati tra il ramo normalizzato e il ramo di scaling, garantendo una transizione fluida senza spostamento della distribuzione.
3. Riduzione (Tapering): Il gate $g$ decade a 0. I parametri del ramo di scaling ($\tilde{\gamma}$) rimangono apprendibili.
4. Folding in Inferenza: Una volta che $g=0$, il livello TaperNorm diventa una semplice mappa lineare/affine ($c \cdot h D\tilde{\gamma}$). Questa operazione può essere matematicamente fusa nelle proiezioni lineari adiacenti (ad es. proiezioni QKV o pesi MLP), rimuovendo l'operazione di normalizzazione esplicita dal grafo di inferenza.

Analisi dell'Ancoraggio della Scala

Il documento identifica una distinzione critica tra i livelli di normalizzazione interni e finali:

• Normalizzazione Finale come Ancoraggio: Il livello di normalizzazione finale agisce come un ancoraggio di scala implicito. Matematicamente, è 0-omogeneo, il che significa che rimuove la componente del gradiente radiale all'output. Senza questo ancoraggio, la minimizzazione della cross-entropy può spingere il modello ad aumentare le magnitudini dei logit indefinitamente ("logit chasing") per ridurre la perdita, portando a instabilità.
• Alternativa Esplicita: Quando la normalizzazione finale viene rimossa (nelle impostazioni "All-Taper"), gli autori introducono una funzione di perdita per scala a target fisso ($L_{aux}$). Questa perdita ausiliaria penalizza le deviazioni delle scale dei token da un target fisso $s_{tgt}$, fornendo una forza di ripristino radiale esplicita che contrasta la tendenza al "logit chasing".

Vengono testate due configurazioni sperimentali:

1. Internal-Taper: I normalizzatori dei blocchi interni vengono ridotti a mappe lineari, ma la normalizzazione finale rimane attiva (ancoraggio implicito).
2. All-Taper: Tutti i normalizzatori, incluso quello finale, vengono ridotti. Questo richiede la funzione di perdita per scala a target fisso per mantenere la stabilità (ancoraggio esplicito).

Contributi Chiave

1. Rimozione della Normalizzazione a Cancellazione: Un metodo per ridurre i moduli RMSNorm/LayerNorm a riscalaggi lineari/affini fissi, consentendo l'eliminazione delle statistiche per token in inferenza tramite folding dei pesi.
2. Teoria dell'Ancoraggio della Scala in Output: Un'analisi che dimostra come il ruolo primario e distinto del livello di normalizzazione finale sia la rimozione dei gradienti radiali. Il documento mostra che senza questo ancoraggio, la cross-entropy induce una crescita illimitata dei logit, e che una funzione di perdita per scala a target fisso può efficacemente sostituire questo ruolo stabilizzante.
3. Validazione Empirica: Esperimenti sul pre-addestramento di TinyStories e sul fine-tuning di GPT-2 mostrano che la normalizzazione interna può essere rimossa con aumenti minimi della perdita di validazione. L'approccio raggiunge fino a 1.18× di throughput superiore nella decodifica autoregressiva con cache KV dopo la fusione degli scaling interni.

Risultati

• Pre-addestramento (TinyStories):
  • Internal-Taper (+aux): Ha raggiunto perdite di validazione entro lo 0.7–1.5% rispetto alla baseline per tutte le dimensioni del modello (da 1M a 30M parametri).
  • All-Taper (+aux): Ha raggiunto perdite entro lo 0.8–1.8% rispetto alla baseline, dimostrando che un modello completamente privo di normalizzazione è fattibile con un ancoraggio di scala esplicito.
  • Dinamiche: I modelli senza ancoraggio di scala hanno mostrato una rapida crescita delle magnitudini dei logit ("logit chasing"). La funzione di perdita per scala a target fisso ha soppresso con successo questo comportamento, allineando le magnitudini dei gradienti con la baseline ancorata.
• Throughput di Inferenza:
  • Su una NVIDIA H100, la configurazione Internal-Taper (fusa) ha raggiunto un throughput di token al secondo 1.15–1.18× superiore rispetto alla baseline RMSNorm standard. La configurazione Internal-Taper (non fusa) (operazioni di scaling esplicite) ha raggiunto un speedup di 1.10–1.14×.
• Fine-tuning (GPT-2):
  • I metodi proposti Internal-TaperLN e All-TaperLN hanno raggiunto perdite di cross-entropy competitive con l'approccio "FakeLN" (Baroni et al., 2025) e i modelli Vanilla standard sui dataset OpenWebText e The Pile.
  • A differenza dei metodi di rimozione sequenziale, TaperNorm utilizza un singolo gate globale per rimuovere tutti i livelli selezionati simultaneamente.

Significato e Affermazioni

Il documento afferma che i livelli di normalizzazione interni e finali svolgono ruoli distinti nei transformer pre-norm. Mentre i normalizzatori interni possono essere sostituiti da mappe lineari indipendenti dal campione apprese con degradazione delle prestazioni trascurabile, il livello di normalizzazione finale svolge una funzione unica come ancoraggio della scala in output.

Il significato di questo lavoro risiede in:

1. Efficienza: Fornisce una via pratica per rimuovere le operazioni di normalizzazione computazionalmente costose dal grafo di inferenza, producendo guadagni misurabili nel throughput (fino a 1.18×) senza riaddestramento da zero.
2. Insight Teorico: Chiarisce che la stabilità fornita dal livello di normalizzazione finale è dovuta in gran parte alla sua rimozione dei gradienti radiali. Questo insight permette la progettazione di modelli "completamente privi di normalizzazione" sostituendo l'ancoraggio implicito con una funzione di perdita esplicita per scala a target fisso.
3. Robustezza: L'approccio mantiene le dinamiche standard pre-norm nelle prime fasi dell'addestramento, transitando a mappe lineari solo successivamente, il che appare più stabile rispetto ai metodi che tentano di addestrare senza normalizzazione fin dall'inizio.

Gli autori rimangono modesti riguardo alla portata, notando che gli esperimenti sono limitati a modelli fino a 30M parametri e a stack hardware specifici. Riconoscono che la robustezza agli outlier di attivazione nei regimi da miliardi di parametri e l'applicabilità ad altri obiettivi oltre alla cross-entropy rimangono domande aperte per lavori futuri.