RiverMamba: A State Space Model for Global River Discharge and Flood Forecasting

Thank you for your time and appreciating the detailed descriptions in the paper, the thorough ablation studies and the promising results.

The only machine learning model among the baselines is a LSTM. Space filling curves might be well-suited for Mamba, but it is not clear whether this is the right approach for a LSTM. (Graph) convolutional networks would be a more interesting baseline.

There seems to be some misunderstanding about the LSTM baseline. To train the LSTM, we followed the same protocol as originally proposed in Nearing et al (2024), which considers only temporal context but does not include any spatial connections. The space filling curves are thus not used in combination with the LSTM. As pointed out, an LSTM would also not be able to handle such long curves. The baselines are described in more details in the suppl. (page 17). In addition, we note that LSTM is the most popular backbone for river discharge forecasting. Recent GNN models for hydrology still use LSTM in their backbone for the temporal modeling e.g. (Gauch et al., "Towards Deep Learning River Network Models", EGU General Assembly 2025, Vienna, Austria, 27 Apr–2 May 2025, EGU25-9768, 2025) and many studies use LSTM to resolve the spatial modeling e.g. (Yu. et al., "Enhancing long short-term memory (LSTM)-based streamflow prediction with a spatially distributed approach", Hydrol. Earth Syst. Sci. 28, 2024). Note that real-world gridded data like the GloFAS dataset has a spatial resolution of 7000 × 3200 (suppl. A.7 page 7) and we consider 4 time steps as input. A spatio-temporal graph would consist of 7000 × 3200 × 4 = 89,600,000 nodes, which would be computationally very expensive. To the best of our knowledge, such a GNN baseline does not exist since GNNs are only proposed in the literature at local scales or coarse resolution, i.e., basin level analysis. Note that we compare to a Transformer architecture with flash attention in Figure 6 and Tables 6 and 7 of the supp.

Limited reproducibility (no code provided)

Thank you for the comment. The code will be made publicly available upon publication as we mention in the NeurIPS Paper Checklist regarding ``Open access to data and code'' and in the suppl. on page 53.

Related work uses GNNs to model the spatio-temporal relationships. Why did you opt for space filling curves combined with Mamba?

The main motivation is that we aim toward developing a model that operates at grid scale and achieves a high accuracy. Mamba avoids the computational cost of the GNNs (see above clarification about GNNs) and Transformers (see Section F.1 in suppl. material). The space-filling curves allow the model to learn the spatial relations adaptively without the complexity of a graph structure.

Dear Reviewer ugtP,

We sincerely appreciate your feedback. Your recognition that our work is very solid and interesting is highly encouraging for us.

In the paper, we compared to 6 baselines: LSTM, Climatology, Persistence (Section. 4.1), GloFAS (Section. 4.2), and Transformer with Flash-attention and Mamba2 (Supplementary. Section F.1). Since there is no GNN or CNN-based model that we can directly compare to, adding an additional GNN or CNN baseline requires building a new model from scratch and adapting the input data. This is beyond the scope of this paper. We believe that the comparison to the state-of-the-art baselines is sufficient to validate the proposed approach.

Thank you for the detailed and clear summary you have provided about this work. We are glad that you found our work novel and the evaluation thorough.

The model produces only point forecasts without uncertainty estimates. These, however, provide important additional information decision makers can use in the operational flood forecasting context.

As mentioned on page 9, extending the model to output a probabilistic forecast is an interesting future direction.

The ablation studies mentioned in the main text and summarized in Table 2 leave out one possible combination preventing assessment of any separate impact of one hyperparameter from possible synergistic effects. Specifically, sub-tables (a) and (b) lack results for using only the recency factor or LOAN exclusively in the forecast block.

There might be some misunderstanding about Table 2. Table 2 represents actually 3 distinctive tables (a), (b), and (c). Table 2 (b) evaluates exclusively the impact of LOAN. We will revise it to make it clearer.

