$K^2$VAE: A Koopman-Kalman Enhanced Variational AutoEncoder for Probabilistic Time Series Forecasting

Reply to W1. Why were CRPS and (NMAE) chosen as the main evaluation metrics? Why not use other evaluation criteria?

Thank you for your valuable comments. In probabilistic forecasting, evaluation metrics including CRPS (Continuous Ranked Probability Score), CPRS_sum, NMAE (Normalized Mean Absolute Error), NMAE_sum, NRMSE (Normalized Root Mean Square Error), and NRMSE_sum are widely adopted as standard performance metrics [1] [2] [3]. From a comprehensive evaluation perspective, CRPS and its aggregated variant CRPS_sum quantify the distributional approximation quality, while NMAE/NMAE_sum or NRMSE/NRMSE_sum assess point estimation accuracy. Typically, one metric from each category is selected for evaluation. The ProbTS benchmark [3], recognized as a comprehensive framework for probabilistic forecasting models, employs CRPS and NMAE as primary evaluation criteria. This benchmark configuration, having undergone extensive empirical validation across numerous algorithms, was adopted in our study to ensure methodological consistency and equitable comparison with baseline models.

[1] Kollovieh, Marcel, et al. "Predict, refine, synthesize: Self-guiding diffusion models for probabilistic time series forecasting." Advances in Neural Information Processing Systems 36 (2023): 28341-28364.

[2] Rasul, K., Seward, C., Schuster, I., and Vollgraf, R. Autoregressive denoising diffusion models for multivariateprobabilistic time series forecasting. In ICML, volume139, pp. 8857–8868, 2021a.

[3] Zhang, Jiawen, et al. "ProbTS: Benchmarking point and distributional forecasting across diverse prediction horizons." Advances in Neural Information Processing Systems 37 (2024): 48045-48082.

Reply to W2. The conclusion of the paper exceeds two lines and needs to be adjusted to meet the final format requirements of ICML.

Thanks! We will modify this issue.

Reply to W3. For input token embeddings, the conventional patch partitioning method was not used, and no detailed explanation was provided, which is confusing.

1. The conventional patching strategy, as introduced in PatchTST [4], aims to uniformly partition multivariate time series across channels to facilitate subsequent channel-wise independent modeling. The dimensional transformation process operates as: $X \in \mathbb{R}^{B \times N \times T} \to X^\prime \in \mathbb{R}^{B \times N \times n \times p} \to X^{patch} \in \mathbb{R}^{(B \times N) \times n \times p} \to X^{embedding} \in \mathbb{R}^{(B \times N) \times n \times d}$, where $n$ denotes the number of patches, $p$ denotes the patch size. This architecture collapses the batch dimension $B$ and variate dimension $N$ into a single axis, effectively enforcing univariate modeling at the patch level.
2. In contrast, our adopted patching mechanism maps multivariate subsequences to high-dimensional representations compatible with Koopman operator theoretic frameworks, enabling systematic state transition modeling. The modified transformation pipeline is formalized as: $X \in \mathbb{R}^{B \times N \times T} \to X^\prime \in \mathbb{R}^{B \times N \times n \times p} \to X^{patch} \in \mathbb{R}^{B \times n \times (N \times p)} \to X^{embedding}\in \mathbb{R}^{B \times n \times d}$, where $n$ denotes the number of patches, $p$ denotes the patch size. Our architecture collapses the patch dimension $p$ and variate dimension $N$ into a single axis.

[4] Nie, Yuqi, et al. "A Time Series is Worth 64 Words: Long-term Forecasting with Transformers." The Eleventh International Conference on Learning Representations.

Replay to W4. Although the authors analyzed the overall performance, key components, and efficiency of K2VAE , there is a lack of analysis on parameter sensitivity. It is hoped that the authors can add an analysis of key parameters.

Thanks for your valuable comments. Following your suggestions, we systematically evaluated three critical hyerparameters: Patch Size $p$, and the dimensionality $d$ of hidden layers in the Measurement Function, and the number of hidden layers $l$ in the Measurement Function. Specific experimental results are presented in the following Tables :

https://anonymous.4open.science/r/K2VAE-D957/sensitivity.md

