MUSTANG: Multi-Variable Spatio-Temporal Meta-Learning for Water Data Imputation

Abstract

River monitoring stations record multiple hydro-environmental variables over a common river-network topology. While their sampling frequencies and temporal dynamics differ substantially, shared riverine drivers imply that densely observed hydrological variables can inform sparsely sampled water-quality targets; the practical goal is therefore to impute a data-scarce target variable by integrating several related, better-observed sources. In practice, a single source variable rarely serves as a perfect proxy due to distinct dynamics, making the integration of multiple sources essential. However, in spatio-temporal imputation, most existing transfer learning and domain adaptation methods focus on the single-source-to-target setting and do not directly address the multi-source-to-target problem. To solve this problem, we propose Mustang, a multi-variable spatio-temporal meta-learning framework for hydrological data imputation. Specifically, we first develop a spatio-temporal conditional diffusion model that encodes a directed river-network prior into spatial attention, which improves the inductive bias of diffusion models for data imputation. Second, to capture intrinsic cross-variable dependencies, we equip this imputation model with a meta-learning algorithm that allows knowledge transfer from multiple source variables, enabling reliable target-variable learning despite heterogeneous dynamics and scarce observations. We further test robustness under simulated extreme-sparsity regimes, demonstrating that meta-learning enables stable target-variable imputation with very limited observed target values. Experiments on real-world hydrological time series show that Mustang consistently improves target-variable imputation across diverse missing-data scenarios, demonstrating the benefit of jointly modeling shared river-network structure and multi-source transferable dynamics. Our code implementation is available here. https://github.com/Coropt/MUSTANG.