There is a partially incorrect statement about Flash-attention in the appendix (line 261) claiming linear complexity; while this accurately describes the space complexity improvement, Flash-attention's time complexity remains quadratic in the input length, which undermines the claimed efficiency comparison.

Thank you for the note. We will correct it and make it more clear if we refer to runtime or memory. In Fig. 6 in the appendix, we compare both.

The sentence spanning the lines 137-140 is overly long and complex, making it difficult to follow on first reading.

We will revise the sentence.

The relationship between the number of hindcast layers and the chosen temporal input dimension is not clearly explained in the main text, though the former seems to be dependent on the latter.

Thank you for this comment. In our implementation, we chose T=4 as for the GloFAS operational system and a temporal down-sampling of 2. Consequently, we defined 3 layers to encode the input (the first layer doesn't use down-sampling). For example, when T=8 and down-sampling is 2, one would need 4 hindcast layers. The Table below shows the temporal resolution for each layer w.r.t. the input temporal dimension T.

Figure 4's layout is somewhat ambiguous -- it is difficult to discern whether the first serialisation layer belongs to the hindcast block, and including it within the outer shaded region would clarify this organizational structure.

Thank you for this suggestion. We will modify the figure accordingly.

While the appendix describes how the authors set up and trained a version of Google's LSTM, it does not explicitly clarify how this relates to the LSTM results shown in the figures versus results obtained directly from Google's system.

All results shown in experiments are based on the trained Google's LSTM on our dataset except the results shown in suppl. section J pages 23-28 which are obtained directly from the published Google reforecast. We will make this clear in the paper.

Some typos were found during reading the paper.

Thank you for the careful reading. We will fix the typos and check the paper for any other ones.

The model demonstrates significant practical value by effectively detecting very rare floods that might be missed by existing LSTM approaches as impressively shown in Figure 47 for example. However, the model's spatial generalization capabilities appear limited based on the results in appendix section J, where performance on spatially out-of-sample ungauged basins does not significantly differ from the LSTM baseline and regresses quickly with increasing return periods.

As described in the suppl. section J, there are plausible reasons behind this observation:

We do not use nowcasting, i.e. weather forecast at $t=0$, as input. We expect that this would substantially improve the results at early lead time and it is the main reason for the close performance at early lead time. Other differences are: The Google model (Nearing et al. 2024) was trained on ($\sim5680$) stations compared to ($3366$) for RiverMamba. The ungauged streamflow forecast becomes better when the number of stations increases. There are also differences in the input initial conditions, i.e., Google's LSTM uses precipitation estimates from the NASA Integrated Multi-satellite Retrievals for GPM (IMERG) early run as input. In addition, Google's LSTM uses an ensemble of three separately trained LSTMs (Nearing et al. 2024). The results in Section J are thus not directly comparable, but we added them for completeness. We expect that using nowcasting and IMERG as input and an ensemble would improve the results of RiverMamba on ungauged basins further. For the other results in the paper, we used the same training and evaluation setup for all methods to ensure a fair comparison.

Does weighting of no-flood events equally in equation (8) to the ones with a one-year return period not skew the results to the no-flood case? Weighting floods not just by their return period but also by augmenting it with an additive flood offset (let's say 1) might help the flood scarcity problem a bit, which also has been mentioned in the conclusion.

As suggested, we conducted an additional experiment. To this end, we trained a model with weighted flood events by their return periods + flood offset of 1. The F1 results are shown below for the reanalysis dataset and different return periods:

Adding an offset does not improve the results.

Could you provide computational time comparisons to better position the model in the landscape of operational forecasting along Google's LSTM and GloFAS?

Neither Google (Nearing et al. 2024) nor GloFAS (Harrigan et al. 2023) provided the compute time for the operational forecast. The inference time for RiverMamba is reported in supp. Fig. 6. In the table below, we report the inference time (seconds) w.r.t. the number of input points for our model with 4 days as a hindcast (first row), a trained version of Google's LSTM with 4 days as a hindcast (second row), and a Google's LSTM version as in Nearing et al. (2024) with one year hincast (third row). We use one A100 GPU for all runs. All machine learning approaches are very fast. We expect that GloFAS is by several magnitudes slower, which is a practical advantage of machine learning approaches for this task.

Is there an architectural reason for not training on quantile regression, thereby providing a way of estimating the uncertainty of the model's predictions?

There is no architectural reason that prevents the training using a quantile loss. An ensemble of different models can be trained with different quantiles. We agree that an ensemble and estimating the uncertainty of the model are an interesting direction for future work, but we focused on deterministic forecast to thoroughly analyze and evaluate the design of the network architecture. We think that the results demonstrate the strong capabilities of RiverMamba for large scale river discharge forecast. Using an ensemble not only increases the computational cost for running all ablation studies, but it also requires a different evaluation framework, e.g, adding Continuous Ranked Probability Score (CRPS) and spread/skill ratio metrics. Given that we have already 54 pages for the suppl. material, adding an ensemble and uncertainty estimation is out of scope, but an interesting future research direction.

Could you evaluate the model's performance with larger temporal gaps between training and test data (e.g., 5-10 years) to assess robustness to climate change impacts on flood drivers?

In the Table below, we report the evaluation for each year in the test dataset separately. The model is trained over Europe for the years 1979-2018 and shown are averaged (KGE $\mid$ F1) for the return periods 1.5-20 years for both reanalysis and GRDC data.

Note that the evaluation is highly dependent on the ratio of the flood events that happened during a specific year and on the dataset. On the observational data the modeling of the discharge time series (KGE) decreases with time while the model stays robust regarding flood detection (F1). This is probably due to anthropogenic intervention, as mentioned in the manuscript. Reanalysis data has more data points to train on them and the model shows a good generalization ability with a larger gap between training and testing.

I would like to thank the authors for addressing all noted points in my review, clarifying how the model hyperparameters regarding the number of layers were chosen and for running the additional tests. The experiments on the proposed offset were insightful.

I would like to clarify my question regarding the ablation studies in Table 2, as my initial phrasing may have been unclear. My concern is that for both Table 2a and 2b, a specific experiment is missing that would allow for a complete assessment of individual versus synergistic effects. To fully understand the source of the performance gains and rule out potential synergistic effects, I believe two specific ablations are needed:

• In Table 2a, a result for using only the recency factor in the objective function.
• In Table 2b, a result for using the LOAN layer in the forecast block only.

These additions would make it possible to isolate the individual contribution of each component.

I would like to thank the authors for addressing all noted points in my review, clarifying how the model hyperparameters regarding the number of layers were chosen and for running the additional tests. The experiments on the proposed offset were insightful.

We appreciate your insightful feedback, which has significantly improved the clarity and completeness of the paper. We are very glad that all noted points were addressed.

I would like to clarify my question regarding the ablation studies in Table 2, as my initial phrasing may have been unclear. My concern is that for both Table 2a and 2b, a specific experiment is missing that would allow for a complete assessment of individual versus synergistic effects. To fully understand the source of the performance gains and rule out potential synergistic effects, I believe two specific ablations are needed:

• In Table 2a, a result for using only the recency factor in the objective function.
• In Table 2b, a result for using the LOAN layer in the forecast block only.

These additions would make it possible to isolate the individual contribution of each component.

Thank you for this further clarification. We indeed misunderstood your question. As suggested, we conducted the additional experiments as shown in the tables:

Table (a): Objective function

Table (b): Location embedding

Using only the recency factor $\hat{u}$ (Table (a) third row) improves the results on both KGE and F1 metrics. In Table (b) third row, using LOAN layer in the forecast blocks increases the F1 score while decreases KGE. The combination of layers in both hindcast and forecast blocks balances the metrics.

We will include these additional ablations in the final version for completeness.

Thank you for the detailed review and the thoughtful feedback. We appreciate your time and we are glad that you found our approach interesting and the contribution strong.

It is unclear how the LSTM baseline is trained and used.

We describe the details of the LSTM baseline in Section G.3 of the suppl. As in Nearing et al. (2024), we do not use spatial connections for the LSTM. The sequence length for the LSTM is 14 days, which performed better than using a longer sequence length. We will state this clearly in the revised version.

The LSTM baseline from Nearing et al. (2024).

The original model from Nearing et al. (2024) is not publicly available. We thus built the encoder-decoder LSTM from the neuralhydrology repository and followed the model structure as described in Nearing et al. (2024). The LSTM model is trained similarly to Nearing et al. (2024) but on a gauged setting (Section G.3 in Suppl.). For a fair comparison of the methods, we also use the same input to the LSTM model and RiverMamba, which is different to Nearing et al. (2024). We will clarify this in the revised manuscript and in the tables, in particular in Table 3. We do not think that a model should achieve a median $R^2>0.7$. The metric is affected by many factors like the length of the evaluation set (length of the extrapolation), lead time, input data, number of gauged points, resolution, region of study, etc.

Different sequence lengths can be used if you feed in the states and current runoff to the model compared to a model that has to simulate everything just from weather inputs.

Thank you for this clarification. We will integrate this into the manuscript to make it clearer.

The motivation for a global receptive field.

We will revise the formulation since the term "global" might be misleading. As mentioned in the suppl. material on page 16, we split the Earth into smaller domains with 311K nearest points for each. We thus do not connect points from the Amazon to points in Siberia. However, we give the network the possibility to consider a very large spatio-temporal context. For instance, river networks like Amazon can be very large and providing a model that can handle a very large spatio-temporal context is thus very intuitive.
Finally, we present a general backbone model that is not limited to medium-range river discharge forecast and specific input variables or task. For instance, the model can be used for seasonal discharge forecasting. Furthermore, weather forecast and climate modeling generally require a large receptive field to preserve the dynamic over lead time and RiverMamba is a promising backbone for such tasks.

The model relies on ERA5-Land which is not available in real-time but has a $\sim30$ day lag.

Please kindly note that ECMWF states: "The ERA5-Land dataset is available for public use for the period from 1950 to 5 days before the current date." ... "The ERA5-Land-T version delivers non-checked close to Near-Real-Time (NRT) daily updates. ERA5-Land-T is synchronized with the close to NRT daily updates provided by the ERA5 climate reanalysis (ERA5T).", so there is no risk of having a 30-day lag. In the RiverMamba framework, ERA5-Land can be also replaced by an analysis or ERA5 similar to GloFAS. We will revise the manuscript to make this issue clear.

L44. Not entirely true. a) The Google Model runs globally at 300k locations. b) Yang et al. already have a global gridded model.

