Breaking Computational Bottlenecks in Land Surface Modelling with Shifted-Window Transformers

The development of hyper-resolution land surface modelling poses significant computational challenges. Detailed water balance assessments, ensemble-based uncertainty quantification, and climate scenario exploration all require running physics-based models like VIC, Noah-MP, and CLM at continental scales with high spatial resolution, long temporal spans, and multiple parameter configurations. The computational cost becomes prohibitive. Machine learning surrogates have recently emerged as potential solutions; however, existing LSTM and CNN approaches have fundamental architectural problems. Sequential processing prevents parallel computation, limited receptive fields miss long-range dependencies, and most approaches only predict single variables, which restricts comprehensive hydrological analysis.

We present a shifted-window transformer framework that simultaneously predicts multiple land surface fluxes (runoff, evapotranspiration, and soil moisture) while maintaining computational efficiency at continental scales. The hierarchical attention mechanism captures both local temporal patterns through windowed self-attention and global temporal context through shifted-window operations. This eliminates recurrent bottlenecks. We adapt vision transformers for hydrological regression by tokenizing meteorological sequences temporally, using relative position biases to encode lag-dependent hydrological relationships, and designing multi-task regression heads that preserve both nonlinear interactions and direct physical drivers.

We demonstrate the approach by emulating the VIC model across India's 76,390 land grid cells at 6 km resolution, spanning diverse climate regimes. Training uses sparse spatial sampling with only a small fraction of available locations. This allows us to evaluate how well the surrogate generalizes VIC's process behaviors to the newer, unseen regions and parameter configurations. We test multiple variants, including autoregressive formulations that incorporate previous timestep outputs, and benchmark everything against LSTM baselines to isolate the contributions of the architecture.