Adaptive Information Masking Dropout for Scalable Uncertainty Estimation in Geospatial Foundation Models

Abstract

Geospatial foundation models are increasingly used to reason over large, heterogeneous Earth observation data. However, their predictions are only useful for high-stakes spatial decision-making when accompanied by reliable uncertainty estimates. Estimating uncertainty at scale remains challenging. Deep ensembles require training and maintaining multiple models, while Monte Carlo (MC) Dropout avoids this cost but incurs inference overhead that grows linearly with the number of stochastic forward passes. Existing adaptive dropout methods improve upon uniform masking, but typically rely on per-instance iterative optimization, limiting throughput in large-scale geospatial applications. We propose Adaptive Information Masking (AIM), a post-hoc dropout method that replaces per-instance optimization with a single offline spectral decomposition. AIM uses ZCA whitening to decorrelate feature channels during calibration and generates spatially non-uniform dropout masks from saliency in the decorrelated feature space. This design preserves the practical advantages of MC Dropout while eliminating the repeated optimization required by existing adaptive approaches. Across geospatial foundation models, including Prithvi and DOFA, and downstream tasks such as landslide segmentation and land-cover classification, AIM consistently outperforms existing methods. AIM is well-suited for large-scale distributed uncertainty estimation. By decoupling calibration from stochastic inference, AIM enables calibration to be performed once and reused across all subsequent predictions. We exploit this property by performing GPU-based calibration and CPU-parallel inference on Apache Spark, reproducing single-node uncertainty estimates exactly while achieving substantial distributed speedups in scenarios where iterative adaptive baselines are computationally impractical.