a) We will rephrase the sentence. We wanted to highlight here the limited spatial modeling of the LSTM.

b) Please see L91-L93 and ref [60], we already discuss the preprint version by Yang et al. We will update the reference and text based on the published version. Thank you for the note.

L89 There is a quite some literature missing.

Thank you for this note. We will briefly discuss the mentioned methods in the related works. Kindly note that the first reference [1] is already mentioned in the manuscript.

L121. Why is RiverMamba not including any information from the nowcast timestep?

RiverMamba intentionally excludes the analysis data (t=0) from the input, and instead starts forecasting from t+1 using only past data up to t–1. The rationale behind this design is to ensure broader applicability, since many weather forecast systems especially ML models provide only lead time $t>0$ and we wanted to make the model more generic so it can work without nowcasting or a specific type of forecast (IFS-HRES). However, adding nowcasting to the model is straightforward and can be added optionally. We will mention this explicitly in the manuscript.

L123. Why another shift by 2 days for CPC and 1 day for GloFAS and ERA5?

We will revise this sentence and remove the word "shift" since we do not shift data. To our knowledge, CPC data are available with 2 days lag. Consequently, we only consider data until t-2.

Figure 4 appears before Fig 3.

We will correct it.

The dimensions in Eq 2.

Thank you for this note. We will correct it. PyTorch handles the dimensions for different operations, i.e., $\sigma$ and $\mu$ will be duplicated $K$ times along the last dimension to match $**X**$.

