Gridded Transformer Neural Processes for Large Unstructured Spatio-Temporal Data

Many thanks for your review and for the concerns you have raised, in particular with regards to providing additional clarity into how we arrived at our conclusions given the results. We address each of them below.

“The experiment results are not analyzed in detail enough to substantiate the claims.”

Thank you for this valuable feedback. While our original approach of unified discussion aimed for conciseness—given that the takeaways across experiments are consistent—we agree that this structure could be more reader-friendly.

In the revised version, we:

1. Added a focused discussion after each experiment that explicitly connects those specific results to our claims.
2. Maintain the unified discussion section, but introduce it as a summary of the findings across the three experiments.

“The claim that introducing ViT and Swin Transformer into the transformer backbone… feels somewhat weak.”

We agree that our text could better articulate our key contributions. To clarify: our innovation is not the mere use of ViT and Swin Transformer architectures, but rather developing a novel framework that enables these powerful grid-based architectures to process unstructured data. We revised the manuscript to emphasise this distinction.

“{(x_{c, n}, z_{c, n})}n=1^N_c should be {(z{c,n})}_n=1^N_c”

Our notation ${(x_{c, n}, z_{c, n})}$ is intentional. The spatial locations $x_{c, n}$ of the context tokens are essential to our model’s operation in two key ways:

1. They are used in the Swin Transformer to construct appropriate window partitions.
2. They are required in the grid decoder for spatial interpolation.

We will add a clarifying note to the caption of Figure 2 to make this dual purpose explicit.

“e_s(x_{c,n,c},x_{c,n,x}) should be e_s(x_{c,n,s},y_{c,n,s})”

Great spot, thank you! We have corrected this typo.

“GNNs… Could the authors comment on them”

We are pleased that you brought up the connections between transformer-based architectures and GNNs, specifically for handling unstructured data. The two different approaches are intimately related, and we will do our best to provide an overview of this relationship, and the limitations thereof.

1. Computational complexity challenges.
   • Fully-connected GNNs, like regular transformers, face quadratic complexity with respect to the number of nodes (number of observations)
   • While sparse connectivity (used in GraphCast and GenCast) helps with gridded data, it poses challenges for unstructured data. In particular, to induce sparse connectivity we would require dynamically computing $k$ nearest neighbours for each node, which a straightforward implementation of demands quadratic computational complexity.
2. Architectural equivalence. Our approach can be mapped to a GNN equivalent as follows:
   • The pseudo-token grid encoder = GNN connecting observations to the nearest nodes on some latent grid of pseudo-tokens with directed edges.
   • The Swin Transformer processing = GNN with fixed sparse connectivity between nodes within each window and staggered GNN updates.
   • The grid decoder = GNN connecting a latent grid of pseudo-tokens to target prediction locations with directed edges.
3. Practical considerations.
   • While theoretically very similar, empirical evidence suggests that transformer-based implementations achieve superior computational efficiency compared to their GNN counterparts (e.g. see discussion in second paragraph of Section 2 of [1]).
   • Furthermore, the widespread adoption of transformer architectures has led to transformer-specific operations, such as attention mechanisms, being highly optimised for modern GPU architectures [2] and seamlessly integrable, making transformers an efficient and convenient choice from a practical perspective.

If you think it will add value to the paper, we are happy to add a discussion of these connections in the Appendix.

[1] Lang et al. AIFS - ECMWF’s Data-Driven Forecasting System. arXiv preprint arXiv:2406.01465 (2024).

[2] Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv preprint arXiv:2307.08691 (2023).

“Could the authors provide an explanation for this difference in time cost”

The difference in computational cost can be attributed to the cross-attention operation being more computationally expensive than a kernel-interpolation operation. This is due to cross-attention operations requiring additional overhead costs, such as those associated with passing tokens pointwise through an MLP within the cross-attention layer.

Dear Reviewer VoCT,

We appreciate that you are likely to be reviewing multiple other papers; however, as we are approaching the end of the discussion period (less than one week remaining), we would greatly appreciate your feedback on our rebuttals. Your additional insights would be valuable in helping us improve the paper further.

Kind regards, The authors of paper 9931

We thank the reviewer for their comments, and address each of the concerns below. As we note, many of the concerns and questions are already addressed in the paper. We therefore hope the reviewer might consider increasing their score.

“What are the queries, keys, and values in the cross-attention operation in the grid encoder”

Thank you for your clarifying question. The cross-attention operation in our grid encoder is precisely defined in Equations 3 and 4, where:

