2025-11-25T22:34:18.624435

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

Inférence Autorégressive Efficace pour les Modèles Probabilistes Transformer

Informations Fondamentales

  • ID de l'article: 2510.09477
  • Titre: Efficient Autoregressive Inference for Transformer Probabilistic Models
  • Auteurs: Conor Hassan, Nasrulloh Loka, Cen-You Li, Daolang Huang, Paul E. Chang, Yang Yang, Francesco Silvestrin, Samuel Kaski, Luigi Acerbi
  • Classification: stat.ML cs.LG
  • Date de publication: 10 octobre 2025 (prépublication arXiv)
  • Lien de l'article: https://arxiv.org/abs/2510.09477

Résumé

Les modèles d'inférence probabiliste amortie basés sur Transformer (tels que les Processus Neuraux, les Réseaux Pré-ajustés et les Modèles Fondamentaux Tabulaires) excellent dans les prédictions marginales uniques. Cependant, de nombreuses applications pratiques, allant de l'interpolation de signaux à la prédiction multi-colonnes tabulaire, nécessitent de capturer les dépendances entre prédictions dans une distribution conjointe cohérente. Les architectures purement autorégressive peuvent générer efficacement de telles distributions, mais sacrifient la capacité de conditionnement d'ensemble flexible qui rend ces modèles puissants en méta-apprentissage. Inversement, la méthode standard pour obtenir des distributions conjointes à partir de modèles basés sur des ensembles nécessite un réencodage coûteux de l'ensemble de conditions augmenté à chaque étape autorégressive. Cet article introduit le Tampon Autorégressif Causal, qui préserve les avantages des deux paradigmes. La méthode découple l'encodage du contexte de la mise à jour de l'ensemble de conditions, le modèle traitant le contexte une seule fois et le mettant en cache, tandis qu'un tampon dynamique capture les dépendances entre cibles. Sur des fonctions synthétiques, des signaux EEG, des modèles cognitifs et des données tabulaires, la méthode atteint une accélération de l'échantillonnage conjoint jusqu'à 20 fois tout en correspondant à la précision de prédiction des bases de référence fortes.

Contexte et Motivation de la Recherche

Problème Central

Les modèles probabilistes basés sur Transformer existants font face à un goulot d'étranglement d'efficacité fondamental : lors de la génération de distributions conjointes, l'ensemble de conditions complet doit être réencodé à chaque étape autorégressive. Spécifiquement :

  1. Limitations des modèles avec conditionnement d'ensemble: Les Processus Neuraux (NP), les Réseaux Pré-ajustés (PFN) et autres modèles excèlent dans les prédictions marginales, mais lors du déploiement autorégressif, ils nécessitent un réencodage répété du contexte, entraînant une complexité computationnelle de O(K(N+K)²)
  2. Insuffisances des modèles purement autorégressifs: Bien que computationnellement efficaces, ils manquent de capacité de conditionnement d'ensemble flexible, limitant les applications dans les tâches de méta-apprentissage

Importance

La prédiction de distributions conjointes est cruciale dans plusieurs applications clés :

  • Dépendances temporelles dans l'interpolation de signaux
  • Corrélations de caractéristiques dans la prédiction multi-colonnes tabulaire
  • Dépendances séquentielles dans la modélisation de données comportementales
  • Évaluation de vraisemblance conjointe dans la sélection de modèles bayésiens

Limitations des Méthodes Existantes

  1. Déploiement autorégressif TNP-D: Nécessite un réencodage de l'ensemble de conditions croissant à chaque étape
  2. TNP-A: Nécessite de traiter des ensembles de cibles répétés lors de l'entraînement et de l'inférence, avec des frais généraux computationnels énormes
  3. TNP-ND: Limité aux distributions gaussiennes multivariées, capacité d'expression réduite

Contributions Principales

  1. Proposition du mécanisme Tampon Autorégressif Causal: Découple l'encodage du contexte du conditionnement d'ensemble de la prédiction séquentielle, permettant un échantillonnage conjoint et une évaluation de vraisemblance efficaces
  2. Conception d'une stratégie d'entraînement unifiée: Utilise l'attention masquée et l'apprentissage par curriculum de taille de tampon, permettant à un modèle unique d'apprendre les deux modes opératoires avec un coût supplémentaire minimal
  3. Vérification de l'applicabilité générale: Réalise une accélération de l'échantillonnage conjoint jusqu'à 20 fois sur les TNP/PFN et les Modèles Fondamentaux Tabulaires, tout en maintenant une précision de prédiction comparable
  4. Optimisation de la complexité théorique: Réduit la complexité computationnelle de O(K(N+K)²) à O(N²+NK+K²)

Détails de la Méthode

Définition de la Tâche

Étant donné un ensemble de contexte C = {(xₙ, yₙ)}ᴺₙ₌₁ et un ensemble de cibles T = {(xₘ, yₘ)}ᴹₘ₌₁, l'objectif est d'apprendre la distribution de prédiction p_θ(y₁:ₘ|x₁:ₘ; C), où θ représente les paramètres du modèle.

