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

Inferencia Autorregresiva Eficiente para Modelos Probabilísticos Transformer

Información Básica

  • ID del Artículo: 2510.09477
  • Título: Efficient Autoregressive Inference for Transformer Probabilistic Models
  • Autores: Conor Hassan, Nasrulloh Loka, Cen-You Li, Daolang Huang, Paul E. Chang, Yang Yang, Francesco Silvestrin, Samuel Kaski, Luigi Acerbi
  • Clasificación: stat.ML cs.LG
  • Fecha de Publicación: 10 de octubre de 2025 (preimpresión arXiv)
  • Enlace del Artículo: https://arxiv.org/abs/2510.09477

Resumen

Los modelos de inferencia probabilística amortizada basados en Transformer (como Procesos Neurales, Redes Ajustadas Previas y Modelos Fundamentales Tabulares) demuestran un excelente desempeño en predicciones marginales únicas. Sin embargo, muchas aplicaciones prácticas, desde la interpolación de señales hasta la predicción de múltiples columnas tabulares, requieren capturar distribuciones conjuntas coherentes que modelan dependencias entre predicciones. Las arquitecturas puramente autorregresivas pueden generar eficientemente tales distribuciones, pero sacrifican la capacidad flexible de acondicionamiento por conjuntos que hace que estos modelos sean potentes en meta-aprendizaje. Por el contrario, el método estándar para obtener distribuciones conjuntas de modelos basados en conjuntos requiere una recodificación costosa del conjunto de condiciones aumentado completo en cada paso autorregresivo. Este artículo introduce el búfer autorregresivo causal, que retiene las ventajas de ambos paradigmas. El método desacopla la codificación de contexto de las actualizaciones del conjunto de condiciones, permitiendo que el modelo procese el contexto una sola vez y lo almacene en caché, mientras que un búfer dinámico captura las dependencias entre objetivos. En funciones sintéticas, señales EEG, modelos cognitivos y datos tabulares, el método logra aceleraciones de hasta 20 veces en la velocidad de muestreo conjunto mientras mantiene la precisión de predicción comparable a líneas base sólidas.

Antecedentes de Investigación y Motivación

Problema Central

Los modelos probabilísticos basados en Transformer existentes enfrentan un cuello de botella de eficiencia fundamental: cuando es necesario generar una distribución conjunta, se debe recodificar todo el conjunto de condiciones en cada paso autorregresivo. Específicamente:

  1. Limitaciones de los modelos acondicionados por conjuntos: Los Procesos Neurales (NP), Redes Ajustadas Previas (PFN) y modelos similares sobresalen en predicción marginal, pero requieren recodificación repetida del contexto durante el despliegue autorregresivo, resultando en una complejidad computacional O(K(N+K)²)
  2. Insuficiencias de los modelos puramente autorregresivos: Aunque son computacionalmente eficientes, carecen de capacidad flexible de acondicionamiento por conjuntos, limitando su aplicación en tareas de meta-aprendizaje

Importancia

La predicción de distribuciones conjuntas es crítica en múltiples aplicaciones clave:

  • Dependencias temporales en interpolación de señales
  • Correlaciones de características en predicción de múltiples columnas
  • Dependencias secuenciales en modelado de datos de comportamiento
  • Evaluación de verosimilitud conjunta en selección bayesiana de modelos

Limitaciones de Métodos Existentes

  1. Despliegue autorregresivo TNP-D: Requiere recodificación del conjunto de condiciones en crecimiento en cada paso
  2. TNP-A: Requiere procesamiento de conjuntos de objetivos repetidos tanto en entrenamiento como en inferencia, con gastos computacionales enormes
  3. TNP-ND: Limitado a distribuciones gaussianas multivariadas, con capacidad expresiva reducida

Contribuciones Principales

  1. Propone el mecanismo de búfer autorregresivo causal: Desacopla la codificación de contexto de acondicionamiento por conjuntos de la predicción secuencial, permitiendo muestreo conjunto eficiente y evaluación de verosimilitud
  2. Diseña una estrategia de entrenamiento unificada: Utiliza enmascaramiento de atención y aprendizaje por currículo de tamaño de búfer, permitiendo que un único modelo aprenda ambos modos de operación con costo adicional mínimo
  3. Verifica aplicabilidad amplia: Logra aceleración de muestreo conjunto de hasta 20 veces en TNP/PFN y Modelos Fundamentales Tabulares, manteniendo precisión de predicción comparable
  4. Optimiza complejidad teórica: Reduce la complejidad computacional de O(K(N+K)²) a O(N²+NK+K²)

Explicación Detallada del Método

Definición de Tarea

Dado un conjunto de contexto C = {(xₙ, yₙ)}ᴺₙ₌₁ y un conjunto de objetivos T = {(xₘ, yₘ)}ᴹₘ₌₁, el objetivo es aprender la distribución de predicción p_θ(y₁:ₘ|x₁:ₘ; C), donde θ son los parámetros del modelo.

