Forecasting with Memory
Why we need retrieval augmented forecasting and recent advances in the field
Introduction
Forecasting models often perform well on benchmark datasets but face significant challenges in real-world applications. One reason is data imbalance: rare but critical events such as extreme weather, demand spikes, or market disruptions are sparsely represented in training data.
Another challenge is the limited scale of forecasting datasets. Unlike vision or language models trained on massive corpora, forecasting models typically learn from relatively small datasets.
These limitations have led to research in Retrieval-Augmented Forecasting (RAF), an approach inspired by retrieval-augmented generation in large language models. Instead of relying solely on model parameters, RAF retrieves relevant historical examples and uses them as additional context during prediction. In this article, I explore a recent study that combines time series embeddings with diffusion-based forecasting, allowing models to retrieve and learn from similar historical patterns when generating predictions.
Retrieval-Augmented Forecasting with Diffusion Models
A recent study, Retrieval-Augmented Time Series Diffusion Model (RATD) extends this idea to diffusion-based forecasting models.
The framework first retrieves historical sequences that resemble the current forecasting scenario. It then uses those examples to guide the diffusion model during forecast generation.
The intuition is simple: if similar patterns have occurred before, their future trajectories may provide useful clues about what happens next.
Key Innovations in RATD
Embedding and Retrieval
RATD learns a latent representation for each time series using a pretrained encoder. These embeddings capture higher-level temporal characteristics, allowing sequences with similar underlying dynamics to be located near one another in the embedding space.
Retrieval is then performed using nearest-neighbor search on these embeddings, returning the top-k historical examples that are most relevant to the forecasting query.
Reference Modulation Attention (RMA)
Retrieved examples are not equally useful. Some may contain valuable information about future behavior, while others may only share partial similarities with the input sequence.
To address this, RATD introduces Reference Modulated Attention (RMA). Rather than directly feeding retrieved sequences into the forecasting model, RMA extracts latent features from the retrieved references and selectively injects them into the diffusion denoising process. This allows the model to emphasize informative patterns while suppressing irrelevant information.
RMA is particularly important because time series diffusion models lack the explicit conditioning signals commonly used in image generation. The retrieved references effectively act as external guidance.
Results and Future Directions
RATD was evaluated on multiple benchmark datasets, including electricity, weather, exchange rate, and wind forecasting tasks. Across most settings, it outperformed both transformer-based and diffusion-based baselines.
Despite these gains, retrieval-augmented diffusion forecasting introduces additional computational overhead from both nearest-neighbor retrieval and iterative diffusion sampling.
Future research will likely focus on improving retrieval efficiency and developing hybrid architectures that combine retrieval mechanisms with large-scale forecasting foundation models.