Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation

Sintesi Tecnica: I Transformer come Sperimentatori Bayesiani In-Context

Formulazione del Problema
Il documento affronta la sfida della progettazione di esperimenti adattivi per stimare l'Effetto Medio del Trattamento (ATE) in contesti sequenziali a due bracci con covariate. La strategia ottimale per tali esperimenti è l'allocazione di Neyman "oracle", che assegna le probabilità di trattamento in base alle varianze residue incondizionate per braccio sconosciute, al fine di minimizzare la varianza dello stimatore dell'ATE. Implementare questo processo in modo ottimale richiede uno sperimentatore bayesiano che aggiorni sequenzialmente le credenze non parametriche sui momenti (medie e secondi momenti) dei potenziali esiti e traduca tali incertezze a posteriori in probabilità di assegnazione del trattamento.

Un collo di bottiglia critico negli spazi di covariate ad alta dimensionalità è la sconosciuta regolarità (smoothness) delle funzioni degli esiti potenziali. Prior eccessivamente regolari inducono bias, mentre prior troppo irregolari convergono lentamente. Inoltre, l'impiego di un processo completo di aggiornamento e progettazione bayesiana in tempo reale è computazionalmente intensivo, richiedendo un'ingegneria specifica per ogni esperimento. Gli autori indagano se questo complesso processo di inferenza e allocazione sequenziale possa essere "ammortizzato" tramite l'apprendimento in-context, specificamente attraverso architetture Transformer, per creare una policy che mappi direttamente le storie sperimentali in probabilità di assegnazione ottimali senza il calcolo esplicito della computazione bayesiana ad ogni step.

Metodologia
Gli autori propongono un framework centrato su uno "sperimentatore bayesiano in-context", implementato come una policy Transformer addestrata per imitare un "insegnante di Neyman posteriore bayesiano".