1. Queries: the pseudo-tokens, each representing the token value at a fixed pseudo-location on a grid.
2. Keys: The set of context tokens for which the pseudo-token is amongst the nearest $k$ neighbours.
3. Values: The same set of context tokens used as keys.

Importantly, we do not need explicit positional encodings for the pseudo-tokens because:

1. Each pseudo-token corresponds to a fixed spatio-temporal location in the grid.
2. The pseudo-tokens are treated as learnable parameters.
3. Through optimisation, the model implicitly learns the appropriate spatial-temporal relationships, effectively encoding positional information into the pseudo-token values themselves.
4. This design choice simplifies the architecture while maintaining the model’s ability to capture spatio-temporal relationships in the data.

“Does this mean each pseudo-token only attends to one context point that is closest to its grid location?”

The parameter $k=1$ (discussed on line 289) actually implies that each pseudo-token attends to the set of context tokens for which it is the closest pseudo-token, not the other way around. To illustrate with a simple 1-D example, suppose that ‘$o(n)$’ denotes the $n$-th context token and ‘$x(m)$’ denotes the $m$-th pseudo-token, and we have the following structure of context and pseudo-tokens:

$o(1) o(2) x(1) x(2) o(3) o(4) x(3)$

Then, the attention relationships are as follows:

1. Pseudo-token $x(1)$ will attend to $(o(1), o(2))$ (as it’s the nearest pseudo-token to both).
2. Pseudo-token $x(2)$ attends to $(o(3))$.
3. Pseudo-token $x(3)$ attends to $(o(4))$.

This mechanism is visualised in Figure 3 of our paper. It ensures that each pseudo-token integrates information from all relevant context tokens in its neighbourhood, not just a single closest point.

“...where do the target tokens Z_t come from?”

You are correct: they are embeddings of the target inputs $x_t$. We discuss this in lines 170-171 in Section 2.2.3.

“...more ablation studies about choosing k in both the encoder and decoder…”

We appreciate this suggestion, and agree. We have performed a set of additional experiments for the weather station experiment which explores this in greater depth. These are shown in Appendix E.2.1 - Analysis of influence of nearest-neighbour encoding and decoding (highlighted in red), alongside a short discussion about the results.

While we agree that it is interesting to perform these ablations across different datasets and tasks, these experiments are expensive to run and we only have limited time available to provide the results. This is the reason we did not run them on the multi-modal experiment too. However, the influence of these parameters depending on the task we are considering is something we are willing to explore further.

Dear Reviewer zSqu,

We appreciate that you are likely to be reviewing multiple other papers; however, as we are approaching the end of the discussion period (less than one week remaining), we would greatly appreciate your feedback on our rebuttals. Your additional insights would be valuable in helping us improve the paper further.

Kind regards, The authors of paper 9931

Many thanks for your review. We appreciate your questions, and we answer both below. If there is anything else we can clarify then please let us know. If not, and you feel that the contributions of our work are strong enough, then we encourage you to please raise your score to recommend acceptance.

“...have you considered using large-scale meteorological models like Pangu…”

As we discuss throughout the paper, a primary motivation for this research is that the current fleet of large-scale meteorological models like Pangu are fundamentally limited to gridded data structures. This constraint prevents us from evaluating these models on unstructured data, which is what we use in our experiments, and so we cannot use them as baselines.

This limitation is particularly significant for weather and climate modelling, where raw observations rarely conform to a regular grid. Our work addresses this challenge in two key ways:

1. We develop novel methods to transform unstructured data into a structured grid of pseudo-tokens, making them compatible with established architectures like ViT and Swin Transformer.
2. We introduce techniques to evaluate predictions at arbitrary target locations, not just the gridded locations.

These innovations enable our approach to bridge the gap between irregularly sampled real-world observations and grid-based deep learning methods. .

“...the model’s complexity might still be high…”

We appreciate the concern about model complexity. While our model’s computational requirements do scale with dataset size, this scaling is both predictable and manageable. Specifically:

1. The computational complexity is dominated by our grid processing components (ViT/Swin Transformer)---the same architectures successfully deployed in production-scale meteorological models.
2. Crucially, our model’s complexity scales linearly with dataset size: see lines 295-296 and line 333 for the computational complexity of the grid encoder and grid decoder, and the computational complexity of the grid processor is dependent on the number of pseudo-tokens, which can be freely chosen but generally scales linearly with the dataset size. This allows for precise resource planning when scaling to larger datasets.
3. These architectures have demonstrated robust performance on ultra-large datasets when provided with appropriate computational resources (e.g., Aurora [1], AIFS [2], Pangu [3], ClimaX [4]).
4. This linear scaling relationship is particularly advantageous compared to potential alternatives that might exhibit quadratic or higher-order scaling behaviours.

