Retrieval-Augmented Diffusion Models for Time Series Forecasting

Thank you very much to reviewer zh2V for the careful reading and consideration, as well as for acknowledging the innovative aspects of our approach .

Q1: I am wondering whether there is any previous work on time series RAG?

A1: In the related work section, I mentioned that previous studies have combined time series prediction with RAG[1][2]. Our proposed method differs from these prior works in several key aspects:

Different Utilization of References: Previous methods either directly fused features using existing cross-attention modules or concatenated them for fusion. In contrast, we introduce a novel RMA(Reference Modulated Attention) module for feature fusion.

Distinct Research Motivation: We believe that time series diffusion models may offer advantages over transformer networks in certain prediction tasks. However, these advantages are not always pronounced (as mentioned in your subsequent question). We identify the lack of guidance mechanisms as a major barrier to achieving greater performance with diffusion models. Due to this absence, diffusion models must directly learn the mapping from Gaussian distribution to data distribution, which can be challenging. We propose an enhanced retrieval-guided mechanism and aim to provide new insights for future research on time series diffusion models, particularly in the study of guidance mechanisms.

Differences in Experimental Completeness: Previous studies did not provide comprehensive evaluations of public datasets, or experiment details, making it difficult to assess model performance. for the sake of comprehensive experimental comparison, we present here the MSE results of our replicated method under our experimental settings, demonstrating that our approach shows noticeable advantages.

Clearly, RATD exhibits significant performance advantages, indicating that our method can better utilize references for prediction.

In summary, while there are similarities between prior works and our method to some extent, there are significant differences as well. We believe our approach advances the integration of RAG with time series prediction tasks.

Q2: I am wondering whether there is any justification in choosing diffusion models for time series forecasting given its potential extra computational cost.

A2: The diffusion model is a generative framework and the transformer is a specific model structure. Two methods follow different technical paths, making it challenging to assess their relative performance. To be specific, many existing efforts focus on enhancing performance by updating the architecture of transformers [3], while related works on diffusion models concentrate on updating the framework as a whole [4]. However, I believe time series diffusion models hold greater research potential. Firstly, time series diffusion models are currently constrained by inference costs, preventing the direct adoption of powerful time series transformer models as their denoising network. In the future, if transformer models with significantly lower inference costs and high performance emerge, the performance of diffusion models could have substantial improvements. Secondly, current time series diffusion models lack clear and effective guidance mechanisms. This work is among the few that address this issue, and I believe that appropriate guidance mechanisms could greatly enhance the performance of diffusion models (as demonstrated in the experimental section of this paper). Further research into guidance mechanisms may activate the potential of time series diffusion models in the future.

Q3: In section 4.2, how is DR different from DR′ if C represents all categories? Is n a hyperparameter for DR′ that also needs to be introduced in section 4.2?

A3: In fact, we employed two different methods for constructing retrieval datasets, primarily to correspond to two types of time series datasets: one smaller and lacking clear category labels, such as wind, electric...; and one with complete category annotations, often larger, such as MIMIC-IV. For the former, we could only use the first method to construct the retrieval dataset, while for the latter, both methods were applicable. If all category data from the latter case were applied, as you mentioned, then there would be no difference between the two construction methods.

Additionally, you are correct in noting that the sample counts for each category should also be included in the formula. Thank you for your correction, and I will review it again.

Q4: ...would there be memory concerns if you are working with long-horizon time series forecasting / long context windows?

A4: Your concerns are reasonable, and we have considered them. Therefore, we used a linear layer (shown on the right side of Figure 3) to reduce the length of the concatenated sequence, which facilitates subsequent computations. We will emphasize this point in the final version of the paper.

[1]Yang S, Eisenach C, Madeka D. MQ-ReTCNN: Multi-Horizon Time Series Forecasting with Retrieval-Augmentation[J]. 2022. [2]Jing B, Zhang S, Zhu Y, et al. Retrieval-based time series forecasting[J]. arXiv preprint arXiv:2209.13525, 2022. [3] Zhang Y, Yan J. Crossformer: Transformer utilizing cross-dimension dependency for multivariate time series forecasting[C]//The eleventh international conference on learning representations. 2022. [4] Rasul K, Seward C, Schuster I, et al. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting[C]//International Conference on Machine Learning. PMLR, 2021: 8857-8868.

