Adaptive inference and function vectors in deep transformers

Problema
I Transformer sono emersi come un substrato general-purpose per l'apprendimento di correlazioni complesse in diversi domini, eppure i loro meccanismi interni per l'apprendimento in-context (ICL) rimangono misteriosi. Mentre il lavoro empirico ha identificato due motivi primari per l'ICL — circuiti di attenzione che implementano algoritmi specifici (ad es. copia, discesa del gradiente) e la formazione di "vettori di funzione" (rappresentazioni interne del task) — la comprensione teorica di questi ultimi rimane limitata. Nello specifico, non è chiaro come i transformer profondi utilizzino i vettori di funzione per eseguire l'inferenza, e quali ruoli computazionali specifici ricadano sulla profondità e sui blocchi feedforward (MLP) in questo processo. Le teorie esistenti si concentrano spesso su circuiti basati sull'attenzione, lasciando inesplorati i meccanismi abilitati dalle interazioni multi-livello tra attenzione e MLP.

Metodologia
Gli autori propongono un quadro teorico che modella un transformer profondo come un sistema interagente a campo medio che implementa un'inferenza bayesiana distribuita. Il modello è soggetto a vincoli architettonici sulla comunicazione (dimensione del canale) $d$, località e profondità ($L$).

1. Modello Teorico: Il framework postula che un transformer elabori un insieme di token accoppiati da una variabile di contesto latente $\theta$. L'inferenza è mediata da "vettori di funzione" ($\Phi_\ell$), che agiscono come variabili di stato compresse.

   • Ad ogni livello $\ell$, una non-linearità locale (MLP) elabora un token e lo stato corrente $\Phi_\ell$ per produrre un embedding $\lambda_{\ell+1}$.
   • Un'operazione di tipo attention aggrega questi embedding attraverso i token per formare una nuova statistica $\phi_{\ell+1}$.
   • Lo stato viene aggiornato tramite somma diretta: $\Phi_{\ell+1} = \Phi_\ell \oplus \phi_{\ell+1}$.
   • Si ipotizza che i blocchi MLP agiscano come "router" (scelta di quali informazioni trasmettere) e decoder, mentre l'attenzione aggrega le informazioni.

2. Strategia di Ottimizzazione: Gli autori inquadrano la selezione degli embedding come un problema di programmazione dinamica. La strategia ottimale al livello $\ell$ minimizza la perdita attesa al livello $L$, trattando efficacemente la scelta dell'embedding come una strategia di design sperimentale per l'inferenza adattiva.

3. Casi Analitici: La teoria è analizzata sotto due distribuzioni a priori per la variabile latente $\theta$:

   • Prior Gaussiana: Dove la variabile latente segue una distribuzione normale standard.
   • Prior ad Albero (Tree Prior): Una struttura gerarchica in cui $\theta$ è costruito in base a un percorso attraverso un albero binario bilanciato. Questa struttura testa il vantaggio della profondità nel risolvere variabili latenti a scale sempre più fini.

4. Validazione Empirica: Le previsioni sono testate utilizzando un'architettura transformer a attenzione lineare vincolata con tre configurazioni:

   • Modello A: Un modello profondo ($L=M$) con dimensione del canale $d=1$ e blocchi MLP a ogni livello.
   • Modello B: Un modello a singolo livello con dimensione del canale $M$ e MLP identici.
   • Modello C: Un modello a singolo livello con dimensione del canale $M$ e senza MLP (solo attenzione lineare).
   • Perturbazione Causale: Viene condotto un esperimento di "key-patching" in cui matrici di key da un contesto "donatore" vengono inserite in un contesto "target" a livelli specifici per testare il meccanismo di routing.

Contributi Chiave

• Teoria dei Vettori di Funzione: Il documento fornisce una teoria formale che definisce i vettori di funzione come rappresentazioni dello stato interno che comprimono le informazioni del contesto per la predizione. Dimostra che questi vettori permettono al modello di inferire variabili di contesto latenti a scale sempre più fini attraverso i livelli.
• Ruolo di Profondità e MLP: La teoria prevede che, per prior non gaussiani e gerarchici, la profondità e i blocchi MLP siano essenziali per implementare strategie di inferenza adattiva che superano gli approcci non adattivi a singolo livello. Al contrario, per le prior gaussiane, la profondità non offre alcun vantaggio rispetto a un singolo livello con un budget di comunicazione totale equivalente.
• Mappatura Meccanicistica: Il lavoro mappa i componenti teorici all'architettura del transformer, assegnando all'MLP il ruolo di "router" che seleziona le statistiche da misurare e al meccanismo di attenzione il ruolo di aggregazione di tali statistiche.
• Conferma Sperimentale: Esperimenti numerici confermano che i transformer profondi con MLP (Modello A) possono apprendere strategie di inferenza adattiva su prior ad albero gerarchici, superando significativamente le strategie non adattive ottimizzate a singolo livello (Modello B) e i modelli di attenzione lineare (Modello C).

Risultati

• Prior Gaussiana: Nel caso di una prior gaussiana, l'errore quadratico medio (MSE) è indipendente dalla specifica realizzazione del vettore di funzione e dipende solo dalla strategia di misurazione. La teoria e gli esperimenti mostrano che un modello profondo (Modello A) non performa meglio di un modello a singolo livello con capacità di canale equivalente (Modello B).
• Prior ad Albero: Per la prior gerarchica ad albero, la teoria prevede che una strategia adattiva (misurare una coordinata alla volta in base ai risultati precedenti) produca una perdita MSE che scala esponenzialmente con la profondità ($\alpha^M$), mentre le strategie non adattive scalano come una potenza di $M$.
  • Gli esperimenti mostrano che il Modello A raggiunge un MSE significativamente più basso rispetto al Modolo B, confermando che il modello apprende una strategia di inferenza adattiva.
  • Il Modello C (senza MLP) performa in modo simile alla baseline teorica non adattiva, confermando la necessità degli MLP per il meccanismo di routing adattivo.
• Perturbazione Causale: Gli esperimenti di key-patching rivelano che l'inserimento di key da un contesto donatore (foglia $\sigma'$) in un contesto target (foglia $\sigma$) causa un forte aumento della perdita se i percorsi divergono precocemente nell'albero e l'intervento avviene nelle fasi iniziali della rete. Ciò conferma che lo stato interno del modello (vettore di funzione) evolve adattivamente in base al percorso specifico del contesto.

Significatività
Il documento sostiene di riconciliare due prospettive sull'apprendimento in-context: la visione dell'ICL come inferenza bayesiana su una variabile latente e la visione meccanicistica dell'ICL come esecuzione algoritmica (ad es. discesa del gradiente). Dimostrando che i transformer profondi possono implementare l'inferenza bayesiana come un procedimento distribuito, adattivo e per livelli mediato da vettori di funzione, gli autori dimostrano che la profondità e i blocchi MLP abilitano una classe di algoritmi di ICL più ricca di quella precedentemente descritta (come la discesa del gradiente precondizionata). Questo lavoro colloca i vettori di funzione all'interno di un framework a livello di circuito, suggerendo che l'organizzazione dei modelli da miliardi di parametri favorisce questa struttura perché permette un'inferenza efficiente e adattiva sotto vincoli di comunicazione. Gli autori osservano che, sebbene la loro teoria utilizzi l'aggregazione lineare, essa rispetta l'scambiabilità e può essere sufficiente per scenari in cui il recupero specifico token-a-token dell'attenzione softmax non è il driver primario dell'ICL.