Unlocking Feature Learning in Gated Delta Networks at Scale

Sintesi Tecnica: Sbloccare l'Apprendimento delle Caratteristiche nelle Gated Delta Networks su Scala

Problematica
La scalabilità dei Large Language Models (LLM) affronta due colli di bottiglia primari: l'elevato costo computazionale dell'addestramento di modelli più grandi e la complessità quadratica della standard self-attention dei Transformer. Le architetture ricorrenti lineari, come le Gated Delta Networks (GDN), sono emerse come alternative efficienti che mantengono un'inferenza in tempo lineare pur catturando dipendenze a lungo raggio. Tuttavia, esiste una lacuna critica nel far scalare efficacemente questi modelli. Sebbene la Maximal Update Parametrization (µP) abbia abilitato con successo il trasferimento zero-shot degli iperparametri (specificamente dei tassi di apprendimento) per i Transformer standard e per alcuni modelli lineari, la sua applicazione alle Gated Delta Networks rimane inesplorata. La sfida principale risiede nella dinamica dello stato ricorrente dell'architettura, che viene aggiornato tramite meccanismi di gating dipendenti dai dati e regole delta a matrice completa, piuttosto che attraverso le ricorrenze scalari o le standard strutture feedforward ipotizzate dalle esistenti derivazioni µP. Il lavoro precedente sui Modelli a Spazio di Stato Strutturati (SSM) come Mamba ha dimostrato che la vanilla µP fallisce per gli SSM diagonali, ma l'uso da parte di GDN di aggiornamenti dello stato a matrice completa rende inapplicabili quelle specifiche correzioni.

Metodologia
Gli autori affrontano questa lacuna derivando rigorosamente la formulazione µP per le Gated Delta Networks. La metodologia consiste nel propagare le stime delle dimensioni delle coordinate attraverso l'intero passaggio in avanti (forward pass), inclusi i meccanici di gating unici e le transizioni di stato ricorrenti, per determinare le necessarie regole di scaling per l'inizializzazione, i moltiplicatori forward e i tassi di apprendimento.

Passaggi analitici chiave includono:

1. Analisi del Forward Pass: Gli autori derivano le dimensioni delle coordinate delle caratteristiche proiettate, dello stato latente $S_t$ e del readout. Identificano che la L2-normalizzazione di query e key, combinata con gli aggiornamenti di scrittura rank-one, risulta in specifici scalamenti della varianza. Crucialmente, propongono di inserire un moltiplicatore $\sqrt{d}$ prima dello strato RMSNorm per garantire che l'input allo strato di normalizzazione mantenga una dimensione di coordinata $\Theta(1)$, evitando la perturbazione dello scaling del gradiente necessaria per l'apprendimento delle caratteristiche.
2. Analisi del Meccanismo di Gating: Il documento analizza i parametri di gating non lineari $\alpha_t$ (decadimento) e $\beta_t$ (forza di scrittura). A differenza dei pesi standard, questi sono derivati da pre-attivazioni scalari che coinvolgono vettori di peso addestrabili ($W_\alpha, W_\beta$) e parametri scalari ($a_{log}, b$). L'analisi rivela che questi componenti di gating non seguono le leggi di scaling standard di µP a causa dei loro ruoli specifici nell'equazione di aggiornamento dello stato.
3. Analisi del Backward Pass e Derivazione del Tasso di Apprendimento:
   • AdamW: A causa della normalizzazione dei gradienti da parte del secondo momento di AdamW, l'entità dell'aggiornamento effettivo è coerente tra le diverse classi di pesi, semplificando i requisiti del tasso di apprendimento.
   • SGD: Sotto il regime di plain SGD, le magnitudo dei gradienti grezzi variano significativamente tra le diverse classi di pesi. Gli autori derivano che le matrici di peso di gating ($W_\alpha, W_\beta$) richiedono uno scaling del tasso di apprendimento di $\Theta(1/\sqrt{d})$, e i parametri di gating scalari ($a_{log}, b$) richiedono uno scaling di $\Theta(\sqrt{d})$. Questi scalamenti divergono dal setting standard di µP dove i tassi di apprendimento sono tipicamente uniformi o seguono una semplice regola $1/n$.

Contributi Chiave

1. Derivazione Teorica: Il documento fornisce la prima formulazione µP completa per le Gote Delta Networks. Deriva le stime delle dimensioni delle coordinate attraverso l'intero forward pass e stabilisce varianze di inizializzazione e moltiplicatori forward principati per tutte le classi di pesi.
2. Nuove Regole di Scaling: Gli autori scoprono che i componenti di gating richiedono scalamenti non standard del tasso di apprendimento sotto SGD:
   • Le matrici di peso di gating ($W_\alpha, W_\beta$) richiedono uno scaling di $\Theta(1/\sqrt{d})$.
   • I parametri di gating scalari ($a_{log}, b$) richiedono uno scaling di $\Theta(\sqrt{d})$.
3. Validazione Empirica: Gli autori pre-addestrano modelli linguistici GDN attraverso diverse larghezze (che vanno da 256 a 1563 dimensioni nascoste) utilizzando sia gli ottimizzatori AdamW che SGD.
   • Risultati: Sotto la proposta configurazione µP, il tasso di apprendimento ottimale si trasferisce zero-shot attraverso tutte le larghezze del modello per entrambi gli ottimizzatori. Al contrario, la Standard Parametrization (SP) e l'originale configurazione µP falliscono nel trasferirsi, con i tassi di apprendimento ottimali che cambiano significativamente all'aumentare della larghezza del modello.
   • Stabilità: Le configurazioni proposte consentono un trasferimento stabile del tasso di apprendimento, confermando le derivazioni teoriche.

Significatività e Rivendicazioni
Il documento sostiene che la sua analisi risolve il problema aperto della parametrizzazione delle Gated Delta Networks per l'apprendimento delle caratteristiche su scala. Dimostrando che la derivata formulazione µP permette il trasferimento zero-shot degli iperparametri, il lavoro valida la correttezza della derivazione teorica e la sua utilità pratica. Gli autori sottolineano che questo è il primo lavoro a derivare e validare una parametrizzazione coerente con µP specificamente per le Gated Delta Networks, distinguendosi dal lavoro precedente sugli SSM diagonali. I risultati suggeriscono che una corretta parametrizzazione è essenziale per scalare i modelli ricorrenti lineari e che i meccanismi di gating specifici in GDN richiedono regole di scaling uniche che differiscono sia dai Transformer standard che da altre architetture lineari. Gli autori sperano che queste derivazioni possano guidare la ricerca futura sulle leggi di scaling di altre architetture lineari o ibride.