Efficient Autoregressive Inference for Transformer Probabilistic Models

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Hassan, Conor, Loka, Nasrulloh, Li, Cen-You, Huang, Daolang, Chang, Paul E., Yang, Yang, Silvestrin, Francesco, Kaski, Samuel, Acerbi, Luigi
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866915946070802432
author Hassan, Conor
Loka, Nasrulloh
Li, Cen-You
Huang, Daolang
Chang, Paul E.
Yang, Yang
Silvestrin, Francesco
Kaski, Samuel
Acerbi, Luigi
author_facet Hassan, Conor
Loka, Nasrulloh
Li, Cen-You
Huang, Daolang
Chang, Paul E.
Yang, Yang
Silvestrin, Francesco
Kaski, Samuel
Acerbi, Luigi
contents Set-based transformer models for amortized probabilistic inference and meta-learning, such as neural processes, prior-fitted networks, and tabular foundation models, excel at single-pass marginal prediction. However, many applications require joint distributions over multiple predictions. Purely autoregressive architectures generate these efficiently but sacrifice flexible set-conditioning. Obtaining joint distributions from set-based models requires re-encoding the entire context at each autoregressive step, which scales poorly. We introduce a causal autoregressive buffer that combines the strengths of both paradigms. The model encodes the context once and caches it; a lightweight causal buffer captures dependencies among generated targets, with each new prediction attending to both the cached context and all previously predicted targets added to the buffer. This enables efficient batched autoregressive sampling and joint predictive density evaluation. Training integrates set-based and autoregressive modes through masked attention at minimal overhead. Across synthetic functions, EEG time series, a Bayesian model comparison task, and tabular regression, our method closely matches the performance of full context re-encoding while delivering up to $20\times$ faster joint sampling and density evaluation, and up to $7\times$ lower memory usage.
format Preprint
id arxiv_https___arxiv_org_abs_2510_09477
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Efficient Autoregressive Inference for Transformer Probabilistic Models
Hassan, Conor
Loka, Nasrulloh
Li, Cen-You
Huang, Daolang
Chang, Paul E.
Yang, Yang
Silvestrin, Francesco
Kaski, Samuel
Acerbi, Luigi
Machine Learning
Set-based transformer models for amortized probabilistic inference and meta-learning, such as neural processes, prior-fitted networks, and tabular foundation models, excel at single-pass marginal prediction. However, many applications require joint distributions over multiple predictions. Purely autoregressive architectures generate these efficiently but sacrifice flexible set-conditioning. Obtaining joint distributions from set-based models requires re-encoding the entire context at each autoregressive step, which scales poorly. We introduce a causal autoregressive buffer that combines the strengths of both paradigms. The model encodes the context once and caches it; a lightweight causal buffer captures dependencies among generated targets, with each new prediction attending to both the cached context and all previously predicted targets added to the buffer. This enables efficient batched autoregressive sampling and joint predictive density evaluation. Training integrates set-based and autoregressive modes through masked attention at minimal overhead. Across synthetic functions, EEG time series, a Bayesian model comparison task, and tabular regression, our method closely matches the performance of full context re-encoding while delivering up to $20\times$ faster joint sampling and density evaluation, and up to $7\times$ lower memory usage.
title Efficient Autoregressive Inference for Transformer Probabilistic Models
topic Machine Learning
url https://arxiv.org/abs/2510.09477