Towards Learning Universal, Regional, and Local Hydrological Behaviors via Machine-Learning Applied to Large-Sample Datasets

In plain words

Predicting how much water a river carries is hard when you need an answer for a whole region and not for one river. The traditional approach tunes a separate model for each river. Those models are accurate for the river they were tuned on, but they often do poorly when moved elsewhere. A single model tuned on many rivers at once has usually been worse than the separate ones.

The authors train one neural network, a type of network that remembers past inputs, on daily weather records from 531 US river basins. Their change is a modified memory cell, called the Entity-Aware-LSTM. Each basin's fixed traits, such as its area, slope, climate and soil, go only into the cell's input gate, which decides how much new information is written into memory. The traits therefore decide which parts of the network a basin uses, while the weather drives the rest of the cell.

The result is one model that, in the paper's comparison, predicts river flow better on average than the calibrated hydrology models it was tested against. The learned input gates also give a map in which basins with similar behaviour use similar parts of the network. The pitfalls below mark where the paper is vague and where its code differs.

TL;DR

A single LSTM-type network is trained on 531 CAMELS basins with daily meteorological forcings and 27 static catchment attributes §2.4. In the EA-LSTM the attributes only drive a static input gate eq. 7 §2.2. The model is trained with a basin-normalised NSE loss, NSE* eq. 13 §2.3. The author-run comparison gives it a mean validation NSE of 0.674 (single model) and 0.705 (8-model ensemble) on 447 basins, above all benchmark models listed tab. 3. The official code uses Adam and several settings the paper does not state; see Paper vs code.

Problem

Hydrological models calibrated for many basins at once degrade relative to models calibrated for each basin separately §1 [abstract]. A regional model has to capture how basins differ in geology, climate, soil and topography, so that different attributes lead to different hydrological behaviour §1. An earlier LSTM study on rainfall-runoff used only meteorological inputs and ignored the catchment attributes §1. The paper asks whether one data-driven model trained on many basins can match or beat per-basin calibrated models, and whether it can learn catchment similarity from the attributes §1.

Method

The paper compares three networks §2.6:

• LSTM without static inputs: weather only.
• LSTM with static inputs: the 27 attributes are concatenated to the weather inputs at every time step.
• EA-LSTM with static inputs: the attributes only enter the input gate; the weather drives the rest of the cell.

Each network is trained with two losses, MSE and NSE* §2.6, with eight random seeds each. The predictions are averaged into an ensemble §2.6. Six settings times eight seeds give 48 trained models §2.6.

1. Inputs. Five daily Maurer forcings per basin and day: precipitation, minimum and maximum air temperature, shortwave radiation and vapour pressure §2.4. The 27 static attributes are listed in app. A tab. A1. They are derivable from remote sensing or CONUS-wide products §2.4.
2. Samples. Each sample is a 270-day window: the target day and the 269 days before it §2.6. The target is one daily discharge value §2.6.
3. Standard LSTM. Gates and states are given by eq. 1–6 §2.1. The first hidden and cell states are zero §2.1.
4. EA-LSTM. The input gate depends on the static attributes only and is computed once per basin eq. 7. The forget gate, cell input and output gate use the weather and the previous hidden state eq. 8–10. The cell and hidden states are updated as in the standard LSTM eq. 11–12 §2.2. The static input gate has the same shape as the cell state §2.2 ealstm.py.
5. Read-out. The last hidden state goes through dropout (rate 0.4) and then one fully connected layer with a single output §2.6 main.py.
6. Loss. NSE* averages a per-basin weighted squared error. The weight of a sample is 1 / (s(b) + ε)², where s(b) is the standard deviation of the basin's discharge over the training period, and ε = 0.1 §2.3 eq. 13.
7. Evaluation. The metrics are NSE, the α- and β-NSE decomposition, peak-flow bias (FHV), mid-segment slope bias (FMS) and low-flow bias (FLV) tab. 1 §2.6.1. The benchmark runs come from other groups, not from the authors' own training §2.5.
8. Robustness. Gaussian noise with σ = 0.1 to 1.0 is added to the standardised static attributes: 50 random draws per basin and σ value §2.6.2. The sensitivity of the output to each attribute is the mean absolute derivative, computed analytically with respect to the static inputs §2.6.2 eq. 14–15 morris.py.
9. Embedding. Each basin's input gate vector (256 values) is clustered with k-means for k chosen by the mean silhouette score. The clusters are compared with k-means on the raw attributes via the variance explained in 13 hydrological signatures §2.6.3. UMAP projects the 256-dimensional embedding to 2D §2.6.3.

Key equations. Standard LSTM eq. 1–6:

\begin{aligned} i[t] &= \sigma(W_i x[t] + U_i h[t-1] + b_i) \\ f[t] &= \sigma(W_f x[t] + U_f h[t-1] + b_f) \\ g[t] &= \tanh(W_g x[t] + U_g h[t-1] + b_g) \\ o[t] &= \sigma(W_o x[t] + U_o h[t-1] + b_o) \\ c[t] &= f[t] \odot c[t-1] + i[t] \odot g[t] \\ h[t] &= o[t] \odot \tanh(c[t]) \end{aligned}

