Gridded Transformer Neural Processes for Spatio-Temporal Data
Authors: Matthew Ashman
Cristiana Diaconu
Eric Langezaal
Adrian Weller
Richard E. Turner
Venue: ICML 2025 spotlightposter (2025)
Keywords: neural process
probabilistic machine learning
transformer
spatio-temporal data
spatio-temporal modelling
translation equivariance
Primary area: probabilistic_methods
TL;DR: We develop gridded transformer neural processes (gridded TNPs), a new member of the TNP family that enables the use of efficient attention mechanisms to tackle problems involving large, unstructured data.

Abstract: Effective modelling of large-scale spatio-temporal datasets is essential for many domains, yet existing approaches often impose rigid constraints on the input data, such as requiring them to lie on fixed-resolution grids. With the rise of foundation models, the ability to process diverse, heterogeneous data structures is becoming increasingly important. Neural processes (NPs), particularly transformer neural processes (TNPs), offer a promising framework for such tasks, but struggle to scale to large spatio-temporal datasets due to the lack of an efficient attention mechanism. To address this, we introduce gridded pseudo-token TNPs which employ specialised encoders and decoders to handle unstructured data and utilise a processor comprising gridded pseudo-tokens with efficient attention mechanisms. Furthermore, we develop equivariant gridded TNPs for applications where exact or approximate translation equivariance is a useful inductive bias, improving accuracy and training efficiency. Our method consistently outperforms a range of strong baselines in various synthetic and real-world regression tasks involving large-scale data, while maintaining competitive computational efficiency. Experiments with weather data highlight the potential of gridded TNPs and serve as just one example of a domain where they can have a significant impact.

Gridded Transformer Neural Processes for Spatio-Temporal Data
Matthew Ashman *
Cristiana Diaconu"
Eric Langezaal 12
Adrian Weller" Richard E. Turner"
Abstract
building such models is non-trivial: the data are typically heterogeneous, irregular in space and time, and multi-modal, with each modality potentially sampled at different spatio- temporal locations. A field that has greatly benefitted from these recent advancements, and will serve as the motivating real-world use case throughout this work, is medium range weather and environmental forecasting. Here, a new gener- ation of models have improved performance and reduced computational costs (Bodnar et al., 2025; Lam al., 2023; Price et al., 2025; Bi et al., 2023; Nguyen et al., 2023; Chen et al., 2023b;a). These models operate on environmental variables that are regularly structured in space and time, allowing them to leverage architectures from the vision and language community, such as the Vision Transformer (ViT; Dosovitskiy al. 2020). Swin Transformer (Liu et al., 2021), and Perceiver (Jaegle al., 2021). They trained and deployed on data from computationally intensive sci- entific simulation and analysis techniques, which integrate observational data (e.g. from weather stations, ships, buoys, etc.) with simulation data to provide the best estimate of the atmosphere's 's state. They are currently not trained on observational measurements directly.
Effective modelling of large-scale datasets is essential for many domains, yet exist- ing approaches often impose rigid constraints on the input data, such requiring them to lie on fixed-resolution grids. With the rise of founda- tion models, the ability to process diverse, het- erogeneous data structures is becoming increas- ingly important. Neural processes (NPs), particu- larly transformer neural processes (TNPs), offer a promising framework for such tasks, but struggle to scale to large spatio-temporal datasets due to the lack of an efficient attention mechanism. To address this, we introduce gridded pseudo-token TNPs which employ specialised encoders and de- coders to handle unstructured data and utilise a processor comprising gridded pseudo-tokens with efficient attention mechanisms. Furthermore, we develop equivariant gridded TNPs for applications where exact or approximate translation equivari- ance is a useful inductive bias, improving accu- racy and training efficiency. Our method consis- tently outperforms range of strong baselines in various synthetic and real-world regression tasks involving large-scale data, while maintaining com- petitive computational efficiency. Experiments with weather data highlight the potential of grid- ded TNPs and serve as just one example of a domain where they can have a significant impact.
We are now the cusp of second generation of such models that will integrate unstructured observational data alongside or instead of analysis data altogether, improv- ing accuracy and reducing computational costs. Achieving this is challenging, as such models need to be able to 1) handle data at arbitrary input locations, and 2) effectively combine modalities that are necessarily recorded the same locations. Only few works have attempted this more ambitious problem (Allen et al., 2025; McNally et al., 2024; Xu et al., 2025; Xiao et al., 2024), and the optimal archi- tectures remain unclear. This presents fertile ground for impactful research, with findings easily transferable to other spatio opatio-temporal domains. For instance, the same techniques could be used to model real-life systems governed by partial differential equations (PDEs), by pre-training synthetic PDE simulations, followed by fine-tuning real-world un- structured observations such wind tunnel measurements.
1. Introduction
The proliferation of data from in situ sensors, remote ob- servations, and scientific computing models is transforming spatio-temporal modelling. This has led to surge of in- terest from the machine learning community to develop new tools and models to support these efforts. However,
'Equal contribution University of Cambridge University of Amsterdam Alan Turing Institute. Correspondence to: Matthew Ashman <imca39@cam.cc.u Cristiana Diaconu <cdd43@cam.ac.kk>
What framework most suitable for spatio-temporal tasks containing both structured and unstructured data? Here we advocate for neural processes (NPs; Garnelo al. 2018a;b) family of meta-learning models that map from datasets of arbitrary size and structure predictions over outputs
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Gridded Transformer Neural Processes for Spatio-Tempore Data
arbitrary locations. NPs support probabilistic treatment of observations, enabling them to output uncertainty estimates that are crucial for downstream decision making. Moreover, NPs are flexible-une unlike other models that only perform forecasting, they can solve more general state estimation problems, including forecasting, data fusion, data interpo- lation and data assimilation. Although early versions of NPs were limited, recent developments (Kim et al., 2019; Gordon et al., 2020; Nguyen & Grover, 2022; Ashman et al., 2024a; Feng et al., 2023; Bruinsma et al., 2021; Ashman et al., 2024b) have significantly improved their effective- ness, especially for small-scale spatio-temporal regression problems. NPs are now being used in a wide-range of ap- plications, from climate downscaling, data assimilation and sensor placement (Vaughan et al., 2022: Andersson et al., 2023: Chen et al., 2024; Niu et al., 2024), to tasks as diverse as molecular property prediction (García-Ortegóón et al., 2024). Among NP variants, the convolutional conditional neural process (ConvCNP; Gordon et al., 2020) is the lead- modelling. The ConvCNP
tive distributions at arbitrary spatio-temporal locations. Remarkably, we find that this outperforms full attention. 4.
We adapt computationally expensive methods proposed in the literature (Ashman al., 2024a) to efficiently incorporate translation equivariance into gridded TNPs, improving learning efficiency. Additionally, to account for symmetry imperfections often present in real- world data, we build on Ashman et al. 2024b to achieve approximate translation equivariance with gridded TNPs.
5. We empirically evaluate our model on a range of synthetic and real-world spatio-temporal regression tasks, demonstrating both the ability to i) maintain strong performance on large spatio-tempont
datasets and ii) handle multiple sources of unstructured data effectively, all while maintaining low computational complexity.
2. Background
ing approach for employs a set convolution operation to project unstructured data onto a regular grid through kernel-based interpolation. While effective, the success of transformers over convolu- tional neural networks (CNNs; LeCun et al., 1998) suggests potential for improvement. Indeed, recent developments in transformer- based NPs (TNPs; Nguyen & Grover 2022; Feng et al. 2023; Ashman et al. 2024a) have shown impres- sive performance on small-scale problems. However, unlike the aforementioned large-scale environmental models for gridded data, TNPs are yet to fully leverage efficient atten- tion mechanisms. As result of the quadratic computational complexity of full attention, they have been unable to scale to complex ppatio-temporal datasets. This is because such efficient techniques require structured-more specifically, gridded data, limiting their immediate applicability.
We consider a supervised learning setting with input and out- put spaces X' = RR. and y R Dv. A dataset consists of input-output pairs yx X x 3 with context and target sets De, D, containing [Del = Ne and [Del= Po|=No points. Let X. E (X)) Y, E (J/)*== and X, E (2) Nr, Y, E (2)) denote the inputs and outputs for D. and D.. We denote a single task as g = (D., D.) = ((x., Yo), (X,, Y,)..
2.1. Neural Processes
Neural processes (NPs; Garnelo et al. 201 2018;;b) can be viewed as neural-nettorkk-bssek mappings from context sets De to predictive distributions at target locations X,, p(- X,, D2). In this work, we restrict our attention to conditional NPs (CNPs; Garnelo al. 2018a), which, unless deployed autoregressivell, target marginal predic- tive distributions by assuming that the predictive densities factorise: (YY Dz) = I P(y,, Dz). We de- note all parameters of a CNP by 0. CNPs are trained in a meta-learning fashion, in which the expected predic- tive log-probability is maximised OML = arg max, Lvn. where (0) = Extc) En=,, log po (y Xx.,: Dz)]. In the limit of infinite tasks and model capacity, the global maximum is achieved if and only if the model recovers the ground-truth predictive distributions (Proposition 3.26 by Bruinsma, 2022). For real-world datasets, with only a finite number of training tasks, we approximate this expecta- tion with an average over tasks. In Appendix A we present a unifying construction for CNPs involving three compo- nents: the encoder e: X y + Z, which encodes each (e,,yyeen) DD into a token representation Ze,n E Z, the processor p: (U% 2n) X 22 which processes the set of context tokens and the target input x; to obtain a target dependent token 2, E Z, and the decoder d: 2 d- Py, which maps from the target token to the predictive distribu- tion over the output at that target location. Py denotes the
Motivated by the approach taken by the ConvCNP, we pursue straightforward solution: encoding the input data onto structured grid before passing it through transformer-based architecture. We introduce attention- based mechanisms for encoding the unstructured data, as well as for processing the resulting structured grid, giving rise to gridded TNPs-a general-purpott tool for spatio-temporal state estimation. Our core contributions are:
1. We develop a novel attention-based grid encoder, in- spired by the idea of 'pseudo-token (Jaegle et al., 2021; Feng et al., 2023; Lee et al., 2019), which outperforms existing encoders for opatio-temport data.
2. Equipped with the Above-mentionns pseudo-toker grid encoder, we enable TNPs to use efficient attention mech- anisms, utilising advancements such as the ViT (Dosovit- skiy et al., 2020) and Swin Transformer (Liu et al., 2021).
3. We develop an efficient k-nearest- neiihbost- attention- based grid decoder, facilitating the evaluation of predic-
2

---

