Training Infinitely Deep and Wide Transformers

Sintesi Tecnica: Addestramento di Transformer Infinitamente Profondi e Larghi

Enunciato del Problema
Sebbene i Transformer siano diventati l'architettura dominante nel machine learning moderno, la comprensione teorica delle loro dinamiche di addestramento rimane limitata, in particolare per architetture profonde con molti livelli e teste di attenzione. I quadri teorici esistenti fanno spesso affidamento sul regime del Neural Tangent Kernel (NTK), che assume che i pesi scalino come $1/\sqrt{\text{larghezza}}$, portando a un "addestramento pigro" in cui la rete rimane vicina all'inizializzazione e l'apprendimento delle caratteristiche è congelato. In alternativa, i regimi di campo medio per reti poco profonde o ResNet (modellate come Equazioni Differenziali Ordinarie, ODE) non catturano pienamente la complessità dei Transformer, dove il meccanismo di attenzione accoppia l'evoluzione di multiple distribuzioni di token. È necessario un quadro matematico rigoroso per analizzare l'addestramento basato su gradienti nel limite congiunto di profondità infinita (numero di livelli) e larghezza infinita (numero di teste di attenzione), tenendo conto dell'evoluzione simultanea delle distribuzioni di token e delle distribuzioni dei parametri.

Metodologia
Il documento sviluppa un quadro di campo medio per i Transformer estendendo la teoria del trasporto ottimo condizionato e dei flussi gradiente di Wasserstein, precedentemente applicata alle ResNet, al contesto dei Transformer.

1. Rappresentazione Duale di Campo Medio: Il modello tratta i Transformer come infinitamente profondi (parametrizzati da una variabile di profondità continua $s \in [0,1]$) e infinitamente larghi (parametrizzati da una distribuzione di probabilità $\rho$ sui parametri di attenzione). Crucialmente, il quadro impiega due tipi di rappresentazioni misurali:

   • Distribuzioni di Token: L'evoluzione dei token attraverso i livelli è modellata come una misura di probabilità $\mu(s)$ sullo spazio dei token.
   • Parametri di Attenzione: I parametri a ogni livello sono rappresentati come una misura $\rho(\cdot|s)$ sullo spazio dei parametri $\Theta$ (matrici query, key, value).

2. Formulazione di PDE Neurale: A differenza delle ResNet, che possono essere viste come il controllo di un'ODE neurale su vettori di caratteristiche, l'addestramento dei Transformer in questo quadro corrisponde al controllo di una Equazione Differenziale alle Derivate Parziali (PDE) Neurale. L'evoluzione dei token è governata da un'equazione di continuità in cui il campo di velocità è determinato dall'operatore di attenzione di campo medio, che accoppia le distribuzioni di tutti i token.
   $$\partial_s \mu(s) + \text{div}(\mu(s) \Phi_{\rho(\cdot|s)}[\mu(s)]) = 0$$
   Gli autori riformulano questa prospettiva PDE in un'ODE a valori in uno spazio di Banach che agisce su mappe di flusso $\Lambda$, che trasportano la distribuzione iniziale dei token. Questa riparametrizzazione evita la mancanza di una struttura differenziabile nello spazio delle misure e fornisce una regolarità potenziata per l'analisi.

3. Dinamiche di Addestramento tramite Flusso Gradiente di Wasserstein Condizionato: Il processo di addestramento è modellato come un flusso gradiente sullo spazio delle distribuzioni di parametri equipaggiato con la metrica del Trasporto Ottimo Condizionato (COT). Questa metrica impone che la distribuzione marginale dei parametri attraverso i livelli sia preservata durante l'addestramento. Gli autori derivano il gradiente di Wasserstein condizionato del rischio di addestramento utilizzando l'analisi di sensibilità aggiuntiva, ottenendo una formula esplicita che coinvolge variabili aggiuntive governate da ODE all'indietro.