Thank you very much to reviewer uLZU for acknowledging some advantages of our proposed method. We understand you may have concerns regarding the overall novelty of the entire paper. We will address this point.

Q1: The paper lacks innovation as the retrieval-enhanced diffusion method has already been proposed in the field of image generation...

A1： The innovation of this method is reflected in the following aspects：

• We for the first time utilize RAG on time series diffusion models. In text2image diffusion models, cross-modal guidance mechanisms are well-established, and retrieval-enhanced methods emphasize using semi-parametric approaches to enhance generation efficiency. However, time series diffusion models lack a universally recognized effective guidance mechanism. We are the first to provide a promising RAG-based guidance mechanism for the diffusion generation of time series, which represents a pioneering advancement in time series diffusion model research. Additionally, while existing work utilizes retrieval mechanisms for time series prediction, previous efforts have not adequately analyzed the use of references nor extensively evaluated and compared model performance in experiments. Our novel framework addresses these gaps in previous research and significantly advances the study of retrieval-enhanced time series prediction methods.
• Our method introduces a novel attentional mechanism for time series diffusion, called RMA (Reference Modulated Attention, L167-180), designed to enhance the guidance provided by references in the generation process. While the overall structure is not complex, our proposed RMA is effective and similarly groundbreaking. Additionally, we designed extra ablation experiments to demonstrate the effectiveness of RMA. |Dataset | Exchange| Electricity | Wind | | :-----| :---- | :---- |:---- | | CSDI | 0.077 | 0.379 | 1.066 | | CSDI+Linear| 0.075 | 0.316 | 0.932 | | CSDI+Cross Attention| 0.028 | 0.173 |0.829 | | CSDI+RMA|0.013 | 0.151 | 0.784 |

Where CSDI represents the most basic network framework, CSDI+Linear denotes the approach where inputs and references are concatenated via a linear layer and fed into the network together, CSDI+CrossAttention signifies the use of cross attention to fuse features from inputs and references, and finally, CSDI+RMA, which incorporates an additional RMA. Clearly, RMA is highly effective in leveraging references for guidance.

Regarding the issue of dependency on pre-trained encoders, we discussed this in Section 5.3 (line 240). In the embedding-based retrieval mechanism, different encoders (feature extractors) do indeed introduce performance differences. However, these differences are not decisive. This variability may stem from the effectiveness of the methods we selected and the quality of the retrieved references. This also indicates that our method exhibits high robustness and does not rely excessively on the choice of pre-trained encoders.

Q2: There is a typo in the last sentence of the second paragraph from Section 5.2. “.t”

A2: I apologize for the typo. It was an oversight on our part, and I will review the paper again. Thank you for bringing it to my attention.

Q3: What would be the result if the retrieved sequences were directly used for a simple weighted average for the final generation output...

A3: Your suggestion is very interesting! Indeed, we paid close attention to this issue during our experimental phase. If the references are too similar to the ground truth (GT), the model's predictions can overly rely on the references, thereby failing to learn the true conditional distribution. We assessed this visually during experiments, such as in our analysis in Figure 4 from the paper. Following your advice, we conducted experiments on two datasets.

Specifically, "Reference (Closest)" refers to the closest reference. The results demonstrate that on datasets with strong periodic patterns, the references and predicted results are indeed very similar, whereas, on more complex datasets(wind), the references cannot replace the prediction at all.

Q4: I am uncertain whether the non-diffusion time-series prediction methods used in the author's comparative experiments include retrieval-enhanced modules...

A4: We did not compare our method with the previous RAG-based time series methods in the paper because these papers did not provide comprehensive evaluations on commonly used time series datasets or publicly available code. The implementation details in these papers were also very limited, making replication of their methods challenging. Displaying replicated results directly in the main text could potentially lead to unnecessary misinterpretation. Nevertheless, as you mentioned, for the sake of comprehensive experimental comparison, we present here the MSE results of our replicated method under our experimental settings, demonstrating that our approach shows noticeable advantages.

