TS-RAG: Retrieval Augmented Generation for Time Series Forecasting

The paper, authored by Yixiong Xiao, Congxi Xiao, Shuangli Li, and Jingbo Zhou (Baidu, Inc.), introduces TS-RAG, a novel retrieval-augmented generation (RAG) framework designed specifically for time series forecasting. The work addresses the gap between the proven success of RAG in natural language processing (NLP) and its limited application in time series forecasting.

The authors observe that while deep learning models, particularly transformer-based architectures, have shown impressive performance in time series forecasting, the application of retrieval-augmented generation (RAG) in this domain remains limited. They note that "since RAG has proven effective in enhancing the capabilities of large language models by incorporating relevant external information, retrieving similar time series sequences as references might also improve accuracy in time series forecasting tasks."

However, the authors identify two key challenges to extending RAG to time series:

1. Model capability limitations: "Most time series models are constrained by limited training data, smaller parameter scales, and a lack of the extensive generative capabilities found in large language models. Simply concatenating reference sequences into the prompt, as done in language models, may not yield the expected results."

2. Retrieval efficiency: "The efficiency of similarity retrieval, particularly when using traditional methods like Dynamic Time Warping (DTW), can be computationally expensive, requiring substantial time to identify similar sequences from large datasets."

TS-RAG leverages RAG to enhance forecasting performance by introducing specially designed reference tokens to effectively fuse information from the input sequence with that from retrieved similar sequences, enabling a more robust capture of complex temporal dynamics. Additionally, by employing vector search techniques for efficient similarity retrieval, TS-RAG significantly reduces the computational burden associated with traditional methods like Dynamic Time Warping (DTW).

Problem Definition: Given a multivariate time series x(i)1:L ∈ R(C×L) (where C is the number of channels and L is the sequence length), the objective is to predict future values x(i)L+1:L+T ∈ R(C×T). The authors assume access to a historical database D from which a similarity-based retrieval function R identifies the most relevant sequences r(i) = R(x(i), D) = r(i,1),..., r(i,N). The forecasting task is reformulated as x̂(i)L+1:L+T = fθ(x(i)1:L, r(i)).

Embedding Process: Each channel is divided into non-overlapping patches of length p. The key innovation is that each retrieved sequence r(i,n) is associated with a learnable reference token vref(n)(c). These tokens act as compressed representations of historical patterns and are prepended to the patchified input. The final input embedding is computed as x(i)embed = Wx x̃(i)patch + bx + PE(x̃(i)patch), where Wx and bx are learnable parameters and PE(·) denotes positional encoding.

Reference Fusion: The framework integrates information using two attention mechanisms:

1. Self-Attention: Captures intrinsic temporal dependencies within the input sequence.

2. Cross-Attention: Aligns retrieved sequences with the input sequence, allowing the model to selectively extract useful information.

The final hidden representations pass through a feed-forward network, and the model outputs the forecast leveraging both self-learned representations and retrieved historical patterns to improve predictive accuracy.

Datasets: Six widely used benchmarks were used: ETTh1, ETTh2, ETTm1, ETTm2, Electricity, and Weather, spanning energy consumption, electricity load forecasting, and meteorological measurements. The input sequence length was set to 96 time steps, with predictions for the next 96, 192, 336, or 720 time steps.

Implementation Details: All experiments used PyTorch on an NVIDIA V100 GPU with 32GB memory, adopting the same hyperparameter settings as iTransformer. Training used the Adam optimizer (β1=0.9, β2=0.999) with a learning rate of 0.001 with cosine decay. A Temporal Convolutional Network (TCN) was trained to retrieve similar sequences.

The paper reports that "TS-RAG-CM achieves the best or near-best performance across all six datasets (ECL, ETTh1, ETTh2, ETTm1, ETTm2, Weather), demonstrating superior predictive accuracy. Specifically, it achieves the lowest average MSE (0.310) and MAE (0.348), outperforming all other models. TS-RAG (which does not model cross-channel dependencies, unlike TS-RAG-CM which considers channel dependency) ranked third overall yet achieves results comparable to iTransformer, which explicitly incorporates channel mixing, highlighting the significant contribution of RAG in a channel-independent setting."

