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.

翻译：神经过程（NP）是一类快速发展的模型，旨在直接建模随机过程的后验预测分布。该类模型最初被设计为高斯过程（GP）的可扩展替代方案——由于GP受限于$O(n^3)$的时间复杂度，而最先进的现代NP模型虽能与GP媲美，但其注意力机制仍存在$O(n^2)$的计算瓶颈。我们提出变换器神经过程-核回归（TNP-KR），一种可扩展的神经过程模型，具有以下特点：（1）核回归块（KRBlock），一种简单、可扩展且参数高效的变换器块，其复杂度为$O(n_c^2 + n_c n_t)$，其中$n_c$和$n_t$分别表示上下文点和测试点的数量；（2）基于核的注意力偏置；（3）两种新型注意力机制：扫描注意力（SA），一种基于内存高效扫描的注意力机制，其与核偏置结合可使TNP-KR具备平移不变性；以及深度核注意力（DKA），一种采用Performer风格的注意力机制，可隐式融入距离偏置，进一步将复杂度降至$O(n_c)$。这些改进使得两种TNP-KR变体能够在单块24GB GPU上，于1分钟内对超过100万测试点完成包含10万上下文点的推理。在涵盖元回归、贝叶斯优化、图像补全和流行病学等基准测试中，采用DKA的TNP-KR几乎在所有基准上均优于其对应的Performer版本，而采用SA的TNP-KR则取得了最先进的结果。