Learning the Kalman Gain: An End-to-End Deep Learning–Kalman Filter Hybrid Framework for Hydrological State Updating

Hydrological data assimilation (DA) is commonly implemented with Kalman-type filters whose performance depends strongly on prescribed (often time-invariant) process and observation error covariances. In real catchments, however, model errors are non-stationary and state-dependent, making covariance tuning difficult and poorly transferable across events and forecast horizons. Pure deep learning models can be flexible but may drift from process constraints and provide limited interpretability for state corrections.

We propose a differentiable deep learning–Kalman filter hybrid DA framework that learns a time-varying Kalman gain inside the recursive loop of a process-based hydrological model. Specifically, we preserve the Kalman-style update structure while an LSTM-based gain module ingests model states and innovations and outputs an assimilation gain for state updating at each time step. The coupled system (physical model + neural gain) is implemented end-to-end and trained via backpropagation through time, enabling adaptive corrections without manual covariance calibration.

We evaluate the framework using hourly data and benchmark against an optimized Unscented Kalman Filter (UKF). The proposed method matches UKF performance at short lead times but shows increasing advantages for longer horizons, consistent with improved control of error accumulation under non-stationary errors. Results demonstrate that the proposed method achieves superior forecast accuracy at the 24-hour lead time (NSE ≈ 0.75), surpassing the UKF benchmark. Crucially, even though both models were optimized primarily for short-term updates, KalmanNet exhibits superior stability in extended rollouts.The results suggest that learning the assimilation gain within a physically based model provides a robust pathway for hydrological DA under complex, state-dependent error dynamics while preserving process constraints through explicit model equations.