Eq 2. Why GELU?

We were motivated by two things: first, GELU avoids the dying ReLU problem and improves optimization. Second, it has been shown that a learned activation function initialized as ReLU, ends up in a smooth variant shape similar to GELU (Teney et al., "Do We Always Need the Simplicity Bias? Looking for Optimal Inductive Biases in the Wild", CVPR'25). We conducted an additional experiment (see Table below). The model is trained over Europe for the years 1979-2018 and tested on the validation set (2019-2020). Shown are (KGE $\mid$ F1) averaged for the return periods 1.5-20 years on reanalysis data.

Using GELU performs slightly better.

L236ff: The terminology is a bit off.

Thank you for pointing out the imprecise usage of terminology. We agree that a return period e.g., 1.5 years does not necessarily imply actual flooding in all regions, particularly in highly regulated or flood-resilient areas. A high return period event simply reflects statistical rarity in streamflow magnitude, and should not be equated with a flood event without additional context (e.g., thresholds, inundation). In the revision, we will clarify it and emphasize that the return period is used as a proxy indicator of hydrological extremity, rather than a direct definition of flood severity.

L238. where do these ranges come from?

The return periods are used in GloFAS and operated by ECMWF.

L247. I couldn’t find the results on the ablation study of their loss weighting.

The ablation study regarding the weighting is provided in Table 2 (a).

Suppl. L190. In many operational settings you compute return periods separately in the observed space and in model climatology.

Thank you for this insightful comment. You are correct that in many operational settings, thresholds are computed from the model’s own climatology (e.g., long-term reforecasts), and what matters is whether the model correctly identifies exceedance relative to its own statistical distribution, rather than matching the exact values of the reference. As described in Suppl. L190–L194, we calculated return period thresholds separately from both the GRDC observations and the GloFAS reanalysis. This allows us to evaluate each dataset relative to its own climatology. However, we didn't compute return periods from the trained ML models reforecast as this would require generating a long reforecast climatology to fit a statistical extreme value distribution. This was beyond the scope of the current study. We acknowledge that this approach penalizes model that has a systematic bias but still captures the correct return period, but a systematic bias is also not desirable. We will discuss this issue.

Suppl. Tab 5. Is the loss only computed for those p in P that are associated to a GRDC station or also on all other points (here then on dis24 from GloFAS)?

For GRDC finetuning, P=245954 is the number of input points per sample. The loss is computed only on 3366 points where GRDC observations are available. We do not use the dis24 GloFAS here to compute the loss. We will make this clearer in the revision.

Suppl F2. How can Glofas be dropped here?

We only drop GloFAS from the input and keep the data as a target to train the model. During inference, we still use the last step of GloFAS-Reanalysis and add it to $\Delta **X**$ to generate the discharge. This ensures consistency within the table and isolates the impact of input GloFAS-Reanalysis on the model. RiverMamba still works without taking GloFAS as input. This highlights that the model is more than a post-processor of discharge data and can in fact be used as a backbone for the hydrological modeling. However, if we want to drop GloFAS completely, we need to change the objective function, e.g., by predicting the absolute value of river discharge or the change of discharge w.r.t. climatology.

Fig 13 in the Supplementary.

The GRDC observational time series, which are originally left-labeled, were explicitly converted to right-labeled time series to ensure a temporal consistency with the right-labeled predictions from the GloFAS simulations, RiverMamba, LSTM baseline, and Google reforecast dataset. The F1 scores reported in Fig. 13 are based on synchronized detection windows between model predictions and observations. We will highlight the label alignment step in the revision to avoid potential confusion. Regarding the observed differences between our Fig. 13 and Fig. 3 in Nearing et al. (2024), we note that although both figures present F1 scores by lead time and return period, they differ in several aspects like thresholding strategy, data sampling, and station filtering (evaluation domain).

I would like to thank the authors for their answers. All clarification where on point, helpful, and clear. I do have to admit that I feel like my main issue will still only be partially addressed in the revised manuscript --- and I would like to see an explicit commitment from the authors to clean up the exposition. Either way, I think the manuscript will become much better with the changes the authors promised.

I do have some very minor points that came up :

We describe the details of the LSTM baseline in Section G.3 of the suppl...

I am aware of that and did read it before making my statement. The problem was and still is that the description in Section G.3 is not understandable. As an indication for this take it that 3 out of the 4 reviews asked for a clarification in this regard.

The original model from Nearing et al. (2024) is not publicly available.

This is in contrast to the statement that I found in the paper. There it is written that all models are available: „Fully functional trained models can be found at https://doi.org/10.5281/zenodo.10397664 (ref. 45).“ What is going on here?

For a fair comparison of the methods, we also use the same input to the LSTM model and RiverMamba, which is different to Nearing et al. (2024). [...] We do not think that a model should achieve a median $R^2 > 0.7$. ... [emphasis mine]

I think the input is the main factor why the performances are so different. In your setup, both, RiverMamba and the LSTM, are post processors for Glofas. Hence, the result can be considerably different from models that directly process meteorological inputs. I must have overseen this in the manuscript (and/or it goes back to my comment on Section G.3), but your explanation here was very insightful!

We will revise the formulation since the term "global" might be misleading. As mentioned in the suppl. material on page 16, we split the Earth into smaller domains with 311K nearest points for each. We thus do not connect points from the Amazon to points in Siberia. However, we give the network the possibility to consider a very large spatio-temporal context. For instance, river networks like Amazon can be very large and providing a model that can handle a very large spatio-temporal context is thus very intuitive.

Please make this part of the main paper and reformulate. The actual method that you are using is very different (and actually still very nice!) from the impression that readers get (or at least that I got) when reading the current manuscript.

Please kindly note that ECMWF states: ...

I am not sure why I should note this. Isn’t the statement exactly what I mentioned? Did I miss something?

Thank you for your time and the recognition of the novelty of this work. We are glad that you found our work creative and addresses a critical gap for future research.

When evaluation metrics such as R², KGE, and F1-score are first introduced in the main text, it would be helpful to briefly define them or, at the very least, direct readers to where the definitions can be found (e.g., in the appendix).

Thank you for the suggestion. We described the metrics in the Supplementary due to the limited space for the main text. We will follow your suggestion and direct the readers to where the definitions can be found.

The presentation of results in Table 1 is confusing and potentially misleading. The reported values for R², KGE, and F1-score all exceed 1, which is unconventional and suggests a transformation has been applied. Upon checking the appendix, it appears that KGE has been multiplied by 100, and it is reasonable to assume the same was done for R² and F1-score—though this is not clearly stated. Rather than burying this detail in the appendix, I strongly recommend that the authors explicitly mention this scaling in the main text, ideally in the caption or title of Table 1. This kind of transformation is non-standard, and without proper explanation, it risks confusing readers and misrepresenting the model’s performance.

Thank you for this comment. We agree with you about the metrics definition. Following your comment, we will remove the transformation and report the metrics in their conventional values, to be consistent with the values reported in the figures. Note that we still use the transformation in the rebuttal for consistency with the paper.

Dear Reviewer yDqB,

Thank you again for your time and reviewing. We hope that the response has resolved your concerns. Please let us know if there are still any open issues.