Research map / Architectures and efficiency / Attention and transformers
Gridded Transformer Neural Processes for Spatio-Temporal Data
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces gridded pseudo-token transformer neural processes that encode unstructured spatio-temporal data onto a grid so efficient attention (Swin/ViT) can be used, with translation-equivariant variants.
Problem
Transformer neural processes cannot scale to large spatio-temporal datasets because full attention is quadratic and efficient attention needs gridded inputs, while real data are often unstructured or multi-modal.
Approach
Encode unstructured context points onto a regular grid of pseudo-tokens via a nearest-neighbour cross-attention grid encoder, process the grid with Swin Transformer or ViT, and decode at arbitrary targets with nearest-neighbour cross-attention; add (approximate) translation equivariance via relative positional encodings on the grid.
Key results
- On the t2m station experiment, Swin-TNP with PT-GE at 64x128 reached 1.819 test log-likelihood versus 1.535 for ConvCNP and 1.344 for PT-TNP.
- On the multi-modal wind speed experiment, Swin-TNP m-PT-GE reached 8.603 log-likelihood versus 6.143 for ConvCNP and 3.988 for PT-TNP.
- Pseudo-token grid encoding outperformed kernel-interpolation grid encoding, and nearest-neighbour cross-attention decoding outperformed full cross-attention.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use the pseudo-token grid encoder instead of kernel-interpolation grid encoding (ablation) | gridded TNP · t2m station prediction | log-likelihood: 1.683 -> 1.819 | Helped |
| use the pseudo-token grid encoder instead of kernel-interpolation grid encoding (ablation) | gridded TNP · multi-modal wind speed | log-likelihood: 7.603 -> 8.603 | Helped |
| use patch encoding before projecting onto the latent grid (ablation) | ViTNP · t2m station prediction | log-likelihood: 1.628 -> 1.734 | Helped |
| use patch encoding before projecting onto the latent grid (ablation) | ViTNP · multi-modal wind speed | log-likelihood: 5.371 -> 7.288 | Helped |
| use Swin Transformer processing instead of ViTNP processing (ablation) | gridded TNP · t2m station prediction | log-likelihood: 1.808 -> 1.819 | Helped |
| use Swin Transformer processing instead of ViTNP processing (ablation) | gridded TNP · multi-modal wind speed | log-likelihood: 7.288 -> 8.603 | Helped |
| use nearest-neighbour cross-attention instead of full cross-attention in the grid decoder (ablation) | Swin-TNP PT-GE · t2m station prediction | log-likelihood: 1.636 -> 1.819 | Helped |
| use separate pseudo-token grids for each modality instead of one shared grid (ablation) | Swin-TNP PT-GE · multi-modal wind speed | log-likelihood: 8.073 -> 8.603 | Helped |
| increase grid size from 64 × 128 to 192 × 384 (ablation) | Swin-TNP PT-GE · t2m station prediction | log-likelihood: 1.819 -> 2.053 | Helped |
| increase grid size from 24 × 60 to 48 × 120 (ablation) | Swin-TNP m-PT-GE · multi-modal wind speed | log-likelihood: 8.603 -> 9.383 | Helped |
| add strict translation equivariance (ablation) | Swin-TNP · t2m station prediction | log-likelihood: 1.819 -> 1.895 | Helped |
| add approximate translation equivariance instead of strict translation equivariance (ablation) | Swin-TNP (T) · t2m station prediction | log-likelihood: 1.895 -> 1.926 | Helped |
Surprising findings
- Nearest-neighbour cross-attention in the decoder outperformed full cross-attention in predictive performance, not just cost.
Limitations
- Scalability of pseudo-token count M tied to memory constraints
- Strict equivariance can be overly restrictive for real data, motivating approximate equivariance
Experimental setup
Datasets: ERA5 reanalysis synthetic Gaussian process datasets (SE kernel, SKI) EAGLE fluid dynamics dataset
Benchmarks: GP regression test log-likelihood ERA5 t2m station prediction ERA5 multi-modal wind speed interpolation
Metrics: test log-likelihood forward pass time
Baselines: CNP PT-TNP ConvCNP ViTNP Swin-TNP (T)
Builds on: transformer neural processes ConvCNP Perceiver pseudo-tokens Swin Transformer translation equivariant TNP
Closest papers by meaning
- Translation Equivariant Transformer Neural Processes (ICML 2024 · similarity 0.89)
- Spectral Transformer Neural Processes (NeurIPS 2026 · similarity 0.87)
- On the Integration of Spatial-Temporal Knowledge: A Lightweight Approach to Atmospheric Time Series Forecasting (NeurIPS 2025 · similarity 0.85)
- CALM-PDE: Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs (NeurIPS 2025 · similarity 0.85)
- Approximately Equivariant Neural Processes (NeurIPS 2024 · similarity 0.84)
- iTransformer: Inverted Transformers Are Effective for Time Series Forecasting (ICLR 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.