🤖 machine learning

Retrieval-Augmented Foundation Models for Water Level Prediction in the Everglades

This paper proposes a retrieval-augmented framework that enhances the performance of pre-trained time-series foundation models for Everglades water level forecasting by incorporating analogous historical hydrological episodes, demonstrating significant improvements in long-horizon predictions and extreme event management.

Original authors: Rahuul Rangaraj, Jimeng Shi, Rajendra Paudel, Giri Narasimhan, Yanzhao Wu

Published 2026-06-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Rahuul Rangaraj, Jimeng Shi, Rajendra Paudel, Giri Narasimhan, Yanzhao Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the Everglades as a giant, complex bathtub. Sometimes it's nearly empty (drought), and sometimes it's overflowing (flood). Keeping track of the water level is crucial for protecting wildlife, preventing floods, and managing the local economy.

For a long time, scientists tried to predict this water level using two main tools:

Physics Models: Like trying to calculate the water level by measuring every single drop of rain and every pipe in the system. It's accurate but incredibly slow and requires constant tweaking.
Statistical Models: Like guessing the future based only on what happened in the last few days. It's fast, but if something weird happens (like a massive hurricane), these models often get confused because they've never seen that pattern before.

Recently, a new type of "super-smart" computer brain called a Time-Series Foundation Model (think of it like a super-advanced weather forecaster that has read millions of history books) was introduced. These models are great, but they have a blind spot: they only look at the immediate past (the last few days or weeks) to guess the future.

The Problem: The "Short Memory" Issue

Imagine you are trying to predict a massive flood. If you only look at the last three days of weather, you might see a few light drizzles and think, "No big deal." But you miss the fact that a huge storm system is building up three weeks away, or that a similar massive flood happened exactly 10 years ago under these exact conditions.

The standard "super-smart" models are like people with short-term memory loss; they forget the big historical patterns that could help them predict rare, extreme events.

The Solution: The "Retrieval-Augmented" Approach

The authors of this paper came up with a clever fix called Retrieval-Augmented Forecasting (RAF).

Think of it like this:

The Old Way: You ask a student to predict tomorrow's water level based only on what they ate for breakfast this morning.
The New Way (RAF): Before the student makes a guess, you give them a time machine. You say, "Look at the last 100 days of history. Now, search through a giant library of past records and find the top 5 days in history that look exactly like today."

Once the student finds those "look-alike" days, they also look at what happened after those days in the past. Did the water rise? Did it crash? By studying these historical analogs (similar past situations), the student can make a much smarter guess about the future.

How They Did It

The researchers built a system with three main steps:

The Library (Knowledge Base): They gathered a massive archive of daily water data from the Everglades, including rain, pump flows, gate openings, and water levels from 2020 to 2024.
The Search Engine (Retriever): They created two ways to find the best "look-alike" days from the library:
- Statistical Similarity: Finding days that look mathematically identical to today.
- Mutual Information: Finding days that share the same "hidden patterns" or relationships, even if they don't look exactly the same on the surface.
The Predictor (Forecaster): They fed these "look-alike" days into the super-smart foundation model (called Chronos) to help it make the final prediction.

What They Found

The results were impressive, especially for the things that matter most:

Better Long-Term Guesses: The new method was significantly better at predicting water levels weeks in advance compared to the standard models.
Extreme Event Superpower: The biggest win was during extreme events (like floods or droughts). When the water levels were doing something wild or unusual, the "Retrieval" method saved the day. It was much better at spotting these rare, dangerous situations because it could say, "Hey, this looks just like that big flood we had in 2022!"
No Re-training Needed: The best part is that they didn't have to re-teach the computer model how to be smart. They just gave it better reference materials (the retrieved history) to work with.

The Bottom Line

This paper shows that if you want to predict water levels in complex places like the Everglades, you shouldn't just rely on what happened yesterday. You need to look back at history to find similar moments. By giving AI access to these "historical twins," we can make much more accurate predictions, especially when nature throws us a curveball.

The authors have made their code and data public, so other scientists can use this "time machine" approach for their own water systems, not just the Everglades.

Technical Summary: Retrieval-Augmented Foundation Models for Water Level Prediction in the Everglades

1. Problem Statement

Accurate water level forecasting in the Everglades is critical for flood mitigation, drought management, and biodiversity conservation. While recent Time-Series Foundation Models (TSFMs) (e.g., TimeGPT, TimesFM, Chronos) have demonstrated strong performance on generic tasks through pre-training on massive datasets, their effectiveness in domain-specific hydrological applications remains limited.

The primary challenge identified is that TSFMs typically rely solely on a recent lookback window for inference. In hydrological systems, this window often contains insufficient signals for predicting rare or extreme events (e.g., floods), which occur infrequently and provide little historical context for the model to learn from. Consequently, even advanced data-driven models struggle to anticipate low-frequency, high-impact events without task-specific retraining, which is often impractical when experimental settings or domains change.