Architecture du Modèle

Composants Principaux

  1. Encodeur de contexte rC: Traite les paires de contexte, utilisant l'auto-attention multi-têtes bidirectionnelle, mettant en cache les paires clé-valeur à chaque couche
  2. Encodeur de tampon rB: Utilise l'auto-attention multi-têtes strictement causale sur le préfixe du tampon
  3. Décodeur de cible rtgt: Interroge le contexte mis en cache et le préfixe de tampon visible par attention croisée

Paramétrisation de la Distribution de Prédiction

p_θ(y*₁:K|x*₁:K; C) = ∏ᴷₖ₌₁ p_θ(y*ₖ|rtgt(x*ₖ, [rC(C), b₁:ₖ₋₁]))

où bₖ = rB((xₖ, yₖ), rC(C), b₁:ₖ₋₁)

Conception des Masques d'Attention

Implémente quatre exigences clés :

  • (R1) Immuabilité du contexte: Encodé une seule fois et mis en cache en lecture seule
  • (R2) Causalité stricte du tampon: Le token j ne peut accéder qu'aux positions <j
  • (R3) Flux d'information unidirectionnel du contexte: Pas d'écritures arrière vers C
  • (R4) Les cibles accèdent au contexte mis en cache et au préfixe de tampon visible

Points d'Innovation Technique

1. Conception Découplée

  • Cache de contexte statique: Encodé une seule fois, réutilisé plusieurs fois
  • Tampon dynamique: Mise à jour incrémentale, capture les dépendances entre cibles

2. Curriculum d'Entraînement

  • 50% des cibles accèdent uniquement au contexte
  • 50% des cibles accèdent au contexte + préfixe de tampon de longueur aléatoire
  • Assure que le modèle fonctionne bien dans différents états de tampon

3. Modes d'Inférence Efficaces

  • Échantillonnage autorégressif: Pré-remplissage du contexte, décodage séquentiel des cibles
  • Évaluation de vraisemblance conjointe: Calcul de toutes les probabilités conditionnelles en une seule passe avant
  • Échantillonnage par batch: Partage du cache de contexte, états de tampon indépendants

Configuration Expérimentale

Ensembles de Données

  1. Fonctions synthétiques:
    • Processus Gaussien (GP): Noyaux RBF, Matérn-3/2, Matérn-5/2
    • Fonction en dents de scie: Non-gaussienne, dérivées discontinues
  2. Données EEG: 11 520 essais, 122 sujets, 7 canaux pertinents, 256 points temporels
  3. Modèle d'Inférence Causale Multisensorielle: Données d'expériences de localisation audiovisuelle, 15 participants
  4. Données Tabulaires: Ensembles de données UCI (consommation électrique, émissions de turbines à gaz, partage de vélos)

Métriques d'Évaluation

  • Vraisemblance logarithmique moyenne: Évaluation de la qualité de prédiction
  • Temps horloge murale: Temps d'exécution réel pour l'échantillonnage, l'évaluation de vraisemblance, les étapes d'entraînement
  • RMSE de vraisemblance marginale logarithmique: Précision des tâches de sélection de modèles

Méthodes de Comparaison

  • TNP-D-Ind: Prédiction indépendante, rapide mais sans modélisation de dépendances
  • TNP-D-AR: Déploiement autorégressif, expressif mais nécessite réencodage
  • TNP-ND: Distribution conjointe gaussienne multivariée, expressivité limitée
  • TNP-A: Modélisation purement autorégressive, lente en entraînement et échantillonnage

Détails d'Implémentation

  • Optimiseur: Adam, taux d'apprentissage 1×10⁻⁴
  • Architecture: Transformer 6 couches, 4 têtes d'attention, dimension 128
  • Tête de prédiction: Modèle de mélange gaussien à 20 composantes
  • Taille de tampon: K=16 (expériences principales)

Résultats Expérimentaux

Résultats Principaux

Efficacité Computationnelle

  • Échantillonnage autorégressif: 3-20 fois plus rapide que TNP-A et TNP-D-AR
  • Évaluation de vraisemblance: Comparable à TNP-A, K fois plus rapide que TNP-D-AR
  • Vitesse d'entraînement: 4-12 fois plus rapide que TNP-A, comparable aux bases de référence les plus rapides

Précision de Prédiction

Ensemble de DonnéesTNP-D-ARTNP-AMéthode (K=16)Méthode (K=1)
GP2.570.802.512.56
Dents de scie1.05-0.431.001.09
EEG-Int0.510.460.520.54
EEG-For1.07-0.040.851.21

Études d'Ablation

  • Impact de la taille de tampon: K=1 équivaut à l'autorégression standard, K=16 montre une légère baisse de performance mais une accélération massive
  • Noyaux Triton personnalisés: Fournissent une accélération significative avec de grands batches
  • Motifs d'attention: Même avec FlashAttention désactivé, TNP-A reste plusieurs ordres de grandeur plus lent que les autres méthodes