Clearly, RATD exhibits significant performance advantages, indicating that our method can better utilize references for forecasting.

Q5: In Formula 1, the symbol D is used ambiguously, representing both distance and feature dimensions.

A5: I apologize for the mistake in using symbols. I will correct the paper and remove any potentially misleading parts.

We are glad to receive your reply! We are also pleased that most of your concerns have been addressed.

Regarding your concerns about the paper's novelty, I would like to provide some additional clarifications. In addition to the lack of an effective guiding mechanism of exisiting time series diffusion models, another core issue is the limited size or quality of existing datasets. Specifically, real-world time series datasets are either too small in scale or highly imbalanced, which may not meet the high-quality data requirements of training diffusion models[1]. Our approach offers a general solution by leveraging useful information from the dataset during the generation process progressively, rather than expending significant resources to augment the existing dataset. In other words, our method makes the most of the limited dataset, thereby addressing the aforementioned core issue to some extent.

Additionally, a critical challenge in time series forecasting is modeling both the periodic and trend components of the time series simultaneously [2]. Our approach provides substantial assistance in direct and explicit modeling of these temporal components by utilizing guidance from references.

Overall, our method focuses on and addresses these core issues, making it both innovative and practical, which have been recognized by other reviewers.

Regarding the supplementary experiments for A4, the performance advantages of our method stem from three aspects:

• Advantages of Diffusion Model Framework: Unlike methods that directly use a transformer, the multi-stage framework of diffusion models reduces the difficulty of learning from prior distributions to data distributions, which helps the model learn more accurate conditional distributions.
• Iterative Guidance Mechanism: Our effective guidance mechanism allows the utilization of conditions (i.e., references) during the diffusion process iteratively (T =100 in our experiment), whereas previous methods only performed feature fusion once. In other words, diffusion models can leverage references as guidance more effectively.
• Better Feature Fusion Module: The proposed RMA (Retrieval-Modulated Attention) module performs better in feature fusion compared to the previous work. We designed an extra experiment to prove it (results are shown in the previous response A1). Our experiments found that compared to the basic cross-attention-based method (used by ReTime) and linear fusion methods (used by MQ-ReTCNN), RMA can integrate an extra information matrix (representing correlations between time and feature dimensions), thus realizing the fusion of feature more effectively.

In summary, the first two advantages stem from the diffusion model framework itself, while the last advantage comes from the new module we proposed. By integrating the three components, our proposed method achieved better experimental results.

Once again, thank you for your reply. We are so honored that you find our work interesting and promising. If you have any further questions, please feel free to ask. Thank you!