Arquitectura del Modelo

Componentes Principales

  1. Codificador de contexto rC: Procesa pares de contexto, utilizando autoatención multi-cabeza bidireccional, almacenando en caché pares clave-valor en cada capa
  2. Codificador de búfer rB: Aplica autoatención multi-cabeza estrictamente causal al prefijo del búfer
  3. Decodificador de objetivo rtgt: Consulta el contexto almacenado en caché y el prefijo de búfer visible mediante atención cruzada

Parametrización de Distribución de Predicción

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

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

Diseño de Máscaras de Atención

Implementa cuatro requisitos clave:

  • (R1) Contexto inmutable: Codificado una sola vez y almacenado en caché como solo lectura
  • (R2) Búfer estrictamente causal: El token j solo puede atender posiciones <j
  • (R3) Flujo de información unidireccional desde contexto: Sin escritura de bordes en C
  • (R4) El objetivo atiende al contexto almacenado en caché y al prefijo de búfer visible

Puntos de Innovación Técnica

1. Diseño Desacoplado

  • Caché de contexto estático: Codificado una sola vez, reutilizado múltiples veces
  • Búfer dinámico: Actualización incremental, captura dependencias entre objetivos

2. Currículo de Entrenamiento

  • 50% de objetivos atienden solo al contexto
  • 50% de objetivos atienden al contexto + prefijo de búfer de longitud aleatoria
  • Asegura que el modelo funcione bien en diferentes estados de búfer

3. Modos de Inferencia Eficiente

  • Muestreo autorregresivo: Precarga contexto, decodifica secuencialmente objetivos
  • Evaluación de verosimilitud conjunta: Pase directo único calcula todas las probabilidades condicionales
  • Muestreo por lotes: Comparte caché de contexto, estados de búfer independientes

Configuración Experimental

Conjuntos de Datos

  1. Funciones Sintéticas:
    • Procesos Gaussianos (GP): núcleos RBF, Matérn-3/2, Matérn-5/2
    • Función Diente de Sierra: no gaussiana, derivada discontinua
  2. Datos EEG: 11,520 ensayos, 122 sujetos, 7 canales relevantes, 256 puntos de tiempo
  3. Modelo de Inferencia Causal Multisensorial: Datos de experimento de localización audiovisual, 15 participantes
  4. Datos Tabulares: Conjuntos de datos UCI (consumo de energía, emisiones de turbina de gas, uso compartido de bicicletas)

Métricas de Evaluación

  • Verosimilitud logarítmica promedio: Evalúa calidad de predicción
  • Tiempo de reloj de pared: Tiempo de ejecución real para muestreo, evaluación de verosimilitud, pasos de entrenamiento
  • RMSE de verosimilitud marginal logarítmica: Precisión en tareas de selección de modelos

Métodos de Comparación

  • TNP-D-Ind: Predicción independiente, rápida pero sin modelado de dependencias
  • TNP-D-AR: Despliegue autorregresivo, expresivo pero requiere recodificación
  • TNP-ND: Distribución conjunta gaussiana multivariada, expresividad limitada
  • TNP-A: Modelado completamente autorregresivo, lento en entrenamiento y muestreo

Detalles de Implementación

  • Optimizador: Adam, tasa de aprendizaje 1×10⁻⁴
  • Arquitectura: Transformer de 6 capas, 4 cabezas de atención, dimensión 128
  • Cabeza de predicción: Modelo de mezcla gaussiana de 20 componentes
  • Tamaño de búfer: K=16 (experimentos principales)

Resultados Experimentales

Resultados Principales

Eficiencia Computacional

  • Muestreo autorregresivo: 3-20 veces más rápido que TNP-A y TNP-D-AR
  • Evaluación de verosimilitud: Comparable a TNP-A, K veces más rápido que TNP-D-AR
  • Velocidad de entrenamiento: 4-12 veces más rápido que TNP-A, comparable a línea base más rápida

Precisión de Predicción

Conjunto de DatosTNP-D-ARTNP-AEste Método (K=16)Este Método (K=1)
GP2.570.802.512.56
Diente de Sierra1.05-0.431.001.09
EEG-Int0.510.460.520.54
EEG-For1.07-0.040.851.21

Experimentos de Ablación

  • Impacto del tamaño de búfer: K=1 equivale a autorregresión estándar, K=16 muestra ligera disminución de rendimiento pero aceleración significativa de velocidad
  • Núcleos Triton personalizados: Proporcionan aceleración significativa en lotes grandes
  • Patrones de atención: Incluso con FlashAttention deshabilitado, TNP-A sigue siendo órdenes de magnitud más lento que otros métodos

Análisis de Casos

