Gridded Transformer Neural Processes for Large Unstructured Spatio-Temporal Data

TL;DR

This paper introduces gridded pseudo-token transformer neural processes that efficiently model large-scale unstructured spatio-temporal data, such as weather observations, outperforming existing methods in accuracy and scalability.

Contribution

It proposes a novel architecture combining specialized encoders, decoders, and gridded pseudo-tokens to enable scalable attention mechanisms for unstructured data in weather forecasting.

Findings

Outperforms strong baselines on synthetic and real-world tasks

Maintains computational efficiency at large scale

Demonstrates potential for improved weather modeling pipelines

Abstract

Many important problems require modelling large-scale spatio-temporal datasets, with one prevalent example being weather forecasting. Recently, transformer-based approaches have shown great promise in a range of weather forecasting problems. However, these have mostly focused on gridded data sources, neglecting the wealth of unstructured, off-the-grid data from observational measurements such as those at weather stations. A promising family of models suitable for such tasks are neural processes (NPs), notably the family of transformer neural processes (TNPs). Although TNPs have shown promise on small spatio-temporal datasets, they are unable to scale to the quantities of data used by state-of-the-art weather and climate models. This limitation stems from their lack of efficient attention mechanisms. We address this shortcoming through the introduction of gridded pseudo-token TNPs which employ specialised encoders and decoders to handle unstructured observations and utilise a processor containing gridded pseudo-tokens that leverage efficient attention mechanisms. Our method consistently outperforms a range of strong baselines on various synthetic and real-world regression tasks involving large-scale data, while maintaining competitive computational efficiency. The real-life experiments are performed on weather data, demonstrating the potential of our approach to bring performance and computational benefits when applied at scale in a weather modelling pipeline.