让概率模型高效生成联合预测,速度提升20倍。
Efficient Autoregressive Inference for Transformer Probabilistic Models
- 用缓存机制只编码一次上下文,后续预测逐步添加到因果缓冲区。
- 在合成函数、脑电数据等任务上,速度比传统方法快20倍,内存降低7倍。
- 适合需要快速联合采样和密度评估的场景,如贝叶斯推断与表格回归。
用于近似概率推断和元学习的基于集合的Transformer模型(如神经过程、先验拟合网络、表格基础模型)在单次通过的边缘预测上表现优异。然而,许多应用需要多个预测的联合分布。纯自回归架构虽能高效生成联合分布,但牺牲了灵活的集合条件建模能力。从基于集合的模型获取联合分布需在每个自回归步骤重新编码完整上下文,效率低下。本文提出一种因果自回归缓冲区,结合两种范式的优点:仅编码一次上下文并缓存;轻量级因果缓冲区捕获生成目标间的依赖关系,每次新预测同时关注缓存的上下文和所有先前预测结果。该方法实现高效的批量自回归采样与联合预测密度评估。训练通过掩码注意力融合基于集合与自回归模式,开销极小。在合成函数、脑电时间序列、贝叶斯模型比较和表格回归任务中,性能接近全上下文重编码方案,同时实现最高20倍的联合采样与密度评估加速,内存消耗降低最高7倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。