En la tarea de inferencia causal multisensorial:

  • Selección de modelos: RMSE LML de 3.56, cercano a 3.47 de TNP-D-AR
  • Predicción de datos: Verosimilitud logarítmica promedio de -2.76, comparable a todas las líneas base sólidas
  • Correlación con valores reales: R²=1.00 (LML), R²=0.92 (ΔLML)

Trabajo Relacionado

Procesos Neurales y Redes Ajustadas Previas

Este método funciona como componente modular que puede integrarse en arquitecturas NP/PFN existentes. Complementa trabajo previo enfocado en escalabilidad de conjuntos de contexto, abordando eficiencia de muestreo conjunto autorregresivo.

Modelos Probabilísticos Transformer

Construye sobre la tendencia de enmarcar inferencia bayesiana como tareas de aprendizaje en contexto, aprovechando variantes NP y PFN basadas en Transformer.

Modelos Fundamentales Tabulares

Se integra naturalmente con modelos como TabPFN y TabICL, proporcionando módulo complementario para predicción conjunta eficiente.

Estimación de Densidad Conjunta Autorregresiva

Relacionado a TNP-A pero con diferencias clave: TNP-A usa repetición de objetivos tanto en entrenamiento como inferencia, mientras este método solo la requiere en evaluación de verosimilitud.

Conclusiones y Discusión

Conclusiones Principales

  1. Avance en Eficiencia: Integra exitosamente la eficiencia de Transformer autorregresivo en marco NP/PFN
  2. Mantenimiento de Rendimiento: Mantiene precisión de predicción mientras acelera significativamente la velocidad
  3. Aplicabilidad Amplia: Valida efectividad del método en múltiples dominios y tareas

Limitaciones

  1. Escalado de Longitud de Búfer: Aún contiene término O(K²) cuando K aumenta, actualmente usa incrustaciones de posición fijas
  2. Desviación de Calidad en Búferes Largos: Posible degradación de calidad comparado con autorregresión exacta que recodifica en cada paso
  3. Ocupación de Memoria: Requiere mantener caché de contexto y estado de búfer

Direcciones Futuras

  1. Mejora de Codificación de Posición: Usar RoPE o ALiBi para soportar secuencias más largas
  2. Decodificación Especulativa: Estrategias de inferencia adaptativa inspiradas en procesos draft-verify
  3. Ajuste Fino Eficiente en Parámetros: Usar adaptadores o LoRA para agregar funcionalidad de búfer a modelos preentrenados

Evaluación Profunda

Fortalezas

  1. Innovación Fuerte: Resuelve ingeniosamente el compromiso entre acondicionamiento por conjuntos y eficiencia autorregresiva
  2. Fundamento Teórico Sólido: Proporciona análisis de complejidad claro y derivaciones matemáticas
  3. Experimentación Integral: Cubre datos sintéticos, datos reales, múltiples dominios de aplicación
  4. Optimización de Ingeniería: Incluye optimizaciones de bajo nivel como núcleos CUDA personalizados
  5. Reproducibilidad: Proporciona detalles de implementación detallados y código de código abierto

Deficiencias

  1. Rango de Aplicabilidad: Principalmente aplicable a secuencias de objetivos de longitud media, secuencias muy largas aún enfrentan desafíos
  2. Análisis Teórico: Carece de límites teóricos en análisis de error de aproximación de búfer
  3. Experimentos de Comparación: No compara con mecanismos de atención eficientes más recientes (como atención lineal)

Impacto

  1. Valor Académico: Proporciona nuevas perspectivas para inferencia eficiente en modelos probabilísticos
  2. Valor Práctico: Reduce significativamente costo computacional de predicción conjunta, haciendo aplicaciones prácticas viables
  3. Escalabilidad: El método tiene buena generalidad, aplicable a múltiples variantes de Transformer

Escenarios Aplicables

  • Aplicaciones que requieren muestreo conjunto frecuente (como cuantificación de incertidumbre)
  • Tareas de predicción secuencial con contexto a gran escala
  • Escenarios con requisitos altos de inferencia en tiempo real
  • Modelado conjunto de datos multimodales

Referencias

Las referencias principales incluyen:

  • Garnelo et al. (2018): Artículo original de Procesos Neurales
  • Nguyen & Grover (2022): Procesos Neurales Transformer
  • Müller et al. (2022): Redes Ajustadas Previas
  • Bruinsma et al. (2023): Procesos Neurales Condicionales Autorregresivos
  • Jingang et al. (2025): Modelo Fundamental Tabular TabICL

Evaluación General: Este es un artículo de investigación de alta calidad que demuestra excelencia en innovación teórica, verificación experimental e implementación de ingeniería. El método resuelve exitosamente un cuello de botella de eficiencia importante en modelos probabilísticos, con amplias perspectivas de aplicación y valor académico.