Multi-Task Learning as a Step Toward Building General-Purpose Hydrological Forecasting Systems

Streamflow and soil moisture are two critical variables in the hydrological cycle, linked through infiltration, runoff generation, and groundwater recharge. Traditional forecasting approaches often treat them independently, overlooking interdependencies and limiting predictive skill. This study investigates Multi-Task Learning (MTL) for daily prediction of both variables using the CAMELS dataset derived from the CARAVAN archive. Two task-aware architectures—a Transformer-based and a Mamba-based model—were developed, where shared representations are combined with task identifiers to distinguish outputs for streamflow and soil moisture. Both models were trained to produce 7-day forecasts with balanced sampling across tasks. Results across more than 600 U.S. basins show that MTL models achieve accuracy comparable to or slightly better than single-task baselines, demonstrating the feasibility of learning shared representations. By confirming that one model can distinguish between processes and predict multiple targets without loss of skill, this study establishes a proof of concept for unified, multi-variable forecasting systems.