POET: Partially Observed Earth Transformer with High-Dimensional Position Embedding

Yuxuan Wang, Haixu Wu, Jiafeng Lin, Huakun Luo, Shiyu Wang, Zhou Ye, Jianmin Wang, Mingsheng Long

Keywords: deep learning, machine learning, earth system

Abstract: The Earth system is integral to every aspect of human life, and accurately forecasting the system states is vital in many domains. Current sensing technology can only obtain partial observations of the Earth, such as meteorological factors collected by multiple weather stations or flood monitoring in different river locations. In this paper, we focus on forecasting physical quantities into the future based on partial observations of scattered stations, recorded as high-dimensional time series. While Transformers are well-suited for processing 1D natural language or 2D vision data, their attention mechanism may struggle to learn higher-dimensional dependencies in Earth data. To advance data-driven Earth modeling, we present Partially Observed Earth Transformer, short as POET, which captures the 3D dependencies underlying the Earth system observations alternately from the temporal, spatial, and variate views. To tackle the position-insensitivity of the attention mechanism, we apply attention with a novel High-dimensional Position Embedding (HiPE) strategy that meticulously encodes the geographical bias of each Earth observation. HiPE not only effectively integrates the off-the-shelf prior knowledge into attention but also automatically discovers the latent relation in the high-dimensional system. In a set of empirical studies, POET achieves consistent state-of-the-art forecasting skills in weather, flood and air quality, across both global and regional Earth systems.

Primary Area: applications to computer vision, audio, language, and other modalities

Submission Number: 11685