Field-Space Attention for Structure-Preserving Earth System Transformers

We introduce Field-Space Attention, a novel, scalable, interpretable, and flexible attention module designed for Earth system machine learning models. The key concept involves computing attention directly within physical space on the HEALPix sphere. This approach ensures that all intermediate states remain as globally defined geophysical fields rather than as abstract latent tokens. This field-centric design maintains the physical meaning of internal representations, renders layer-wise updates interpretable, and offers a simple interface for integrating scientific constraints and prior knowledge throughout the network (see Figure). Field-Space Attention is based on a fixed, non-learned, multiscale, spherical decomposition. It learns structure-preserving deformations that coherently couple information across coarse and fine scales. This enables global context without sacrificing local detail.

We demonstrate the module's effectiveness in representative Earth system learning experiments on spherical grids. We focus on global near-surface temperature super-resolution on a HEALPix grid using ERA5 reanalysis data and benchmark it against widely used Vision Transformer and U-Net–style baselines. Our Field-Space Transformer model trains more stably, converge faster, achieve strong accuracy with substantially fewer parameters, and yield physically interpretable intermediate fields.

By keeping computation in field space and explicitly separating scales, Field-Space Attention is particularly well-suited for high-resolution Earth system modeling. It supports scale-aware inductive biases, principled cross-scale consistency, and the efficient coupling of large-scale dynamics with fine-scale variability. These properties position Field-Space Attention as a compact building block for next-generation, high-resolution Earth system prediction and generative modeling. This includes downscaling, spatiotemporal forecasting, infilling, and data assimilation under stronger physical constraints.