Tapered Language Models

Riassunto Tecnico: Modelli Linguistici Tapered (Tapered Language Models)

Problematica
I moderni modelli linguistici (LM), inclusi i Transformer, le varianti ricorrenti e le architetture basate sulla memoria, condividono una comune impostazione strutturale predefinita: una pila di strati identici in cui i parametri sono allocati uniformemente attraverso la profondità. Nonostante le crescenti prove suggeriscano che gli strati contribuiscano in modo non uniforme all'output finale — dove gli strati iniziali spesso catturano pattern sintattici superficiali e gli strati successivi raffinano i flussi residui piuttosto che trasformarli — questa allocazione uniforme rimane lo standard. Il documento postula che questo approccio "taglia unica" non rifletta l'asimmetria dell'importanza degli strati, potenzialmente sprecando la capacità dei parametri negli strati successivi dove è meno utilizzata.

Metodologia
Gli autori introducono i Tapered Language Models (TLM), un principio architettonico in cui una componente portatrice di parametri viene progressivamente ridotta (tapered) in modo monotono attraverso la profondità sotto un budget totale fisso.

• Componente Target: Lo studio istanzia questo principio sulla larghezza intermedia dell'MLP ($d_{ff}$). Gli MLP sono stati scelti perché dominano il conteggio dei parametri nei moderni LM ed espongono la larghezza come un singolo asse di variazione pulito, indipendente dal modulo di mixing dei token.
• Schedules di Tapering: Invece di allocazioni discrete per blocchi, gli autori propongono schedule di decadimento fluide e continue per $d_{ff}(l)$ attraverso gli strati $l \in \{0, \dots, L-1\}$. Vengono valutate tre schedule:
  1. Lineare: Decadimento a tasso costante.
  2. Coseno: Segue una curva a metà coseno, decrescendo lentamente alle estremità e bruscamente nel mezzo.
  3. Sigmoide: Concentra la transizione in una stretta banda attorno al punto medio.
• Vincoli: Il tapering è rigorosamente preservante del budget. La larghezza media per strato è vincolata a essere uguale alla larghezza uniforme di base ($\frac{1}{L}\sum d_{ff}(l) = d_{ff}^{baseline}$), garantendo che il conteggio totale dei parametri e i FLOP (sia di training che di inferenza) rimangano identici al baseline uniforme.
• Setup Sperimentale: Il principio è valutato su quattro architetture distinte (Standard Transformer, Gated Attention, Hope-attention e Titans) e tre scale di modello (440M, 760M e 1.3B di parametri). Lo studio utilizza un setup di training fisso (tokenizer Llama 3, sequenza di lunghezza 4K) e confronta i modelli tapered con i baseline uniformi su perplexity di validazione in-distribution, perplexity out-of-distribution (WikiText, LAMBADA) e otto benchmark di ragionamento comune senso (commonsense reasoning).

Risultati Chiave

1. La Direzione dell'Allocazione Conta: Un esperimento motivazionale su un Transformer da 440M ha partizionato gli strati in gruppi: iniziali, medi e tardivi. Allocare più capacità agli strati iniziali ("Wider-early") ha migliorato la perplexity di 0.32 punti rispetto al baseline uniforme. Al contrario, allocare più capacità agli stroli tardivi ("Wider-late") ha degradato le prestazioni di oltre 1.0 punto, dimostrando che la direzione della capacità è critica.
2. Schedule e Ratio Ottimali: Una ricerca (sweep) di ratio di larghezza e schedule su un modello da 440M ha rivelato che la schedule Coseno ha costantemente superato le schedule Lineare e Sigmoide. La configurazione ottimale è stata individuata con un rapporto iniziale-finale di 1.50 $\to$ 0.50, raggiungendo una perplexity di validazione di 14.44 (un miglioramento di 1.84 punti rispetto al baseline uniforme di 16.28).
3. Generalizzazione: L'applicazione della configurazione ottimale fissa (Coseno, 1.5/0.5) ai modelli da 760M e 1.3B su tutte e quattro le architetture ha migliorato costantemente la perplexity e l'accuratezza nei task di ragionamento comune senso. I guadagni sono stati osservati indipendentemente dal fatto che il modulo di token-mixing utilizzasse softmax attention, gated attention o meccanismi ricorrenti/basati sulla memoria.
4. Insight Meccanicistico: L'analisi dei transformer pre-addestrati ha rivelato che gli output degli MLP diventano progressivamente più allineati con il flusso residuo esistente all'aumentare della profondità (misurato tramite similarità del coseno). Ciò suggerisce che gli MLP successivi rinforzano il contenuto esistente piuttosto che computare nuove feature. Il tapering allinea l'architettura a questo pattern riducendo la capacità dove è meno utilizzata (strati successivi) e riallocandola agli strati iniziali dove è più ortogonale e utile.

Significatività e Claim
Il paper stabilisce l'allocazione della capacità consapevole della profondità (depth-aware capacity allocation) come un semplice lever di design agnostico rispetto all'architettura per i modelli linguistici. Gli autori sostengono che il tapering della larghezza dell'MLP offre un miglioramento delle prestazioni "gratuito" — potenziando la perplexity e l'accuratezza nei task downstream senza aumentare il numero di parametri o il costo computazionale.

La significatività risiede nel sfidare l'uniformità ereditata dallo chassis del Transformer. I risultati suggeriscono che il principio del tapering riguardi l'allocazione generale dei parametri attraverso la profondità piuttosto che proprietà specifiche dei meccanismi di attenzione. Sebbene gli autori notino che la specifica schedule ottimale (Coseno, 1.5/0.5) sia stata derivata da un modello da 440M e possa non essere l'ottimo globale per tutte le scale, i guadagni costanti attraverso diverse architetture e dimensioni indicano che il principio è robusto. Il paper conclude che l'allocazione consapevole della profondità è una "leva nascosta" evidente che può essere applicata ai modelli di fondazione oltre il linguaggio, inclusi i transformer visivi e multimodali.