When ancient numerical demons meet physics-informed machine learning: adjoint-based gradients for implicit differentiable modeling
Differentiable models are a recently-developed method of combining physics-based models with deep learning techniques, which have enabled much clearer understanding of environmental processes while maintaining predictive performance similar to the state-of-the-art deep learning models. However, there are numerical errors associated with the standard explicit numerical scheme framework, and especially its dependence on sequential calculations. The adjoint method for gradient calculation enables the use of implicit solvers in differentiable modeling, helping mitigate these numerical errors and further improve model prediction performance. This work removed the barrier for complex implicit schemes to enrich differentiable modeling in hydrology.
The adjoint model demonstrated in this work particularly improved low-flow, high-flow, and groundwater simulations, meaning flood peaks and droughts can be more accurately predicted. Additionally, other hydrological problems that require implicit solutions can also benefit from this method, such as groundwater and shallow water equations. More abstractly, the capacity of differentiable models to outperform pure deep learning models in low-flow and high-flow metrics at the median proves that physics-based restrictions (and physical interpretability) and state-of-the-art performance are not mutually exclusive, and pure deep networks are not necessarily the performance ceiling of environmental models (although we do still expect them to be close to optimal). In fact, for rarely observed events or spatial extrapolation, the addition of physical equations may potentially overcome data limitations.
Recent advances in differentiable modeling, a genre of physics-informed machine learning that trains neural networks (NNs) together with process-based equations, have shown promise in enhancing hydrological models’ accuracy, interpretability, and knowledge-discovery potential. Current differentiable models are efficient for NN-based parameter regionalization, but the simple explicit numerical schemes paired with sequential calculations (operator splitting) can incur numerical errors whose impacts on models’ representation power and learned parameters are not clear. Implicit schemes, however, cannot rely on automatic differentiation to calculate gradients due to potential issues of gradient vanishing and memory demand. Here we propose a “discretize-then-optimize” adjoint method to enable differentiable implicit numerical schemes for the first time for large-scale hydrological modeling. The adjoint model demonstrates comprehensively improved performance, with Kling–Gupta efficiency coefficients, peak-flow and low-flow metrics, and evapotranspiration that moderately surpass the already-competitive explicit model. Therefore, the previous sequential-calculation approach had a detrimental impact on the model’s ability to represent hydrological dynamics. Furthermore, with a structural update that describes capillary rise, the adjoint model can better describe baseflow in arid regions and also produce low flows that outperform even pure machine learning methods such as long short-term memory networks. The adjoint model rectified some parameter distortions but did not alter spatial parameter distributions, demonstrating the robustness of regionalized parameterization. Despite higher computational expenses and modest improvements, the adjoint model’s success removes the barrier for complex implicit schemes to enrich differentiable modeling in hydrology.