2. Methodology

The authors propose a Retrieval-Augmented Forecasting (RAF) framework designed to enrich the input context of pre-trained foundation models without requiring task-specific retraining. The framework operates in three stages:

2.1 Knowledge Base Construction

A domain-specific multivariate hydrological dataset from Everglades National Park is curated. The data is split chronologically:

Knowledge Base ( $D$ ): The first 85% of the data serves as a retrieval pool containing pairs of lookback windows ( $X_{past}$ ) and corresponding future windows ( $Y_{future}$ ).
Test Set ( $T$ ): The remaining 15% is used for evaluation.

2.2 Knowledge Retrieval

Given a current input window, a retriever identifies high-relevance historical analogs from the knowledge base to serve as contextual information. The paper evaluates two distinct retrieval strategies:

Similarity-based Retrieval: Uses a T5 encoder to embed lookback windows into a fixed-length vector space (512 dimensions). Similarity is measured via Euclidean distance between the embedding of the current input and candidate historical windows.
Mutual Information (MI)-based Retrieval: Leverages information theory to quantify statistical dependencies between the current input and historical candidates. MI captures both linear and nonlinear relationships, serving as a relevance signal to identify contexts most informative about the current system state.

The top- $k$ highest-scoring samples are retrieved, forming a context set $C$ .

2.3 Retrieval Augmentation Strategies

The retrieved context is integrated with the original input using three distinct strategies:

Strategy A (Averaging-First): The top- $k$ retrieved contexts are averaged pointwise to create a single representative context ( $\bar{C}$ ), which is then concatenated with the original input.
Strategy B (Augmenting-First): Each of the $k$ retrieved contexts is individually concatenated with the original input to form $k$ separate augmented inputs. The forecaster generates $k$ predictions, which are then averaged for the final output.
Strategy C (Long-Context Concatenation): All $k$ retrieved contexts are directly concatenated with the original input to form a single, extended sequence.

2.4 Forecaster

The framework utilizes Chronos, a pre-trained time-series foundation model, as the forecaster. The RAF approach (Sim-RAF and MI-RAF) is model-agnostic but demonstrated here using Chronos without fine-tuning, relying on zero-shot inference enhanced by the retrieved context.

3. Experimental Setup

Dataset: Daily records from October 16, 2020, to December 31, 2024, comprising 38 variables (water levels, rainfall, PET, gate/pump flows) across five key monitoring stations in the Everglades.
Task: Predict water levels at five stations over horizons of 7, 14, 21, and 28 days using a 100-day lookback window.
Baselines: Compared against task-specific deep learning models (DLin, TimeMixer, KAN, iTransformer) and other foundation models (TimeGPT, TimesFM, Timer, Chronos).
Metrics: Mean Absolute Error (MAE), Root Mean Squared Error (RMSE), and the Symmetric Extremal Dependence Index (SEDI) for evaluating extreme event prediction.

4. Key Results

Performance Superiority: The retrieval-augmented methods (Sim-RAF and MI-RAF) consistently outperformed all baselines, including the strongest pre-trained model (Chronos).
Long-Horizon Gains: The benefits of retrieval augmentation were more pronounced for long-term predictions.
- For a 7-day horizon, Sim-RAF improved MAE and RMSE by 4.1% and 5.1% over Chronos, respectively.
- For longer horizons, Sim-RAF achieved 7.0% (MAE) and 5.8% (RMSE) improvements, while MI-RAF showed even larger gains of 13.5% (MAE) and 7.4% (RMSE).
Extreme Events: The framework yielded disproportionately larger gains during extreme events, as measured by SEDI. This indicates that retrieving analogous historical episodes effectively compensates for the lack of immediate historical signals during rare hydrological events.
Spatial Variability: Performance varied across stations, with the NP205 station showing the lowest performance, attributed to weaker water-level correlations with other stations.

5. Significance and Claims

The paper claims to provide empirical evidence that analog-based retrieval can significantly benefit pre-trained time-series foundation models in environmental science. Specifically:

Generalizability: The proposed RAF pipeline is basin-agnostic and can be applied to other hydrological systems given time-series data and historical archives, without requiring domain-specific retraining.
Decision Support: By improving long-horizon and extreme-event forecasts, the method offers practical insights for environmental decision-making in complex, non-stationary systems like the Everglades.
Methodological Insight: The study highlights the strengths and failure modes of foundation models when applied to hydrology, demonstrating that enriching input context with retrieved historical analogs is a viable strategy to overcome the limitations of short lookback windows in data-sparse or extreme regimes.

The authors emphasize that while the evaluation is specific to the Everglades, the framework's design allows for transferability to other hydrological contexts, offering a new direction for enhancing water level prediction in ecologically sensitive systems.