EA-LSTM eq. 7–12: the input gate uses the static inputs x_s and is not indexed by time:

\begin{aligned} i &= \sigma(W_i x_s + b_i) \\ f[t] &= \sigma(W_f x_d[t] + U_f h[t-1] + b_f) \\ g[t] &= \tanh(W_g x_d[t] + U_g h[t-1] + b_g) \\ o[t] &= \sigma(W_o x_d[t] + U_o h[t-1] + b_o) \\ c[t] &= f[t] \odot c[t-1] + i \odot g[t] \\ h[t] &= o[t] \odot \tanh(c[t]) \end{aligned}

NSE* loss eq. 13, with B basins, N days per basin, and ε = 0.1:

\mathrm{NSE}^* = \frac{1}{B} \sum_{b=1}^{B} \sum_{n=1}^{N} \frac{(\hat{y}_n - y_n)^2}{(s(b) + \epsilon)^2}

The EA-LSTM forward pass, one basin-sample at a time:

# EA-LSTM forward pass for one sample  [§2.2] [eq. 7-12]
# Inputs: weather (T, 5) standardised; attributes (27,) standardised
# Output: one predicted value per sample

def ea_lstm_forward(weather, attributes, params):
    # Static input gate: time-independent, computed once  [eq. 7]
    input_gate = sigmoid(params.W_i @ attributes + params.b_i)
    # Hidden and cell states start at zero  [§2.1]
    hidden = zeros(H)
    cell = zeros(H)
    for day in range(T):
        # Gates driven by the weather and the previous hidden state  [eq. 8-10]
        forget = sigmoid(
            params.W_f @ weather[day] + params.U_f @ hidden + params.b_f)
        cell_input = tanh(
            params.W_g @ weather[day] + params.U_g @ hidden + params.b_g)
        output = sigmoid(
            params.W_o @ weather[day] + params.U_o @ hidden + params.b_o)
        # Cell update: the static gate scales what is written  [eq. 11]
        cell = forget * cell + input_gate * cell_input
        # Hidden state read-out  [eq. 12]
        hidden = output * tanh(cell)
    # Dropout on the last hidden state, then one linear unit  [§2.6]
    return linear(dropout(hidden, rate=0.4))

Training the NSE* model as the code runs it. The paper does not state these settings; they are from the official code.

# Training loop as the official code runs it  [code: main.py#L335-L420]
# Not stated in the paper. Samples from all basins are shuffled together.

def train(model, samples, seed):
    # Adam with learning rate 1e-3  [code: main.py#L335] [code: main.py#L52]
    optimizer = adam(model.parameters(), learning_rate=1e-3)
    for epoch in range(1, 31):
        # Learning rate drops at epoch 11 and epoch 21  [code: main.py#L344]
        if epoch == 11:
            set_learning_rate(optimizer, 5e-4)
        if epoch == 21:
            set_learning_rate(optimizer, 1e-4)
        for batch in shuffled_batches(samples, batch_size=256):
            # Mean over the batch  [code: papercode/nseloss.py#L52]
            loss = nse_star_loss(model(batch.weather, batch.attributes),
                                 batch.target, batch.basin_std)
            loss.backward()
            # Clip the global gradient norm to 1  [code: main.py#L416-L417]
            clip_grad_norm(model.parameters(), max_norm=1)
            optimizer.step()
        # Keep the checkpoint of each epoch  [code: main.py#L354-L355]
        save_checkpoint(model, epoch)

Hyperparameters

Results

Author-reported. The benchmark comparison covers 447 basins that every benchmark model ran on. The EA-LSTM and LSTM are trained on the 1999–2008 calibration period and tested on the 1989–1999 validation period §2.6. The main claim is that the EA-LSTM outperforms the benchmark models on overall NSE §3.2.

Single-model and ensemble results on the 531 basins tab. 2:

Benchmark comparison on the 447 common basins, validation period tab. 3. The ± values are as printed in the paper; it does not say what they measure:

Setup: the EA-LSTM and LSTM rows are the paper's own models; the benchmark rows come from earlier studies and are re-used as given §2.5. The α-NSE is best near one, the β-NSE near zero, and the three bias metrics near zero tab. 3 footnotes.