These empirical findings substantiate that $K^2$VAE maintains robust performance across different hyperparameter configurations. To achieve best performance, we recommand a group of stable hyperparameters for both short-term probabilistic forecasting and long-term probabilistic forecasting:

1. Patch Size:
   • Short-horizon forecasting tasks achieve peak performance with patch_size $p$: 8
   • Long-horizon forecasting benefits from extended context capture with patch_size $p$: 24
2. Network Architecture:
   • Number of hidden layers $l$: 2-3 layers yield optimal performance-efficiency balance
   • Dimensionality of hidden layers $d$ : 256 provides sufficient representational capacity

Thanks again for the Reviewer auu7's valuable opinion!

Reply to Q1: Why can the proposed method solve the problem of error accumulation well?

Thank you for your valuable comments. In $K^2$VAE, the KoopmanNet models time series in linear dynamical systems, where uncertainties are represented as deviations from the linear system’s predictions. Then KalmanNet works in a two-phase operation:

1. Prediction Phase: Outputs state predictions and covariance matrices of uncertainty distributions for the linear system. $\hat{z} _{k} = A z _{k-1} + B u _k$ $\hat{\mathrm{P}} _k = A\mathrm{P} _{k-1}A^T + Q$

2. Update Phase: Integrates observations to compute the Kalman gain $K_k$, refining predictions and covariance via: $K _k = \hat{\mathrm{P}} _k H^T(H\hat{\mathrm{P}} _k H^T + R)^{-1}$ $z _k = \hat{z} _{k} + K _k(\hat{x} _k^H - H \hat{z} _{k})$ $\mathrm{P} _k = (I - K _k H) \hat{\mathrm{P}} _k$

The prediction and the update phase fuse the information from Integrator ($u_k$) and KoopmanNet ($\hat{x}^H _k$) to effectively eliminate error accumulation in the prediction $\hat{z} _k \to z _k$ and the covariance matrix $\hat{\mathrm{P}} _k \to \mathrm{P} _k$. To demonstrate KalmanNet’s effectiveness, we include an ablation study of it in https://anonymous.4open.science/r/K2VAE-D957/ablations.md. The experiments demonstrate the importance of the KalmanNet in multiple PTSF tasks, particularly in tasks of L=336 and 720.

Reply to Q2: According to my understanding, this paper constructs a probabilistic time series model based on Koopas. This should further clarify. Why is the experiment not compared with Koopas?

We first explain the differences between $K^2$VAE and Koopa in the following points:

1. Koopa applies Koopman Theory to address the non-stationarity issue in point forecasting tasks. $K^2$VAE utilizes the KoopmanNet to make It easier to model the uncertainties. Koopa is a model completely based on the Koopman Theory and our proposed $K^2$VAE benefits from the Koopman Theory, Kalman Filter, and VAE.
2. Koopa is a deterministic model while $K^2$VAE is a non-deterministic model. Koopa is a MLP based model while $K^2$VAE adopts the generative architecture VAE, and excels at modeling the distributions over the target horizon.
3. Koopa assumes that the time series can be accurately modeled based on the pure Koopman Structures, so that it utilizes multi-scale MLP to simulate the Koopman process, trying to model an unbiased linear dynamical system in the measurement space. While $K^2$VAE is a model tailored for PTSF tasks, which starts from analyzing the ``bias'' in the linear dynamical system modeled by KoopmanNet, using a KalmanNet to model these uncertainties.
4. A key design in Koopa's is that it utilizes multiple MLP layers to continuously models the residual parts. In our asssumptions, such residual parts may be uncertainties which are hard to model. Therefore, $K^2$VAE not only designs a KoopmanNet to linearize the time series, but also models the uncertainties through the subsequent KalmanNet.

We include Koopa in our baselines. Specifically, we equip it with a Gaussian distribution, which is consistent with other point-forecasting baselines such as FITS, iTransformer and PatchTST. The experimental results are shown in https://anonymous.4open.science/r/K2VAE-D957/comare_with_koopa.md . In both long-term and short-term prediction tasks, Koopa lags behind $K^2$VAE.

Reply to Q3: Is the proposed method for more accurate prediction of long time series or for effective uncertainty modeling over longer periods of time? if the former, you can add any probabilistic model with koopa.