1. L'Insegnante Bayesiano:

   • Modello: L'insegnante mantiene credenze bayesiane non parametriche sulle medie dei bracci ($\mu_w$) e sui secondi momenti ($q_w$) utilizzando prior di serie gaussiane. Tali prior sono indicizzate da un parametro di regolarità $\beta$.
   • Adattamento Gerarchico: Per gestire la regolarità sconosciuta, l'insegnante impiega un'estensione bayesiana gerarchica in cui $\beta$ è trattato come una variabile casuale con un prior su un intervallo compatto. Ciò consente al posteriore di concentrarsi sulla classe di regolarità appropriata man mano che i dati si accumulano.
   • Allocazione: L'insegnante aggiorna le credenze utilizzando una quasi-verosimiglianza gaussiana semplificata (perdita al quadrato) e assegna il prossimo unità di trattamento secondo una propensione di Neyman posteriore: $e_t(x) \propto \sqrt{\bar{v}_{1,t-1}(x)} / (\sqrt{\bar{v}_{0,t-1}(x)} + \sqrt{\bar{v}_{1,t-1}(x)})$, dove $\bar{v}$ è la media posteriore della varianza residua.
   • Garanzie Teoriche: Il documento dimostra che, sotto deboli condizioni di regolarità, il posteriore dell'insegnante si contrae al tasso minimax non parametrico (adattato alla regolarità reale) e che il design risultante converge all'allocazione di Neyman oracle, garantendo l'efficienza semiparametrica per lo stimatore finale dell'ATE (tramite l'Inverse Probability Weighting aumentato, AIPW).

2. Espressività e Costruzione del Transformer:

   • Architettura: Gli autori dimostrano costruttivamente che un Transformer può rappresentare la mappatura dell'insegnante.
     • L'Attenzione come Statistiche Sufficienti: I blocchi di auto-attenzione aggregano i token storici per calcolare le statistiche empiriche delle risposte e i prodotti covarianza-vettore, agendo efficacementamente come statistiche sufficienti per il modello di serie gaussiane.
     • FFN come Discesa del Gradiente: Le reti feed-forward (FFN) per token implementano i passi di discesa del gradiente proiettata per risolvere il problema di regressione ridge mascherata a rango finito corrispondente all'aggiornamento della media posteriore bayesiana.
     • Mascheramento Dinamico: L'architettura utilizza una maschera di dimensione effettiva dinamica che attiva un numero di coordinate di base ($J$) dipendente dalla dimensione del campione ($N$) e dal prior di regolarità. Ciò impone il compromesso bias-varianza: $J$ è scelto in modo che la matrice di Gram empirica sia ben condizionata e che la precisione del prior non domini.
     • Mixture-of-Experts (MoE): Una porta MoE approssima il posteriore gerarchico sulle classi di regolarità. Essa pesa diversi "esperti" (ciascuno corrispondente a un livello di regolarità $\beta_k$) in base a punteggi di verosimiglianza marginale approssimativi derivati dalle statistiche basate sull'attenzione.
   • Generalizzazione tramite ERM: Il documento analizza l'apprendimento di questa policy tramite la minimizzazione del rischio empirico supervisionato (ERM). Il Transformer è addestrato su traiettorie di pre-training finite per imitare il comportamento input-output dell'insegnante (mappatura storia-propensione). La teoria garantisce che il modello addestrato possa approssimare il comportamento dell'insegnante fino a errori di approssimazione e stima, senza dover recuperare esplicitamente lo stato posteriore interno o eseguire i passi di gradiente durante l'inferenza.

Contributi Chiave

1. Design Bayesiano e Adattamento della Regolarità: Gli autori definiscono un insegnante di Neyman posteriore bayesiano che si adatta alla regolarità degli esiti sconosciuta tramite un prior gerarchico. Dimostrano che questo design raggiunge tassi di contrazione vicini all'oracle e un'efficienza dell'ATE ottimale.
2. Espressività del Transformer: Forniscono una prova costruttiva del fatto che i Transformer possono rappresentare la mappa di aggiornamento e design bayesiano. Nello specifico, mostrano come i meccanismi di attenzione e i blocchi di discesa del gradiente proiettata possano realizzare la regressione ridge mascherata a rango finito e come una porta MoE possa implementare l'adattamento della regolarità.
3. Teoria della Generalizzazione: Il documento stabilisce un limite di generalizzazione ERM, dimostrando che il pre-training supervisionato su traiettorie finite consente al Transformer di apprendere la mappa di sperimentazione adattiva, imitando il comportamento storia-propensione dell'insegnante.
4. Validazione Empirica: Gli esperimenti su dati sintetici mostrano che:
   • Un singolo Transformer, addestrato solo su etichette di miscela bayesiana (senza una testa MoE esplicita), raggiunge tassi di minimax adattivi alla regolarità su livelli di regolarità non visti.
   • Il Transformer di design riproduce il comportamento di allocazione di Neyman bayesiano, rispondendo direzionalmente alle varanze residue specifiche per braccio, nonostante non apprenda esplicitamente i momenti degli esiti nei suoi stati latenti.
   • La policy appresa migliora la precisione della stima dell'ATE (riducendo l'MSE) rispetto alla randomizzazione uniforme e si avvicina alle prestazioni dell'allocazione di Neyman oracle, anche sotto rapporti di varianza out-of-distribution.

Risultati e Significato
Il documento sostiene che i Transformer possono fungere da "sperimentatori bayesiani in-context ammortizzati". Imparando direttamente la mappatura dalla storia sperimentale alla propensione del trattamento, questi modelli evitano la necessità di risolvere un intero problema di design bayesiano al momento dell'impiego.

La significatività risiede nella dimostrazione che:

• L'ammortizzazione è fattibile: Una complessa inferenza bayesiana sequenziale e un design ottimale possono essere distillati in una policy a rete neurale fissa.
• L'adattività emerge: Il modello può adattarsi alla regolarità delle funzioni e alle strutture di varianza sconosciute senza un'ingegneria esplicita del parametro di regolarità o delle pipeline di stima dei momenti.
• L'efficienza è preservata: Le policy apprese raggiungono un'efficienza statistica comparabile all'allocazione di Neyman oracle, validando l'uso dei Transformer per un design sperimentale statisticamente fondato.

Gli autori segnalano i limiti, affermando che lo studio attuale è limitato a esperimenti a due bracci e ambienti sintetici, e che l'estensione di tale approccio a spazi di azione più ricchi, esiti ritardati e log del mondo reale rimane una direzione futura. L'impatto più ampio è inquadrato come un miglioramento dell'efficienza dei dati in domini come la salute digitale e i servizi online.