Significance reported by the authors against the two best hydrological models §3.2: against the HBV upper ensemble, p = 1.9e-4 for the single model and p = 6.2e-11 for the ensemble. Against the basin-calibrated mHM, p = 4.3e-6 and p = 1.0e-13 respectively. Effect sizes (Cohen's d) are 0.22 and 0.40 against HBV, and 0.26 and 0.45 against mHM §3.2.

Robustness (validation period, standardised attributes) §3.3 fig. 6: the median NSE falls from 0.71 without noise to 0.48 at σ = 1. The 1st percentile falls from 0.13 to -5.87 and the 5th from 0.34 to -0.94. Lower percentiles are hit harder than the median.

Embedding §3.4 fig. 9 fig. 10 fig. 11 fig. 12: the mean silhouette is highest at k = 5 for the EA-LSTM embedding and at k = 6 for the raw attributes. Within-cluster variance of 12 of 13 hydrological signatures is reduced as much as or more than with raw attributes; the exception is the day of the year when half the annual flow occurs (hfd-mean) §3.4. Both the raw-attribute and the embedding clusters correspond to different geographic regions of the US §3.4 fig. 11.

Paper vs code

Repository at commit d11815824793737f28f125d55477e42f8900a78e, Apache-2.0 [code: LICENSE]. The paper links it in its code availability statement. The file headers name an earlier version of the paper with a different title, so search on the arXiv id or the author names rather than the title ealstm.py.

• Optimizer. The paper says "stochastic gradient descent" §2.6. The code uses Adam main.py.
• Loss normalization. Eq. 13 sums over basins and days and divides by B eq. 13. The code takes the mean over all samples in the batch nseloss.py. This is a constant factor, not a change in the relative weights, but it matters for the gradient clipping, which the paper does not mention.
• ε term. The text calls ε "added to the variance" §2.3. Eq. 13 and the code add it to the standard deviation s(b) eq. 13 nseloss.py.
• Units in the loss. The code standardises discharge with the global constants above datautils.py and divides the loss weight by the basin's raw standard deviation, computed before normalisation datasets.py. I read this as equivalent to eq. 13 up to a constant factor; it was not run.
• Negative predictions. Predictions below zero are set to zero at evaluation main.py. The paper does not say so.
• Missing discharge. Training samples with negative (missing) discharge or NaN are dropped datasets.py. The paper does not say so.
• Checkpoint. Evaluation loads the epoch-30 weights main.py.
• Robustness count. The paper reports 531 × 10 × 50 = 265,500 evaluations §2.6.2. The code also runs σ = 0 once per basin, giving 531 × 501 main.py.
• NSE in robustness. Observations below zero are masked out before computing the NSE main.py. The paper does not say so.

No discrepancies found in the EA-LSTM gate equations, the gate order (forget, output, cell input), or the 27 attributes and 5 forcings: the code matches the paper there ealstm.py datautils.py.

Pitfalls

• Settings the paper omits. Learning rate, schedule, batch size, epochs, gradient clipping, forget-gate bias, initialization and the optimizer are only in the code (see Hyperparameters). An implementation that follows the paper alone will not match the reported numbers. The paper also gives no training hardware or run time.
• Reversed split. The model is trained on 1999–2008 and evaluated on 1989–1999 §2.6. It is not a forward-in-time test.
• Warm-up. Each window reaches 269 days before the target day, so the first targets of a period need forcings from before that period datasets.py.
• Hyperparameter search. The grid covers hidden size, dropout, sequence length and layer count app. B. The final set is the one with the highest median NSE on the calibration data of held-out basins in 4-fold cross-validation, not on the validation period app. B. The grid lists 156 hidden states; 256 was used.
• Sensitivity method. The paper calls the Morris method a global sensitivity measure §2.6.2. The computed quantity is the mean absolute local gradient of the output with respect to the standardised attributes, averaged over days ranking.ipynb eq. 14–15. That is a local derivative measure; it does not sample the input space. The gradients are taken with respect to the standardised inputs and output, so the values are not in physical units.
• Table and text disagree. The text gives 0.23 as the single-model mean NSE of the LSTM without static inputs trained with MSE §3.1. Table 2 gives 0.24 tab. 2.
• Static means static. The attributes are assumed constant in time, which the paper calls a strong assumption §4. The gate is fixed per basin, so it cannot follow land-cover change.
• Comparison is not one-to-one. Benchmark outputs were produced by other groups for different model setups. The FUSE runs come from personal communication and may not reflect the final FUSE model §2.5 footnote. The benchmark outputs were to be released on HydroShare with a DOI still to be added §2.5 footnote [code and data availability statement]. Those results cannot be rechecked from the paper alone.
• Interpretation of the ranking. The EA-LSTM single model has a lower mean NSE than the LSTM with static inputs (0.674 vs 0.685 on 447 basins) tab. 3. The paper calls the EA-LSTM's cost "small" and worth the interpretability §3.1.
• Cluster and embedding details. The paper does not give the silhouette range, the k-means initialisation, or the UMAP parameters. The code uses k = 3 to 14, k-means with 200 initialisations for the final clusters, and UMAP with 10 neighbours clusterutils.py clusterutils.py embedding.ipynb.
• Costs. 48 trained models §2.6, each with 30 epochs in the code, and 265,500 forward passes for the robustness test §2.6.2.

Lineage

• Builds on: the LSTM §2.1 and the rainfall-runoff LSTM of Kratzert et al. 2018 §1 §2.1. The paper says the standard LSTM is described in that earlier work.
• Compared against: SAC-SMA, VIC, FUSE, HBV and mHM, individually and regionally calibrated §2.5. The discussion relates the EA-LSTM to the MPR approach, which also links parameters to catchment attributes, but with a fixed model structure §4.