# Scalable Spatiotemporal Inference with Biased Scan Attention Transformer Neural Processes

arXiv: https://arxiv.org/abs/2506.09163 · HTML: https://mlsanity.com/paper/2506.09163/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 43 references that are in the map.
- First posted on arXiv 10 Jun 2025.

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Efficient Attention: Attention with Linear Complexities](https://mlsanity.com/paper/1812.01243/) (2018)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro](https://mlsanity.com/paper/1912.11554/) (2019)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [Self-attention Does Not Need O(n2) Memory](https://mlsanity.com/paper/2112.05682/) (2021)
- [Convolutional Conditional Neural Processes](https://mlsanity.com/paper/1910.13556/) (2019)
- [Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling](https://mlsanity.com/paper/2207.04179/) (2022)
- [Leveraging redundancy in attention with Reuse Transformers](https://mlsanity.com/paper/2110.06821/) (2021)
- [Convolutional conditional neural processes for local climate downscaling](https://mlsanity.com/paper/2101.07950/) (2021)
- [Latent Bottlenecked Attentive Neural Processes](https://mlsanity.com/paper/2211.08458/) (2022)
- [Graph Neural Processes for Spatio-Temporal Extrapolation](https://mlsanity.com/paper/2305.18719/) (2023)
- [LazyFormer: Self Attention with Lazy Update](https://mlsanity.com/paper/2102.12702/) (2021)
- [Environmental Sensor Placement with Convolutional Gaussian Neural Processes](https://mlsanity.com/paper/2211.10381/) (2022)
- [Translation Equivariant Transformer Neural Processes](https://mlsanity.com/paper/2406.12409/) (2024)
- [Memory Efficient Neural Processes via Constant Memory Attention Block](https://mlsanity.com/paper/2305.14567/) (2023)
- [Approximately Equivariant Neural Processes](https://mlsanity.com/paper/2406.13488/) (2024)

## Built on by (most cited first, 1 of 1)

- [Spectral Transformer Neural Processes](https://mlsanity.com/paper/2605.09498/) (2026)

## Abstract (arXiv)

Neural Processes (NPs) are a rapidly evolving class of models designed to directly model the posterior predictive distribution of stochastic processes. While early architectures were developed primarily as a scalable alternative to Gaussian Processes (GPs), modern NPs tackle far more complex and data-hungry applications spanning geology, epidemiology, climate, and robotics. These applications have placed increasing pressure on the scalability of these models, with many architectures compromising accuracy for scalability. In this paper, we demonstrate that this trade-off is often unnecessary, particularly when modeling fully or partially translation-invariant processes. We propose a versatile new architecture, the Biased Scan Attention Transformer Neural Process (BSA-TNP), which introduces Kernel Regression Blocks (KRBlocks), group-invariant attention biases, and memory-efficient Biased Scan Attention (BSA). BSA-TNP is able to: (1) match or exceed the accuracy of the best models while often training in a fraction of the time, (2) exhibit translation invariance, enabling learning at multiple resolutions simultaneously, (3) transparently model processes that evolve in both space and time, (4) support high-dimensional fixed effects, and (5) scale gracefully, running inference on over 1M test points and 100K context points in under a minute on a single 24GB GPU. Code is provided as part of the `dl4bi` package.