$K^2$VAE focuses on probabilistic forecasting, which is in line with ``effective uncertainty modeling over longer periods of time.'' Combining with the answers to Q1 and Q2, $K^2$VAE employs KalmanNet to achieve better uncertainty modeling. Based on the generative model VAE, $K^2$VAE accurately describes the variational distribution in the measurement space constructed by KoopmanNet, thus enhancing the ability to construct the target distribution of the forecasting horizon. On the other hand, Koopa is constructed for the deterministic long-term point forecasting tasks. Its characteristic lies in constructing a multi-scale MLP structure and modeling the time series from the perspective of decomposition. And $K^2$VAE does not directly apply the structure of Koopa. Instead, it has a greater inclination towards the application of the original Koopman Theory. For the part of uncertainties, $K^2$VAE models it with the KalmanNet. We will discuss the differences from Koopa in the Related Works and will include Koopa as an additional baseline (as elaborated in Q2) in the paper.

Reply to Other Problems

Figure 3 is redrawn in https://anonymous.4open.science/r/K2VAE-D957/README.md .Thanks again for the valuable comments! May I ask if our responses have resolved your questions? Or we can have further discussions!

Reply to W1. Koopman theory requires that the measurement function maps inputs to several observable variables to construct a linear dynamical process in that space. Traditional methods typically choose fixed basis functions to meet this requirement, but this paper adopts a learnable measurement function. How can we ensure that the constructed space is within a linear dynamical system?

Thank you for your valuable comments. Koopman Theory requires that the measurement function possess sufficient projection capabilities to map nonlinear inputs into a high-dimensional measurement space and model them using a linear dynamical system. Traditional methods would use some fixed functions, such as polynomial bases, because these functions have strong high-order fitting capabilities. Similarly, in deep learning, the universal approximation theorem guarantees that multi-layer perceptron models (MLP) have strong fitting capabilities, theoretically able to approximate any continuous function to arbitrary precision. $K^2$VAE, as a deep learning model, also adopts multi-layer perceptron fitting for measurement functions. KoopmanNet can construct a linear dynamic system because it designs the Measurement Function $\psi$ and Decoder $\psi^{-1}$ in a fully symmetric manner and controls the outputs of the linear system constructed by the Koopman Operator $\mathcal{K}$ close to the contextual time series in the original space by using reconstruction loss $\mathcal{L} _{rec}$. Additionally, the Koopman Operator is also linear, which aligns with the assumptions of Koopman Theory.

Reply to W2. In the ablation experiments for the components of KoopmanNet in Table 3, why does using the local Koopman operator lead to numerical instability?

This is because we use one-step eDMD to accelerate the efficient computation of $\mathcal{K}_{loc}$, as shown in the following steps.

$X^{P^\ast} _{back} = \left[x^{P^\ast} _1,x^{P^\ast} _2,\cdots,x^{P^\ast} _{n-1}\right]$ $X^{P^\ast} _{fore} = \left[x^{P^\ast} _2,x^{P^\ast} _3,\cdots,x^{P^\ast} _{n}\right]$ $\mathcal{K} _{loc} = X^{P^\ast} _{fore} (X^{P^\ast} _{back})^{\dagger}$

By using a one step offset and then calculating the pseudo-inverse, we estimate $\mathcal{K} _{loc}$. This process is prone to matrix computation errors when the difference between the contextual length and the horizon length is significant (e.g., 96 and 720). And this phenomenon often occurs when the model is just beginning to be optimized, and the space constructed by the measurement function is not good enough. Theoretically, though the Moore-Penrose Pseudo Inverse always exists, the corresponding computational method relies on Singular Value Decomposition (SVD), which can be numerically unstable and lead to computation failures. Therefore, we need to use a learnable $\mathcal{K} _{glo}$ to ensure the computational stability at the beginning of training.

Reply to W3. There is a lack of sensitivity analysis regarding the parameters, which obscures understanding the behavior of the proposed method in different hyperparameter settings.

Thanks for your valuable comments. Following your suggestions, we systematically evaluated three critical hyerparameters: Patch Size $p$, and the dimensionality $d$ of hidden layers in the Measurement Function, and the number of hidden layers $l$ in the Measurement Function. Specific experimental results are presented in the following Tables :

https://anonymous.4open.science/r/K2VAE-D957/sensitivity.md