1. Comparison with Direct Sequence Concatenation: The authors tested whether NLP-style direct concatenation of retrieved sequences into the prompt is effective for time series. Results showed that TS-RAG and TS-RAG-CM consistently outperform Direct Sequence Concatenation: "At 96 length, TS-RAG reduces MSE by 12.1%, while TS-RAG-CM improves it by 14.2%. For 192 length, TS-RAG achieves a 13.8% reduction, and TS-RAG-CM improves it by 17.6%. At 336 length, TS-RAG lowers MSE by 14.0%, and TS-RAG-CM further reduces it by 17.4%. The largest improvement is observed at 720 length, where TS-RAG decreases MSE by 16.2% and TS-RAG-CM achieves an 18.2% reduction. This demonstrated that direct sequence concatenation is not as effective in time series forecasting models as it is in NLP tasks because time series models lack the strong generalization and generative capabilities of LLMs, highlighting the necessity of the reference fusion framework proposed in TS-RAG."

2. Reference Number Analysis: The authors found that TS-RAG and TS-RAG-CM achieve the best performance when using a single reference sequence. In contrast, incorporating two or four reference sequences generally leads to higher MSE and MAE, because retrieving too many similar sequences may introduce redundant or conflicting patterns, making it harder for the model to extract relevant information.

3. Retrieval Approach Comparison: TCN-based retrieval achieved the best forecasting performance with lower MSE and MAE compared to DTW-based and Euclidean Distance (ED) retrieval. Regarding efficiency, TCN-based retrieval is slightly slower than ED retrieval (0.004s vs. 0.002s), but it is still significantly faster than DTW-based retrieval (14.27s), making it a more practical choice for real-time applications.

The authors conclude that TS-RAG effectively integrates retrieved historical sequences into the forecasting process and addresses key limitations of existing deep learning models by incorporating external information through reference tokens, facilitating a more structured fusion of complex temporal dependencies. The framework consistently outperforms state-of-the-art models, highlighting its potential to enhance forecasting performance in both common and rare-event scenarios, particularly when training data is scarce or lacks diversity. This work bridges the gap between retrieval-augmented methods and time series forecasting, showcasing how ideas from natural language processing (NLP) can be effectively adapted to the time series domain.

Improvements to AI systems based on TS-RAG:

1. Add retrieval-augmented forecasting to any time series model. Instead of relying solely on the input window, the system retrieves historically similar time series segments from a database and fuses them via learnable reference tokens prepended to the input embedding. This lets the model leverage recurring patterns, rare events, and domain-specific trends that are not visible in the current input alone.

2. Replace NLP-style prompt concatenation with a dedicated reference fusion mechanism. Directly concatenating retrieved sequences (as done in LLM RAG) fails for smaller time series models. The improved system uses self-attention on the input sequence and cross-attention between input tokens and reference tokens, allowing the model to selectively extract useful information from retrieved sequences without overwhelming or confusing the predictor.

3. Use efficient vector-search retrieval instead of DTW. The system embeds time series windows (e.g., via a temporal convolutional network) into a vector space, enabling sub-linear approximate nearest-neighbor search. This cuts retrieval time from 14 seconds (DTW) to milliseconds, making retrieval-augmented forecasting feasible for real-time and large-scale applications.

4. Automatically calibrate the number of references to avoid pattern conflict. The improved system can identify that a single retrieved reference often yields the best accuracy, while multiple references inject redundant or conflicting patterns. The system can adaptively choose the optimal reference count per prediction, or be trained to weight references by relevance, improving forecast stability.

5. Improve forecasts in data-scarce or low-diversity regimes. By pulling similar historical sequences from a database, the system gains access to external temporal patterns that act as pseudo-training examples, reducing overfitting and improving accuracy when the training set is small or lacks diversity.

6. Enhance channel-aware forecasting. Extend the framework with a channel dependency mechanism (as in TS-RAG-CM) so the system models correlations between variables while still benefiting from retrieved references. This yields the best average MSE and MAE across benchmarks (e.g., 0.310 MSE, 0.348 MAE), outperforming state-of-the-art models like iTransformer.

7. Provide a plug-and-play component for existing transformers. The reference token embedding and cross-attention layer can be inserted into any patch-based transformer or linear forecasting model, enabling those models to become retrieval-augmented without redesigning their core architecture.

What the improved AI system can do:

• Given a multivariate time series input, retrieve the most similar historical windows from a database in milliseconds.

• Fuse those windows into the prediction process via cross-attention over learnable reference tokens, producing forecasts that are more accurate than any single-model baseline.

• Handle long-horizon forecasting (96, 192, 336, 720 steps) with consistently lower error than non-retrieval models.

• Work well in channel-independent settings (improving robustness) and even better when channel dependencies are modeled explicitly.

• Operate in real-time environments where DTW-based retrieval is too slow.

• Automatically ignore redundant retrieved candidates, preventing performance degradation from too many references.

• Adapt to rare-event forecasting by recalling analogous past events from the database, even when the training data does not contain them.