[1] Bodnar et al. Aurora: A Foundation Model of the Atmosphere. arXiv preprint arXiv:2405.13063 (2024).

[2] Lang et al. AIFS - ECMWF’s Data-Driven Forecasting System. arXiv preprint arXiv:2406.01465 (2024).

[3] Bi et al. Accurate medium-range global weather forecasting with 3D neural networks. Nature (2023).

[4] Nguyen et al. ClimaX: A foundation model for weather and climate. ICML (2023).

Dear Reviewer 7UA7,

We appreciate that you are likely to be reviewing multiple other papers; however, as we are approaching the end of the discussion period (less than one week remaining), we would greatly appreciate your feedback on our rebuttals. Your additional insights would be valuable in helping us improve the paper further.

Kind regards, The authors of paper 9931

We sincerely appreciate your thorough review and detailed feedback. However, we believe some aspects of our paper’s scope and contributions have been misunderstood. Below, we delineate both what our paper does and does not aim to accomplish.

What This Paper Is

1. Enables Transformers to Efficiently handle Unstructured Spatio-Temporal Data
   • Novel method to apply highly efficient ViT and Swin Transformers to non-gridded spatio-temporal data.
   • Allows efficient transformer architectures to be applied to irregularly sampled data.
2. An Extension of TNPs
   • Addresses quadratic complexity limitations associated with regular TNPs.
   • Enables efficient attention mechanisms to be used.
   • Improves scalability.
3. A General-Purpose Framework
   • Parameterises dataset-to-distribution mappings. That is, models that map directly from datasets—a collection of input-output pairs in the case of supervised learning—to predictive distributions at arbitrary input locations.
   • Applicable to multiple tasks (e.g. data assimilation, regression, classification).
   • Provides probabilistic predictions crucial for many tasks.
4. A New Method which is Evaluated Using ERA5
   • ERA5 is chosen for convenience and accessibility. As mentioned, it is widely used for evaluating machine learning models that are similar to ours, e.g. [1]-[7].
   • Used as ground-truth for evaluation, not as a predictive comparison.
   • Demonstrates practical applicability and scalability of our method.

What This Paper Is Not

1. The Paper Does Not Aim to Provide an Alternative to ERA5
   • Our goal is not to replace ERA5, but rather we utilise it as a datasource for evaluating our model.
   • ERA5 is widely used for evaluating machine learning methods that are similar to ours, e.g. [1]-[7].
2. The Paper Is Not Weather-Specific
   • Our model is domain-agnostic.
   • Weather experiments serve as validation, not the only application of our methods.
   • Custom architectures would be needed for specific domain applications.
3. The Paper Is Not a Competitor to GraphCast/GenCast
   • These models address structured data forecasting, where data lies on a regular grid or mesh.
   • Our work extends transformer-based models in general (specifically TNPs).
   • We focus on enabling efficient-attention-based architectures, such as ViT and Swin Transformer, for unstructured data.
4. The Paper Is Not Focused on Forecasting
   • We address general prediction tasks: given a dataset of observations at some locations, predict the observations at a different set of locations.
   • One can formulate specific problems by choosing which locations and quantities to evaluate.
   • Our scope includes prediction at any spatio-temporal location.
   • Forecasting is just one possible application.

References

[1] Gordon et al. Convolutional Conditional Neural Processes. ICLR (2020).

[2] Foong et al. Meta-Learning Stationary Stochastic Process Prediction with Convolutional Neural Processes. NeurIPS (2020).

[3] Scholz et al. Sim2Real for Environmental Neural Processes. arXiv preprint arXiv:2310.19932 (2023)

[4] Holderrith et al. Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes. ICML (2021).

[5] FNP: Fourier Neural Processes for Arbitrary-Resolution Data Assimilation. NeurIPS (2024).

[6] Andersson et al. Environmental sensor placement with convolutional Gaussian neural processes. Environmental Data Science 2 (2023).

[7] Bruinsma et al. Autoregressive Conditional Neural Processes. ICLR (2023).

[8] Vaughan et al. Convolutional conditional neural processes for local climate downscaling. Geoscientific Model Development (2022).

[9] Rußwurm et al. Geographic Location Encoding with Spherical Harmonics and Sinusoidal Representation Networks. ICLR (2024).

Dear Reviewer XaBC,

We appreciate that you are likely to be reviewing multiple other papers; however, as we are approaching the end of the discussion period (less than one week remaining), we would greatly appreciate your feedback on our rebuttals. Your additional insights would be valuable in helping us improve the paper further.

Kind regards, The authors of paper 9931