Efficient Autoregressive Inference for Transformer Probabilistic Models
Hassan, Loka, Li et al.
Transformer-based models for amortized probabilistic inference, such as neural processes, prior-fitted networks, and tabular foundation models, excel at single-pass marginal prediction. However, many real-world applications, from signal interpolation to multi-column tabular predictions, require coherent joint distributions that capture dependencies between predictions. While purely autoregressive architectures efficiently generate such distributions, they sacrifice the flexible set-conditioning that makes these models powerful for meta-learning. Conversely, the standard approach to obtain joint distributions from set-based models requires expensive re-encoding of the entire augmented conditioning set at each autoregressive step. We introduce a causal autoregressive buffer that preserves the advantages of both paradigms. Our approach decouples context encoding from updating the conditioning set. The model processes the context once and caches it. A dynamic buffer then captures target dependencies: as targets are incorporated, they enter the buffer and attend to both the cached context and previously buffered targets. This enables efficient batched autoregressive generation and one-pass joint log-likelihood evaluation. A unified training strategy allows seamless integration of set-based and autoregressive modes at minimal additional cost. Across synthetic functions, EEG signals, cognitive models, and tabular data, our method matches predictive accuracy of strong baselines while delivering up to 20 times faster joint sampling. Our approach combines the efficiency of autoregressive generative models with the representational power of set-based conditioning, making joint prediction practical for transformer-based probabilistic models.
academic
Inferenza Autoregressiva Efficiente per Modelli Probabilistici Transformer
I modelli di inferenza probabilistica ammortizzata basati su Transformer (come i Processi Neurali, le Reti Predittive Precedenti e i Modelli Fondamentali Tabulari) eccellono nelle previsioni marginali singole. Tuttavia, molte applicazioni pratiche, dalla interpolazione di segnali alla previsione multi-colonna tabulare, richiedono di catturare distribuzioni congiunte coerenti che modellino le dipendenze tra previsioni. L'architettura puramente autoregressiva può generare efficientemente tali distribuzioni, ma sacrifica la capacità di condizionamento flessibile dell'insieme che rende questi modelli potenti nell'apprendimento meta. Al contrario, il metodo standard per ottenere distribuzioni congiunte dai modelli basati su insiemi richiede una ricodifica costosa dell'intero insieme di condizionamento aumentato ad ogni passo autoregressivo. Questo articolo introduce il buffer autoregressivo causale, che preserva i vantaggi di entrambi i paradigmi. Il metodo disaccoppia la codifica del contesto dall'aggiornamento dell'insieme di condizionamento, con il modello che elabora il contesto una volta e lo memorizza nella cache, mentre un buffer dinamico cattura le dipendenze tra target. Su funzioni sintetiche, segnali EEG, modelli cognitivi e dati tabulari, il metodo raggiunge accelerazioni fino a 20 volte nel campionamento congiunto mantenendo al contempo l'accuratezza predittiva di forti baseline.
I modelli probabilistici basati su Transformer affrontano un collo di bottiglia di efficienza fondamentale: quando è necessario generare una distribuzione congiunta, l'intero insieme di condizionamento deve essere ricodificato ad ogni passo autoregressivo. Nello specifico:
Limitazioni dei modelli con condizionamento su insiemi: I Processi Neurali (NP), le Reti Predittive Precedenti (PFN) e modelli simili eccellono nelle previsioni marginali, ma durante il dispiegamento autoregressivo richiedono ricodifiche ripetute del contesto, risultando in una complessità computazionale O(K(N+K)²)
Insufficienze dei modelli puramente autoregressivi: Sebbene computazionalmente efficienti, mancano della capacità di condizionamento flessibile su insiemi, limitando l'applicazione nei compiti di apprendimento meta
Propone il meccanismo del buffer autoregressivo causale: Disaccoppia la codifica del contesto del condizionamento su insiemi dalla previsione sequenziale, realizzando campionamento congiunto e valutazione della verosimiglianza efficienti
Progetta una strategia di addestramento unificata: Utilizza maschere di attenzione e apprendimento curricolare della dimensione del buffer, consentendo a un singolo modello di apprendere entrambe le modalità operative con costo aggiuntivo minimo
Verifica l'applicabilità diffusa: Realizza accelerazioni fino a 20 volte nel campionamento congiunto su TNP/PFN e modelli fondamentali tabulari, mantenendo al contempo accuratezza predittiva comparabile
Ottimizzazione della complessità teorica: Riduce la complessità computazionale da O(K(N+K)²) a O(N²+NK+K²)
Dato un insieme di contesto C = {(xₙ, yₙ)}ᴺₙ₌₁ e un insieme di target T = {(xₘ, yₘ)}ᴹₘ₌₁, l'obiettivo è apprendere la distribuzione predittiva p_θ(y₁:ₘ|x₁:ₘ; C), dove θ sono i parametri del modello.
Codificatore di contesto rC: Elabora coppie di contesto utilizzando auto-attenzione multi-testa bidirezionale, memorizzando nella cache le coppie chiave-valore ad ogni livello
Codificatore di buffer rB: Utilizza auto-attenzione multi-testa rigorosamente causale sul prefisso del buffer
Decodificatore di target rtgt: Interroga il contesto memorizzato nella cache e il prefisso del buffer visibile tramite attenzione incrociata
Impatto della dimensione del buffer: Con K=1 equivale all'autoregressione standard, con K=16 leggera riduzione delle prestazioni ma accelerazione significativa della velocità
Kernel Triton personalizzati: Forniscono accelerazione significativa con batch di grandi dimensioni
Modalità di attenzione: Anche disabilitando FlashAttention, TNP-A rimane più lento di ordini di grandezza rispetto ad altri metodi
Questo metodo funziona come componente modulare integrabile nelle architetture NP/PFN esistenti. Complementare ai lavori precedenti focalizzati sulla scalabilità dell'insieme di contesto, questo articolo affronta l'efficienza del campionamento congiunto autoregressivo.
Costruisce sulla tendenza di inquadrare l'inferenza bayesiana come compiti di apprendimento contestuale, sfruttando varianti NP e PFN basate su Transformer.
Correlato a TNP-A ma con differenze chiave: TNP-A utilizza ripetizione dei target sia nell'addestramento che nell'inferenza, mentre questo metodo la richiede solo nella valutazione della verosimiglianza.
Bruinsma et al. (2023): Autoregressive Conditional Neural Processes
Jingang et al. (2025): Modello Fondamentale Tabulare TabICL
Valutazione Complessiva: Questo è un articolo di ricerca di alta qualità che eccelle nell'innovazione teorica, verifica sperimentale e implementazione ingegneristica. Il metodo risolve con successo un importante collo di bottiglia di efficienza nei modelli probabilistici, con ampi prospettive di applicazione e valore accademico.