[1] Schramowski P, Brack M, Deiseroth B, et al. Safe latent diffusion: Mitigating inappropriate degeneration in diffusion models[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 22522-22531.

[2]Shen L, Chen W, Kwok J. Multi-Resolution Diffusion Models for Time Series Forecasting[C]//The Twelfth International Conference on Learning Representations. 2024.

Thank you very much to reviewer FgUZ for recognizing our proposed method. Following your suggestion, we have conducted additional experiments and provided further explanations of the model's structure. We hope to receive your approval.

Q1: The experimental setup and comparisons are insufficient and not aligned with existing studies, making it difficult to verify the performance and correct implementation of this work.

A1: We extended the horizon length in our experiments, which is currently a reasonable practice as many popular forecasting methods [1][2] evaluate results by predicting horizon lengths of 96 and above. Of course, we agree with your point, so we supplemented the experiments according to the settings of CSDI, and the results are as follows.

We compared CSDI, TimeGrad, and RATD under the experimental setting of CSDI (with a prediction window length of 24). It is noteworthy that in this experimental setup, the differences in results among different methods were narrowed. This may be due to the reduced difficulty of the task with a smaller prediction window. From the extra experiment, our method still exhibits obvious advantages.

W2: Insufficient Ablation Tests: While Table 3 indicates the importance of choosing good retriever embeddings, there are no systematic ablation tests of the architecture developed (Figure 3).

A2: Thank you very much for your suggestions! I truly appreciate that you took the time to read our paper thoroughly and provided very constructive feedback! To validate the ideas you mentioned, we conducted additional ablation experiments and performed preliminary tests on three datasets. The results are shown in the table below:

Where CSDI represents the most basic network framework, CSDI+Linear denotes the approach where inputs and references are concatenated via a linear layer and fed into the network together (following [6]), CSDI+CrossAttention signifies the use of cross attention to fuse features from inputs and references (following [5]), and finally, CSDI+RMA, which incorporates an additional Reference Modulated Attention (RMA).

Through our experiments, we found that compared to the basic Cross-attention-based method, RMA can integrate an edge information matrix (representing correlations between time and feature dimensions) more effectively. The extra fusion is highly beneficial in experiments, guiding the model to capture relationships between different variables. In contrast, linear-based methods concatenate inputs and references initially, which prevents the direct extraction of meaningful information from references, resulting in comparatively modest performance.

W3: There is no information on GPU memory usage and computation efficiency...

A3: As mentioned in our conclusion and limitation section, the retrieval process incurs additional computational costs. To minimize this cost, we pre-retrieve references for all samples during training and record the corresponding indices (as noted in line 215 of the paper). This means that during actual training and inference, we only need to consider the additional cost introduced by RMA. Since you are concerned about this aspect, we provide a rough overview of GPU memory usage and computational efficiency (all the experimental settings are the same as the paper, we complete the comparison on the wind dataset).

In general, compared to CSDI, our method has a slight disadvantage in terms of memory consumption and computational efficiency. However, this disadvantage is not significant and does not pose serious challenges to applications.

Q4: More comprehensive and detailed comparisons with existing studies...

A4: We have provided supplementary experimental results above in hopes of addressing your concerns. Thank you.

Q5: Systematic ablation tests to illustrate your specific architecture designs...

A5: Similarly, we hope our ablation experiments have addressed your concerns. If you have any further questions, please feel free to reply. Thank you.

[1]Liu Y, Hu T, Zhang H, et al. iTransformer: Inverted Transformers Are Effective for Time Series Forecasting[C]//The Twelfth International Conference on Learning Representations. [2]Nie Y, Nguyen N H, Sinthong P, et al. A Time Series is Worth 64 Words: Long-term Forecasting with Transformers[C]//The Eleventh International Conference on Learning Representations. [3]Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems, 34:24804–24816, 2021. [4]Rasul K, Seward C, Schuster I, et al. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting[C]//International Conference on Machine Learning. PMLR, 2021: 8857-8868. [5]Yang S, Eisenach C, Madeka D. MQ-ReTCNN: Multi-Horizon Time Series Forecasting with Retrieval-Augmentation[J]. 2022. [6]Jing B, Zhang S, Zhu Y, et al. Retrieval based time series forecasting[J]. arXiv preprint arXiv:2209.13525, 2022.

Hello! Here we provide more comprehensive experimental results to address any concerns you may have about the completeness of our experiments. Our experimental findings are as follows, with additional experiments on all datasets mentioned in CSDI using identical experimental setups to CSDI(CRPS means CRPS-sum here).

From the results, it can be observed that our method performs better across all datasets. To quantify this more explicitly, we have calculated the average performance improvement （API） of our method compared to CSDI under both experimental settings.

It appears that our method will indeed achieve better performance on more complex datasets. Notably, our method exhibits higher performance gains in more complex experimental setups, aligning with our motivation to tackle more intricate time series forecasting tasks.

Additionally, we have supplemented results from disintegration experiments on six of these datasets.

As mentioned earlier, CSDI represents the basic network framework, CSDI+Linear denotes the approach where inputs and references are concatenated via a linear layer and fed into the network together, CSDI+CrossAttention signifies the use of cross attention to fuse features from inputs and references, and finally, CSDI+RMA, which incorporates an additional Reference Modulated Attention (RMA).

For further analysis of RMA performance, we have also supplemented additional qualitative and quantitative experiments to compare the similarity between reference and prediction results. Due to our current inability to upload images or links, we present quantitative experimental results in the table below. Specifically, we use the Mean Squared Error (MSE) metric to measure the similarity between the reference and predicted results (where smaller values indicate greater similarity). Additionally, to assess whether the model effectively utilizes explicit information from the reference rather than noise, we employ the STL( Seasonal-Trend decomposition using LOESS) method to evaluate the similarity of different components. The results are as follows:

In some complex time series datasets (Solar), trend components are more important for predictions, while for ECG sequences (Mimic-IV), periodicity is more crucial. Due to enhanced feature integration capabilities, our method adaptively captures the most important information for both cases, leading to better prediction results. In summary, our approach efficiently captures information from the reference and provides beneficial guidance on periodicity and trends for the generated time series.

Finally, we sincerely hope that the additional experiments we propose can address your concerns regarding the completeness of the experiments! If you have any further questions, please let us know!

We thank Reviewer AGBG for the thorough and valuable feedback. We are glad that the reviewer found that the proposed model is effective and our paper is easy to read. The main concern of the reviewer is that the architecture of RATD lacks justification. Please see below for our responses to your concerns.

Q1: The proposed RATD lacks justification for architecture.

A1: We here justify the effectivness of our architecture from following aspects:

• We for the first time introduce the concept of RAG to time series diffusion models. In other words, and introduces a new guiding mechanism for time series diffusion models. Describing time series data directly is challenging, and currently, there is no universally recognized effective guidance mechanism that guarantees the effectiveness of diffusion models in time series conditional generation tasks.
• The reference-augmented mechanism we propose is highly effective. Our model framework is based on CSDI (mentioned in Line 167), and building upon CSDI, our approach achieves superior performance. We have listed the performance improvement ratios (MSE) in the table below.

• The RMA (Reference Modulated Attention, L167-180) we propose is also robust and effective. In terms of robustness, the performance does not rely excessively on pre-trained encoders（from line 250 on paper). Additionally, we designed extra ablation experiments to demonstrate the effectiveness of RMA (Metric: MSE). |Dataset | Exchange| Electricity | Wind | | :-----| :---- | :---- |:---- | | CSDI | 0.077 | 0.379 | 1.066 | | CSDI+Linear| 0.075 | 0.316 | 0.932 | | CSDI+Cross Attention| 0.028 | 0.173 |0.829 | | CSDI+RMA|0.013 | 0.151 | 0.784 |

