Surface water quality prediction via an MLA-Mamba hybrid neural network with GRPO optimization

Why predicting river health matters

Rivers and lakes are our drinking water sources, irrigation supplies, and wildlife habitats. Yet their quality can change quickly when pollution washes in from farms, factories, or cities. Authorities often find out only after damage is done. This study explores how modern artificial intelligence can act like a smart early-warning system, forecasting changes in water quality days in advance so that managers have time to respond.

Old tools, new problems

For decades, scientists have tried to predict water quality using mathematical formulas and traditional statistics. These methods either simulate chemistry and flow in great detail or fit past measurements with relatively simple curves. Both approaches struggle with the messy reality of rivers, where weather, upstream discharges, and biological activity all interact in complex, nonlinear ways. They often miss sudden pollution spikes or fail to account for how conditions at one monitoring station ripple downstream to another. As a result, forecasts can be too crude for confident decision-making.

Teaching a neural network to read a river

The authors propose a new deep learning model, called MLA-Mamba, designed specifically for this tangle of space and time. Instead of looking at a single sensor in isolation, the model ingests a week’s worth of hourly data from multiple monitoring stations, together with supporting information such as water temperature, flow, and acidity. It then learns to predict four key indicators that signal organic pollution and nutrient loading: a chemical oxygen demand index (CODMn), ammonia (NH3–N), total phosphorus (TP), and total nitrogen (TN). The model combines two specialized components. One focuses on patterns over time, spotting cycles, slow drifts, and delayed effects. The other looks across space, learning how upstream and neighboring stations move together. By fusing these views, the network builds a richer picture of how water quality evolves.

Capturing both time trends and upstream influence

Inside the MLA-Mamba framework, the "Mamba" module concentrates on the temporal story. It scans long sequences of measurements, using ideas from state-space models and modern recurrent networks to hold onto information from days earlier without being overwhelmed. This helps it recognize seasonal patterns and lingering impacts from past disturbances. In parallel, a "Multi-Head Local Attention" module weighs how strongly each station relates to the others at a given moment, with a built-in bias toward nearby sites in the same river reach. If an upstream station suddenly records a jump in ammonia, the attention mechanism can quickly shift focus to that signal when predicting conditions downstream. A multi-task setup lets the model learn all four water quality indicators together, so that shifts in one pollutant can inform expectations for the others.

Smarter training for noisy environmental data

Training such a network on real-world sensor records is tricky: data are noisy, gaps occur, and standard optimization methods can get stuck. To tackle this, the researchers introduce a custom training strategy they call Gradient Reparameterization Optimization (GRPO). GRPO adjusts how fast each parameter in the network learns based on how its gradient behaves over time, roughly speeding up in stable directions and slowing down when updates start to oscillate. It also enforces a minimum step size so learning does not stall on flat parts of the error surface. The team further uses dropout not only to prevent overfitting but also to estimate uncertainty, by running the model multiple times and examining how much its predictions vary. This yields confidence bands around each forecast, giving managers a sense of how trustworthy any particular prediction is.

Putting the model to the test

The authors evaluate MLA-Mamba on several years of hourly data from two river stations in China, one upstream of the other. The model takes the previous seven days of data and predicts the next three days. It is compared against eight alternatives, ranging from classic statistical methods to modern deep learning architectures such as long short-term memory (LSTM) networks, convolutional–recurrent hybrids, and Transformer models. Across all four indicators and both locations, MLA-Mamba consistently delivers the lowest prediction errors. In many cases, it reduces typical errors by 10–20 percent relative to strong deep learning baselines. When parts of the model are disabled in controlled tests—removing the spatial attention, substituting a standard LSTM for the Mamba module, turning off the GRPO optimizer, or training each indicator separately—performance noticeably degrades. This shows that each ingredient contributes to the gains.

What this means for protecting water resources

In plain terms, the study demonstrates that a tailored hybrid neural network can make more accurate and reliable short-term forecasts of river pollution than today’s standard tools. By simultaneously tracking multiple pollutants across multiple stations, and by quantifying how certain it is about its own predictions, the MLA-Mamba framework could underpin early-warning systems that trigger inspections or temporary controls before thresholds are breached. While the approach still depends on good-quality monitoring data and must be tested across more rivers and extreme events, it offers a promising route toward smarter, data-driven management of surface waters.

Citation: Wei, R., Chen, H. & Wang, H. Surface water quality prediction via an MLA-Mamba hybrid neural network with GRPO optimization. Sci Rep 16, 5845 (2026). https://doi.org/10.1038/s41598-026-36229-3

Keywords: water quality forecasting, river pollution, deep learning, spatio-temporal modeling, environmental monitoring