Gridded Transformer Neural Processes for Spatio-Temporal Data

Thank you for your review. We address your concerns below.

Similarity to Aardvark - While our method shares similarities with Aardvark, we highlight key differences:

• 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 and a comparison between the PT-GE and KI-GE for this dataset here.
• 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.
• 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.
• Training Objective: Aardvark optimises a latitude-weighted RMSE, while we optimise the predictive log-likelihood, enabling probabilistic modelling, crucial for chaotic systems.
• Decoder Design: Aardvark uses a lightweight convolutional decoder for station forecasts, while we introduce a k-nearest-neighbour attention-based grid decoder.

Both 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.

Moreover, as the reviewer noted, our method allows for efficient incorporation of (strict and approximate) translation equivariance, further enhancing performance and generalisability.

Other Strengths and Weaknesses

Aardvark 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.

"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.

Other Comments or Suggestions

Thank you, we will correct this in a revised version.

Questions

The 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).

To address your question, we ran a set of additional experiments with larger grid sizes for the ConvCNP.

• 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.
• 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.

We 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).

[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.

We thank the reviewer for their detailed review and address their concerns below.

Claims and evidence

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.

• “Additional experiments on OOD regions”-This is addressed by the wind experiment, but we welcome suggestions for additional experiments.
• “More direct comparisons with standard TNPs OOD”-Table 2 provides this comparison, showing that plain Swin-TNP fails on OOD Europe.

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).

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).

Weaknesses and Questions

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.

That 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.

“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.

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. 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.

Regarding 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.

Scaling Laws

Due 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].

[1] Ashman, M., Diaconu, C., Weller, A., Bruinsma, W.P., & Turner, R.E. (2024). Approximately Equivariant Neural Processes.

[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.

[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.

[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.

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.

Claims and Evidence

3. 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:

• Swin-TNP uses nearest-neighbour cross-attention (NN-CA) and achieves a log-likelihood of 1.819.
• 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).

This supports our claim that NN-CA outperforms full cross-attention.

We 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:

• Swin-TNP with PT-GE (1.819) vs Swin-TNP $\cancel{NN}$ (i.e. full attention) with PT-GE (1.636)
• Swin-TNP with KI-GE (1.683) vs Swin-TNP $\cancel{NN}$ (i.e. full attention) with KI-GE (1.544)

Similarly, in Appendix G.1 Table 3, we observe the same pattern in the GP experiment for a grid size of $32 \times 32$:

1. $l=0.5$: Swin-TNP with PT-GE (0.844) vs Swin-TNP $\cancel{NN}$ with PT-GE (0.837)
2. $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.
3. $l=0.5$: Swin-TNP with KI-GE (0.844) vs Swin-TNP $\cancel{NN}$ with KI-GE (0.834)
4. $l=0.1$: Swin-TNP with KI-GE (0.723) vs Swin-TNP $\cancel{NN}$ with KI-GE (0.716)

Other Strengths and Weaknesses---”Comparing attention mechanisms"

This 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.

Furthermore, 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).

Other Strengths and Weaknesses---”Comparing grid sizes"

Thank 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$.

Table 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.

We 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.

Moreover, 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.

If the reviewer feels that we have sufficiently addressed their concerns, we would greatly appreciate a reconsideration of the score.