Gridded Transformer Neural Processes for
Data
space of distributions over the output space V
Transformer Neural Processes Transformer neural pro- cesses (TNPs) use transformer- architectures as the processor in the CNP construction described in Section 1. First, each context point "xc.n.e... E De and target imput xtn E X are encoded to obtain an initial set of context and target tokens, Ze E IRNNxxDs and Z E R NNN.DD, xD; Trans- formers are then used to process the union Z = ZZ us- ing series of MHSA and MHCA operations. keeping only the output tokens corresponding to the target inputs. These processed target tokens are then mapped to predictive distri- butions using the decoder. The specific transformer- based architecture is unique to each TNP variant, generally consist- ing of MHSA operations acting on the context tokens or pseudo-token representation of the context-and MHCA operations updating the target tokens given the context to- kens. We provide diagram of two popular TNP variants in Appendix A: the regular TNP Nguyenn& & Grover, 2022) and the induced set transformer NP (ISTNP; Lee et al. 2019).
2.2. Transformers and Transformer Neural Processes
Transformers can be viewed as general set functions (Lee et al., 2019), making them ideal for NPs, which must ingest datasets. This section briefly overviews transformers and their integration into CNPs, leading to the TNP family.
Self-Attention and Cross-Attention Broadly speaking, transformer-based architectures consist of two operations: multi-head self-attention (MHSA) and multi-head cross- attention (MHCA). Informally, the MHSA operation up- dates set of tokens using the same set, whereas the MHCA operation updates one set of tokens using a different set. More formally, let Z E NR denote set of N D-- dimensional input tokens. The MHSA operation updates this set of tokens for n = I.....NA NN as
The application of TNPs to large datasets is impeded by the use of MHSA and MHCA operations acting on the entire set of context and target tokens. Even when pseudo-tokens are used in pseudo-token TNPs (PT-TNPs), the number of pseudo-tokens M required for accurate predictive inference generally scales with the complexity of the dataset, and the MHCA operations between the pseudo-tokens and the con- text and target tokens quickly becomes prohibitive as the size of the dataset increases. This motivates the use of effi- cient attention mechanisms, as used in the ViT (Dosovitskiy et al., 2020) and Swin Transformer (Liu et al., 2021).
H Zn +cat cat (( On (Zn; 2m mm Wv. n=1) Wo. )
(L)
Here, V. E RD and Wo E RHDxxD; are the value and projection weight matrices, where и denotes the number of 'heads', and 0% is the attention mecha- transformed
nism. This is most often a inner-product between pairs of tokens: On (Z,, Zm)) = softmax (z W, WT k J J ) where Wo, E
DD DDo and W K.A E DRDxxDo are the query and key matrices. The MHCA operation updates one set of tokens, Z1 E xD, using another set of tokens, Zy E RR N2xxx D., in similar manner for n =1, I,...;N.. N1:
2.3. Translation Equivariance in TNPs
A useful inductive bias in spatio-temporal modelling, where the data are roughly stationary, is translation equivariance (TE): upon translating the data in space or time, the predic- tions of the model should translate accordingly. Previous works have already considered such symmetries in TNP ar- chitectures. Ashman et al. (2024a) introduced the translation equivariant TNP (TF-TNP) denoted by TNP (7), where T represents the d-dimensional group of translations (with d the dimensionality of the data). Here, translation equiv- ariance is achieved by restricting the attention mechanism Oh,n. On (Z,,, Am X,, xm) to only depend on differ- ences, rather than absolute input locations, resulting in the TE-MHSA and TE-MHCA operations. These are defined analogously to Equations 1 and 2, with (th,n.m defined as:
N₂
I H ah n 2m 2z.m WV.. Wo. (2)
11,a cat
T
MHSA / MHCA operations are combined with layer- normalisation and point-wise MLPs to obtain MHSA / MHCA blocks. Unless stated otherwise, we adopt the order used by Vaswani et al. (2017) (detailed in Appendix G).
Sreudo-Token-Beee Transformers Transformers based on pscudo-tokens, first introduced by Jaegle et al. 2021 with the Perceiver, remedy the quadratic computational complexity induced by the standard transformer by condensing the set of N tokens, Z E RND,,, into a smaller set of M < N 'pseudo-tokens', U Mx..
(b Wo..WW H 2m Xn..)
using the MHCA operation. The processor acts on these pseudo-tokens (instead of the original set), reducing the computational complexity from ((N2) cost of MHSA on the original set-to ONN + the cost of MHCA between the original set and the pseudo-tokens, followed by MHSA on the pseudo-tokens.
Ohhmm.m ΣM ethh (x W Wa. WK. X
W X
where Ph: : R x R. + R is a learnable function, param- eterised by an MLP, and the initial tokens {z,, W. only depend on output values. TE-MHCA is defined analogously. A more thorough treatment is provided in Appendix B.
3

---

Gridded Transformer Neural Processes for Opatio-Temporrs Data
Strictly equivariant architectures introduce useful inductive biases but be overly restrictive for modelling complex real-world data that are rarely perfectly symmetric. To ad- dress this, Ashman et al. (2024b) propose the approximately equivariant NP, which also generalises beyond translation to other symmetry groups. This is achieved through a simple architecture modification, by introducing additional fixed basis functions that enable the model learn symmetry- breaking features in data-driven manner. Here, we focus on the approximately TE-TNP, denoted as TNP (T), repre- senting approximate equivariance with respect to T(d).
the first goal. However, most of these approaches (as dis- cussed in Section 1) are based efficient transformer-based architectures that can only be applied on structured, grid- ded data. Moreover, they primarily focus on forecasting', hence constraining predictions to the same spatial locations as the inputs. Some methods partially address this limita- tion by handling missing data (Nguyen et al., 2023), but they still assume a gridded data structure, hence not sat- isfying the second desideratum. An alternative that can address both points is using graph neural networks (GNNs; Bronstein et al. 2021) as backbone, but attempts in the weather community proved that GNNs are harder to scale than ceaaffommerbbasee approaches leading to worse perfor- mance (e.g. 36.7M parameters for GraphCast (Lam et al., 2023) compared to 1.3B for Aurora (Bodnar et al., 2025)).
Ashman et al. (2024b) show approximately equivariant TNPs outperform their non-equivariant counterparts, but they suffer from large time and space complexity due to the need pass pairwise distances through an MLP when evaluating the weights of the translation equivariant atten- tion mechanism. For this reason, their application has been limited to relatively small datasets.
NPs arise promising framework due to their ability to model stochastic processes, which can be evaluated at any target location, and to flexibly condition on unstructured data. Indeed, our work reflects a more general trend towards more flexible methods, reflected in works such as Aardvark (Allen et al., 2025) and FuXi-DA (Xu et al., 2025), both end-to-end weather prediction models that handle both un- structured and structured data, as well as AtmoRep (Lessig et al., 2023) and Prithvi WxC (Schmude et al., 2024). In Appendix C, we discuss these methods in detail and provide an extended comparison between our approach and GNNs.
3. Related Work
Transformers for Point Cloud Data A closely related research area is point cloud data modelling (Tychola al., architectures employed for
2024), with transformer-bee variety of tasks (Lu et al., 2022). To handle large datasets, the use of efficient attention mechanisms has been explored. Several notable approaches use voxelisation, whereby un- structured point clouds are encoded onto a structured grid (Mao al., 2021; Zhang et al., 2022). prior to employing efficient architectures that operate on grids. Our pseudo- token grid encoder draws inspiration from the voxel-based set attention (VSA) of He et al. (2022), which also cross- attends local neighbourhoods of unstructured tokens onto a structured grid of pseudo-tokens, However, VSA uses the same initial pseudo-token values for all grid locations, as the 'cloud' in which points exist has no unobserved information. Their use of fixed set of initial values manifests a specific implementation in which all tokens attend to the same set of -heudo-tokens. In contrast, spatio-temporal problems may involve potentially unobserved, fixed topographical infor- mation, such as elevation, land use and soil type. To enable the model to capture this, we employ different pseudo-token values for each grid location.
4. Gridded Transformer Neural Processes
While TNPs have shown promising performance on small to medium-sized datasets, they are unable to scale to large spatio-temporal data. To address this, we consider the use of efficient attention mechanisms that have proved effective for gridded spatio-temporal data (Bodnar et al., 2025; Lam et al., 2023; Price et al., 2025; Nguyen et al., 2023; Bi et al., 2023). Such methods are not immediately applicable with- out first structuring the unstructured data. We achieve this by drawing upon methods developed in point cloud modelling notably the VSA (He et al., 2022) and develop the pseudo- token grid encoder: an effective attention- based method for encoding unstructured data onto a grid.
We provide an illustrative diagram of our proposed ap- proach in Figure 1, which decomposes the processor p: (Un 2n) X' + Z into three parts: the grid encoder, pge: U 2n + ZM, which embeds the set { ((xe,n Ze,n 1 into tokens [ at gridded loca- tions {vm) i 2. the grid processor, Pap: zd + z.M. which transforms the token values; and 3. the grid decoder, Pad: M x X + 2, which maps the tokens onto the loca- tion x,. Here, the latent space Z = Zyoken X X. We discuss
Models for Structured Weather Data Our method is motivated by the need develop models that can 1) scale to massive spatio-temporal datasets and 2) flexibly han- dle unstructured data. Here, unstructured data refers to, for example, observations that do not lie on a regular grid (c.g., off-the-grid weather station data, solutions to PDEs on unstructured meshes) or modti-modal data collected at varying spatial locations. Recent advances, particularly in weather modelling, have made significant progress on
Predicting the entire gridded state at time f + given history of previous gridded state(s).
4

---

Gridded Transformer Neural Processes for Opatio-Tempone Data
((Ke,n Nesaa ] J
(X Ran) m
{( V Up 1 mmml
((Ver Mnn mum
Point-wise encoding
Grid encoding
Grid processing
Encodkat
Processor
Decoder
X
Point-wisc decoding
Context observation
Grid decoding
X
Context token
X
Pscudo-token
X
Processed pseudo-token
Target location
Target token
((S
((V
Up
B
Target predictive distribution
Figure 1. Illustration of the complete gridded TNP pipeline, We highlight the encoder (blue), processor (red) and decoder (green).
the choices for these components in this section.
Unstructured
Cross-atlend to grid
Structured
Grid Encoder: the Pseudo-Token Grid Encoder Let 2n RDP= be the token representation of the input-output pair (x) after point-wise embedding. We introduce of I Ma pseudo-tokens U IR x.- MD. D corresponding grid locations V E M. xD, (i.e. pseudo-token umi...,m.. E RD is associated with
Context Initial pscudo-tokens
D'Dumm' tokens
Pseudo-tokens
location V mmD...,m.. IR PR)"). For ease of reading, we re-
Figure 2. An illustration of the pseudo-token grid encoder in 2-D. For efficient implementation of cross attention the pscudo- token grid, we pad sets of neighbourhood tokens with 'dummy' tokens, that each neighbourhood has the same cardinality.
place the product I1 Ma with and the indexing nota- tion mi..... ... ,, with m. The pseudo-token grid encoder obtains pseudo-token representation U IR MxD: of De on the grid V by cross-attending from each set of tokens (Ze,n mx) each initial pseudo-token uh:
Grid Processor: Efficient Attention Transformers While we can choose any architecture for processing the pseudo-token grid, including CNNs, we focus on transformer architectures employing efficient atten- tion mechanisms consider the VIT (Dosovitskiy et al., 2020) and Swin Transformer (Liu al., 2021). Typically, ViT uses patch encoding to downsample grid of tokens to coarser grid of tokens, upon which MHSA operations are applied. We explore both patch encoding and directly encoding to smaller grid via the grid encoder, which, when used in gridded TNPs, we collectively refer to as ViTNPs. The Swin Transformer has the advantage of being able to operate with finer grid of eseudo-tokens as only local neighbourhoods of tokens attend each other at any one time. We refer its in gridded TNPs as Swin-TNP.
Un + MHCA (u%, [Ze,n (eEm(ym)))) Vm E M. (4)
Here, 9T denotes the index set of input locations for which V. is amongst the & nearest grid locations. practice, we found that & = suffices. While this opera- tion seems computationally intensive, we apply two tricks make it computationally efficient. First, note that pro- vided the pseudo-token grid is regularly spaced, the set of nearest neighbours for all grid locations can be found in O(N)). Second, the operation described in Equation 4 can be performed in parallel for all m E M by 'padding' each set (ne with max; [9(vv;; k)] [ k)
"dummy tokens' and applying appropriate masking, resulting in computational complexity of O (M max; [ot(v;; k) Restricting the neighbourhood size can further reduce this. We illustrate the pseudo-token grid encoding in Figure 2. In our experiments, we compare the performance of the eseudo-token grid encoder with simple intrel-intrroolation onto grid, a popular method used by the ConvCNP. We pro- vide detailed deseription this approach in Appendix D.
the Cross-Attention Grid Decoder In PT-TNPs, all pscudo-tokens U cross-attend to all target tokens Z,, with computational complexity of O(MN,,, which is prohibitive for large N, and M. We propose nearest-neighbour cross-attention, in which the pseudo-
Grid Decoder:
5

---

Gridded Transformer Neural Processes for Spatio-Tempordl Data
tokens {um) m me)((xe..ik) attend the target token Ze,ni
with unstructured data, the need to dynamically compute this quantities cannot be overcome, as the locations of the points change at each pass. However, if we instead operated on uniformly discretised domain, we can simply maintain a fixed matrix of attention weighting biases, closely resem- bling the relative positional encodings (Shaw et al., 2018) commonly used in transformers. Fortunately, this ex- actly the case in the Swin-TNP processor, where the initial, unstructured data has already been projected onto a fixed, structured grid of pseudo-tokens of window size. Given that attention is only performed among these structured pseudo- tokens, we can efficiently implement the TE attention mech- anism by introducing of a set of learnable relative positional encodings added to the pre-softmax -ormmanx--ooalseddd attention weights. This comes with negligible computational cost.
Zr,0 f MHCA (zf,w0 (mmttst(s...)))))
Here, M(x kc) denotes the index set of the lo grid locations that are closest to xt.n. 2 With equally spaced grid locations, this can be computed for all target tokens in (N)), and the operation in Equation 5 can be parallelised with com- plexity of (k:Ne), leading to a significant computational reduction for lo < M. We also found that nearest-neighhour outperformed full coss-attentionn in practice, likely due to its beneficial inductive bias for spatio-temporal data.
4.1. Handling Multiple Data Modalities
Spatio-temporal datasets often include multi-modal data. For structured data (i.e. gridded data, with modalities shar- ing the same input locations). this is typically addressed by concatenating the modes along the channel dimension. In contrast, unstructured data-where different modalities different ap-
For the pseudo-token grid encoding and grid decoding, we must compute relative positional encodings dynamically using the TE-MHCA mechanism proposed by Ashman et al. (2024a). However, due to the use of the nearest-neighhbur cross cross-attention mechanisms, the number of pairwise compu- tations scales only linearly with the number of context and target points. This allows for an efficient and scalable imple- mentation of the gridded Swin-TNP (7) model. The encoder and decoder equations are analogous to Equations 4 and 5, but we replace the MHCA operations with TE-MHCA. Fi- nally, we relax strict equivariance through an approximately TE model San--Sii--NNN (T). This can be implemented fol- lowing Ashman et al. (2024b) by adding a few learnable basis functions to the inputs of the TE components of the architecture. We provide more details in Appendix G.
can be recorded distinct locations-requit proach. While existing frameworks mentioned in Section 3 generally lack support for this setup. our method can be extended to such scenarios. One approach, which we call the single pseudo-token grid encoder, is to use modality- specific encoders and apply the aseudo-token grid encoder to the union of tokens. Formally, let De = U De.e de- note the context dataset partitioned into specifi smaller datasets. For each source and context datapoint, we obtain Ze,n.s Nenn.)) where e X x Ve. + Z denotes the source-specific point-wise encoder and V. the output space for source s. The set of context tokens is obtained Ze {(2.,n.)) proach, called the multi pseudo-token grid encoder, we first form a separate pseudo-token grid for each modal- ity. For each set of source-specific context tokens Z,,= = (es (x..,,ss, mEr,, we obtain gridded pseudo-token representation U, E MxD. We then pass their wisnt-wise concatenation through some function to obtain a unified gridded pseudo-token representation. Both methods can be applied analogously to kernel-interpolattonn grid encoding.
5. Experiments
In second ap-
We evaluate our gridded TNPs on synthetic and real-world regression tasks involving large datasets. We show con- sistent improvements over baselines, especially as task complexity and dataset size increase, while maintaining low computational complexity. Throughout, compare gridded TNPs with two different grid encoders (GEs) the pernel-interpolation grid encoder (KI-GE) and our pseudo-token grid encoder (PT-GE)-and two different grid processors Seresssor--Sinnnn Transformer and VIT. We also compare to the following baselines: the CNP (Garnelo et al., 201 8a), the PT-TNP using the induced set-transformer architecture (Lee al., 2019; Ashman al., 2024a)², and the ConvCNP (Gordon et al., 2020). We do not compare to the regular ANP or TNP (Kim et al., 2019; Nguyen & Grover, 2022) as they are unable to scale to the considered context set sizes. Details on experiments, architectures, datasets and hardware can be found in Appendices F and G. We also provide a public implementation of gridded TNPs in the repository
4.2. Introducing Translation Equivariance
To improve training efficiency and generalisation capabili- ties, we also consider incorporating translation equivariance into gridded TNPs, as a useful inductive bias for spatio- temporal modelling- As mentioned in Section 2.3, the TNP (T) of Ashman et al. (2024a) cannot be naively applied to large data because of the need to pass pairwise distances through an MLP when computing the attention weights, re- sulting in high time and space complexity. When working
'oor all experiments, we use the maximum number of pseudo- tokens allowed by memory constraints.
For dimension d= di)( we select the nearest c/dd( neighbours per dimension. Appendix E. provides more details.
6

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
http://github..co//// hwwp:///.oooooooooaaongg//ddddaa
(including the use of mean aggregation grid encoder, as used by FuXi-DA (Xu et al., 2025)), an ablation of the grid size, and the predictive means of selected models an example dataset with é = 0.1, illustrating gridded TNPs' superior ability to capture complex data. We also provide a smaller- scale experiment comparing gridded TNPs to the standard TNP (which cannot scale to the larger-scale version due to the lack of efficient attention). We show that in regimes where the standard TNP is applicable, gridded TNPs are able to recover its performance. These additional results demonstrate that the improvements in efficiency of gridded TNPs do not come at the expense of predictive performance.
To illustrate the generality of our framework, we include an additional study in Appendix G.6 using the scaee EAGLE fluid-dy fuid-dynamics dataset (Janny et al., 2023). The re- sults show that gridded TNPs can be effectively and directly applied-without task-specific modifications- PDE mod- elling on irregular meshes.
5.1. Meta-Learning Gaussian Process Regression
We begin with two synthetic 2-D regression tasks with datasets drawn from a Gaussian process (GP) with squared- exponential (SE) kernel approximated using structure ker- nel interpolation (SKT) (Wilson & Nickisch, 2015). Each dataset contains 1.1 X 10 datapoints (N,- = 1 x 10 con- text and N₂- = x 10 target points), with inputs sampled uniformly from - We consider two SE. kernel length- scales: f= 0.1 and l = 0.5, implying roughly 10, 000 and 500 'wiggles', respectively, in each dataset. While both se- tups are non-trivial for typical NP synthetic experiments, we expect the former to be intractable for most it the model to both assimilate finely grained information and modulate its predictions accordingly. In Figure 3, we plot the test log- ligllikeissoos against the time taken for single forward pass (FPT) for two ViT variants (with and without patch encoding), and Swin-TNP. We compare them to two strong baselines: the PT-TNP with M = 256 pseudo-tokens, and the ConvCNP with the same grid sizes as the Swin-TNP. Training and inference for all models is performed on one NVIDIA GeForce RTX 2080 Ti.
5.2. Real-World Weather Data
We perform two experiments on data from the ERA5 reanal- ysis by the European Centre for Medium-Range Weather Forecasts (ECMWF; Hersbach et al. 2020). We train on data between 2009-2017, validate on 2018 and test on 2019, Af- ter outlining the experimental setups, we present key conclu- sions from aggregated results, highlighting the consistency across experiments. Training and inference are performed using single NVIDIA A A10 80GB with 32 CPU cores.
Combining Weather Station Observations with Struc- tured Reanalysis This experiment tests our models' ability to combine sparse, unstructured observations with struc- tured, on-the-grid data. For this, we use skin temperature (skt) and 2m temperature 2t)m We construct each con- text dataset by combining the 22 at a random subset of 9, 957 weather station locations (proportion sampled from Uu,,.,s) with a coarsened 180 x 360 core-cooresoonding to a grid spacing of 1 of skt values. The targets are the 2m values at all 9, 957 weather station locations.
CoasUNNP 61
VITNP K16 92
VITNP, KI 22
Sain-INP KI
PT-INP, Af 206
VITNP, PT
VITNI,
Soin-NNP 6T61
0.85
Combining Multiple Sources of Unstructured Wind Speed Observations We also evaluate the ability of grid- ded TNPs to model sparsely sampled eastward (u) and northward (v) components of wind at three different pres- sure levels: 700hPa, 850hPa and 1000hPa (surface level). This dataset is considered unstructured, with input loca- tions that differ across the six variables. We perform spatio- temporal interpolation over a latitude / longitude range of [25", 40,449) [-125, -66%,, which corresponds to the con- tiguous US, spanning four hours. The proportion of the 543, 744 observations used as the context dataset is sampled from 0...5. The target set size is fixed at 135, 936.
Text Mootlo---llllle
0.84
0.72
0,83
0.70
0.68
0.81
0,565
EPT (ms)
EPT (mz)
Figure 3. Test log-likelihood vs. forward pass time (FPT) for the GP datasets. Baselines are hatched. The considered grid sizes are 64% x 64 and 32x x 32, shown as 64 and 32. For ViTNPs, we include results with and without patch encoding, the former indicated by the -> symbol in-between the pre- and post-patch-encode grid sizes. KI/ / PT: kermel-interpolation// / pseudo-token grid encoding.
The conclusions of the two weather experiments, empirically supported by the results in Table (where we bold the results that support each conclusion), are as follows:
Our results show that gridded TNPs outperform all baselines (hatched markers in Figure 3) especially on the more com- plex datasets, where they significantly surpass the strongest CNnyCNP baseline-while maintaining competitive com- putational complexity. Further, by comparing the darker to the lighter markers, we observe that the PT-GE outperforms the KI-GE. We show in Appendix G. 1 table of all results
1. Gridded TNPs outperform all baselines while maintain-
'skt is the Earth's 's surface temperature, directly output by atmospheric models, while t22 is found by interpolating between the lowest model temperature and skt (Owens & Hewson, 2018).
7

---

Gridded Transformer Neural Processes for Opatio-Temporal Data
Table Test log- likelihood (1) for the 2m station prediction (left) modti wind speed (right) experiments. The standard errors the oogliikeihhod are all below 0.010/0.020 For the station experiment, NN indicates that NN-CA is not employed. For the wind speed experiment, m- and s- stand for multi and single GE, FPT: forward pass time for a batch size of eight in ms.
(a) 2m station prediction experiment.
(b) Multi-modal wind speed experiment.
Model
GE
Grid size
Log-lik.
FPT
Model
GE
Grid size
Log-lik.
FPT
1)
CNP
0.636
32
1)
CNP
-1.593
33
PT-TNP
M 256
1.344
230
PT-TNP
M = 61
3.988
166
ConvCNP NN SetConv
64 x 128
1.535
96
Con NN SetConv
24 x 60
6.143
210
Swin-TNP PT-GE
64 x 128
1.819
127
Swin-TNP m-PT-GR
24 x 60
8.603
375
2) ViTNP KI-GE
48 x 96
1.628
167
2)
VITNP m-KI-GE VITNP m PT- 24 x - 12 30
12 30
5.371
349
VITNP KI-GE 144 x 288 48 x 96
1.734
171
7.288
392
3) VITNP KI-GE 144 x 288 1888-444 48 x 96
1.734
171
3)
VITNP m- KI-GE 24 x 60 + 12 x 30
6,906
372
VITNP PT-GE 144 x 288 +48 x 96
1.808
215
VITNP m-PT-GE 21 x 60 -> 12 30
7.288
392
Swin- TNP KI-GE
64% x 128
1.683
121
TNi m-KI-GE Swin-TNP m-PT-GR
24 x 60
7.603
355
Swin- PT-GE
64 x 128
1.819
127
24 x 60
8.603
375
4) ViTNP PT-GE
144 288 + 18 x 96
1.808
215
4)
ViTNP m-PT-GE 21
60 12 30
7.288
392
Swin-TNP PT-GE
64 x 128
1.819
127
Swin-TNP M-PT-GE
24 x 60
8.603
375
5) Swin- TNP PT-GE
64 x 128
1.819
127
6)
Swin-TNP S-PT-GE
24 x 60
8,073
364
Swin-TNP PT-GE
64 x 128
1.636
144
Swin-TNP m-PT-GE
21 x 60
8.603
375
7) ConvCNP NN SerConv ConvCNP SetConv
64 x 128
1.535
96
7)
COnvCNP NN SetConv
24 x 60
6.143
210
192 x 384
1.689
74
ConvCNP SetConv
48 x 120
7,841
64
Swin- TNP PT-GE Swin-TNF PT-GE
64 x 128
1.819
127
Swin-TNP m-PT-GE
24 60
8.603
375
192 x 384
2.053
306
Swin-TNP m-PT-GE
48 120
9.383
369
ing relatively low computational complexity.
pendix G.5 we conduct further experiments using larger grid show that
2. Performing patch encoding prior to projecting onto the latent grid gives better performance in ViTNPs.
sizes for the strongest Maseline-ConvCNN even at its largest configuration, ConvCNP underperforms relative to our best Swin-TNP variant and typically requires more parameters.
3. Pseudo-token approaches outperform kernel interpola- tion in the grid encoder, especially for coarse grids.
4. Swin-TNP achieves performance comparable to or better than that of ViTNP at a lower computational cost.
Incorporating Translation Equivariance Finally, we study the effect of translation equivariance with the effi- cient Swin-TNI (T') and the approximately equivariant Swin-TNP (T). The architectures are analogous to the non- equivariant Swin-TNF TNP. For the wind speed experiment, we also show the ability of Swin-TNP (?') to generalise spa- tially, by testing on a region unseen during training (i.e. we train on the US, and test on Europe). The results in Ta- ble show that 1) equivariant models outperform their non- equivariant counterpart, 2) they generalise spatially, and 3) relaxing strict equivariance further boosts performance.
5. In the grid decoder, nearest-neighbour cross-attention both reduces the computational cost and improves pre- dictive performance relative to full pross-attention, as highlighted in the t.2m station experiment.
For the multi- experiment, the use of the multi (m-) grid encoders outperform the use of single (s-) ones.
7. When varying the grid size, even the smaller Swin-TNP outperforms the larger ConvCNP baseline.
In Appendices G.2 and G.3 we show per-experiment results, as well as predictive errors for an example dataset for the Swin-TNP, ConvCNP, and PT-TNP, indicating that Swin- TNP also has the most accurate predictive uncertainties. For the 22m station prediction experiment, we provide results for the Swin-TNP using either solely t2m or skt measure- ments, showing that the model leverages both information sources. Additionally. we experiment with different num- bers of nearest neighbours in the encoder and decoder, as well as with richer context sizes, to evaluate the robustness of the model under varying configurations.. Finally, in Ap-
6. Conclusion
This paper introduces gridded TNPs, an extension to the family of TNPs which enables the use of efficient attention- based transformer architectures--uch as the VIT and Swin Transformer- and thus facilitates the application of TNPs large-scale spatio-temporal data. We also propose method to leverage gridded TNPs to address the scalability issues of translation equivariant TNPs, and incorporate a
8

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
Table Test log- log-likelihood (f) for Swin-TNP (T) and (T), compared the non-equivariant Swin-TNP. For the t2m station experiment
we report results for Swin-TNP PT-GE 64 x 128, while for the wind speed experiment, we use Swin-TNP m-PT-GE 24 x 60. FPT: forward pass time for batch size of eight in ms.
(a) 2 stations.
(b) Wind speed.
US
Europe
Model
Log-lik. (1)
FPT
Log-lik. (1)
Log-lik. (T)
FPT
Swin-TNP (2)
1.895
125
9.972
10.087
373
Swin-TNP (T)
1.926
142
10.086
10.429
467
Swin-TNP 1.819
127
8.603
-13.265
375
useful inductive bias for spatio-temporal tasks. Gridded TNPs consist of three components: the grid encoder, the grid processor, and the grid decoder. Through a comprehen- sive study, we identify the optimal design choices for each component. The best-performing gridded TNP variant in- corporates: 1. our novel pseudo-token grid encoder to move unstructured data onto structured grid pseudo-tokens, 2. the Swin Transformer to process this grid, 3. the pseu- do-token grid decoder equipped with the nearest-neighbour cross-atted cess-attention evaluate arbitrary locations, 4. and ap- proximate translation equivariance for efficient training and generalisation capabilities. We show that gridded TNPs out- perform strong baselines on large-scale synthetic and real- world regression tasks, with context sets with over 100,000 datapoints. This work marks step toward building flexible architectures for large, unstructured spatio-temporal data, with applications including, but not limited to, weather and environmental, as well as PDE modelling.
References
Allen, A., Markou, S., Tebbutt, W., Requeima, J., Bruinsma,
W. P. Andersson, T. R., Herzog, M., Lane, N. D., Chantry, M., Hosking, J. S., et al. End-to-end data-driven weather prediction. Nature, pp. 1 1-8, 2025.
Andersson, T. R., Bruinsma, W. P., Markou, S., Requeima, J., Coca-Castro, A., Vaughan, A., Ellis, A.-L., Lazzara, M. A., Jones, D., Hosking, S., et al. Environmental sensor placement with convolutional Gaussian neural processes. Environmental Data Science, 2:e32, 2023.
Ashman, M., Diaconu, C., Kim, J., Sivaraya, L.,, Markou, S., Requeima, J., Bruinsma, P., and Turner, Translation Equivariant Transformer Neural Processes. In Forty-first International Conference Machine Learn- ing, 2024a.
Ashman, M., Diaconu, C., Weller, A., Bruinsma, W., and Turner, R. Approximately Equivariant Neural Processes. Advances in Neural Information Processing Systems, 37: 97088-97123, 2024b.
Acknowledgements
CD is supported by the Cambridge Trust Scholarship. EL acknowledges financial support for this conference from
Bi, K., Xie, L., Zhang, H., Chen, X., Gu, X., and Tian, Q. Accurate medium-range global weather forecasting with 3d neural networks. Nature, 619(7970):333-536 2023. doi: 10. 10.10386415886003-0666555
the Amsterdam ELLIS Unit, Qualcomm and Ørpheus AI. AW acknowledges support from a Turing AI fellowship un- der grant EP/VO25279/1 and the Leverhulme Trust via CFI. RET is supported by Google, Amazon, ARM, Improbable, EPSRC grant HP/T005386/1, and the EPSRC Probabilistic Al Hub (ProbAI, EP/Y028783/I). We thank Anna Vaugh- aun, Stratis Markou and Wessel Bruinsma for their valuable feedback.
Bodnar, C., Bruinsma, W. P., Lucic, A., Stanley, M., Allen, A., Brandstetter, J., Garvan, P., Riechert, M., Weyn, J. A., Dong, H., Gupta, K., Thambiratnam, K., Archibald, A. T., Wu, C.-C., Heider, E.,, Welling, M., Turner, R. E., and Perdikaris, P. A foundation model for the Earth sys- tem. Nature, 2025, doi: 10.1038/s4 1104441586-000 5586-025-09005--1
Impact Statement
Bronstein, M. M., Bruna, J., Cohen, T., and Veličković, P. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges. arXiv preprint arXiv:2104.13477 2021.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Bruinsma, W. P. ral Processes. PhD thesis, Department of Engi- neering, University of Cambridge, 2022. URL https://wwet www.siepository... wam.oosittiroam.mmmo...mmmmmmm.. cam. ac, uk/ handle/1810/354385
Convolutional Conditional Neu-

