提出新型神经过程模型,实现超大规模上下文推理与高精度预测。
Transformer Neural Processes - Kernel Regression
- 采用核回归模块与新型注意力机制,降低计算复杂度至O(nc)。
- 支持10万上下文点、百万测试点的快速推理,单卡1分钟内完成。
- 适用于图像补全、贝叶斯优化等任务,性能超越现有方法。
神经过程(NPs)是一类快速发展的模型,旨在直接建模随机过程的后验预测分布。最初作为高斯过程(GPs)的可扩展替代方案提出,后者受O(n³)运行时复杂度限制。尽管现代最准确的NPs常可媲美GPs,但仍因注意力机制存在O(n²)瓶颈。我们提出Transformer神经过程-核回归(TNP-KR),具备:(1) 核回归块(KRBlock),一种简单、可扩展且参数高效的Transformer块,复杂度为O(nc² + ncnₜ),其中nc和nt分别为上下文点和测试点数量;(2) 基于核的注意力偏置;(3) 两种新注意力机制:扫描注意力(SA),一种内存高效扫描式注意力,结合核偏置可实现平移不变性;深度核注意力(DKA),类似Performer的注意力机制,隐式包含距离偏置,进一步将复杂度降至O(nc)。这些改进使两种TNP-KR变体均能在单个24GB GPU上对超过10⁵个上下文点和超过10⁶个测试点实现推理,耗时不足一分钟。在元回归、贝叶斯优化、图像补全及流行病学等多个基准测试中,采用DKA的TNP-KR几乎在所有任务上优于其Performer基线,而采用SA的TNP-KR达到当前最佳性能。
原文摘要 · Abstract (English)
Neural Processes (NPs) are a rapidly evolving class of models designed to directly model the posterior predictive distribution of stochastic processes. Originally developed as a scalable alternative to Gaussian Processes (GPs), which are limited by $O(n^3)$ runtime complexity, the most accurate modern NPs can often rival GPs but still suffer from an $O(n^2)$ bottleneck due to their attention mechanism. We introduce the Transformer Neural Process - Kernel Regression (TNP-KR), a scalable NP featuring: (1) a Kernel Regression Block (KRBlock), a simple, extensible, and parameter efficient transformer block with complexity $O(n_c^2 + n_c n_t)$, where $n_c$ and $n_t$ are the number of context and test points, respectively; (2) a kernel-based attention bias; and (3) two novel attention mechanisms: scan attention (SA), a memory-efficient scan-based attention that when paired with a kernel-based bias can make TNP-KR translation invariant, and deep kernel attention (DKA), a Performer-style attention that implicitly incoporates a distance bias and further reduces complexity to $O(n_c)$. These enhancements enable both TNP-KR variants to perform inference with 100K context points on over 1M test points in under a minute on a single 24GB GPU. On benchmarks spanning meta regression, Bayesian optimization, image completion, and epidemiology, TNP-KR with DKA outperforms its Performer counterpart on nearly every benchmark, while TNP-KR with SA achieves state-of-the-art results.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。