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
Effiziente autoregressives Inferenz für Transformer-Wahrscheinlichkeitsmodelle
Auf Transformer basierende amortisierte probabilistische Inferenzmodelle (wie neuronale Prozesse, vorhergehend angepasste Netzwerke und tabellarische Grundmodelle) zeigen hervorragende Leistungen bei einzelnen Randvorhersagen. Viele praktische Anwendungen – von Signalinterpolation bis zur Vorhersage mehrerer Spalten – erfordern jedoch kohärente Gelenkverteilungen, die Abhängigkeiten zwischen Vorhersagen erfassen. Rein autoregressives Architekturen können solche Verteilungen effizient generieren, opfern aber die flexible Mengenkonditionalität, die diese Modelle im Meta-Learning stark macht. Umgekehrt erfordert die Standardmethode zur Gewinnung von Gelenkverteilungen aus mengenbasierten Modellen eine teure Neukodierung des gesamten erweiterten Bedingungssatzes bei jedem autoregressiven Schritt. Dieses Paper führt kausale autoregressives Puffer ein, die die Vorteile beider Paradigmen bewahren. Die Methode entkoppelt die Kontextkodierung von der Bedingungssatzaktualisierung – das Modell verarbeitet den Kontext einmal und speichert ihn zwischen, während ein dynamischer Puffer Zielabhängigkeiten erfasst. Bei synthetischen Funktionen, EEG-Signalen, kognitiven Modellen und tabellarischen Daten erreicht die Methode bis zu 20-fache Beschleunigung beim gemeinsamen Sampling bei gleichzeitiger Beibehaltung der Vorhersagegenauigkeit starker Baselines.
Bestehende auf Transformer basierende Wahrscheinlichkeitsmodelle sehen sich einem grundlegenden Effizienzengpass gegenüber: Bei der Generierung von Gelenkverteilungen muss der gesamte Bedingungssatz bei jedem autoregressiven Schritt neu kodiert werden. Konkret:
Einschränkungen mengenkonditionalter Modelle: Neuronale Prozesse (NPs), vorhergehend angepasste Netzwerke (PFNs) und ähnliche Modelle sind hervorragend bei Randvorhersagen, erfordern aber bei autoregressiver Bereitstellung wiederholte Neukodierung des Kontexts, was zu einer Rechenkomplexität von O(K(N+K)²) führt
Unzulänglichkeiten rein autoregressiver Modelle: Obwohl rechnerisch effizient, fehlt ihnen die flexible Mengenkonditionalität, was ihre Anwendung bei Meta-Learning-Aufgaben einschränkt
Einführung des kausalen autoregressiven Puffermechanismus: Entkopplung der Kontextkodierung von Mengenkonditionalität und Sequenzvorhersage für effizientes gemeinsames Sampling und Wahrscheinlichkeitsbewertung
Entwurf einer einheitlichen Trainingsstrategie: Verwendung von maskierter Aufmerksamkeit und Puffergröße-Curriculum-Learning, das es einem einzelnen Modell ermöglicht, beide Betriebsmodi mit minimalen zusätzlichen Kosten zu erlernen
Validierung breiter Anwendbarkeit: Erreichung von bis zu 20-facher Beschleunigung beim gemeinsamen Sampling auf TNPs/PFNs und tabellarischen Grundmodellen bei Beibehaltung vergleichbarer Vorhersagegenauigkeit
Optimierung der theoretischen Komplexität: Reduzierung der Rechenkomplexität von O(K(N+K)²) auf O(N²+NK+K²)
Gegeben ein Kontextsatz C = {(xₙ, yₙ)}ᴺₙ₌₁ und ein Zielsatz T = {(xₘ, yₘ)}ᴹₘ₌₁ besteht das Ziel darin, die Vorhersageverteilung p_θ(y₁:ₘ|x₁:ₘ; C) zu erlernen, wobei θ die Modellparameter sind.
Kontextkodierungsmodul rC: Verarbeitet Kontextpaare mit bidirektionaler Multi-Head-Selbstaufmerksamkeit und speichert Schlüssel-Wert-Paare bei jedem Layer
Pufferkodierungsmodul rB: Verwendet strikte kausale Multi-Head-Selbstaufmerksamkeit für das Pufferprefix
Zieldekodermodul rtgt: Fragt zwischengespeicherten Kontext und sichtbares Pufferprefix durch Kreuzaufmerksamkeit ab
Diese Methode fungiert als modulare Komponente, die in bestehende NP/PFN-Architekturen integriert werden kann. Sie ergänzt frühere Arbeiten, die sich auf die Skalierbarkeit von Kontextsätzen konzentrieren, während dieses Paper die Effizienz des autoregressiven Gelenksampling adressiert.
Verwandt mit TNP-A, aber mit Schlüsselunterschieden: TNP-A verwendet Zielwiederholung sowohl beim Training als auch bei der Inferenz, während diese Methode dies nur bei der Wahrscheinlichkeitsbewertung benötigt.
Bruinsma et al. (2023): Autoregressive Conditional Neural Processes
Jingang et al. (2025): TabICL tabellarische Grundmodelle
Gesamtbewertung: Dies ist ein hochqualitatives Forschungspaper mit hervorragenden Leistungen in theoretischer Innovation, experimenteller Validierung und technischer Implementierung. Die Methode löst erfolgreich einen wichtigen Effizienzengpass in Wahrscheinlichkeitsmodellen und hat breite Anwendungsperspektiven und akademischen Wert.