---

Gridded Transformer Neural Processes for Spatio-Tempordl Data
Bruinsma, W. P., Requeima, J., Foong, A. Y. K., Gordon, J., and Turner, R. E. The Gaussian Neural Process. In Proceedings of the 3rd Symposium on Advances in Ap- proximate Bayesian Inference, 2021.
Garcia-Ortegón, M., Seal, S., Rasmussen, C., Bender, A., and Bacallado, S. Graph neural processes for molecules: evaluation on docking scores and strategies to improve generalization. Journal of Cheminformatics, 16(1):115, 2024.
Chen, K., Han, T., Gong, J., Bai, L., Ling, F., Luo, J.-J., Chen, X., Ma, L... Zhang, T., Su, R., et al. FengWu: Push- ing the Skillful Global Medium--Range Weather Forecast Beyond 10 Days Lead. arXiv preprint arXiv:2304.0294h. 2023a.
Gardner, J., Pleiss, G., Weinberger, K. Q. Bindel, D., and Wilson, A. G. GPyTorch: Blackbox Matrix-Matrix Gaus- sian Process Inference with GPU Acceleration. Advances in Neural Information Processing Systems, 31, 2018.
Chen, K., Ye, P., Chen, H., Han, T., Ouyang, W., Chen, T., BAI, L., et al. Fnp: Fourier neural processes for rrbittaty-resolution data assimilation. Advances in Neu- ral Information Processing Systems, 37:137847-137872. 2024.
Garnelo, M., Rosenbaum, D., Maddison, C., Ramalho, T.,
Saxton, D., Shanahan, M., Teh, Y. W., Rezende, D., and Eslami, S. A. Conditional Neural Processes. In Interna- tional Conference on Machine Learning, pp. 1704-1713. PMLR, 2018a.
Chen, L., Zhong, X., Zhang, F., Cheng, Y., Xu, Y., Qi, Y., and Li, H. FuXi: A cascade machine learning forecasting system for 15-day global weather forecast. NPJ Climate and Atmospheric Science, 6(1):190, 2023b.
Garnelo, M.. Schwarz, J., Rosenbaum, D., Viola, F, Rezende, D. J., Eslami, S., and Teh, Y. W. Neural Pro- cesses. arXiv preprint wxXiv1807.01622 2018b.
Dao, T. 2ltetonttn Faster Attention with Better Parallelism and Work Partitioning. In Proceedings of the International Conference on Learning Representations (ICLR), 2024.
Gilmer, J., Schoenholz, S. S., Riley, F., Vinyals, O., and Dahl, G. E. Neural Message Passing for Quantum Chem- istry. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1263-1272 PMLR, 06-11 Aug 2017.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. FlashAt- tention: Fast and Memory-
Exact Attention with AO-aeeressss. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 16344-16359, Curran Associates, Inc., 2022.
Gordon, J., Bruinsma, W. P., Foong, A. Y., Requeima, J., Dubois, Y., and Turner, R. E. Convolutional conditional neural processes. In Proceedings of the International Conference on Learning Representations (ICLR), 2020.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn,
He, C., Li, R., Li. S.,, and Zhang, L. Voxel Set Transformer: A Set-to-Set Set Approach to 3D Object Detection from Point Clouds. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 8417-8427, 2022.
D.. Zhai, X., Unterthiner, T., Dehghani, M. Minderer, M., Heigold, G., Gelly, S., et al. An Image is Worth 16x 16 Words: Transformers for Image Recognition at Scale. In International Conference on Learning Representations, 2020.
Hersbach, H., Bell, B., Berrisford, P., Hirahara. S., Horányi, A., Muñoz-Sabater, J., Nicolas, J., Peubey, C., Radu, R., Schepers, D., Simmons, A., Soci, C., Abdalla, S., Abel- lan, X., Balsamo, G., Bechtold, P., Biavati, G., Bidlot, J., Bonavita, M., De Chiara, G., Dahlgren, P., Dee, D., Diamantakis, M., Dragani, R., Flemming, J., Forbes, R., Fuentes, M., Geer, A., Haimberger, L., Healy, S., Hogan, R. J., H6lm, E., Janisková, M., Keeley, S., Laloyaux, P., Lopez, P., Lupu, C., Radnoti, G., de Rosnay, P., Rozum, I., Vamborg, F., Villaume, S., and Thépaut, J.-N. The ERA5 global reanalysis. Quarterly Journal of the Royal Meteorological Society, 146(730): 1999)))::999-0000 2049, 2020.
Dunn, R. H., Willett, K. M., Thorne, P. W., Woolley, F. V., Durre, I., Dai, A., Parker, D. E., and Vose, R. S. HadISD: quality-controlledd global synoptic report database for selected variables at long-term stations from 1973-2011. Climate of the Past, 865):1649-1679, 2012. doi: 10.5194/ cp-8- 1049-2224- URL httpss:/s//ppt copernicus. org/articces//////////o htt.aaaiiiccee/////////222/22
Feng, L., Hajimirsadeghi, H., Bengio, Y., and Ahmed, M. O. Latent Bottlenecked Attentive Neural Processes. In The Eleventh International Conference on Learning Repre- sentations, 2023.
Feng. L., Tung, F., Hajimirsadeghi, H., Bengio, Y., and Ahmed, M. O. Memory Efficient Neural Processes via Constant Memory Attention Block. In Forty-fire Interna- tional Conference on Machine Learning, 2024.
Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., and Carreira, J. Perceiver: General Perception with Iterative Attention. In International Conference on Ma- chine Learning, pp. 4651 1664. PMLR, 2021.
10

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
Janny, S., Benetteau, A., Nadri, M., Digne, J., Thome, N., and Wolf, C. Eagle: Large-scale learning of turbulent fluid dynamics with mesh transformers, In International Conference on Learning Representations (ICLR), 2023.
Mao, J., Xue, Y., Niu, M., Bai, H., Feng. J., Liang, X., Xu, H., and Xu, C. Voxel Transformer for 3D Object Detection. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 3144-3153, IEEE Com- puter Society, 2021.
Kim, H., Mnih, A., Schwarz, J., Garnelo, M., Eslami, A., Rosenbaum, D., Vinyals, O., and Teh, W. Attentive Neural Processes. In International Conference on Learn- ing Representations, 2019.
McNally, A., Lessig, C., Lean, P., Boucher, E., Alexe, M., Pinnington, E., Chantry, M., Lang, S., Burrows, C., Chrust, M., et al. Data driven weather forecasts trained and initialised directly from observations. arXiv preprint arXiv:2407. 15586, 2024.
Kovachki, N., Li, Z., Liu, B.. Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces with applications to PDEs. Journal of Machine Learning Re- search, 24(89):1-97, 2023.
Nguyen, T. and Grover, A. Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Model- ing. In International Conference on Machine Learning. pp. 16569-16594. PMLR, R, 2022.
Citimarriin--nntte D., Meanti, G., Pagliana, N., Verri, A., Mazzino, A., Rosasco, L., and Seminara, A. Physics informed machine learning for wind speed prediction. Energy, 268:12 2023.
Nguyen, T., Brandstetter, J., Kapoor, A., Gupta, J. K., and Grover, A. ClimaX: A foundation model for weather and climate. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, vol- ume 202 of Proceedings of Machine Learning Research, pp. 25904-25938, PMLR, 23-29 Jul 2023.
Lam, R.. Sanchez-Gonzalez, A., Willson, M., Wirnsberger, P., Fortunato, M., Alet, F., Ravuri, S., Ewalds, T., Eaton- Rosen, Z., Hu, W. et al. Learning skillful medium-range global weather forecasting. Science, 382(6677): 1416- 1421, 2023.
Niu, R., Wu, D., Kim, K., Ma, Y., Watson-Parris D., and Yu, R. Multi-Fidelity Residual Neural Processes for Seal- able Surrogate Modeling. In Forty-first International Conference on Machine Learning, 2024.
LeCun, Y., Bottou, I.., Bengio, Y., and Haffner, P. Gradient- Based Learning Applied to Document Recognition. Pro- ceedings of the IEEE, 86(1 (::27))):2822223 1998.
Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh,
Owens, R. and Hewson, T. ECMWF Forecast User Guide. Technical report, ECMWF, Reading, 05/2018 2018.
Y. W. Set Transformer: A Framework for Attention- Based Innuutaaion-Invariann Neural Networks. In In- lernational Conference on Machine Learning, pp. 3744- 3753. PMLR, 2019.
Pagliardini, M., Paliotta, D., Jaggi, M., and Fleuret, F. Fast attention over long sequences with dynamic sparse flash attention. In Thirty-sevnnth Conference on Neural Infor- mation Processing Systems, 2023.
Lessig, C., Luise, I., Gong, B., Langguth, M., Stadtler, S., and Schultz, M. AtmoRep: A Stochastic Model of Atmosphere Dynamics Using Large Scale Representation Learning. In Proceedings of the NewrlPS 2023 Workshop on Tackling Climate Change with Machine Learning: Blending New and Existing Knowledge Systems, 2023.
Price, I., Ganchez-Conzalez, A., Alet, F., Andersson, T. R., El- A., Masters, D., Ewalds, T., Stott, J., Mohamed, S., Battaglia, P., et al. Probabilistic weather forecasting with machine learning. Nature, 637(804)):84-90, 2025.
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Ronneberger, O., Fischer, P., and Brox, T. U-Net: Convo- lutional Networks for Biomedical Image Segmentation. In Medical Image Computing and Computer-Assisted Intervention 2010: 18th International Confer- ence, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234-241. Springer, 2015.
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. In Proceedings of the 7th International Con- ference Learning Representations (ICLR), 2019.
RuBwurm, M., Klemmer, K., Rolf, E.,, Zbinden, R., and Tuia, D. Geographic Location Encoding with Spherical Harmonics and Sinusoidal Representation Networks. In The Twelfth International Conference on Learning Repre- sentations, 2024.
Lu, D., Xie, Q., Wei, M., Gao, K., Xu, L., and Li, J. Trans- formers in 3D Point Clouds: A Survey. arXiv preprint DxXiv:2055007417 2022.

---