Étude de Cas

Dans la tâche d'inférence causale multisensorielle :

  • Sélection de modèles: RMSE LML de 3.56, proche du 3.47 de TNP-D-AR
  • Prédiction de données: Vraisemblance logarithmique moyenne de -2.76, comparable à toutes les bases de référence fortes
  • Corrélation avec les vraies valeurs: R²=1.00 (LML), R²=0.92 (ΔLML)

Travaux Connexes

Processus Neuraux et Réseaux Pré-ajustés

La méthode proposée fonctionne comme un composant modulaire pouvant s'intégrer aux architectures NP/PFN existantes. Elle complète les travaux antérieurs se concentrant sur l'évolutivité de l'ensemble de contexte, en s'attaquant à l'efficacité de l'échantillonnage conjoint autorégressif.

Modèles Probabilistes Transformer

S'appuie sur la tendance à encadrer l'inférence bayésienne comme des tâches d'apprentissage contextuel, exploitant les variantes NP et PFN basées sur Transformer.

Modèles Fondamentaux Tabulaires

S'intègre naturellement aux modèles comme TabPFN et TabICL, fournissant un module complémentaire pour la prédiction conjointe efficace.

Estimation de Densité Conjointe Autorégressive

Connexe à TNP-A mais avec des différences clés: TNP-A utilise la répétition de cibles lors de l'entraînement et de l'inférence, tandis que cette méthode ne l'exige que lors de l'évaluation de vraisemblance.

Conclusions et Discussion

Conclusions Principales

  1. Percée d'efficacité: Intègre avec succès l'efficacité du Transformer autorégressif dans le cadre NP/PFN
  2. Maintien de la performance: Préserve la précision de prédiction tout en augmentant considérablement la vitesse
  3. Applicabilité générale: Valide l'efficacité de la méthode dans plusieurs domaines et tâches

Limitations

  1. Extension de la longueur de tampon: Présente toujours un terme O(K²) lorsque K augmente, utilisant actuellement des plongements de position fixes
  2. Dérive de qualité pour longs tampons: Peut présenter une dégradation de qualité par rapport à l'autorégression exacte avec réencodage à chaque étape
  3. Consommation mémoire: Nécessite de maintenir un cache de contexte et des états de tampon

Directions Futures

  1. Amélioration du codage de position: Utiliser RoPE ou ALiBi pour supporter des séquences plus longues
  2. Décodage spéculatif: Stratégies d'inférence adaptatives s'inspirant des processus draft-verify
  3. Ajustement fin efficace en paramètres: Utiliser des adaptateurs ou LoRA pour ajouter des capacités de tampon aux modèles pré-entraînés

Évaluation Approfondie

Points Forts

  1. Innovation forte: Résout élégamment le compromis entre conditionnement d'ensemble et efficacité autorégressive
  2. Fondation théorique solide: Fournit une analyse de complexité claire et des dérivations mathématiques
  3. Expérimentation complète: Couvre données synthétiques, données réelles, multiples domaines d'application
  4. Optimisation d'ingénierie: Inclut des optimisations bas niveau comme les noyaux CUDA personnalisés
  5. Reproductibilité: Fournit des détails d'implémentation détaillés et s'engage à ouvrir le code source

Insuffisances

  1. Portée d'application: Principalement adapté aux séquences de cibles de longueur modérée, les séquences très longues restent problématiques
  2. Analyse théorique: Manque d'analyse des limites théoriques de l'erreur d'approximation du tampon
  3. Expériences comparatives: N'a pas comparé avec les mécanismes d'attention efficaces les plus récents (comme l'attention linéaire)

Impact

  1. Valeur académique: Fournit de nouvelles perspectives pour l'inférence efficace dans les modèles probabilistes
  2. Valeur pratique: Réduit significativement les coûts computationnels de la prédiction conjointe, rendant les applications pratiques possibles
  3. Extensibilité: La méthode possède une bonne généralité, applicable à diverses variantes de Transformer

Scénarios d'Application

  • Applications nécessitant un échantillonnage conjoint fréquent (comme la quantification d'incertitude)
  • Tâches de prédiction séquentielle avec contexte à grande échelle
  • Scénarios avec des exigences élevées d'inférence en temps réel
  • Modélisation conjointe de données multimodales

Références

Les principales références incluent :

  • Garnelo et al. (2018): Article original sur les Processus Neuraux
  • Nguyen & Grover (2022): Transformer Neural Processes
  • Müller et al. (2022): Prior-Fitted Networks
  • Bruinsma et al. (2023): Autoregressive Conditional Neural Processes
  • Jingang et al. (2025): Modèles Fondamentaux Tabulaires TabICL

Évaluation Générale: Cet article de recherche de haute qualité excelle dans l'innovation théorique, la vérification expérimentale et l'implémentation d'ingénierie. La méthode résout avec succès un goulot d'étranglement d'efficacité important dans les modèles probabilistes, possédant des perspectives d'application larges et une valeur académique significative.