Integrating multi-task learning with a differentiable physics constrained framework for hydrological forecasting

Why smarter water forecasts matter

Farmers, city planners and disaster managers all rely on knowing how water moves through soil and into the air. Yet predicting key pieces of this puzzle—such as how wet the ground is at different depths and how quickly water evaporates back to the atmosphere—is still surprisingly difficult. Traditional physics-based models struggle with the messy complexity of real landscapes, while modern artificial intelligence can be accurate but often ignores the basic laws of nature. This study introduces a new way to blend both worlds, aiming to make hydrological forecasts that are not only precise but also physically trustworthy.

Linking several water signals at once

Most deep learning systems in hydrology are trained to predict just one thing at a time—for example, only soil moisture or only evaporation from land to air, known as evapotranspiration. But in the real world these variables are tightly connected: the amount of water in the soil influences how much can evaporate, and evaporation in turn changes soil moisture. The authors build a multi-task learning model that predicts several related variables together: soil moisture in three depth layers and evapotranspiration. The model uses a recurrent neural network (an LSTM) that digests a year’s worth of daily weather inputs, such as rainfall, temperature, solar radiation and wind, along with information about the local soil and vegetation. A shared core network learns general patterns, while several smaller output branches specialize in each variable.

Teaching the model to respect water balance

To move beyond a purely data-driven approach, the researchers embed a simple but powerful physical rule: the water balance. Over a given period, incoming water from precipitation must be accounted for by what leaves as runoff and evapotranspiration, plus changes in stored water such as soil moisture and other smaller processes. Instead of treating this equation as an external check, they build it directly into the model’s training process. The neural network’s predictions for soil moisture and evapotranspiration are penalized whenever they violate the water balance, nudging the system toward solutions that conserve water overall.

Letting the network learn the missing pieces

Real landscapes contain many water pathways that are hard to describe exactly, including snowmelt, groundwater exchanges and water intercepted by plants and buildings. These processes vary from place to place and are awkward to encode with fixed formulas. The authors tackle this by inserting a second neural network directly into the water balance equation to represent the “residual” term—the hard-to-capture leftovers. This embedded network receives the same meteorological and land-surface inputs and learns, without direct supervision, to estimate the missing processes so that the overall water budget closes. Because all parts of the system are differentiable, the main predictor and this residual network are trained together in one end-to-end framework.

Testing performance around the globe

The team evaluates their method using the LandBench 1.0 dataset, which combines decades of global weather and land-surface information at daily resolution. First, they compare the multi-task model against four separate single-task networks. Predicting several variables at once proves beneficial: soil moisture predictions improve at both shallow and deeper depths, and evapotranspiration becomes much more consistent in terms of how well the model captures variability and bias. Next, they test three versions of the multi-task setup: one with no physical constraints, one with a simple penalty based on a fixed water-balance formula, and their new differentiable physics-constrained framework. The simple penalty often harms performance, showing that rigid but imperfect physics can mislead learning. In contrast, the differentiable framework usually yields the best scores for soil moisture, especially in deeper layers, and offers modest gains for evapotranspiration.

Benefits for dry regions and rare events

A key strength of the new framework emerges when data are limited or conditions are extreme. When the authors sharply reduce the amount of training data, the unconstrained model’s predictions become scattered and biased, while the physics-constrained version stays closer to reality, particularly for soil moisture at all depths. The approach also improves performance for very low values of soil moisture and evapotranspiration—the kinds of conditions associated with drought. In arid regions such as North Africa and the Middle East, the model better captures how evaporation becomes limited by scarce soil water, a behavior that purely statistical models often miss. High, energy-driven extremes of evapotranspiration remain challenging, highlighting the need to add energy-balance equations in future work.

What this means for future water forecasting

For non-specialists, the main message is that the authors have created a forecasting system that learns from vast amounts of data while still “remembering” that water cannot appear or disappear without cause. By tying a multi-output neural network directly to the water balance and allowing it to learn the fuzzy, hard-to-model parts of the hydrological cycle, they achieve more accurate and more robust predictions across many climates, especially when observations are sparse or conditions are unusual. This kind of differentiable, physics-aware learning offers a promising route toward hydrological tools that scientists, decision-makers and the public can trust to behave sensibly even when the future does not look like the past.

Citation: Yan, Y., Yu, Z., Zhu, J. et al. Integrating multi-task learning with a differentiable physics constrained framework for hydrological forecasting. Sci Rep 16, 13824 (2026). https://doi.org/10.1038/s41598-026-41277-w

Keywords: hydrological forecasting, soil moisture, evapotranspiration, physics-informed AI, multi-task learning