Where CSDI represents the most basic network framework, CSDI+Linear denotes the approach where inputs and references are concatenated via a linear layer and fed into the network together, CSDI+CrossAttention signifies the use of cross attention to fuse features from inputs and references, and finally, CSDI+RMA, which incorporates an additional RMA. Clearly, RMA is highly effective in leveraging references for references.

Q2: What is the main difference between RAG and the proposed work?

A2: In our related work (section 2.2, line 76), we mentioned RAG and positioned our proposed RATD as a new application of RAG in the study of time series diffusion models. Compared to previous time series RAG methods, our approach still holds significant advantages. This advantage stems from the iterative structure of the diffusion model and our proposed Reference Modulated Attention, where references can repeatedly influence the generation process, allowing references to provide a informative guidance for the conditional generation process.

Q3: How to define stability in time-series data? Someone may not agree with stability in Figure 1-(c)...

A3: To validate our conclusions about 'stability,' we designed additional experiments. In these experiments, we calculated the variance of 15 repeated prediction results, which are presented in the table below. By evaluating the variance, we found that our method demonstrates a clear advantage in stability through the use of additional references for guidance.

Q4: Why necessarily the proposed RATD is needed in time-series forecasting? In other words, how to link the relationship between RATD and TSAD?

A4 : As we mentioned in A1, our method provides a new guiding mechanism for time series diffusion models, and it proves to be highly effective in experiments. Furthermore, our method is plug-in, allowing its application to other transformer-based time series diffusion models. The design principles and motivations behind RATD may also offer new insights for future work, particularly in designing novel guiding mechanisms. All of these aspects underscore the ‘necessity’ of our approach. If TSAD refers to Time Series Anomaly Detection, directly applying the retrieval process may not yield optimal results, as the anchors used for retrieval could be anomalous. However, with appropriate design, it might be possible to cleverly mitigate this issue, thereby making retrieval-based methods potentially viable for TSAD.

Q5: There is a typo in the last sentence of the second paragraph from Section 5.2. “.t”