These empirical findings substantiate that $K^2$VAE maintains robust performance across different hyperparameter configurations. To achieve best performance, we recommand a group of stable hyperparameters for both short-term probabilistic forecasting and long-term probabilistic forecasting:

1. Patch Size:
   • Short-horizon forecasting tasks achieve peak performance with patch_size $p$: 8
   • Long-horizon forecasting benefits from extended context capture with patch_size $p$: 24
2. Network Architecture:
   • Number of hidden layers $l$: 2-3 layers yield optimal performance-efficiency balance
   • Dimensionality of hidden layers $d$: 256 provides sufficient representational capacity

Reply to Other Problems

We will carefully examine and correct the normativity of the citation, thank you for the correction!

Thanks again for the Reviewer rXKc7's valuable opinion! We deal with W1-W3, and provide analysis and evidence from the perspective of model design and experiment. Do our reply solve your question? If there are any other questions, we can discuss them further!

Reply to W1. The experimental design of this paper is very comprehensive, evaluating multiple datasets on both long and short step tasks. However, in Table 5, some datasets have the same name but different actual lengths. What is the reason for this?

Thank you for your valuable comments. We follow the evaluation protocol in ProbTS, a well known benchmark for probabilistic forecasting tasks. Some datasets with same names but different suffixes, e.g., Electricity-S and Electricity-L, are different datasets with different lengths and channels. The ETT datasets share the same in short- and long-term probabilistic forecasting tasks. We have summarized the details of datasets in our appendix, we list the table here:

https://anonymous.4open.science/r/K2VAE-D957/datasets_info.md

Reply to W2. The related work section discusses the application of VAE in time series data but seems to omit some classic algorithms, such as D3VAE. It is suggested to include the discussion and experiments of these algorithms.

Thanks for the suggestions. Although $D^3$VAE is closer to a Diffusion model, it does indeed follow the paradigm of VAE. Specifically, $D^3$VAE is a dual-directional variational auto-encoder that combines diffusion, denoising, and factorization. By coupling diffusion probability models, it expands time series data, reduces data uncertainty, and simplifies the inference process. Furthermore, it treats latent variables as multivariate and minimizes total correlation to separate them, thereby enhancing the interpretability and stability of predictions. Both $D^3$VAE and $K^2$VAE aim to better model the uncertainty in the target window through decoupling methods, and we will discuss the similarities and differences in the Related Works section. In terms of experiments, we also provide a detailed comparison with $D^3$VAE, assessing $D^3$VAE in both short- and long-term PTSF tasks:

https://anonymous.4open.science/r/K2VAE-D957/compare_with_d3vae.md

We keep the contexutal length equals to the horizon length to meet the setting of $D^3$VAE. The results show that $K^2$VAE outperforms $D^3$VAE in both short and long step scenarios, and $D^3$VAE also seems less suitable for long-term PTSF tasks.

Reply to W3. Although the ablation studies of key components in KoopmanNet and KalmanNet are discussed, the complete removal of these two modules to analyze their impacts has not been considered.

For KoopmanNet and KalmanNet, they act as two important parts of $K^2$VAE. We further analyze the impact of these two key modules on PTSF tasks, and supplement the ablation experiments using the same datasets in the paper:

https://anonymous.4open.science/r/K2VAE-D957/ablations.md

The results demonstrate that KalmanNet is more critical for LPTSF tasks. When KalmanNet is removed, the arrcuracy breaks down sharply, indicating that the importance of KalmanNet to effectively eliminate the cumulative errors in LPTSF tasks. When removing KoopmanNet, on the other hand, would lead to a performance decline across all the tasks. Without KoopmanNet, the nonlinear time series is hard to model, and the uncertainties in such nonlinear system is also difficult to be captured through KalmanNet, which furtherly declines the model's modeling capabilities. The experiment shows that both KoopmanNet and KalmanNet are critical and indispensable in our design.

Reply to W4. The conclusion of the article needs to be adjusted to meet the final format requirements of ICML, as it currently exceeds two lines.

Thanks for your reminder! We will fix it.

Thanks again for the Reviewer iEFP's valuable opinion! We deal with W1-W4, and provide analysis and evidence from the perspective of model design and experiment. Do our reply solve your question? If there are any other questions, we can discuss them further!