4. Analisi del Neural Tangent Kernel (NTK): Per stabilire la convergenza, il documento indaga l'iniettività dell'NTK per i meccanismi di attenzione. Gli autori derivano condizioni necessarie e sufficienti affinché l'NTK sia definito positivo, collegando questa proprietà all'indipendenza lineare delle funzioni log-sum-exp (funzioni generatrici dei cumulanti) delle distribuzioni di token modulo funzioni affini.

Contributi e Risultati Chiave

• Ben Posizionamento delle Dinamiche Forward e di Addestramento: Il documento stabilisce l'esistenza e l'unicità delle soluzioni per il passaggio in avanti (evoluzione dei token) e per l'evoluzione dell'addestramento (flusso della distribuzione dei parametri). Dimostra che l'evoluzione dei token è caratterizzata da mappe di flusso che soddisfano ODE in spazi di funzioni e che le curve di flusso gradiente esistono e sono uniche nello spazio metrico di Wasserstein condizionato.
• Formulazione Esplicita del Gradiente: Utilizzando l'analisi di sensibilità aggiuntiva, gli autori derivano una formula esplicita per il gradiente di Wasserstein condizionato del rischio di addestramento. Questo coinvolge variabili aggiuntive che soddisfano ODE all'indietro, fornendo una base rigorosa per l'addestramento basato su gradienti nel limite di campo medio.
• Condizioni di Iniettività NTK: Un contributo tecnico centrale è la caratterizzazione dell'iniettività NTK per i meccanismi di attenzione. Il documento dimostra che l'iniettività NTK è equivalente all'indipendenza lineare delle funzioni log-sum-exp modulo funzioni affini. Si dimostra che questa condizione è soddisfatta da distribuzioni di token diverse, tra cui:
  • Distribuzioni discrete che soddisfano una condizione di distinzione (verificata quasi certamente per campioni i.i.d. da misure assolutamente continue).
  • Distribuzioni uniformi su cubi con raggi distinti.
  • Distribuzioni Laplace multivariata simmetriche con matrici di covarianza distinte.
  • Misure gaussiane a due componenti con parametri di separazione distinti.
• Convergenza ai Minimi Globali: Sotto l'assunzione che l'NTK sia positivo all'inizializzazione (specificamente, la parte $V$ dell'NTK) e che la perdita iniziale sia sufficientemente piccola, il documento dimostra che il flusso gradiente converge a un minimizzatore globale con un tasso di convergenza lineare. Questo risultato elimina i minimi locali spurii dal paesaggio di ottimizzazione nel regime specificato.

Significato e Affermazioni
Il documento afferma di fornire una base matematica rigorosa per l'addestramento di Transformer infinitamente profondi e larghi. Estendendo la teoria di campo medio dalle ResNet (ODE neurali) ai Transformer (PDE neurali), colma il divario tra la comprensione teorica delle reti residuali profonde e le dinamiche più complesse dei meccanismi di attenzione.

Gli autori sottolineano che il loro quadro cattura l'accoppiamento di multiple distribuzioni di token attraverso il meccanismo di attenzione, una caratteristica assente nelle analisi standard delle ResNet. L'istituzione di condizioni di iniettività NTK fornisce indicazioni costruttive per l'inizializzazione, garantendo che l'addestramento basato su gradienti converga ai minimi globali senza rimanere intrappolato in minimi locali spurii, a condizione che l'inizializzazione sia sufficientemente vicina alla soluzione ottimale.

Il documento nota modestamente che la sua analisi è limitata all'attenzione di tipo encoder (non causale). Estendere questo quadro alle architetture decoder con mascheramento causale è identificato come una sfida tecnica sostanziale a causa della rottura delle simmetrie utilizzate nelle dimostrazioni. Inoltre, la formulazione attuale mantiene la dimensione della testa fissa mentre lascia che il numero di teste cresca all'infinito; catturare la scalatura accoppiata della dimensione di embedding e del numero di teste rimane un problema aperto.