A5: I apologize for the typo. I will correct the paper again. Thank you for bringing it to my attention.

Q6: Why Table 2 is separated from Table 1 independently?

A6: The main reason we separated Table 1 and Table 2 is due to our use of different strategies for constructing the retrieval databases (as mentioned in line 211). Additionally, most of the popular baselines have not yet been evaluated on the MIMIC-IV dataset, given its later release.

Q7: What can you define as the best retrieval for forecasting?

A7: Generally, a good pre-trained time series encoder can be a good retriever because it can effectively embed the trend of time series. In this paper, we leverage state-of-the-art pre-trained encoder models, and use them for retrieving the most similar anchor time series as reference. Such reference can naturally be the best retrieval for the conditional forecasting, which have been further proved by our empirical results in paper.

We thank Reviewer HjSX for the thorough and valuable feedback. We are glad that the reviewer found that the proposed idea is novel. The reviewer's main concern is the use of pre-trained models. Please see below for our responses to your concerns.

Q1: The additional complexity and computation cost may limit the efficiency of the proposed method.

A1: The demand for pre-trained models does indeed incur additional computational costs. However, the increased training costs do not impose severe constraints on the practical application of the model. The additional training costs stem from two components: pretraining the encoder and the retrieval process. Specifically, pretraining the encoder is straightforward, as we employ the structurally simple yet effective TCN. Typically, TCN can be trained on a dataset within a few hours on a standard single GPU (e.g., Nvidia RTX 3090), thereby avoiding excessively high time costs. Similarly, the entire retrieval process also incurs a few extra hours of computational time. Compared to the time costs required for training diffusion models (ranging from a day to several days), the extra cost is not a serious issue. Furthermore, it is worth noting that our model does not incur excessively high sampling costs (as depicted in Figure 6 of the original text), which is crucial because sampling costs of diffusion models are a more pertinent concern compared to training costs.

Q2: Experimental results do not contain standard deviation, which potentially limits the confidence and significance of the performance superiority.

A2: Thank you very much for your suggestion! Calculating the standard deviation will indeed further assess the stability of the model's generated results. Following your advice, we conducted additional experiments including some popular baselines, and the MSE results are shown in the table below (with the same experimental settings as in this paper, we repeat the test 15 times to calculate the standard deviation ).

Our method also demonstrates superior certainty in the results based on additional experiments.

Q3: Some baselines are missing. For instance, the authors mentioned MQ-ReTCNN and ReTime as retrieval-augmented time series models.

A3: We did not compare our method with the previous RAG-based time series methods in the paper because these papers did not provide comprehensive evaluations on commonly used time series datasets or publicly available codes. The implementation details in these papers were also very limited, making replication of their methods challenging. Nevertheless, as you mentioned, for the sake of comprehensive experimental comparison, we present here the MSE results of our replicated method under our experimental settings, demonstrating that our approach shows noticeable advantages.

RATD exhibits significant performance advantages, indicating that our method can better utilize references for prediction.

Q4: How is the time series encoder Eϕ pre-trained with representation learning tasks (Line148)? Since the embedding quality highly affects the retrieval precision and quality, some analysis and evaluations should be conducted on the time series embedding quality.

A4: We trained TCN for time series prediction tasks for pre-training. In Section 5.3（Influence of Retrieval Mechanism， line 240）of the paper, we extensively discussed the question of which encoder structure to adopt. I believe this discussion holds the same significance as the "discussion on embedding quality" you mentioned because the quality of embeddings is difficult to assess directly and can only be judged by comparing their contributions to experiment improvements. It is worth noting that through comparative experiments, we found that the differences in results brought by different encoders are not significant. This also demonstrates that our approach has strong robustness.

Q5: What is the architecture of the pre-trained encoder Eϕ? How do you encode the multivariate time series into a single embedding? Do you use any pooling operation? More analysis and explanations on the encoder are beneficial for paper clarity.

A5: Your concerns may be solved after reading Section 5.3（Influence of Retrieval Mechanism， line 240）. As we addressed above: we conducted a comprehensive discussion on the structure of the encoder. The TCN architecture supports the encoding of multivariate time series (resulting embeddings are not one-dimensional) without the need for additional pooling layers.