Gridded Transformer Neural Processes for Spatio- Tpatoo-Temporal Data
Ryali, C.,, Hu, Y.-T., Bolya, D., Wei, C., Fan, H., Huang, P.-Y,, Aggarwal, V., Chowdhury, A., Poursaeed, O., Hoff- man, J., Malik, J., Li, Y., and Feichtenhofer, C. Hiera: A Hierarchical Vision Transformer without the Bells-and- Whistles. In Krause, A., Brunskill, E., Cho, K., Engel- hardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learn- ing. volume 202 of Proceedings of Machine Learning Research, pp. 29441 -99444. PMLR, 23-29 Jul 2023.
Wilson, A. and Nickisch, H. Kernel Interpolation for Scal- able Structured Gaussian Processes (KISS-GP). In In- ternational Conference Machine Learning, pp. 1775- 1784. PMLR, 2015.
Xiao, Y., Bai, L., Xue, W., Chen, H., Chen, K., Chen, K., Han, T., and Ouyang, W. Towards Self-contained Data-driven Global Weather Forecasting Framework. In Salakhutdinov, R., Kolter, Z. Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.), Pro- ceedings of the 41st International Conference on Machine Learning, volume 235 Proceedings of Machine Learn- ing Research, pp. 54255-5444 54275. PMLR, 21-27 Jul 2024.
Schmude, J., Roy, S., Trojak, W., Jakubik, J., Civitarese, D. S., Singh, S., Kuehnert, J., Ankur, K., Gupta, A., Phillips, C. E., Kienzler, R., Szwareman, D., Gaur, V., Shinde, R., Lal, R., Silva, A. D., Diaz, J. L. G., Jones, A., Pfreundschuh, S., Lin, A., Sheshadri, A., Nair, U., Anantharaj, V., Hamann, H., Watson, C., Maskey, M., Lee, T. J., Moreno, J. B., and Ramachandran, R. Prithvi WxC: Foundation Model for Weather and Climate. arXiv preprint arXiv:2409. 00.44099.3990 13598, 2024.
Xu, X., Sun, X., Han, W., Zhong, X., Chen, L., Gao, Z., and Li, H. Fuxi-da: A generalized deep learning data assimi- lation framework for assimilating satellite observations. npj Climate and Ammospheric Science, 8(1):156, 2025.
Zhang, C., Wan, H., Shen, X., and Wu, Z. PVT: Point- Voxel Transformer for Point Cloud Learning. Inter- national Journal of Intelligent Systems, 37(12):11185- 12008, 2022.
Shaw, P., Uszkoreit, J., and Vaswani, A. Self-Attention with Relative Position Representations. In Walker, M., Ji, H., and Stent, A. (eds.), Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, Volume 2 (Short Papers), pp. 464 New Orleans, Louisiana, June 2018. Association for Computational Lin- guistics. doi: 18.18653/v 1/N18-20044 2074.
Tu, Z., Talebi, H., Zhang, H., Yang, F., Milanfar, P., Bovik, A., and Li, Y. Maxvit: Multi-axis vision transformer. In Avidan, S., Brostow, G., Cissé, M., Farinella, G. M., and Hassner, T. (eds.), Computer Vision - ECCV 2022, pp. 459 479, Cham, 2022. Springer Nature Switzerland.
ISBN
Tychola, K. A., Vrochidou, E., and Papakostas, G. A. Deep learning based computer vision under the prism of 3D point clouds: a systematic review. The Visual Computer,
pp. 1-, 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, t.., and Polosukhin, I. Atten- tion Is All You Need. Advances in Neural Information Processing Systems, 30, 2017.
Vaughan, A., Tebbutt, W., Hosking, J. S., and Turner, R. E. Convolutional conditional neural processes for local cli- mate downscaling. Geoscientific Model Development, 15 (1):251-268, 2022.
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph Attention Networks. In 6th International Conference on Learning Representations (ICLR), 2018
12

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
A. A Unifying Construction of Conditional Neural Processes
The many variants CNPs differ in their construction of the predictive distribution p(-|x,, Dz). Here, we introduce a construction of CNPs that generalises all variants. CNPs are formed from three components: the encoder, the processor, and the decoder. The encoder, e: X' y + first encodes each (Xe,s) Ve,n) D. into some latent representation, or token, Z Z. The processor, p: (U 2n) X' + 2, processes the set of context tokens e(Dc) {c(Xe,,x Ve,u h together with the target input obtain a target dependent token, E Z Finally, the decoder, d: Z + Py, maps from the target token to the predictive distribution over the output at that target location. Here, Py denotes the set of distributions over V. We illustrate this decomposition in Figure 4.
Encode, e(.) {Ze,n}r }
Process, p(:,x:) 2t
Decode, d(-)
(xe, Yenia
p( 1 Xc)
Figure unifying construction of CNPs, with D. {(Xe,n Ve,a) and Be,n e(Xc,ns Ve..).
We present below several schematics showing the architectures of different members of the CNP family, and detail how they can be constructed following this universal construction.
Original CNP The first architecture is based on (Garnelo et al., 8a) and is the least complex of the CNP variants, using summation the permutation-invarian aggregation. The diagram is shown in Figure 5. The encoder of CNP is MIP, which maps from each concatenated pair (e,,) E Dr some representation Ze,n E RDp The processor sums together these representations, and combines the aggregated representation with the target input using pl p( In: x,) = +(E,. Zo,n, xz). T is often just the concatenation operation. Finally, the decoder consists another MLP which maps from Z₁ = =(En, Ze,n X₂) to the parameter space some distribution over the output space (e.g. Gaussian).
z (1)
ple(D.).z.)
Decoder
2x(D,)
Processor
Point-wise embeddings
Context dependent
Target dependent
6e
dC
Context target dependent
(Wels
(IeeNNs McNN)
Figure 5. A diagram illustrating the architccture of the plain CNP (Gamelo ct al., 2018 First, the context set (Xe.n, Fe.u) and the target tokens Xt.x encoded using point-wise embeddings. These fed into the CNP processor, which performs simple permutation- invariant aggregation of the context tokens. These are then concatenated with the target tokens and fed into the decoder. which outputs the parameters of the specified NP distribution based target representation (in this case, mean variance a Gaussian).
'The space of target dependent tokens does not need to be the same that of context tokens have used Z in both cases for simplicity. It is also possible for Z to be the product of multiple spaces, e.g. Z Zroken N' where we retain information about the input locations.
13

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
ConvCNP Another member of the CNP family is the ConvCNP (Gordon al., 2020), that embeds sets into function space in order to achieve translation equivariance. This can lead more efficient training in applications where such inductive bias is appropriate, such as stationary time-series spatio-temporal regression tasks. We show in Figure 6 schematic of the architecture. The encoder of ConvCNP is simply the identity function. The processor then encodes the discrete function represented by input-output pairs {(x,n, Ye,a) onto a regular grid using the kernel-interpolation grid encoder (KI-GE). It then processes this grid using a CNN, which is afterwards combined with the target location using a kernel-interpolation grid decoder (KI-GD). Letting U (ummm denote the set of values on gridded locations
V {vm)m we can decompose the ConvCNP processor as
Um + Σ ye,u Vege Vege(v Xx,n.)
(6)
U t CNN(U,
(7)
2t f Σ Un Umt/ga(X Vnn.
(8)
Here, gee Vod: R x R + R denote the KI-GE kernel and KI-GD kernel. As with the CNP, the decoder is an MLP
mapping to the parameter space some distribution over the output space.
2 (1)
ple(D.),
Ze (D.)
- 4))
(ony!t)
Decoder
Processor
Cony
z (D2)
Point-wist embeddings
1. Ben) N/p - xen)
e(D)) (D(
Context dependent
Φe
Φ
61
Target dependent
Context and larget dependent
(1e,1 Vc.1)
(aaNNs Me,N))
21
Figure diagram illustrating architecture of ConvCNP (Gordon 2020), First, the context (Xe.n, Vo,n) and target tokens Xt.n are encoded using the identity encoder. These are fed into the ConvCNP processor, which uses KI-GE to project the tokens into function space. These are then evaluated discrete locations using a pre-specified resolution. followed by multiple layers of CNN-based architecture acting upon the discretised signal. To decode at arbitrary locations, KI-GD is used, giving rise to the target token representation. This fed into the decoder, which outputs the parameters of the specified NP distribution (in this case, mean and variance of Gaussian).
There are a number different architectures used for the different members of the TNP family. We provide below diagrams for two members mentioned the main paper, namely the TNP of (Nguyen & Grover, 2022) and the induced
TNPs
14

---

Gridded Transformer Neural Processes for Spatio-Temports Data
transformer (ISTNP) of (Lee al., 2019). For the standard TNP, the encoder consists of MLP mapping from each (Xe,n Ye,n) De some representation (token) Be,u DR The processor begins by embedding the target location in the same space the context tokens, giving ž E IRP It then iterates between applying MHSA operations on the set of context tokens, and MHCA operations from the set of context tokens into the target token:
2 + MHSA(Z,) (Zc) L.
(9)
24 < MHCA (z;; Z,)
Again, the decoder consists of an MLP mapping from 2ς to the parameter space of some distribution over outputs.
μ
z plc(DD),
(D.)
MHCA (s. (D.)
MHSA ( (Z!--1) V/
Decoder
2.
MHCA (2
(3) Z (D.))
Processor
MHSA (3) (2."))
(D.)
MHCA W1 (#., 2. (D.)
Point-wise embeddings
MHSA (Z.)
(D.) (.().
Context dependent
6
6
Target dependent
Context and target dependent
(We,1s Vc.1)
(Ee,N's Me,Nr)
21
Figure diagram illustrating the architecture of the TNP (Nguyen Grover, 2022). First, the context (Xe..) Ve..) and the target tokens Xt.n are encoded using point-wise embeddings to obtain the context set representation Z. and target representation Z, These are fed into the TNP processor, which takes in the union of [Ze, Ze and outputs the token corresponding to the target inputs Z(4) At each layer of the processor, the context set representation first updated through an MHSA layer, which representation through MHCA layer between the target set representation from the previous layer and the updated context representation. Finally, decoder outputs the parameters of specified NP distribution based on the target representation from final layer (in this case, mean and variance of a Gaussian).
used to modulate target set
One of the main limitations of TNPs is the cost of the attention mechanism, which scales quadratically with number of input tokens. Several works (Feng et al., 2023; Lee al., 2019) addressed this shortcoming by incorporating ideas from the Perceiver-style architecture (Jaegle al., 2021) into NPs. The strategy to introduce act information bottleneck between the context and target sets. Provided that M << Ne, where N. is the number of context points, this leads to significant reduction in computational complexity. The architecture we consider in this work is called the induced set transformer NP (ISTNP), and differs from the plain TNP in the calculations performed in the processor. At each layer, the pseudo-token representation first updated through an MHCA operation from the context set to the pseudo-tokens. The updated representation is then used to modulate the context and target sets separately, through separate MHCA operations:
of M 'pseudo-tokens' which
15

---

Gridded Transformer Neural Processes for Spatio-Tenan Temporal Data
U < MHCA(U; Z.)
2, + MHCA(Z;: U) x
(10)
Z. + (ZC U)
Thus, as apparent in Figure 8, the context and target sets never interact directly, but only through the 'pseudo-tokens' The computational cost at each layer reduces from O(N2 N,N,)) in the plain TNP, where N. and N, represent the number of context and target points, respectively. to (MM, + Nt)). This is significant reduction provided that M << No. resulting in apparent linear dependency on Ne. However, in practice, M is not independent of Ne, with more pseudo-tokens needed as the size of context set increases.
Ct
(.))).).
10 .
(1 g
MICCA (2] 3, (,eri
2
Decoder
2
MIICA (2!" gr00g
MHCA (erla, 2!!
MECCA (" pr(0))
Processor
Point-wise embeddings
MECA (2,, D
MHCA (U 2.)
MIICA (4, tri
(b.J 275
depandent
de
40
de
Targat depandent
Contest and target depondent
(Pa1)
(30_,0 Ml,N)
Figure 8. diagram illustrating the architecture of the ISTNP (Lee et al., 2019). As opposed the regular TNP of Nguyen & Grover 2022, the ISTNP uses a 'summarised" representation of the context set through the use of pseudo-tokens. They are first randomly initialised (U))). Then, their representation is updated through cross attention with the context representation from the previous layer (i.e. at layer l: @ MHCA ("(("- z )) This updated set of pseudo-tokens tokens (()) is then used to modulate both the context set representation at the current layer through eross aatenntntn (i.e. Z MHCA(Z (1-1) U()), well the target representation (i.e. T MHCA(Z( U.)). Thus, the context and target set representations do not interact directly, but only through the pseudo-tokens, which act as an information flow bottleneck between the two in order to decrease the computational demands of the TNP.
B. Translation and Approximate Equivariance
We aim to give a brief introduction to Translation Equivariant TNPs (TE-TNPs) (Ashman al., 2024a) and Approximately Equivariant NPs (Ashman et al., 2024b), and refer the reader to the aforementioned papers for more thorough treatment. Pseudo-code for a forward pass through these models is provided in Algorithms 3 and 4 in Ashman et al. (2024b).
B.1. Translation Equivariant TNPs
As mentioned in Ashman et al. (2024a), if the ground-truth stochastic process are interested in modelling is stationary, building translation equivariance into the model leads to more efficient training by reducing the model search space. Moreover, as formally indicated in Theorem 2.2 in Ashman et al. (2024a), this inductive bias also allows NPs to generalise spatially.
In order to achieve translation equivariance, we must ensure that each token zn translation invariant with respect to its
16

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
inputs-this can be achieved by only allowing relative, rather than absolute dependencies on the inputs. More specifically, the attention mechanism only depends on pairwise distances X; Xj, resulting in the translation equivariant MHSA and MHCA operations, which we denote TE-MHSA and TE-MHCA. In addition, unlike the non-equivariant counterpart, we force the token encoder to only depend on output values yn.
The TE-MHSA operation is constructed analogously to Equation 1, but the inputs to the attention mechanism only depend on pairwise differences between the inputs, rather than their absolute value
Zn f- cat (
Σ Wi
Oi (Zns Zen Xn. Xm Zar WW V.h
H
Wo.
(11)
h=1
There various ways in which the translation equivariant attention weighting can be implemented, but we choose to follow Ashman et al. (2024a) and set Ot,,n.m = Orh(Z,, Zm: Xn - Xm), as indicated in Equation 3, as
(0a " "W.. (.,, 2m. Xmm)
Chamm
EM=Ic cP (2. (z.. W... (Wx.a)' "m X. - X.a)
(12)
where Ph: R x R. + R is a learnable function, parameterised by MLP, and {zn N only depend output values. TE-MHCA is defined analogously
N2 On (Z 22,m, X1,n X2,m 78..m
Z1,n cat. (f
T Ww.. V.l A=1
Wo.
(13)
In Ashman et al. (2024a), the authors also propose optional translation equívariant updates to the inputs, but we do not use them in our implementation. Together with layer normalisation and pointwise MLPs, the TE-MHSA and TE-MHCA operations form the TE-MHSA and TE-MHCA blocks. The TE-TNP (TNP (T)) shares identical architecture to the regular TNP, with the MHSA and MHCA blocks replaced by TE-MHSA and TE-MHCA blocks.
Translation Equivariant Gridded TNPs MHCA blocks replaced by TE-MHSA / TE-MHCA blocks:
The architecture of Swin-TNP (T") is similar to that of Swin-TNP, with MHSA/ /
1.
Pseudo-token grid encoder: we follow the definitions from Equation 4, but the pseudo-tokens are shared for each grid No
point, and { [ze.n JWM., only depend on output values {ye.n
1m < TE-MHCA (u
Vm E M.
(14)
2. Grid processor: We use Swin Transformer augmented with fixed set of learnable relative positional encodings. The number of these encodings is equal to the number of grid points multiplied by the number of heads in the attention mechanism. Their introduction comes with negligible computational cost.
3. Pseudo-token grid decoder: similarly to Equation 5 we perform nearest-neighhour cross-attention, crstt blocks with TE-MHCA blocks.
but replace the MHCA
Ztn TE-MHCA (2f,ns (immmme
(15)
B.2. Approximately Equivariant NPs
Although introducing inductive biases such as translation equivariance has been shown to improve sample efficiency and the ability to generalise, real-world data rarely perfectly symmetric. To address this, Ashman al. (2024b) propose Approximately Equivariant Neural Processes, which relax strict equivariance assumptions within the training domain in a data-driven manner, while maintaining equivariance out-of-distributio to allow for good generalisation capabilities.
The main insight towards achieving approximate equivariance is outlined in Theorems 2 and 3 in Ashman et al. (2024b), indicating that any non-equivariant operator between function spaces can be constructed as, or approximated by, an equivariant mapping with additional, fixed functions as input. They use this result to construct the approximately equivariant neural process, which, following the unifying construction in Appendix A, can be constructed as a G-equivariant CNP with encoder X x y >> Z, G-invariant processor p: (U 2n) X - 2, and a decoder d: -> Py. Approximate
17

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
equivariance is achieved by inserting B fixed inputs (4 into the processor, which become additional model parameters that break the G-invariance of the processor:
N.
N, II ply (D)) Kr, (1..,
p(Y,| D.) II ply Xr, D2)
(16)
The authors argue that the number of additional inputs determines the degree to which the model breaks equivariance, with more inputs resulting in more non-equivariant processor. In the limit B + oo we eventually recover any non-equivariant processor.
There are various ways in which these additional inputs can incorporated into G-equívariant architectures. Here, we only outline how they are included into the translation equivariant TNP (7). The authors consider fixed functions on {xe,n) 1 which are summed to the initial context tokens {ze.n 1 These are obtained by performing a Fourier expansion on the context token locations, followed by passing the result through an MLP with two hidden layers of dimension Dz.
The authors introduce two additional techniques to ensure generalisation beyond the training distribution. These are features inside the training domain;
motivated by the idea that we only have information about local symmetry- outside, where no data are available, it is more reliable to assume strict equivariance. To determine the most suitable cquivariant mapping that should be applied outside the training regime, the authors propose the following view: Any approximately equivariant mapping can be decomposed into an equivariant component, which is, in some sense, the best "cquivariant approximation" of the original mapping, and a residual non non non-equivariant component. The latter can then be approximated by an equivariant mapping with additional fixed inputs (according to Theorems and 3). Thus, they propose the following:
1. During training, they set the additional basis functions to with some fixed probability. This allows the models to learn good approximation of the equivariant component of the underlying system when the additional inputs are zero.
2. Forcefully set the additional fixed inputs to zero outside the training domain.
In our implementation of the Swin-TNP (T), we only add additional fixed inputs to the initial context set embeddings- before feeding (Ze,n No into the pseudo-token grid encoder, we perform:
Ze,n = =(Y=,n) + ((xe.))
(17)
where o is implemented with an MLP, and t(x,..) is obtained by performing the Fourier expansion, as mentioned above.
C. Related Work
Models for Structured Weather Data As mentioned in Section 3, recently there has been significant interest from the research community in developing models that are able to scale to massive spatio-temporal datasets-sone of the most successful of which can be seen in weather modelling. The majority operate on structures, gridded data, and are only trained to perform forecasting. Models such as Aurora (Bodnar et al., 2025), GraphCast (Lam et al., 2023), GenCast (Price et al., 2025), Pangu (Bi et al., 2023), FuXi (Chen et al., 2023b) and FengWu (Chen et al., 2023a) share a similar encoder-processoen dreceesorrddeodder to our construction of CNPs presented in Section 2.1, in which the input grid is projected onto some latent space which is then transformed using an efficient form of information propagation (e.g. sparse attention with Swin Transformer (Bodnar et al., 2025; Bi et al., 2023), message passing with GNNs (Price et al., 2025)) before being projected back onto the grid at the output.
At the same time, there has been line of work aiming to develop more flexible methods for spatio-temporal data, either by enabling models to deal with unstructured data, or by adjusting the training task. Perhaps the most relevant works to ours are Aardvark (Allen et al., 2025) and FuXi-DA (Xu et al., 2025), both end-to-end weather prediction models that handle both unstructured and structured data. Aardvark employs kernel interpolation, also known as SetConv (Gordon et al., 2020), to move unstructured data onto grid, followed by a VIT which processes this grid and outputs gridded predictions, whereas FuXi-DA simply averages observations within each grid cell. We provide extensive comparison between the kernel interpolation approach to structuring data and our pseudo-token grid encoder in Section 5, demonstrating significantly better performance. Additionally, in the GP experiment, we evaluate our encoder against FuXi-DA's encoder, with the results in Table 3 showing superior performance.
Another relevant example is Lessig et al. (2023), a task-agnostic stochastic model of atmospheric dynamics, trained to predict randomly masked or distorted tokens. While the task agnosticism of this approach shares similarities with that of
18

---

Gridded Transformer Neural Processes for Spatio-Tempore Data
NPs, it is unclear how this approach extends to settings in which data can exist arbitrary spatio-temporal locations, and they are more limited in their approach to model multiple sources of the input data.
Finally, Schmude et al. (2024) propose Prithvi WxC, foundation model for weather and climate that can be effectively tuned for multiple distinct use cases. To this end, they employ a training objective that combines ideas from masked reconstruction with forecasting, whereby they predict a future state with a variable lead time, based on masked current and past states. While able to adapt to various downstream tasks through fine-tuning (which generally requires additional architectural components), the backbone transformer is only able to operate on gridded data. The architecture is based on a 2D vision transformer that combines approaches from Hiera (Ryali et al., 2023) and Max Vit (Tu et al., 2022). This implies that, as it is, the model is not able to handle data at arbitrary apatooltemporal locations, and would need architectural adaptations to do so. Indeed, our grid encoder could represent a solution to move the initial unstructured data onto structured grid, compatible with the foundation model's backbone.
Connections between our method and GNNs It is well known that regular transformers can be interpreted as GNNs (Bron- stein al., 2021); in particular, fully-connected graph attention networks (GAT; Veličković et al. 2018). This interpretation extends to our construction of the transformer processor. In particular, we can interpret the grid encoder as dynamically constructing graph connecting the tokenised observations to their nearest neighbours on some latent grid of pseudo-tokens with directed edge. The operation of our pseudo-token grid encoder is then single GAT update. The grid processor can be interpreted as a fully-connected GAT when a ViT is used, or a sparsely-connece Finally, similar to the interpretation of the grid encoder, the grid decoder can be interpreted as dynamically constructing a graph connecting the tokenised target locations to their nearest neighbours on the same latent grid of pseudo-tokens with a (reversed) directed edge, and performing single GAT update.
GAT when a Swin Transformer is used.
Given these connections, it reasonable to question why we chose use attention-based updates rather than other forms of GNN updates, such as message passing. Our response is straightforward: empirical evidence suggests that transformer-batt based implementations achieve superior computational efficiency compared to GNN counterparts for which operations cannot be performed in parallel (Veličković et al., 2018). Furthermore, the widespread adoption of transformers has led to transformer- specific operations, such as the attention mechanism, being highly optimised for modern GPU architectures (Dao et al., 2022; Dao, 2024; Pagliardini et al., 2023).
Universal Physics Transformers tions and methodologies as work. UPTs consider a different, but related problem of learning approximate solutions to partial differential equations (PDEs), which can be formulated as learning mappings between function spaces. The broader family of models which achieve this are known as neural operators (Kovachki et al., 2023). In practice, neural operators work with a finite number of function evaluations according to some discretisation scheme. Dependent on the choice of discretisation scheme, the set of function evaluations may not lie on a regularly spaced grid or mesh, and when the number of function evaluations is large we arrive at similar problem to that considered in this work. UPTs use a single message passing layer (Gilmer al., 2017) local neighbourhoods of tokenised representations of the discretised function evaluations into a set of pseudo-tokens. The primary difference between our pseudo-token grid encoder and this approach is that in UPTs the pseudo-tokens are not constrained to lie regular grid; rather, the pseudo-locations are randomly sampled from the original set of discretised locations. Their approach is motivated by the desiderata for the original discretisation structure to be preserved; however, it prevents application of efficient grid-based transformer architectures such as Swin Transformer.
The recently proposed universal physics transformers (UPTs; ?) share similar motiva-
D. Kernel-Interpolation Grid Encoder
Let z,, E Rr denote the token representation of input-output pair (x)) after point-wise embedding. We introduce the set of grid locations V E R-Mx xMM,x.-M,. xD. For ease of reading, we shall replace the product D. Ma with M and the indexing notation mi,..., m,,, with m.
The iereel-interpolatioon grid encoder obtains a pseudo-token representation U E M of D. on the grid V by interpolating from all tokens ( {Ze,n In at corresponding locations (xe,a)n In to all pseudo-token locations:
1m
Σ Eeen Vmn Xc,n) Vm E M.
(18)
Here, 1b: X x X + R the kernel used for interpolation, which we take to be the squared-exponential (SE) kernel when
19

---

Gridded Transformer Neural Processes for Spatio-Tenan Temporal Data
X' --RR, :.:
D2 [ee,n,d - (m,d)
VVe(VVm Xx,n) exp
(19)
Lz
where &a denotes the 'lengthscale" for dimension d. Similar the pseudo-token grid encoder, we can restrict the kernel- interpolation grid encoder interpolate only from sets of token { (ne for which m is amongst the k nearest grid locations. The computational complexity of the kemee-interpolation and pseudo-token grid encoders differ only by scale factor when this approach is used.
E. Nearest-Neighbour Cross-Attention in the Grid Decoder
In this section we provide more details, alongside schematics, of our nearest- neirhst-neighhour pross-attention scheme. As mentioned in the main text, in the grid decoder only allow subset of the gridded pseudo-tokens attend to the target token (those in the vicinity of it). Finding the nearest neighbours is not done in the standard k-nearest neighbours fashion, because we use the same number of nearest neighbours along each dimension of the original data (i.e. latitude and longitude for spatial interpolation; latitude, longitude and time for spatio-temporal interpolation). This is to ensure that we do not introduce specific preferences for any dimension. In the case the data lies on grid with the same spacing in all its dimensions, the procedure becomes equivalent to k-nearest neighbours.
In practice, we specify the total number of neirest-n we want use k. We then compute the number of nearest- neighbours in each dimension by kidim ceil am()) ). where dim(x) represents the dimensionality of the input. For efficient batching purposes, we tend to choose &o = (2n - 1)dim(z), where n E N (i.c. 9 for experiments with loogiue-longitude grids, 27 for experiments with latitude-longitud- time grids). We then find the indices in each dimension of these nearest- neighbours by performing an efficient search that leverages the gridded nature of the data, leading to a computational complexity of (kNN,, with N, the number of target points. When the neighbours go off the grid (i.e. for targets very close to the edges of the grid), only consider the number of viable (i.e. within the bounds of the grid) neighbours.
This procedure is visualised in Figure 9, where we consider both grids with the same spacing along each dimension, well as grids with different spacings. We cover both the case of central target point, as well target point closer to the edges of the grid.
Example 2D
2.0
2.0
3
1.0
1.0
Dimension2
0.5
Dimension 2
0.5
Dimension
Dimonsion
-0.5
-0.5
-1.0
-2
X
-1.5
-1.5
-20
-2.0
Dimension
Dimension
Dimension
Dimension
(a) Same/ / Central
(b) Same Edge
(c) Diff' / Central
(d) Diff Edge
Figure Example of our nearest- neighbous procedure in 2D on 5 x grid for nearest neighbours. We consider four different cases.
Same / Diff refers to whether the grid spacing is the same in the two dimensions or different. Central Edge refers to the position of the target. In the case of an edge target, we do not consider invalid neighbours (i.e. those that are outside the grid bounds).
Accounting for non-Euclidean geometry A lot of our experiments are performed on environmental data, distributed across the Earth. For our purposes, we assume the Earth shows cylindrical geometry, whereby there is no such thing as a grid edge along the longitudinal direction (i.e. a longitude of -880 is the same as 1809). This is not the same for latitude, where one extreme corresponds to the North Pole, and the other one to the South Pole. Thus, we would like to allow for the grid 'roll' around the longitudinal direction when computing the nearest neighbours. In this case, there should be no edge
20

---

Gridded Transformer Neural Processes for Spatio-Temporyy Data
target points along the longitudinal dimension. This procedure is graphically depicted in Figure 10 and it in experiments environmental data.
Non-rolling Dimension
-1
-3
-2
Rolling Dimension
Figure 10. Example of nearest- neirest-neighhours procedure in 2D on a x grid for 9 nearest neighbours. We allow rolling along the horizontal dimension (e.g. longitude), but do allow rolling along the vertical one (i.e. latitude). Hence, the neighbours extend the other side of the grid horizontally, but not vertically. This example corresponds to cylindrical symmetry.
Example for 3D data We also provide examples of the nearest-neiighburss procedure on 3D spaces. These dimensions could represent, for example, latitude, longitude and time, or latitude, longitude and height/pressure levels. Figure we consider case where we do not allow for rolling along any dimension, while in Figure we allow for rolling along one of the dimensions.
Dimension 3
Dimension
v / 2 3 4
1 y /
-2
Dimension 2
Dimension
Dimension I
Dimension J
(a) Central
(b) Edge
Figure 11. Example of our nearest-neighhouss procedure in 3D on a 5 x 9 3 grid for 27 nearest neighbours. We consider two different
cases whether the target is central or near the edge of the grid. In the case an edge target (right), we do not consider invalid neighbours (that are outside the grid bounds).
F. Hardware specifications
For the smaller synthetic GP regression experiment, we perform training and inference for all models on a single NVIDIA GeForce RTX 2080 Ti GPU with 20 CPU cores. For the other two, larger experiments, we perform training and inference for all models on single NVIDIA A100 80GB GPU with 32 CPU cores.
21

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
Non-rolling dimension
Non-rolling dimension
1 シッリン y / 3 7
y 0 2
Rolling dimension
Rolling dimension
1 Lyyo
Non-tolling dimension
Non-nolling dimension
2
(a) Central
(b) Edge
Figure 12. Example of nearest-ncighhouss procedure in 3D 5 x grid for 27 nearest neighbours. We allow for rolling along the second dimension, but not along any of the other ones.
G. Experiment Details
Common optimiser details For all experiments and all models, the Adam optimiser (Loshchilov & Hutter, 2019) with fixed learning rate of 5 10-4 and apply gradient clipping gradients with magnitude greater than 0.5.
Common likelihood details For all experiments and all models, employ Gaussian likelihood parameterised by mean and inverse-softplus variance, i.e. the decoder of each model outputs
Me, log(expo 1) d(z.), p( Dr, xt) = -N(5M(0))
(20)
For the experiment modelling skin and 2m temperature (skt and 22m) with richer context, set minimum noise level of o 2 = 0.01 by parameterising
Pe. log(exp (o. minn)² d(zz), p( De, x.) = (..
(21)
CNP details For the CNPs, we encode each Ye,n) E Dr in RD using MLP with two-hidden layers of dimension D.. We obtain a representation for the entire context set by summing these representations together, z, = Σ.. Z which then concatenated with the target input x,. The concatenation [z,, x₂] is decoded using an MLP with two-hidden layers of dimension D.. We use D = 128 all experiments.
Cony CNP details For the ConvCNP model, we use a U-Net architecture (Ronneberger et al., 2015) for the CNN consisting of 11 layers with input size C. Between the five downward layers we apply pooling with size two. For the five upward layers, we use 2C' input channels and C' output channels, the input channels are formed from the output of the previous layer concatenated with the output of the corresponding downward layer. Between the upward layers apply linear up-sampling to match the grid size of the downward layer. In all experiments, we use C = 128, kernel size of five nine, and stride of one. We use SE kernels for the SetConv encoder and SetConv decoder with learnable lengthscales for each input dimension. The grid encoding modified similarly to the pseudo-tokes grid encoder, whereby we only interpolate from the set of observations for which each grid point is the closest grid point. Unless otherwise specified, we also modify the grid decoding similarly to the pseudo-token grid decoder, whereby we only interpolate from the 3D= nearest points on
22

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
grid to each target location. We resize the output of the SetConv encoder dimension C using MLP with two hidden layers of dimension C. We resize the output of the SetConv decoder using MLP with two hidden layers of dimension C.
Common transformer details For each MHSA MHCA operation, we construct a layer consisting of two residual connections, two layer norm operations, one MLP, together with the MHSA / MHCA operation as follows:
Z Z + CATIAA/MHCAAlayer--ooa (Z))
(22)
Z + Z + MLP(layer-norm,, (2))
All MHSA/ MHCA operations use H = 8 heads, each with Dv = 16 dimensions. We use D₃ = DQK 128 throughout.
PT-TNP details We use an induced set transformer (IST) architecture for the PT-TNPs, with each layer consisting of the following set of operations:
U MHCA-layer[ U: Ze)
2, f- MHCA-layer Zri U)
(23)
Z. < MHCA-layer Zo: U).
In all experiments we use five layers, and encoder / decoder MLPs consisting of two hidden layers of dimension D,-
The ViT architecture consists of optional patch encoding, followed by five MHSA layers. The patch encoding is implemented using a single linear layer. In all experiments we use five layers, and encoder decoder MLPs consisting of two hidden layers dimension Dz.
VIT details
Swin Transformer details Each layer of the Swin Transformer consists of two MHSA layers applied to each window, and shifting operation between them. Unless otherwise specified, use window size of four and shift size of two for all dimensions (except for the time dimension in the final experiment, as the original grid only has four elements in the time dimension). For the second experiment in which the grid covers the entire globe, we allow the Swin attention masks to 'roll' over the longitudinal dimension, allowing the pseudo-tokens near 180 longitude to attend to those near - -880 longitude. In all experiments, we use five Swin Transformer layers (10 MHSA layers in total), and encoder / decoder MLPs consisting of two hidden layers of dimension D.. We found that the use of hierarchical Swin Transformer as used the original Swin Transformer and models such as Aurora (Bodnar et al., 2025)-did did not lead any improvement in performance.
Spherical harmonic embeddings When modelling input data the sphere (i.c. the final two experiments), the CNP, PT-TNP, and gridded TNP models first encode the latitude/ / longitude coordinates using spherical harmonic embeddings following (Rußwurm et al., 2024) using 10 Legendre polynomials. We found this to improve performance in all cases.
Temporal Fourier embeddings When modelling input data through time (i.e. the final experiment), the CNP, PT-TNP, and gridded TNP models first encode the temporal coordinates using a Fourier embedding. The time value is originally provided in hours since 1st January 1970. Following (Bodnar al., 2025), we embed this using Fourier embedding of the following form:
2t 2 for 0 L/2
Emb(t) = [ cos
sin Ae
(24)
X;-
where the \; are log-spaced values between the minimum and maximum wavelength. We set Amin 1 and Amaz 8760,
the number of hours in year. We use L 10.
Great-circle distance For methods using the kernel-interpolation grid encoder (i.e. the ConvCNP and some gridded TNPs), we use the great-cirele distance, rather than Euclidean distance, as the input into the kernel when modelling input data on the sphere. The haversine formula determines the great-circle distance between two points X1 = (41, A2) and X2 = (f2 22), where À and oo denote the latitude and longitude, and given by:
cos(Ap) COS P1 cos $2 (1 cos( (A))
d(xi, X2) 2r aresin
(25)
2
where Ap= =0=02 62 - 11, AA X2- - > and r is taken to be 1.
23

---

Gridded Transformer Neural Processes for Opatio-Tempone Data
G.I. Meta-Learning Gaussian Process Regression
We utilise the GPyTorch software package (Gardner et al., 2018) for generating synthetic samples from GP. As the number of datapoints in each sampled dataset is very large by GP standards (1.1 x 10*), we approximate the SE kernel using structured kernel interpolation (SKI) (Wilson & Nickisch, 2015) with 100 grid points in each dimension. We use an observation noise of o 0.1 for the smaller and larger lengthscale tasks. In Figure 13, we show an example dataset generated using lengthscale of 0.1 to demonstrate the complexity of these datasets. We were unable to compute ground truth log- -og-eiicl values for these datasets without running into numerical issues.
In addition, we also plot the predictive means (Figure 14) and predictive errors (Figure 15) in the form of heatmaps for a number of CNP models on different example dataset.
We train all models for 500, 000 iterations on 160, 000 pre-generated datasets using batch size of eight. For all models, excluding the ConvCNP, we apply Fourier embeddings to each input dimension with L = 64 wavelengths with Xmin = 0.01 and Xmax 12. We found this to significantly improve the performance of all models.
We provide a plot of the test log-likelihood against the forward pass time in Figure 16, where we include both the larger models shown in Figure 3, as well as smaller versions. In Table 3, we provide test log-likelihood values for number of gridded TNPs and baselines for both tasks. We observe that even when increasing the size of the baseline models they still underperform the smaller gridded TNPs. We include results for the Swin-TNP with full attention grid decoding (no nearest-neighbour cross-attention indicated as NN), which fail to model the more complex dataset when using the PT-GE.
Table 3. Test log-likelihood (1) for the synthetic GP regression dataset. NN indicates that NN-CA is not employed. FPT: forward pass time for batch size of eight in ms. Params: number of model parameters in units of M.
Model Grid encoder
Grid size
l 0.5 (t)
l 0.1 (1) FPT
Params
CNP
-0.406
0.112
0.21
PT-TNP
M = 128
0.819
0,558
53
1.50
PT-TNP
M = 256
0.819
0.565
74
1.52
ConvCNP
SctConv
32 32
0.801
0.536
13
2.11
ConVCNP
SctConv
64 64
0.830
0.681
6.70
VITNP
KI-GE
32 32
16 x 16
0.841
0.722
30
1.16
VITNP
PT-GE
32 32 + x 16
0.841
0.721
1.39
ViTNP
KI-GE
16 x 16
0.833
0.711
28
1.09
VITNP
PT-GE
16 x 16
0.840
0.712
29
1.22
VITNP
KI-GE
64 X 64 + 32 x 32
0.842
0.728
44
1.16
VITNP
PT-GE
64 x 64 + 32 x 32
0.836
0.727
56
1.78
VITNP
KI-GE
32 32
0.830
0.725
47
1.09
VITNP
PT-GE
32 32
0.837
0.728
53
1.32
Swin-TNP
KI-GE
32 32
0.844
0.723
39
1.09
Swin-TNP
PT-GE
32 32
0.844
0.723
42
1.32
Swin-TNP
Avg-GE
32 32
0.840
0.723
34
1.22
Swin-TNP
KI-GE
64 64
0.846
0.728
62
1.09
Swin-TNP
PT-GE
64 64
0.847
0.730
69
1.72
Swin-TNP
Avg-GE
64 x 64
0.845
0.725
58
1.62
Swin- TNP NN Swin- Swin-TNP TNP NN
KI-GE
32 x 32
0.834
0.716
45
1.09
PT-GE
32 32
0.837
0.709
48
1.32
ConvCNP For the ConvCNP models, we use regular CNN architecture with C = 128 channels and five layers. We use kernel size of five for the smaller ConvCNP (32 x 32 grid) and kernel size of nine for the larger ConvCNP (64 x 64).
Swin-TNP For the Swin-TNP models, we use a window size of 4 x 4 for the smaller model (32 x 32 grid) and a window size x 8 for the larger model (64 x 64 grid). The shift size is half the window size in each case. We also provide results when a simple average pooling is used for the grid encoder, which is similar to Xu et al. 2025 except that the pooling is performed in token space rather than on raw observations.
24

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
6
2-
2
2
2
6
2
2
Ground-truth.
(b) Swin-T TNP, PT-GE, 64 64 grid.
6
4
2
0
2
52
6
2
0
2
4
6
2
0
4
(c) ConvCNP, 64 x 64 grid.
(d) PT-TNP, M 256.
Figure 13. A comparison between the predictive means of a selection of CNP models on synthetic GP dataset with e = 0.1. The noiseless ground-truth dataset shown in Figure 13a, and the context set is a randomly sampled set N = x 10* noisy observations of this. The colour corresponds to the output value, with the same scale used in each plot. Observe the complexity of the ground-truth dataset, which the Twins'NP's predictive mean resembles. The ConvCNP and PT-TNP's predictive means are notably smoother.
25

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
-1.710
1.872
-1.710
1.872
4-
2
2-
0
0.
2
4-
4-
6
-6
4
2
4
-6
4
2
(a) Ground-truth.
(b) Swin-TNP, PT-GE, 64 64 grid.
-1710
1.872
1.710
1.872
0-
2-
2-
4-
4
6-
6
(c) ConvCNP, 64 grid.
(d) PT-TNP, 256.
Figure 14. A comparison between the predictive means of a selection of CNP models on a synthetic GP dataset with & - 0.1. The noiseless ground-truth dataset in Figure 14a, and context set is a randomly sampled set N; = 10th noisy observations of this. The colour corresponds to the output value, with the same scale used in each plot.
26

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
0.5
0.5
-1.710
1.872
4
2-
2
0
2
2
4
4-
6
-6
4
2
-6
-2
4
(a) Ground-truth.
(b) Swin-TNP, PT-GE, 64 64 grid.
0.5
0.5
0.5
0.5
4
4
2.
22
2
2
4-
4-
6
4
-2
0
-6
4
-2
0
4
(c) ConvCNP, 64 64 grid.
(d) PT-TNP, M 256.
Figure 15. A comparison between the difference between the predictive mean and ground-trute truth for a selection of CNP models on synthetic GP dataset with é= 0.1. The noiseless ground-trund truth dataset shown Figure 15a, and the context set randomly sampled set of N. = 1 x 10 noisy observations of this. The colour corresponds to the prediction error, with the same scale used in each plot.
27

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
ConcNNP
PT-IN2, M- 256
VITNP, 61 32
VIINE, PT 61
VITNE, XL
VIINZ, PT
Swin INT. KI6I
Swi-TNP.
ConvCNP
INP,
0.85
0.84
Test lng likeihlodd
0.72
0.83
0.82
0.70
0.81
0.68
0.80
0.536
0.558
0.565
0.67
10
20
30
40
50
60
70
80
40
50
60
70 80
90
FPT (m65)
FPT (ms)
Figure 16. Test log-likelihood vs. forward pass time (FPT) for the GP datasets. For each model, show the results for large and small (transparent) version. Baselines are hatched. The considered grid sizes are 64 x 64 and 32 x 32, shown as 64 and 32. For ViTNPs, we include results with and without patch encoding, the former indicated by the symbol in-between the pre- and post-patch-enooded grid sizes. KT / PT: kemnel-nterpoolation pscudo-token grid encoding.
G.1.I. SMALL-SCALE
GAUSSIAN PROCESS REGRESSION
We also consider smaller GP regression task with datasets drawn from a GP with SE kernel with lengthscale & 0.1. Each dataset this smaller task has a randomly sized context set, N, ~ (111, 1000), and a fixed sized target set N; 100. The inputs are sampled uniformly in the range [-2, 2] each dimension. The use smaller dataset allows us make comparisons with the regular TNP. In Table 4, we compare the performance of the best performing smaller ViTNP and Swin-TNP gridded TNPs from the paper with the regular TNP, implemented with five MHSA layers, token dimension D. = 128, H = 8 heads and Do Dxv 16. We include the standard error of the mean test log-likelhhoo, which demonstrate that there is no significant difference in performance between the regular TNP and Swin-TNN It should be noted, however, that the regular TNP more computationally efficient than both gridded TNPs for this small-scale dataset. This reflects the suitability gridded TNPs for large-scale datasets, as there is little difference in forward pass time for the small-scale datasets here and the large-scale datasets considered the paper for the gridded TNPs. In contrast, the TNP cannot be implemented on the large-scale datasets considered in the paper due the quadratic computational and memory complexity associated with full attention.
Table 4. Test log-likelihood (1) for the synthetic GP regression dataset. FPT: forward pass time for a batch size of eight in ms. Params:
number of model parameters units of M.
Model Grid encoder
Grid size
Test log- (t)
FPT
Params
TNP VITNP Swin-TNP
0.596 0.02
0.60
PT-GE PT-GE
32 x 32 + 16 x 16
-0.637 + 0.02
22
1.39
32 32
-0.616 + 0.02
33
1.32
G.2. Combining Weather Station Observations with Structured Reanalysis
Inspired by the real-life assimilation of 2m temperature (t.),,, we use the ERA5 reanalysis dataset to extract skin temperature (skt) and 2m temperature (t2m) at 0.25% resolution (corresponding to a 721 x 1440 grid). We then coarsen the skt grid to 180 x 360 grid, corresponding to 19 in both the latitudinal and longitudinal directions. This implics that, because 2m lies on finer grid, essentially becomes an off-the-gric variable with respect to the coarsened grid on which skt lies. In order for the experimental setup better reflect real-life assimilation conditions, +22m values real weather station locations In total, there are 9, 957 such weather station locations, extracted from the HadISD dataset (Dunn et al., 2012). We show their geographical location in Figure 17.
assume to only observe off-the-grid
For each task, we first randomly sample a time point, and then use the entire coarsened skt grid as the on-the-grid context
MMore specifically, because the 2m temperature values come from the gridded ERA5 data, only consider the nearest grid points the station locations valid thi-id locations (i.e. if station has coordinates (11.19", 115.43") latitudelloggitue, consider the grid point (44.25", 115.5")
28

---

Gridded Transformer Neural Processes for Temporal Data
180"
120"W
60"W
0"
6
12
180
60'
60'
N
30"N
10
3
30'S
60"S
60
180
120
60"W
0"
60%
120ºE
180
Figure 17. Station distribution within 1 1 patches. The colour indicates the number of stations within each patch, clipped from a maximum value of 22 15. The distribution far from uniform, with dense station areas continents such North America, Europe and parts of Asia, and a sparser distribution Africa, South America, and in the oceans.
data (64, 800 points), as well as Noft, off-the-grid t 2m context points randomly sampled from the station locations. the experiment from the main paper, Not,, M00,0.3. but also consider the case of richer off-the-grid context sets with Nott.e 21.. in Table The target locations are all the 9, 957 station locations.
We train all models for 300, 000 iterations the hourly data between 2009 2017 with a batch size of eight. Validation is performed on 2018 8 and testing on 2019. The test metrics are reported for 16, 000 data samples. Experiment specific architecture choices are described below.
Input embedding We use spherical harmonic embeddings for the latitude longitude values. These are used in the ConvCNP model as the ConvCNP does not modify the inputs in order maintain translation equivariance (in this case, with respect the great-circle distance).
Grid sizes For the main experiment (with results reported in Table 1), we chose a grid size of 64 x 128 for the ConvCNP and Swin-TNP models, corresponding to grid spacing of 2.8125 in both the latitudinal and longitudinal directions.
In Table we report results for richer context set using a grid size 128 256 for the Swin-Thie models, corresponding grid spacing of 1.41 in both the latitudinal and longitudinal directions. The results for the ConvCNP are for grid size of 64 x 128, maintain smaller gap in parameter count between models.
CNP We use a different deepset for the on- and the off-the-grid data, and the mean as the permutation-invariant function to aggregate the context tokens, i.e. z. Nat.c Σ D Con (Xot,,ee, Yoll..,)) + No.. Σ ENer, Coa (Xon,e,n Nou,c,))
PT-TNP We managed to up M 256 pseudo-tokens without running into memory issues. This shows that if we only use two variables (one on- and one off-the-grid) PT-TNPs do not scale well to large data. We use different encoder for the on- and off-the-gnt -rid data, before aggregating the two sets of tokens into single context set.
Conv CNP For the ConvCNP we use a grid size 64 x 128 for all experiments. We first separately encode both the on- and the off-the-grid the specified grid size using the SetConv. We then concatenate the two and project them to
29

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
dimension of C' 128 before passing through the U-Net (Ronneberger et al., 2015). The U-Net uses kernel size of k: with stride of one.
Swin-TNP
For the Swin-TNP models in the main experiment, use grid size of 64 128, window size of 4 4 and shift size of 2 x 2. For the experiment with richer off-the-grid context (i.e. between 0.25 and 0.5 of the off-the-grid data), use grid size of 128 256, window size of 8 x 8 and shift size of 4 4.
Swin-TNP (T) For the Swin-TNP (T) models we use analogous architectures Swin-TNP. For the pseudo-token grid encoder and decoder, parameterise {pr ]==1 by an MLP with two hidden layers of dimension H and output dimension H, where H is the number of heads used in the multi-head attention mechanism. For the translation equivariant Swin processor, we add a fixed set of learnable relative positional encodings.
Swin-TNP (T) For the Swin-TNP (7)) models we use analogous architectures to Swin-TNP (T). To break equivariance, we first perform a Fourier expansion on the context token location, and pass the result through an MLP with two hidden layers of dimension D2. We use 16 Fourier coefficients, and drop the fixed inputs out during training with a probability of 0.5. We do not zero them out as during training we observe the entire X' domain (i.e., the entire world).
G.2.1, ADDITIONAL RESULTS FOR THE MAIN EXPERIMENT
Aggregated experiment results We provide in Table 5 the results for the 2m station prediction experiment. The figures correspond those presented in Table 1, but are aggregated per experiment, and additionally include the root mean squared enror (RMSE) metric and the number of parameters associated with each model. Moreover, we also provide results for the Swin- TNP using just the 2m station data (Swin-TNP 2t)m), or just the skt gridded data (Swin-TNP (skt)). The superior performance of the Swin-TNP using both sources of information proves that the model manages to capture meaningful relationships between the two variables.
Table Test log- oog-likel (1) and RMSE (4) for the 2m station prediction experiment. The standard errors the log-likelihood are all below 0.010, and of the RMSE below 0.013. NN indicates that NN-CA is not employed. FPT: forward pass time for batch size of eight ms. Params: number model parameters in units of M.
Model
GE
Grid size
Log-lik. f
RMSE 4
FPT
Params
CNP
0.636
3.266
32
0.34
PT-TNP
M 256
1.344
1.659
230
1.57
ConvCNP (NN)
SetConv
64 x 128
1.535
1.252
96
9.36
VITNP
KI-GE
48
1.628
1.197
167
1.14
VITNP
PT-GE
48 x
1.704
1.118
181
1.83
VITNP
KI-GE
144 x 288 + 48 x 96
1.734
1.073
171
1.29
VITNP
PT-GE
144 x 288 - 48 x 96
1.808
1.021
215
6.69
Swin-TNP Swin-TNP
KI-GE
128
1.683
L,157
121
1.14
PT-GE
64 128
1.819
1.006
127
2.29
Swin-TNP (NNN)
KI-GE
64 128
1.544
1.436
137
1.14
Swin-TNI (NN)
PT-GE
128
1.636
1.273
144
2.29
Swin-TNP (skt)
PT-GE
64 128
1.427
1.330
123
2.29
Swin-TNP (t2m)
PT-GE
64 128
1.585
L.599
107
2.24
ConyCNP
SctConv PT-GE
192 x 384
1.689
1.166
74
9.36
Swin-TNP
192 x 384
2.053
0.873
306
10.67
Example predictions We provide in Figure 18 a comparison for an example dataset between the predictive errors (i.e. difference between predicted mean and ground truth) produced by three models: Swin-TNP with PT-GE, ConvCNP, and PT-TNP. The predictions are performed at all station locations. The stations included in the context set are indicated with a black dot. The figures show how Swin-TNP usually produces lower errors in comparison to the baselines, indicated through paler colours. Examples of regions where this most prominent include central US, as well as southern Australia and southern Europe. We also provide in Figure 19 zoomed-in version of Figure 18 focusing on the US region.
30

---

Gridded Transformer Neural Processes for Spatio- Temporal Data
7.5
5.0
2.5
0.0
-2
-5
-7
(a) Swin- TNP crror.
7.5
5.0
2.5
0.0
-2
-5
(b) ConvCNP error.
7.5
5.0
2.5
0.0
-2
-5
-7
(c) PT-TNP emror.
Figure 18. A comparison between the predictive emor-the difference between predictive and ground truth-of the 2m temperature at all weather station locations at 15:00, 28-01-2019. Stations included in the context dataset are shown as black dots (3% of station locations). The mean predictive log-likelshoods (averaged the globe) for these samples 1.611 (Swin-TNP, PT-GE, grid size of 64 x 128), 1.351 (ConyCNP, grid size of 64 x 128), and 1.271 (PT-TNP, M = 256).
31

---

Gridded Transformer Neural Processes for Tentiole Data
5
0
-5
(a) Swin-TNP error.
(b) COnvCNP error.
(c) PT-TNP CITOr.
Figure 19. A comparison between the predictive error of the 2m temperature at the US weather station locations at 15:00, 28-01-2019. Stations included in the context dataset shown black crosses (+e 3% of station locations). The Swin-TNP uses the PT-GE. Both the Swin-TNP and ConvCNP use a grid size of 64 x 128. The PT-TNP uses M = 256 pscudo-tokens. The mean log-likelihoods of this sample for the three models are 1.611, 1.351, and 1.271, respectively.
cbn
0.4
.Vio
Densty
Demiity
Deasity
0.2
0.2
0.2
0.0
0.0
0.0
0 Normalised error
-2
2
Normalised error
Normalised error
(a) Swin-TNPP histogram.
(b) ConvCNP histogram.
(c) PT-TNP histogram.
Figure 20. A comparison between the normalised predictive error- the predictive error divided by the predieted standard deviation of the 2mm at the US weather station locations at 15:00, 28-01-2019. The context contains observations at 3% of station locations. Each plot shows a histogram of the normalised errors based on the predictions at all station locations, alongside an overlaid standard normal distrihution that perfect predictive uncertainties should follow. The mean log-likelihoods of the normalised predictive errors under standard normal distribution for the Swin-TNT (PT-GF, grid size of 64 x 128), ConvCNP (grid size 64 x 128), and PT-TNP (M- 256) are -1.439. -1.490, and -13380, respectively. reference, standard normal distributionn has a negative entropy of -1.419, indicating that the Swin-TNP TNP has the most accurate predictive uncertainties.
Analysis of the predictive uncertainties For the example dataset considered above, Figure 20 shows histograms of the normalised predictive errors, defined as the predictive errors divided by the predictive standard deviations. We compute the mean log-likelihoods under a standard normal distribution, and compare it to the reference negative entropy of the standard normal distribution of - 1.4419, This acts as indicator of the accuracy of the predictive uncertainties produced by the three models we consider: Swin-TNP (PT-GE, grid size 64 x 128), ConvCNP (grid size 64 x 128), and PT-TNP (M = 256). For the dataset considered in Figure 20, we obtain - 1..49 for Swin-TNP, 1.490 for ConyCNP, and - 1.330 for PT-TNP, indicating that, out of the three models, Swin-TNi outputs the most accurate uncertainties.
Analysis of grid size influence We study to what extent increasing the grid size of the models, and hence their capacity, improves their predictive performance. We repeat the experiment for two models: Swin-TNE (with PT-GE) and ConvCNP with grid size of 192 x 384, corresponding to 0.9375 in both latitudinal and longitudinal directions. For the SWin-TNP we use window size of 8 x 8, and a shift size of 4 x 4. For the U-Net architecture within the ConvCNP we use a kernel size of nine. the decoder of the bigger ConvCNP model, attention is performed over the nearest 9 neighbours, whereas for the smaller ConvCNP we use full attention. This makes the FPT of the bigger model smaller that that of the 64 128 model.
The results are shown in Table 6 (and represent a subset of the results shown in Table 5). In comparison to the Swin-TNP and ConvCNP models which use grid size of 64 x 128, both models improve significantly. However, the bigger ConvCNP with a significant gap in both log-likelihood and RMSE.
still underperforms both the small and big variant of
Analysis of influence of nearest-neighhour encoding and decoding A final ablation we perform in this experiment studies the influence of the number of nearest neighbours considered for the encoder and decoder on the performance of the model. Initially, we focus on the effect of nearest neighbour decoding, investigating two models-withh and without full attention at decoding time. More specifically, we compare Swin-TNPP with PT-GE and KI-GE with grid size of 64 x 128, and either perform full attention in the grid decoder, or cross-attention over the 9 nearest neighbours (NN-CA). For the variants with full attention, we evaluate the log- likeliicl at 25% randomly sampled station locations instead of all of them because of memory constraints. The results are shown in Table 7, and indicate that, not only does NN-CA offer a more
32

---

Gridded Transformer Neural Processes for
Data
Table Test log likelihood (f) and RMSE (L) for the 2m station prediction experiment when varying grid size for two models. The standard emrors of the log- leglikkeihood are all below 0.004, and of the RMSE below 0.005, NN signifies that full attention is applied in the decoder. FPT: forward pass time for a hatch size of eight in ms. Params: number of model parameters in units M. Best results for each
configuration (Swin-TNP/ ConvCNP) are bolded.
Model
GE
Grid size
Log-lik. lik. 1
RMSE
FPT
Params
ConvCNP NN
SetConv
64 128
1.535
1.252
96
9.36
ConvCNP
SetConv
192 384
1.689
1.166
74
9.36
Swin-TNP PT-GE
64 128
1.819
1.006
127
2.29
Swin-TNP
PT-GE
192 384
2.053
0.873
306
10.67
scalable decoder attention mechanism, but it also leads to improved predictive performance when applied to spatio-temporal data. We hypothesise this due to the inductive biases it introduces, which are appropriate for the strong spatio-temporal correlations present in the data we used.
Table 7. Test log- ligelikelihodd (1) and RMSE (4) for the 2 station prediction experiment when varying the decoder attention mechanism neighbt cross-attention and full attention (NN, indicating that NN-CA is not employed). The standard errors of the log- lig-iiklihood and RMSE are all below 0.003. FPT: forward pass time for batch size of eight in Params: number of model parameters in units of
M. Best results for each configuration (PT-GE/ KI-GE) are bolded.
Model
GE
Grid size
Log-lik. t
RMSE
FPT
Params
Swin-TNP
KI-GE
64 x 128
1.683
1.157
121
1.14
Swin-TNP (NK)
KI-GE
64 x 128
1.544
1.436
137
1.14
Swin-TNP Swin-TNP (NN)
PT-GE
64 x 128
1.819
1.006
127
2.29
PT-GE
64 x 128
1.636
1.273
141
2.29
Focusing on just the models using the PT-GE, we also study intermediate regimes for the nearest neighbour decoding mechanism with kite = 25, and kisc = 49. Moreover, we also investigate how the models perform with an increased number of nearest neighbours considered during encoding (kiene = 9). The full results are presented in Table 8, where for each model we specify the number of nearest neighbours considered in the encoder (kenc) and decoder (Rite).
Table Test log likelihood (f) and RMSE (1) for the 2m station prediction experiment when varying the number of nearest neighbours considered for the encoder (kase) and decoder (kaes). NN signifies that full attention is applied in the decoder. The standard errors of the log-likelihood and RMSE are all below 0.003. FPT: forward pass time for batch size of eight in ms. Params: number of model parameters in units of M. Best results for cach configuration (PT-GE / KI-GE) are bolded.
Model
GE
kene
Rise
Grid size
Log lik.
RMSE $
FPT
Params
Swin-TNP
PT-GE
1
9
64 x 128
1.819
1.006
127
2.29
Swin-TNP (NK)
PT-GE
1
64 x 128
1.636
1.273
144
2.29
Swin-TNP
PT-GE
25
64 x
1.787
1.038
207
2.29
Swin-TNP
PT-GE
1
49
64 x 128
1.810
1.028
327
2.29
SWin-TNP
PT-GE
9
9
64 x 128
1.857
0.995
264
2.29
We observe that the performance of the models tends to:
Be optimal for lower values of Aace (9)-we believe this is because locality represents a good inductive bias in the task we consider. We also suspect that with sufficient training, the models with different kiec would eventually reach similar performance, but lower value encourages more efficient training and has a lower computational cost.
Slightly improve with increasing ken--the the gridded pseudo-tokens are, on average, modulated by more context points, hence increasing predictive performance. However, this comes at increased computational cost.
G.2.2. ADDITIONAL. RESULTS FOR RICHER CONTEXT SET
In the previous experiment, the model was provided with relatively little context information about the off-the-grid 2m variable, forcing it to learn meaningful relationships between the relatively sparse off off-the-ar the-grid information and rich on-the- grid data (skt t values). order to investigate whether the model still manages to learn these relationships between the off-
33

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
and on-the-grid data (t2m and skt) and exploit the on-the-grid information even in the presence of more off-the-grid data?, also perform experiment with richer context sets. More specifically, the number of off-the-gre context points is sampled according to Nott,e ~ Mi0.28 25.05.. The results are given in Table 9.
For the ConvCNP we evaluated two models-owe one with full decoder attention and one with neiihst- cross-attention (NN-CA). Similarly to the previous section, we found that the latter has a better performance with log-likelihood of 1.705 (NN-CA) as opposed to 1.635 (full attention). As such, Table 9 shows the results for the NN-CA ConvCNP model.
Table Test log- likelihood (f) and RMSE (L) for the 2m station prediction experiment with richer off fer-the-grid context information. The standard errors of both the log- ligliiklihood and the RMSE are all below 0.003, FPT: forward pass time for a batch size of eight in ms.
Params: number of model parameters in units of M.
Model
GE
Grid size
Log-lik, lik, (f)
RMSE (4)
FPT
Params
CNP
0.715
3.056
27
0.34
PT-TNP
M = 256
1.593
1.403
219
1.57
ConvCNP
SetConv
64 x 128
1.705
1.100
21
9.36
VITNP
KI-GE
48 x 96
1.754
1.112
175
1.14
VITNP
PT-GE
48 x 96
1.988
0.932
188
1.83
VITNP
KI-GE
144 x 288 + 48 x 96
1.842
1.046
179
1.29
VITNP
PT-GE
144 288 + x 96
2.242
0.798
221
6.69
Swin-TNP Swin-TNP
KI-GE
128 256
2.362
0.758
174
1.14
PT-GE
128 256
2.446
0.697
208
5.43
Swin-TNP (skt)
PT-GE
128 256
1.501
1.266
178
5.43
Swin-TNP (t2m)
PT-GE
128 256
2.331
0.909
186
5.38
The results are consistent with the findings from the main experiment:
The performances of the baselines (CNP, PT-TNP, and ConvCNP) are significantly worse than the gridded TNP variants considered.
Among the ViT variants, the ones that employ patch encoding before projecting to a 48 x 96 grid outperform the ones that directly encode to a 48 96 grid.
For each gridded TNP variant, the pseudo-token grid encoder (PT-GE) performs better than the kernel-interpolation one (KI-GE).
The variants with a Swin-transformer backbone outperform the ones with a ViT-backbone, even when the latter has more parameters and higher FPT.
Performing nearest-neighhour cross-attention in the decoder as opposed full attention leads to both computational speed-ups, as well as enhanced predictive performance.
In comparison to the main experiment, the gap in performance between Swin-TNP and Swin-TNP 2t2m) is smaller-thi this is expected, given that the context already includes between 25% and 50% of the off- fff-the-grid station locations. However, the gap is still significant, implying that Swin-TNP manages to leverage the on-the-grid data (skt) and exploits its relationship with the target 2.2m to improve its predictive performance.
G.3. Combining Multiple Sources of Unstructured Wind Speed Observations
In this experiment, we consider modelling the eastward (u) and northward (v) components of wind speed at 700hPa, 850hPa and 1000hPa (surface level). These quantities are essential for understanding and simulating large-scale circulation in the atmosphere, for wind energy integration into power plants, or for private citizens and public administrations for safety planning in the case of hazardous situations (Lagomarsino- Onetoo et al., 2023). We obtain each of the six modalities from the ERA5 reanalysis dataset (Hersbach al., 2020), and construct datasets over latitude / longitude range of [25", 49") [-125 -66"], which corresponds to the contiguous US spanning four hours. We show plots of wind speeds at the three pressure levels for single time step in Figure 21.
For each task, we first sample series of four consecutive time points. From this 4 x 96 x 236 grid, we sample a proportion Pe and Pt of total points to form the context and target datasets, where P. ~ M0.00,00.25 and p: = 0.25. The context and
This is achieved by comparing the performance of model that only given off-the-grid context information, with similar model that provided with both off- as well as on-the-grid data.
34

---

Gridded Transformer Neural Processes for Spatio-Tenan Temporal Data
1000hPa
8506Pa
700hPa
Figure 21. Wind-sper speed and direction each of the three pressure levels. 700hPa, 850hPa and 1000hPa, over the contiguous US 15:00
GMT, 08-06 1997. The colours correspond to the magnitude and direction of the wind speed.
1000hPe
1000hPa
1000hPe
850hPa
700hPa
700hPa
700hPs
(a) Swin-TNP error.
(b) ConvCNP error.
(c) PT-TNP emrrr.
Figure 22, A comparison between the predictive error the difference between predictive mean and ground truth of normalised wind speeds for selection of CNP models small region of the US 04:00, 01-01-2019. Each plot consists of 2,400 arrows with length orientation corresponding to the direction and magnitude the wind-speed error at corresponding pressure level. The colour each arrow is given by the HSV values with hue dictated by orientation, and saturation and value dictated by length (i.e. the brighter the colour, the larger the error). For this dataset, the context dataset consists of 5% of the total available observations, and the corresponding predictive log-likelihoods for the Swin-TNP (PT-GE, grid size 1 24 60), ConvCNP (grid size x 24 60) and PT-TNP (M = 64) are 5.84, 3.23 and 2.11.
35

---

Gridded Transformer Neural Processes for Tentio-e Data
1000hPa
1000bPa
1000hPa
APL 1)
0.3
0.3
Dernlit
Deratty
0.2
0.2
0.2
0.1
0.3
0.0
-4
2
-2
6
Normallood IETTT
Normsllead Wrror
Normallend error
850hPs
850hPa
850hPt.
ARL 1)
0.3
0.5
0.3
Deesity
Deeesity
Dessity
0.2
0.2
0.2
0.1
0.1
-4
Normalised DITOC
Normalised error
Normalised orror
700hPe
700hPa
700hPa.
0.4
ARL 1)
ABB 11
ARL 13
0.5
0.3
0.3
Derseity
Deceatty
0.2
0.2
0:2
0.1
-4
6
Normatlsed BETTO
Normallsed error
Normalised error
(a) Swin-TTiY normalised error.
(b) ConvCNP normalised error.
(c) PT-TNP normalised crror.
Figure 23. A comparison between the normalised predictive error- the predictive error divided by the predictive standard deviation of wind speeds for a selection of CNP models on a small region of the US at 04:00, 01-01-2019. Each plot compares histogram of values for both the and v components with standard normal distribution, which perfect predictive uncertainties follow. For this dataset, the context dataset consists of 5% of the total available observations. The mean ligllikelsooods (averaged over pressure levels and the 4 time points) of the uormalised predictive errors under standard normal distribution for the Swin-TNP (PT-GE, grid size 4 x 24 x 60). ConvCNP (grid size 1 x 24 x 60) and PT-TNP (M = 6-1) are -1,425, -1.501 and -1,514, For reference, standard normal distribution has negative entropy of -1.49.. This indicates that the Swin-TNP has the most accurate predictive uncertainties.
36

---

Gridded Transformer Neural Processes for Spatio-Temporay Data
target locations sampled independently for each of the six modalities. All models are trained for 300, 000 iterations on hourly data between 2009- 2017 with batch size of eight. Validation is performed on 2018 and testing on 2019. Experiment specific architecture choices are described below, and full set of results is provided in Table 10.
Input embedding As the input contains both temporal and latitude / longitude information, we use both Fourier embed- dings for time and spherical harmonic embeddings for the latitude/ / longitude values. These are not used in the ConvCNP as the ConvCNP does not modify the inputs to maintain translation equivariance (in this case, with respect to time and the great-circle distance).
Grid sizes We chose a grid size of 4 x 24 x 60 for the ConvCNP and Swin- TNP models, as this corresponds to grid spacing of 19 in the latitudinal direction and around 1 in the longitudinal direction. For the ViTNP, we chose a grid size of 4 x 12 x 30, corresponding to a grid spacing of 2
CNP different deepset is used for each modality, with the aggregated context token for each modality then summed together to form a single aggregated context token, i.e. 2e > Em Ne..
ex(Xe,n Ven,,)..eeeenn)))
PT-TNP For this experiment, we could only use M = 64 pseudo-tokens for the PT-TNP without running into out-of- memory issues. This highlights a limitation in scaling PT-TNPs large datasets. We note that there does exist work that remedies the poor memory scaling of PT-TNPs (Feng al., 2024); however, this trades off against time complexity which itself is a bottleneck given the size of datasets consider. A different encoder is used for each modality, before aggregating the tokens into a single context set of tokens.
ConyCNP For the ConvCNP, use a grid size of 4 x 24 60. Each modality is first grid encoded separately using the SetConv, concatenated together and then to C 128 dimensions before passing through the U-Net.
Swin-TNP For the Swin-TNP models, we use a grid size of 4 24 x 60, window size of 4 x 4 x 4 and a shift size of 0 x 2 2.
Swin-TNP (T) For the Swin-TNP (T) models encoder and decoder, we parameterise [pr H by an MI.P with two hidden layers of dimension H and output dimension H, where H is the number of heads used in the multi-head attention mechanism. For the translation equivariant Swin processor, we add a fixed set learnable relative positional encodings.
analogous architectures to Swin TNN. For the pseudo-toke grid
Swin-TNP (T) For the Swin-TNP (?") models we use analogous architectures to Swin-TNP (T)). To break equivariance, we first perform a Fourier expansion on the context token locations, and pass the result through an MLP with two hidden layers dimension Dz- We use 8 Fourier coefficients, and drop the fixed inputs out during training with a probability of 0.5. We zero them out outside the training region (the US, corresponding to a latitude longitude range of [25", 49º] [-125%; 1. Although minor, the additional memory overhead from the extra basis functions required us slightly reduce the proportion of target points during training, validation, and testing Pt 0.235.
G.4. Incorporating translation equivariance
We provide the full set of results (including RMSE and parameter count) in Table 11. For the t2m station experiment, the standard errors of the log- ligeeiisiooods are all below 0.004 and those of the RMSE below 0.003. For the wind speed experiment when tested on the US region, the standard errors of the log- ligelikelssoods are below 0.010 for the Swin-TNP and Swin-TNP (Z'), and below 0.030 for Swin-TNP (T)). The same holds true when tested on Europe, with the exception of the Swin-TNP which has a standard error of 0.065. The standard errors of the RMSE are below 0.013.
G.S. Additional comparisons to ConyCNP
The ConvCNP is a strong bascline: its FPT is fairly low when we incorporate the NN-CA encoding and decoding mechanisms, it incorporates translation equivariance, and its number of parameters does not increase with the grid size. To perform a comprehensive comparison between Swin-TN- and ConvCNP, we provide additional results for larger grid sizes for both roal-world experiments. The results are illustrated in Table and show that:
37

---

Gridded Transformer Neural Processes for Dpatio-Tempory Data
Table 10. Test log- likelihood (1) and RMSE (L) for the multi-modat wind speed dataset. All grids have grid size of in the time dimension. The standard errors of the log- leelikklisoods are all below 0.02, and of the RMSE below 0.005. NN signifies that full attention is applied in decoder. FPT: forward pass time for a batch size of eight ms. Params: number model parameters in units of M.
Model
GE
Grid size
Log-lik. f
RMSE
FPT
Params
CNP
-1.593
2.536
33
0.66
PT-TNP
M 64
3.988
1.185
166
1.79
ConvCNP NN
SetConv
24 x 60
6.143
0.784
210
14.41
ViTNP
multi KI-GE
12 30
5.371
0.908
349
1.49
VITNP
multi PT-GE
12
7.754
0.651
374
2.36
VITNP
multi KI-GE
24 X
-> 30
6.906
0.718
372
1.55
VITNP
multi PT-GE
24 x 60 + 12 30
7.288
0.681
392
3.25
Swin-TNP
multi KI-GE
24 60
7.603
0.651
355
1.49
Swin-INP
multi PT-GE
24
8.603
0.577
375
3.19
Swin-TNP
single KI-GE
24 60
7.794
0.642
366
1.39
Swin-TNP
single PT-GE
24 x 60
8.073
0.614
364
1.67
ConvCNP
SctConv
48 120
7.841
0.615
64
14.41
Swin-TNP
multi PT-GE
48 120
9.383
0.509
369
6.51
Table 11. Test log- gegliikelihodd (1) for Swin-TNN (T)) and (T), compared to the non-cquivariant Swin-TNP. FPT: forward pass time for a batch size of eight ms. Params: number model parameters units of M.
stations
Wind speed
US
Europe
Model
Log-lik. (f)
RMSE ())
FPT
Params
Log-lik. (f)
RMSE (1)
Log-lik. (1)
RMSE (J)
FPT
Params
Swin-TNP (T)
1.895
0.981
125
1.20
9.972
0.485
10.087
0.483
373
2.09
Swin-TNP (T)
1.926
0.965
142
1.24
10.086
0.469
10.429
0.477
467
2.12
Swin-TNP 1.819
1.006
127
2.29
8.603
0.577
-13.265
3.000
375
3.19
In the station experiment, the ConvCNP with the largest grid size (384 x 768) still lags behind the smaller Swin-TNP (64 x 128), with twice the FPT, and 4x more parameters.
In the multi-mode wind speed experiment, ConvCNP at the native resolution of the grid (96 x 236) still lags behind the translation equivariant version of Swin-TNP (T) with 24 x 60 grid size. Efficiency- wwise the ConvCNP has a lower FPT, but it requires 7x more parameters. ConvCNP (96 x 236) also falls behind the bigger Swin-TNP (48 x 120), and, although has a lower FPT, it requires 2x more parameters.
Table 12. Test log- ligliiklihoodd (f) for the 2 station prediction (left) multi-modal wind speed (right) experiments. The standard emror of the ligeliioddo are all helow 0.010/0.030. For the wind speed experiment, stands for multi GE. FPT: forward pass time for hatch size of eight in ms. Params: number of model parameters in M.
(a) 2m station prediction experiment.
(b) Multi-modal wind speed experiment.
Model
GE
Grid size
TE Log-lik. t FPT Params
Model
Grid size TE Log-lik. 1
FPT Params
ConvCNP SetConv 192 x 384 V
1.689
74
9.36
ConvCNP SetConv 48 x 120 v
7.841
64
14.41
ConvCNP SetCony 288 x 576
/
1.784
149
9.36
ConvCNP SetConv 72 x 180
/
8.478
100
14.41
ConVCNP SetCony 384 x 768 /
1.808
259
9.36
ConvCNP SetConv 96 x 236
/
9.253
142
14.41
Swin-TNP PT-GE 64 x 128
X
1.819
127
2.29
Swin-TNP m-PT-GE 24 x 60
x
8.603
375
3.19
Swin-TNP (T) PT-GE
64 x 128
V
1.895
125
1.20
Swin-TNP (T) m-PT-GE 24 x 60
V
9.972
373
2.09
Swin-TNP PT-GE 192 x 384
X
2.053
306
10.67
Swin-TNP m-PT-GE 48 x 120
x
9.383
369
6.51
G.6. The EAGLE dataset
To show that our approach can be used as a general-purpost tool for spatio-temporal state estimation, we additionally apply to a large-scale fluid dynamics dataset called the EAGLE (Janny et al., 2023). For a thorough description of the data, we refer the reader to Janny et al. (2023). The dataset composed of 1.1 million simulations on an irregular 2D
38

---

Gridded Transformer Neural Processes for Spatio-Temporal Data
mesh, modelling unsteady fluid dynamies caused by moving flow source interacting with nonlinear scene structure. More specifically, EAGLE models the airflow generated by 2D drone moving within a 2D scene with varying floor profile. The dataset comprises of three types of floor geometries (step, triangular, and spline), and for each geometry it contains roughly 200 different scenes, giving total of approximately 1, 200 simulations (two simulations per geometry depending on whether the drone goes to the left or right of the scene). Each simulation contains 990 time steps. The simulated quantities are the velocity field and the pressure field over the entire domain, sampled on an irregular 2D mesh.
The task we address with Swin- TNP spatial interpolation, modelling the pressure and velocity at randomly sampled points in space at randomly sampled time point. The inputs to the Swin-TNP are 1) the location of the drone, the locations of the domain boundaries (open or closed), well as 3) proportion Pa. (40005.0.220 of randomly sampled context points containing the velocity and pressure values at the context locations. Swin-TNP predicts the velocity and pressure at the remaining target locations. To indicate the different types of inputs (i.e. node associated with the drone / with a boundary / within the fluid), we apply a similar strategy as the one for handling multi-modal data, whereby we use input-specific point-wit wise encoders.
We train the models for 1, 000, 000 iterations with batch size of eight. Testing is performed on 80, 000 samples, where a sample is defined as a randomly-chosen time step from one of the trajectories included in the test set. We report the results in Table 13 for two different grid sizes-32 12 and 60 x 36-and using both the PT-GE and the KI-GE. All models use NN-CA in the encoder and decoder with kenc = 1 and litec = 9, a window size of 4 x 4 and a shift size of 2 x 2. Moreover, we provide some example predictions of the velocity field in Figure 24. These results show that our approach can be easily applied to challenging meatio-tempora datasets, and that, consistent with the findings from the other experiments, our proposed pseudo-token grid encoder outperforms the kernel-interpon -nteppiinteroolaiion grid encoder. Future work includes extending the analysis to spatio-temporal forecasting, and comparing to established approaches that are able to deal with irregular meshes.
Table 13. Test log-likelihood (f) for the Eagle experiment. The results are based 80, 000 samples, where sample is one randomly-
chosen time step from one of the trajectories included in the test set (which are also randomly sampled). FPT: forward pass time for a batch size of eight ms. Params: number of model parameters in units of M. Best results are bolded.
Model
GE
Grid size
Log lik. t
Log-lik lik std.
FPT
Params
Swin-TNP
PT-GE
32 x 12
8.153
0.018
1.16
Swin-TNP
KI-GE
32 x 12
7.748
0.018
40
1.41
SWin-TNP
PT-GE
60 x 36
8.319
0.018
40
1.79
Swin-TNP
KI-GE
60 x 36
8.059
0.017
40
1.41
39

---

Gridded Transformer Neural Processes for Tempooal Data
True
Predicted Swin-TNP (PT-GE)
Difference
Predicted Std
10.0
5.0
2.5
True
Predicted Swin-TNP (PT-GE)
Difference
Predicted
10.0
True
Predicted Swin-TNP (PT-GE)
Difference
Predicted Std
10.0
2.5
True
Predicted Swin-TNP (PT-GE)
Difference
Predicted Std
10.0
7.5
True
Predicted Swin-TNP (PT-GE)
Difference
Predicted Stdl
10.0
7.5
Figure 24. Five example predictions for the EAGLE dataset. The first column shows the PDE state, and the coloured dots (white and magenta) represent the context set that is fed to the Swin-TNV TNP. The magenta points denote boundaries, while the white dots are points within the fluid mesh. The second column shows the predictions of the Swin TNP with PT-GE (32 12) terms of predicted velocity (i.e. the model predicts the horizontal and vertical velocities and pressures, and we are plotting / vg)). The third column shows the difference between true and predicted velocities, and last columa shows the predicted standard deviation the velocity v.
40


Reviews:
[{"review_id": "dAHNSJI98i", "reviewer": "Reviewer_wqS4", "summary": "The work introduces a novel approach for modeling large-scale spatio-temporal data without being limited to fixed-resolution grids. Instead of traditional methods that constrain inputs, the authors build on transformer neural processes (TNPs) by developing gridded pseudo-token TNPs. These models use specialized encoders and decoders to manage unstructured, heterogeneous data and employ an efficient attention mechanism based on gridded pseudo-tokens. Additionally, they propose an equivariant variant that leverages translation equivariance as an inductive bias, boosting both accuracy and training efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and for describing our experimental design “sound and well-aligned with the goals of large-scale spatio-temporal modelling.” We address your concerns below.\n\n**Claims and Evidence**\n\n3. We believe there may be some confusion here because, as mentioned in the caption of Table 1, $\\cancel{NN}$ indicates that nearest-neighbour cross-attention is **not** employed (hence full-attention is used). In Table 1:\n- **Swin-TNP** uses nearest-neighbour cross-attention (NN-CA) and achieves a log-likelihood of **1.819**.\n- **Swin-TNP $\\cancel{NN}$** does *not* use nearest-neighbour cross-attention (i.e., uses full cross-attention) and achieves a lower log-likelihood of 1.636 (worse performance).\n\nThis supports our claim that NN-CA outperforms full cross-attention.\n\nWe further validate this claim in Appendix G.2.1—**Analysis of influence of nearest-neighbour encoding and decoding** Table 7--- where we test this claim for the kernel interpolation grid encoder (KI-GE) too:\n- Swin-TNP with PT-GE (1.819) vs Swin-TNP $\\cancel{NN}$ (i.e. full attention) with PT-GE (1.636)\n- Swin-TNP with KI-GE (1.683) vs Swin-TNP $\\cancel{NN}$ (i.e. full attention) with KI-GE (1.544)\n\nSimilarly, in Appendix G.1 Table 3, we observe the same pattern in the GP experiment for a grid size of $32 \\times 32$:\n1) $l=0.5$: Swin-TNP with PT-GE (0.844) vs Swin-TNP $\\cancel{NN}$ with PT-GE (0.837)\n2) $l=0.1$: Swin-TNP with PT-GE (0.723) vs Swin-TNP $\\cancel{NN}$ with PT-GE (0.709) — we will correct a typo that mistakenly indicates 0.109 currently.\n3)  $l=0.5$: Swin-TNP with KI-GE (0.844) vs Swin-TNP $\\cancel{NN}$ with KI-GE (0.834)\n4) $l=0.1$: Swin-TNP with KI-GE (0.723) vs Swin-TNP $\\cancel{NN}$ with KI-GE (0.716)\n\n**Other Strengths and Weaknesses**---”Comparing attention mechanisms\"\n\nThis also addresses your comment from **Other Strengths and Weaknesses** regarding the comparison of different attention mechanisms. While we focused this analysis on our best-performing model (Swin-TNP) due to computational constraints, we believe that the reason NN-CA outperforms full attention is task-dependent rather than model-dependent. As noted in Appendix G.2.1, locality provides a good inductive bias for spatio-temporal tasks, which should hold regardless of the model architecture.\n\nFurthermore, instead of comparing k-nearest-neighbour and full attention for all models (rather than just for the best-performing ones), we chose to focus computational resources on exploring other aspects of the attention mechanism—specifically, the **number** of nearest neighbours used in the encoder ($k_{\\text{enc}}​$) and decoder ($k_{\\text{dec}}$​). The results, presented in Table 8, highlight the importance of the locality inductive bias in our tasks (see the discussion under Table 8 for further details).\n\n**Other Strengths and Weaknesses**---”Comparing grid sizes\"\n\nThank you for your suggestion. Please note we do already include an analysis of the grid size for two models on the t2m experiment: Swin-TNP with PT-GE and ConvCNP for grids of sizes $64 \\times 128$ and $192 \\times 384$.\n\nTable 6 in Appendix G.2.1 shows how the test log-likelihood increases with increased grid size for both models in the t2m experiment. We believe that this trend should generally hold as the model has access to more fine-grained information. However, this comes at an increased computational cost, with the ideal trade-off being experiment-specific.\n\nWe also provide an extensive ablation of grid size for all models in the synthetic GP experiment. Table 3 in Appendix G.1 demonstrates that increasing grid size generally improves performance across models, but the gains are task-dependent. For example, in the synthetic GP experiments, the performance gains with increasing grid size are more pronounced when the task is harder (for $l=0.1$ when the model needs to assimilate finely grained information), which aligns with our expectations.\n\nMoreover, during the rebuttal phase we ran additional experiments for the ConvCNP with increased grid sizes (larger than for Swin-TNP) to address reviewer’s tnKQ comment. The results can be found [here](https://anonymous.4open.science/r/ICML_2025-06DE/swintnp_convcnp_comparison.png).\n\n\nIf the reviewer feels that we have sufficiently addressed their concerns, we would greatly appreciate a reconsideration of the score.", "claims_and_evidence": "1. **Novel Attention-Based Grid Encoder**  \n   This claim is supported by clear empirical evidence. The authors compare the pseudo-token grid encoder (PT-GE) with a kernel-interpolation grid encoder (KI-GE) on both synthetic and real-world tasks. Experimental results (e.g., Figure 3 and Table 1) show that PT-GE consistently yields higher predictive log-likelihoods.  \n\n2. **Enabling Efficient Attention via the Pseudo-Token Grid Encoder**  \n   This claim is supported by clear empirical evidence. By integrating architectures inspired by the ViT and the Swin Transformer, the resulting TNP variants (ViTNP and Swin-TNP) achieve competitive performance with lower computational costs compared to baselines. Experiments reporting prediction quality and forward pass times support this design choice.  \n\n3. **Efficient k-Nearest-Neighbour Attention-Based Grid Decoder**  \n   While the method reduces computational complexity, the claim that it “outperforms full attention” in terms of predictive performance is not consistently supported, making this claim problematic. The authors propose a grid decoder using k-nearest-neighbour cross-attention and claim that it outperforms full cross-attention. However, in the t2m station prediction experiment (Table 1), the variant using nearest-neighbour attention (NN-CA) attains a lower log-likelihood (1.636) compared to the full cross-attention variant (1.819).  \n\n\n4. **Efficient Incorporation of Translation Equivariance**  \n   There is clear and convincing evidence supporting this claim. The paper adapts methods from Ashman et al. to incorporate translation equivariance, and experiments (see Table 2) show improved log-likelihoods in both t2m and wind speed tasks, as well as enhanced spatial generalisation.  \n\n5. **Empirical Evaluation on Large-Scale and Multi-Source Data**  \n   There is clear and convincing evidence supporting this claim. Extensive experiments on synthetic Gaussian process regression and real-world weather datasets compare gridded TNPs with multiple baselines (including ConvCNP and standard TNPs). Results demonstrate improved predictive performance and computational efficiency, confirming the model’s scalability and its ability to handle unstructured, multi-source data.", "methods_and_evaluation_criteria": "1. The use of predictive log-likelihood and forward pass time as evaluation metrics provides a comprehensive view of both the model’s accuracy and its computational efficiency.\n\n2. The choice to benchmark on synthetic Gaussian process regression tasks allows for controlled experiments that test the model under varying levels of complexity.\n\n3. The application to real-world weather data, including t2m station prediction and multi-modal wind speed interpolation, ensures that the evaluation reflects practical challenges and the scalability of the proposed approach.", "theoretical_claims": "Checked. No incorrect proofs.", "experimental_designs_or_analyses": "The experimental designs and analyses are generally sound and well-aligned with the goals of large-scale spatio-temporal modeling.\n\nFor the Gaussian Process Regression, The synthetic experiments are set up using Gaussian processes with two different kernel length-scales (ℓ = 0.1 and ℓ = 0.5), allowing the authors to control the level of complexity and the number of fine-grained features (or \"wiggles\") in the data.\n\nFor the Real-World Weather Data Experiments, the experiments include both t2m station prediction and multi-modal wind speed interpolation. These tasks are highly relevant to weather forecasting applications.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "comparing the effects of different grid sizes, or attention mechanisms independently would strengthen the analysis.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hEm7R4NIot", "reviewer": "Reviewer_t5r6", "summary": "This paper introduces Gridded Transformer Neural Processes (TNPs), a new framework for modeling large-scale spatio-temporal data. Existing approaches, such as Conditional Neural Processes (CNPs), Convolutional CNPs (ConvCNPs), and Transformer Neural Processes (TNPs), struggle with either scalability or handling unstructured data efficiently. The authors propose a novel solution that incorporates pseudo-token grids and efficient attention mechanisms to improve performance on structured and unstructured datasets. \n\nMain algorithmic and conceptual contributions:\n---\n\n1. Gridded Pseudo-Token TNPs: \n- Introduces pseudo-tokens on a structured grid to efficiently process unstructured spatio-temporal data.\n- Uses cross-attention to encode data into a grid before applying transformers.\n\n2. Efficient Attention Mechanisms for Scalability\n- Uses Vision Transformers (ViT) and Swin Transformers to efficiently model spatio-temporal dependencies.\n- Nearest-neighbor cross-attention further reduces computational complexity.\n\n3. Handling Multi-Modal Data\n- Compares single pseudo-token grid encoding (s-PT-GE) vs. multi pseudo-token grid encoding (m-PT-GE) for handling datasets with multiple sources (e.g., satellite data, ground observations).\n\n4. Incorporating Translation Equivariance\n- Introduces translation equivariant TNPs (TE-TNPs) to enforce shift-invariance in spatial patterns.\n- Proposes approximate TE-TNPs, which improve generalization while allowing minor symmetry-breaking adaptations.\n\nMain results and findings:\n---\n1. Synthetic Gaussian Process Regression\n- Gridded TNPs outperform ConvCNPs and PT-TNPs, especially on high-resolution data.\n- Pseudo-token grid encoding (PT-GE) is superior to kernel interpolation (KI-GE).\n\n2. Real-World weather forecasting\n- Temperature Prediction at Weather Stations: Gridded TNPs outperform ConvCNPs by better fusing gridded and unstructured data.\n- Multi-Modal Wind Speed Forecasting: Multi pseudo-token encoding (m-PT-GE) improves accuracy in datasets with multiple data sources.\n- Translation Equivariance: Equivariant Swin-TNP (T) generalizes better, especially when tested on unseen geographic regions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and address their concerns below.\n\n**Claims and evidence**\n\n**Out-of-Distribution (OOD) and Translation Equivariance (TE)** - We perform TE experiments in the station and wind speed experiments. In the former, where we use stations from all regions globally, there are no OOD regions to test on. However, we explicitly study OOD in the multi-modal wind experiment by training on the US and evaluating on both the US (in-distribution, ID) and Europe (OOD). Table 2 shows that plain Swin-TNP fails OOD on Europe (low test log-likelihood), while the strictly TE ($T$) and approximately TE ($\\widetilde{T}$) variants generalise better due to appropriate inductive biases.\n\n- “Additional experiments on OOD regions”-This is addressed by the wind experiment, but we welcome suggestions for additional experiments.\n- “More direct comparisons with standard TNPs OOD”-Table 2 provides this comparison, showing that plain Swin-TNP fails on OOD Europe.\n\n**Comparison between strict and approximate TE-TNP** - This comparison is outside this paper's scope, as it is extensively studied in Ashman et al. [1]. In general, approximate TE is helpful when symmetries are present but some symmetry-breaking features exist (e.g. real-world data).\n\n**Benefits of NN-CA** - As noted in our response to reviewer wqS4, we performed a direct ablation study comparing full attention with NN-CA for the Swin-TNP (with PT-GE and KI-GE) in 1) the t2m station experiment (Table 7) and 2) in the synthetic GP experiment (Table 3).\n\n**Weaknesses and Questions**\n\n**Limited Comparisons with GraphCast, FourCastNet, and Pangu-Weather** - Our work focuses on developing a flexible, general, and efficient approach to spatio-temporal modelling rather than an application-specific method. We compare to fundamental techniques from weather modelling that can handle unstructured data: ViT + KI-GE (Aardvark) and Avg-GE (FuXi-DA). For GraphCast, which relies on GNNs, we discuss connections in the Appendix. FourCastNet and Pangu-Weather assume gridded states, making them incompatible with our setup, where models 1) do not require the full state, 2) can handle off-the-grid observations, and 3) can combine modalities that are not observed at the same locations.\n\nThat said, a future direction is adapting our method to forecasting conditioned on past observations at arbitrary locations, enabling comparisons. However, models like GraphCast required ~4 weeks on 32 TPUs, whereas ours were trained for at most a few days on a single GPU, making a direct comparison unfair at this stage.\n\n**“Hyperparameter Sensitivity: Pseudo-token Count ($M$) and Grid Resolution”** - In the synthetic GP experiment (Table 3) we vary the number of pseudo-tokens in PT-TNP ($M=128$, $M=256$) and test with different grid sizes. We also analyse this in the t2m experiment for Swin-TNP with PT-GE and ConvCNP (Table 6). In the real-world experiments, PT-TNP used the maximum $M$ allowed by memory constraints, and decreasing it would just increase the performance gap to our models.\n\n**Narrow OOD Testing** - We conducted additional preliminary experiments on the PDE-based Eagle dataset [2]. Due to space constraints, we cannot provide a detailed description of the experiment, but we provide example predictions of the fluid velocity and a table comparing Swin-TNP (PT-GE) with Swin-TNP (KI-GE) for two grid sizes [here](https://anonymous.4open.science/r/ICML_2025-06DE). These preliminary results show how our method can easily be applied to any spatio-temporal task, and that PT-GE outperforms KI-GE in this setting too.\n\nRegarding generalisation of TE models—In [3] the authors prove that if a translation equivariant prediction map is equipped with a receptive field (i.e. if the output stochastic processes have only local dependencies), then the prediction map generalises spatially (Theorem 2.2). Thus, TE models should generalise in any domain where translation equivariance is an appropriate inductive bias.\n\n**Scaling Laws**\n\nDue to computational constraints, we did not perform experiments on varying dataset sizes, but we acknowledge that exploring data efficiency and scaling laws is valuable, building on existing studies like [4].\n\n[1] Ashman, M., Diaconu, C., Weller, A., Bruinsma, W.P., & Turner, R.E. (2024). Approximately Equivariant Neural Processes.\n\n[2] Janny, S., B'eneteau, A., Thome, N., Wolf, M.N., Digne, J., & Wolf, C. (2023). Eagle: Large-Scale Learning of Turbulent Fluid Dynamics with Mesh Transformers.\n\n[3] ​​Ashman, M., Diaconu, C., Kim, J., Sivaraya, L., Markou, S., Requeima, J., Bruinsma, W.P., & Turner, R.E. (2024). Translation Equivariant Transformer Neural Processes.\n\n[4] Bodnar, C., Bruinsma, W.P., Lucic, A., Stanley, M., Vaughan, A., Brandstetter, J., Garvan, P., Riechert, M., Weyn, J.A., Dong, H., Gupta, J.K., Thambiratnam, K., Archibald, A.T., Wu, C., Heider, E., Welling, M., Turner, R.E., & Perdikaris, P. (2024). A Foundation Model for the Earth System.", "claims_and_evidence": "The claims made in the paper are generally well-supported by experimental results and conceptual arguments. The authors provide both synthetic and real-world experiments to validate their contributions, and the benchmarks against strong baselines (ConvCNP, PT-TNP, etc.) help substantiate their findings.\n\nClaims that need stronger justification:\n---\n1. Gridded TNPs generalize better to unseen spatial regions (translation equivariance claim).\nThe paper shows some evidence that equivariant models generalize better, but:\n- Additional experiments on out-of-distribution (OOD) regions could strengthen this claim.\n- More direct comparisons with standard TNPs in OOD settings would be useful.\n\n2. Approximate translation equivariance improves generalization while allowing flexibility.\nThe paper claims that relaxing strict TE improves performance, but:\n- The comparison between strict TE-TNP (T) and approximate TE-TNP (T̃) is not fully explored beyond weather data.\n- More systematic analysis of when and why approximate TE works better would clarify this finding.\n\n3. Nearest-neighbor cross-attention (NN-CA) improves both computational efficiency and accuracy.\nThe experiments support the efficiency claim but do not fully isolate the accuracy benefit of NN-CA.\n- A direct ablation study (comparing full cross-attention vs. NN-CA) on a fixed dataset would clarify whether NN-CA improves accuracy beyond just efficiency.", "methods_and_evaluation_criteria": "The proposed methods and benchmarks are generally well-chosen for large-scale spatio-temporal modeling.\nERA5 and Gaussian Process regression are strong datasets for testing model performance.", "theoretical_claims": "Most derivations are reasonable and align with prior literature.", "experimental_designs_or_analyses": "The experimental design and analysis in the paper are generally well-structured and appropriate for evaluating the proposed Gridded Transformer Neural Processes (TNPs). The authors conduct both synthetic and real-world experiments, use strong baselines, and report multiple performance metrics. However, there are areas where further validation or alternative setups could strengthen the findings.\nFor example, lack of comparison with state-of-the-art Transformer forecasting models GraphCast (Google DeepMind), FourCastNet (NVIDIA), Pangu-Weather (Huawei).", "relation_to_broader_scientific_literature": "The paper extends Neural Processes (CNPs, TNPs) by incorporating transformer-based dependencies for spatio-temporal modeling. It builds on Vision Transformers (ViT, Swin) and adapts them using pseudo-token grid encoding, similar to point cloud transformers. Compared to Graph Neural Networks (GNNs) used in weather models (GraphCast, Aurora), Gridded TNPs focus on scalability and parallelization.\n\nFor multi-modal learning, it improves on kernel interpolation (Aardvark Weather, FuXi-DA) by dynamically encoding structured and unstructured data. The translation-equivariant model (TE-TNP) follows Ashman et al. (2024a, 2024b) but improves efficiency using precomputed relative positional encodings. Finally, the work contributes to weather forecasting, aligning with GraphCast, FourCastNet, and Pangu-Weather, but offering a more flexible structured + unstructured data approach.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "Strengths:\n---\n- Scalability: Efficient pseudo-token encoding and nearest-neighbor cross-attention (NN-CA) improve computational efficiency.\n- Multi-Modal Data Handling: Supports both structured (gridded) and unstructured (sensor-based) data.\n- Translation Equivariance: Introduces efficient TE-TNPs and approximate TE (T̃-TNPs) for better generalization.\n- Comprehensive Ablations: Tests different encoders, transformers, and attention mechanisms.\n\nWeaknesses:\n---\n- Limited Comparisons: No direct benchmarking against GraphCast, FourCastNet, or Pangu-Weather.\n- Hyperparameter Sensitivity: No detailed study on pseudo-token count (M) and grid resolution effects.\n- Narrow OOD Testing: Generalization is only tested in weather forecasting, not other spatio-temporal domains.", "other_comments_or_suggestions": "-", "questions_for_authors": "Comparison to State-of-the-Art Weather Forecasting Models:\n---\n- Why does the paper not compare against transformer-based forecasting models like GraphCast, FourCastNet, or Pangu-Weather?\n- If Gridded TNPs were tested against these models, how would their performance compare in terms of accuracy and efficiency?\nIf the proposed method is competitive with state-of-the-art models, it strengthens its contribution to weather forecasting. If not, it may suggest limitations in its practical applicability.\n\nPseudo-Token Grid Hyperparameters:\n---\n- How does the number of pseudo-tokens (M) and grid resolution affect performance?\n- Have experiments been conducted to determine optimal scaling laws for different dataset sizes?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JMwFxR5mWt", "reviewer": "Reviewer_tnKQ", "summary": "The authors present an approach for spatio-temporal modelling. This involves encoding point observations into grid based “pseudo-tokens” using an attention mechanism and processing these pseudo tokens using a Transformer Neural Process approach. A grid decoder also using an attention mechanism allows prediction at arbitrary spatio-temporal locations.\n\nThe authors evaluate this approach on synthetic data and weather prediction tasks and find that it performs favourably compared to baseline approaches. Ablations of their approach are performed allowing us to see the impact of various elements of it, and many details are provided in the extensive appendices. \n\nThe author’s approach combines recent advancements in order to scale effectively to allow spatio-temporal modelling from large datasets. This approach can handle both unstructured “off the grid” and structured “on the grid” data, and can combine multiple modalities of data.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We address your concerns below.\n\n**Similarity to Aardvark** - While our method shares similarities with Aardvark, we highlight key differences:\n- **Methodological advancement, rather than weather-specific**: Aardwark is specifically designed for weather modelling, while our method provides a general framework for spatio-temporal tasks. For example, during the rebuttal we conducted experiments on a different spatio-temporal task - PDE modelling with the Eagle dataset [1] - to show the generality of our approach. Example predictions can be seen [here](https://anonymous.4open.science/r/ICML_2025-06DE/eagle_figures.png) and a comparison between the PT-GE and KI-GE for this dataset [here](https://anonymous.4open.science/r/ICML_2025-06DE/eagle_table.png).\n- **Grid Encoding**: Aardvark uses a kernel-interpolation grid encoder (KI-GE), while we propose a pseudo-token grid encoder (PT-GE), which outperforms KI-GE in our experiments.\n- **Transformer Choice**: Aardvark uses a Vision Transformer (ViT) in the encoder, while we use a Swin Transformer, proven to perform better in high-resolution and dense prediction tasks. This is also supported by our experiments.\n- **Training Objective**: Aardvark optimises a latitude-weighted RMSE, while we optimise the predictive log-likelihood, enabling probabilistic modelling, crucial for chaotic systems.\n- **Decoder Design**: Aardvark uses a lightweight convolutional decoder for station forecasts, while we introduce a k-nearest-neighbour attention-based grid decoder.\n\t\nBoth approaches handle observations at arbitrary locations, but our contributions/modifications (PT-GE, replacing ViT with Swin, and k-nearest-neighbour decoder) improve upon Aardvark’s approach, and, unlike Aardvark, our method supports a probabilistic treatment.\n\nMoreover, as the reviewer noted, our method allows for efficient incorporation of (strict and approximate) **translation equivariance**, further enhancing performance and generalisability.\n\n**Other Strengths and Weaknesses**\n\nAardvark is a modular system for end-to-end weather forecasting, while our work focuses on fundamental research to identify the best approach for spatio-temporal tasks. Rather than directly comparing to Aardvark, which was trained at a much larger scale and in a modular fashion with a focus on forecasting, we propose architectural improvements that could enhance its performance---replacing Aardvark’s encoder (ViT with KI-GE) with our Swin transformer with PT-GE, and its decoder with our k-nearest-neighbour attention-based grid decoder.\n\n\"comparison to a weather predicting approach\" - Please see response to reviewer t5r6 regarding comparisons to other weather models (GraphCast, FourCastNet, Pangu-Weather). Moreover, please note that we compare to the fundamental methods (ViT + KI-GE) used in Aardvark in all experiments and to the grid encoder used in FuXi-DA (Avg-GE) in the GP experiment.\n\n**Other Comments or Suggestions**\n\nThank you, we will correct this in a revised version.\n\n**Questions**\n\nThe FPT for the ConvCNP approach is indeed lower, but this is due to the fact that we incorporated the k-nearest-neighbours into the kernel-interpolation method (in both the encoder and decoder). If we wanted to run the ConvCNP with full attention in the decoder for the $192 \\times 384$ (station experiment) or $48 \\times 120$ (multimodal experiment) grids, we would require more than 80 GB of memory, which is more than what we used for the TNP experiments. Moreover, while FPT is an important efficiency metric, ConvCNP requires significantly more parameters than Swin-TNP (up to 7x).\n\nTo address your question, we ran a set of [additional experiments](https://anonymous.4open.science/r/ICML_2025-06DE/swintnp_convcnp_comparison.png) with larger grid sizes for the ConvCNP.\n- In the station experiment, the ConvCNP with the largest grid ($384 \\times 768$) still lags behind the smaller Swin-TNP ($64 \\times 128$), with twice the FPT, and $4$x more parameters.\n- In the wind experiment, ConvCNP at the native resolution ($96 \\times 236$) still lags behind the TE Swin-TNP with $24\\times 60$ grid size. Efficiency-wise the ConvCNP has a lower FPT, but it requires $7$ times more parameters. ConvCNP ($96 \\times 236$) also falls behind Swin-TNP ($48 \\times 120$), and although has a lower FPT, it requires $2$x more parameters.\n\nWe hope these additional results help you better compare our approach to the strong ConvCNP benchmark, showing that the latter either lags behind in terms of performance and efficiency (as in the t2m experiment), or struggles to match the performance of our strongest TE models even when we increase the grid size to the native data resolution (as in the wind experiment).\n\n[1] Janny, S., B'eneteau, A., Thome, N., Wolf, M.N., Digne, J., & Wolf, C. (2023). Eagle: Large-Scale Learning of Turbulent Fluid Dynamics with Mesh Transformers.", "claims_and_evidence": "Claims made are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "Yes, the synthetic dataset and the weather prediction dataset make sense. The evaluation criteria seem sensible.", "theoretical_claims": "No proofs or theoretical claims.", "experimental_designs_or_analyses": "Both the synthetic data and weather predicting tasks seem sound. Perhaps the weather predicting tasks could be expanded as this seems to be the focus of the paper from the introduction section.", "relation_to_broader_scientific_literature": "The approach used does seem quite similar to that developed by Vaughan et al. (2024) “Aardvark Weather: End-to-end data-driven weather prediction”. In their work, a similar grid encoder and decoder is used that can also deal with off the grid and multimodal data, and a Swin transformer is used to process these grid tokens, and so should be able to efficiently deal with large numbers of grid tokens.\n\nAardvark weather does seem to include many of the key elements of this work such as enabling TNPs to use efficient attention mechanisms, and the use of gridded “pseudo tokens” which aggregate information from the local area and are used as inputs to the transformer component, and similarly a decoder which aggregates information from local pseudo tokens and allows predictions at arbitrary locations.\n\nA difference is that this work uses an attention mechanism to aggregate this local information rather than a kernel interpolation approach. This work also seems to add “translational equivariance” and “approximate translational equivariance”, which is an incorporation of previous work seen in Ashman et al. “Approximately Equivariant Neural Processes” (NeurIPS 2024) and Ashamn et al. “Translation Equivariant Transformer Neural Processes” (ICML 2024).\n\nThe authors have provided significant details of their approach and experiments investigating the impact of different components of the approach, which will be helpful to the broader scientific community.", "essential_references_not_discussed": "All relevant references seem to be included.", "other_strengths_and_weaknesses": "The paper is quite well written and provided many helpful background details. The investigation into their approach / ablations is extensive. The approach seems more generally applicable than many of the spatio-temporal forecasting papers that are cited.\n\nThe baseline grid encoder approach and other TNP approaches seems to be a strong baselines, and it appears that thought has been put into optimizing hyperparameters etc of these baselines.\n\nPerhaps a comparison to a weather predicting approach from the literature would be useful given that the introduction seems to focus on this area quite heavily, however I can understand this is difficult to do. For example some modification of the Vaughan et al. (2024) “Aardvark Weather: End-to-end data-driven weather prediction” approach.\n\nI wonder if the originality compared to Aardvark weather is significant enough. The differences mostly seem to relate to using a local attention based mechanism instead of a \"kernel interpolation\" approach.", "other_comments_or_suggestions": "I believe a crossed out NN should probably be included in row 3 of the multimodal results of Table 1 to indicate that the kernel here is not using the k nearest neighbours approach.", "questions_for_authors": "1. Looking at Table 6, given that:\nThe number of parameters of the ConvCNP approach remains fixed as you change the grid size; The FPT for the fine grid size used by this approach is still much lower than for your approach (and might not change too much if you are using the knn KI); Performance for the ConvCNP approach seems to improve with increasing number of grid points,\ndo you think that using a finer grid could lead to convCNP outperforming your approach while still being competitive on parameters and computation time? Are there practical reasons why this is not feasible?\n\nA response could allow me to evaluate the effectiveness of the proposed method compared to existing approaches.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All reviewers agreed that this paper proposes a novel method that improves the performance of Neural Processes (NP) on high dimensional data.
There is a consensus that the paper is well motivated, well presented and contains extensive experimentation that validates the claims.
The contributions include a method to incorporate attention with pseudo-tokens, comparison between different methods, and improvements to prior NP methods.
Decision: Accept (spotlight poster)