Field-Space Attention for Structure-Preserving Earth System Transformers

TL;DR

The work tackles the challenge of modeling continuous geophysical fields while preserving physical structure by proposing Field-Space Attention, a field-space, fixed multiscale attention mechanism operating directly on spherical field representations. Field-Space Transformers perform structure-preserving deformations on a fixed HEALPix multi-scale decomposition, enabling coherent cross-scale interactions and interpretable intermediate states with improved optimization and parameter efficiency. Applied to global temperature super-resolution on ERA5 data, they achieve faster, more stable convergence with fewer parameters than latent-space Vision Transformer and Transformer-U-Net baselines, while embedding physical priors via scale conservation and spherical-harmonic position embeddings. This field-centric approach offers a strong building block for Earth system foundation models and multi-variable, space-time tasks that require physically grounded, multi-scale reasoning.

Abstract

Accurate and physically consistent modeling of Earth system dynamics requires machine-learning architectures that operate directly on continuous geophysical fields and preserve their underlying geometric structure. Here we introduce Field-Space attention, a mechanism for Earth system Transformers that computes attention in the physical domain rather than in a learned latent space. By maintaining all intermediate representations as continuous fields on the sphere, the architecture enables interpretable internal states and facilitates the enforcement of scientific constraints. The model employs a fixed, non-learned multiscale decomposition and learns structure-preserving deformations of the input field, allowing coherent integration of coarse and fine-scale information while avoiding the optimization instabilities characteristic of standard single-scale Vision Transformers. Applied to global temperature super-resolution on a HEALPix grid, Field-Space Transformers converge more rapidly and stably than conventional Vision Transformers and U-Net baselines, while requiring substantially fewer parameters. The explicit preservation of field structure throughout the network allows physical and statistical priors to be embedded directly into the architecture, yielding improved fidelity and reliability in data-driven Earth system modeling. These results position Field-Space Attention as a compact, interpretable, and physically grounded building block for next-generation Earth system prediction and generative modeling frameworks.

Figures (7)

• Figure 1: Conceptual comparison of convolutional networks, Vision Transformers and the proposed Field-Space Transformer.a, Convolutional models map the input field $x$ to the output $y$ through stacked local convolutions, typically using an encoder-decoder with skip connections. b, A Vision Transformer maps $x$ to patch tokens, transforms them by self-attention in a latent token space, and decodes tokens back to patches to obtain $y$. c, The Field-Space Transformer maintains a field-valued state: it forms multi-scale patches, applies attention across scales, and predicts a residual update $\Delta x_n$ that is added to the current field state $x_n$ to produce $x_{n+1}$.
• Figure 2: Field-Space Attention Block.a, Exemplary The Field-Space Attention Block is applied to a spherical input field $x_0$ in a super-resolution task. Each block predicts a residual update $\Delta x_n$ that is added to the current state $x_n$ to obtain $x_{n+1} = x_n + \Delta x_n$, gradually refining spatial structures. b, Architecture of a single Field-Space Attention block. The field is first tokenized from space to feature tokens, followed by adaptive layer normalization, multi-head attention, and a linear projection. The updated tokens are then mapped back to space by reverse tokenization. A second residual sublayer repeats the pattern using a MLP instead of attention. The proposed multi-scale attention uses a multi-scale tokenization on the HEALPix sphere, followed by multi-head attention, and updates to each scale by inverse tokenization.
• Figure 3: Parameter efficiency of Field-Space Transformer compared to baseline architectures. Root-mean-square error (RMSE) of super-resolution (x1024) of ERA5 global surface temperatures as a function of the number of trainable parameters for our proposed Field-Space Transformer, a standard Vision Transformer (defined on the HEALPix sphere), a multi-scale ViT which uses our multi-scale tokenization, and a convolutional U-Net with attention layers (Trans-U-Net). Our model consistently achieves lower RMSE at substantially smaller model sizes, indicating superior optimization properties. Note that the largest Field-Space Transformer (32.6 M) overfits to the training dataset, while the largest ViT model exhibits unstable training properties. The numeric values are shown in Supp. Table \ref{['tab:comparison']}, while a subset of learning curves is shown in Figure \ref{['fig:losses_models']}.
• Figure 4: Temporal performance and example snapshots of selected models.a Super-resolution snapshots for two days from our Field-Space Transformer (FST, $15.5$ M ), Multi-Scale Vision Transformer (MS-ViT, $54.8$ M), Vision Transformer (ViT, $54.8$ M), and Transformer-U-Net ($99.6$ M). For each model we show the reconstructed temperature field in the left and the corresponding error field $\Delta T$ in the right column. b Time series of the daily mean absolute error (MAE) over the full evaluation period. The Field-Space Transformer attains the lowest and most stable MAE, followed by MS-ViT and ViT, while U-Net shows substantially larger and more variable errors
• Figure 5: Layer-wise residual fields of our Field-Space Transformer. Layer outputs of the $15.5$ M parameter Field-Space Transformer with zoom levels $Z=\{3,5,8\}$ for temperature super-resolution by a factor of $1024$. The panels show the residual fields $r^{(5)}$ and $r^{(8)}$ (zoom levels 5 and 8) across the stack of Field-Space Attention blocks (FSA) and the final scale-constraining layer (SC; bottom). Zoom level 3 is omitted for clarity, since the data are defined on this grid and the learned updates at this scale remain close to zero ($\approx 10^{-8}$). The model gradually builds up structured corrections at the finer scales that match the ground-truth residual fields (bottom row).