Paper: Generative Modeling of Regular and Irregular Time Series Data via Koopman VAEs (ICLR-2024)

Reviewer comment (Presentation):
It seems that 0 should be excluded from the MuJoCo results in Table 10.

Author response:
Thank you for noticing this. We fixed it.

Reviewer comment (Experiments):
I am curious about the role of the predictive loss term. In this paper, a predictive loss term was added to the object function. Therefore, I am curious about how much the predictive loss term affects the performance that is superior to existing baseline models.

Author response:
Indeed, ablating the effect of the predictive loss term is important. This ablation appears in Table 4 (after correcting the typo $\beta \rightarrow \alpha$). The results indicate that the benefit of the predictive loss term can be significant. Notable examples are stocks regular, energy $30\%$, and stocks $50\%$. More generally, removing this loss term results in inferior error estimates in most cases.

Reviewer comment (Novelty):
The novelty of the study is limited. While the authors claim their method is based on the Koopman operator, their model is, in fact, more similar to the dynamic linear model, or Kalman filter.

Author response:
We agree with Reviewer HuLD that the dynamic linear model and Kalman filter are related to our approach. The main difference between these techniques and Koopman-based methods is the domain where the dynamics are investigated. For the dynamic linear model and Kalman filter it is the state space, whereas in Koopman methods it is the observable (function) space. In practice, given inputs from the state space $x_{1:T}$, the dynamic linear model and Kalman filter can be viewed as studying a dynamical system $A$ respecting the equation $x_t = A x_{t-1}$. In contrast, Koopman approaches transform states to the observables domain, and study the dynamics there. Namely, $\psi(x_t) = A \psi(x_{t-1})$, where $\psi$ is a nonlinear transformation. Our framework, KVAE, uses nonlinear encoder and decoder, and it employs linearity in the latent space. Therefore, we believe it aligns better and can be justified more profoundly from a Koopman perspective.

Reviewer comment (Novelty):
One of the claims is that the model can take care of irregular time series, but the capability is simply using a pre-existing model, neural controlled differential equations.

Author response:
Indeed, we use NCDE to handle irregularly sampled time series information, but we do not state our model's ability to handle irregular time series as a contribution. The ability to handle irregular time series is simply a feature of our model. In retrospect, we believe that utilizing NCDE is a strength of our framework as it provides strong results and it could be potentially replaced by any other module that interpolates time series data.

Reviewer comment (Theory):
I know trying to optimize both the prior and posterior in a loss function has become a trend in the deep learning community. However, theoretically speaking, trying to optimize the prior and posterior jointly in the KL loss leads to an ill-posed problem, where a unique solution does not exist. Simply put, it becomes a ping-pong game between the prior and posterior. You can easily show it by computing the parameters of the distributions in a local minima. How do you deal with this ill-posedness?

Author response:
Thank you for raising this great question. In our experience, we did not observe any issues during training or inference due to the mentioned ill-posedness. Furthermore, our ablation study in the revised version Section 5.3 and Table 4 considers a model whose prior is not learned ($\alpha=0$, no GRU). The results show a significant drop in discriminative scores in comparison to KVAE's performance. Thus, while jointly optimizing the prior and posterior may be ill-posed in theory, the strong performance we obtain outweighs this potential shortcoming. If given the opportunity, we would be happy to discuss this limitation in the final version.

Reviewer comment (Theory):
What is \( z_t \) and what is \( y_t \)? Are they different random variables? Based on the paper, it looks like both \( y_t \) and \( z_t \) denote the latent code, meaning that they are the same variable. I understand that the authors used \( y_t \) and \( z_t \) to distinguish between the prior and posterior latent code. But the way it is formulated now is not correct. For example, how do you define \( KL[q(z)\|p(y)] \) in eq. (3)? Shouldn't it be \( KL[q(z)\|p(z)] \) or \( KL[q(y)\|p(y)] \)?

Author response:
Regrettably, there seems to be a typo in our notation. Thank you for pointing this out. We fixed this issue in the revised version, where we always use $z$ for the latent variable.

Reviewer comment (Theory):
What is the probabilistic model for \( p(x_{1:t}|z_{1:t}) \)? Is it a parametric distribution, e.g., normal distribution? How do you compute the log likelihood function?

Author response:
We would like to thank the reviewer for pointing this out. Following a common choice in probabilistic modeling, we assume that $p$ is a Gaussian distribution whose mean is the output of the decoder. Under this assumption, the log likelihood becomes a simple MSE between $x_{1:T}$ and the output of the decoder, 
$\tilde{x}_{1:T}$. 

A similar reasoning applies also for the $L_\text{pred}$ term. We added a new discussion in the revised version Appendix C, describing technical details related to the evaluation of the loss function.

Reviewer comment (Theory):
If the modulus of the eigenvalues of \( A \) is not strictly 1, i.e., \( |\lambda| = 1 \), the system either grows or decays exponentially fast. It should be a hard constraint, not a soft constraint. How do you guarantee this?

Author response:
We constrain the eigenvalues by adding a penalty term $L_\text{eig}$ to the objective function, i.e., it is a soft constraint. In general, one could construct a hard constraint by parametrizing A in the special orthogonal group [1] or by decomposing A to its Koopman representation and limiting the eigenvalues to lie on the unit circle [2]. Furthermore, we would like to note that exponential growth and decay modes are ineffective in our framework since we employ the matrix A to the power of 1. That is, we multiply $\bar{z}_{1:T-1}$ by A only once 

to obtain $z_{2:T}$, and thus, long-term growth and decay effects do not apply. Please see the definition of $z_t$ given by the end of the Prior modeling paragraph in the revised version. 

    [1] "Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group" by Mario Lezcano-Casado, and David Martınez-Rubio. 

    [2] "Physics-informed probabilistic learning of linear embeddings of nonlinear dynamics with guaranteed stability" by Shaowu Pan, and Karthik Duraisamy.

Reviewer comment (Theory):
Based on Eq. (6), \( y_t \) becomes deterministic once \( y_{t-1} \) is observed. Then the probability distribution becomes a delta function, \( p(y_t|y_{1:t}) = p(y_t|y_{t-1}) = \delta (y_t - Ay_{t-1}) \). How do you compute the KL divergence of the Dirac delta distribution?

Author response:
Thank you for noticing this point. We write the explanation with $z_{1:T}$ due to fixing our typo. 

The distribution $p(z_{1:T})$ is defined as $p(z_{1:T}) = p(\bar{z}_{1:T}) \prod_t \delta(z_t - A \bar{z}_{t-1})$. 

In practice, the KL divergence for $p(z_{1:T})$ is always evaluated with $z_{1:T}$ associated with $\bar{z}_{1:T}$. 

Therefore, we approximate $\delta(z_t - A \bar{z}_{t-1})$ with $1$, 

leading to $p(z_{1:T}) \approx p(\bar{z}_{1:T})$. 

The KL divergence can be computed for $p(\bar{z}_{1:T})$ using a closed-form expression given its mean and variance (as well as the mean and variance of $q$). 

We added this discussion to Appendix C in the revised version.

Reviewer comment (Theory):
How do you find the correct initial condition \( y_0 \) to represent \( x_{1:T} \)? As discussed by the authors, \( x_{1:T} \) is transformed to \( y_{1:T} \). Then, since the model is linear, once \( y_0 \) is determined, the rest of the sequence is determined as \( y_t = A^t y_0 \). Hence, it is crucial to find \( y_0 \) that describes \( x_{1:T} \) the best. How do you choose \( y_0 \) and how do you guarantee that the choice of \( y_0 \) is the optimal?

Author response:
We would like to clarify a potential misunderstanding. 

While Eq. (6) specifies a fully linear dynamical system, our approach implements it in an approximative fashion. 

Specifically, we define z_t = A \bar{z}_{t-1}

(after switching $y_t \rightarrow z_t$), and thus, $z_t$ depends on 

$\bar{z}_{t-1}$ and not on $z_1$. In practice, we take $z_1 = \bar{z}_1$.

Reviewer comment (Writing):
Why do you use two different letters, \( y \) and \( z \), for the prior part and the posterior part, respectively? Usually in VAE papers, the latent variable is always \( z \), and we just say \( p(\cdot) \) for prior and \( q(\cdot) \) for posterior. The current notation in the paper might also be okay, but the reviewer wondered if there could be a particular intention to use the two letters.

Author response:
Regrettably, there seems to be a typo in our notation. Thank you for pointing this out. We fixed this issue in the revised version, where we always use $z$ for the latent variable.

Reviewer comment (Experiments):
Although the paper focuses on the generation capability of the models, some more experiments on the reconstruction / inference capability could also be interesting.

Author response:
Thank you for suggesting that. Following your suggestion, we added two new figures to the revised version showing several reconstruction and inference examples. Specifically, we plot in Fig. 15, the reconstructed signals in the regular setting for all datasets (dashed lines) vs. the ground-truth data (solid lines). In addition, we show in Fig. 16, the inferred signals (dashed lines) in the irregular 50\% case. This experiment demonstrates the inference capabilities of our model since half of the samples were removed during training. We plot the complete signals (solid lines) for comparison.  In addition, we want to emphasize that our climate experiments could also be considered as inference test cases. To evaluate our model abilities in the climate setting, we omit some of the signals, which we refer to as "points of interest", and generate them to compare with the ground truth.

Reviewer comment (Experiments):
As stated above, the presence of GRU in the decoder makes it a little difficult to assess the real utility of the linear structure in the prior model. Do you have some observations when you did not use a nonlinear sequence-to-sequence model in the decoder?

Author response:
In general, the predictive loss term improves the results in comparison to models without the predictive loss term (i.e., $\alpha=0$). This is shown in our ablation study in Table 4. In addition, our preliminary results with an MLP decoder show a similar trend. For further details regarding using an MLP decoder, please see the response above.

Reviewer comment (Experiments):
In practice, how linear is the sequence of \( y \)? In the reviewer’s understanding, the linearity of the dynamics of \( y_{1:T} \) is not a hard constraint but rather is imposed in a soft manner as regularization. The reviewer is curious to what extent \( y \) could become linear with such a soft constraint.

Author response:
Indeed, we preferred to promote linearity using a soft constraint instead of a hard constraint. We believe that it provides the learning algorithm with the flexibility to enforce linearity if possible. In addition, if the data is too complex or the encoder is not expressive enough, the predictive loss term will be high, serving as a tool for identifying potential modeling problems. In practice, the datasets we consider in this work admit a high-level of learnable linearity as we show in Table 1 below. Specifically, we report the relative predictive loss error, i.e., 
$$|z_{2:T} - A\bar{z}_{1:T-1}|^2 / |\bar{z}_{1:T-1}|^2$$,
where $z_{2:T}, \bar{z}_{1:T-1}$ are the latent variables and $A$ is the approximated Koopman operator. Importantly, all errors are below $6e-2$, where some errors are one- and two-orders of magnitude below that.

| Data   |   | Regular   | 30\%      | 50\%      | 70\%      |
|--------|---|-----------|-----------|-----------|-----------|
| Sines  |   | $4.14e-3$ | $2.32e-3$ | $4.29-3$  | $1.32e-2$ |
| Stocks |   | $1.11e-3$ | $4.84e-2$ | $5.86e-2$ | $3.35e-4$ |
| Energy |   | $5.79e-4$ | $3.88e-4$ | $2.23e-3$ | $6.73e-3$ |
| MuJoCo |   | $1.81e-2$ | $9.22e-3$ | $1.35e-4$ | $8.20e-3$ |

Reviewer comment (Experiments):
A large literature of sequential VAEs for time series data generation is omitted e.g. [1,2,3,4], despite a large number of baselines being used in the experiments section. Considering there is heavy development in this area, it would be useful to compare KVAE to these methods.

Author response:
Thank you for bringing these works to our attention. We agree with the reviewer that sequential VAEs are actively developed in recent years, on a broad range of applications. In our work, we focused on comparing with baseline generative approaches that target the same benchmarks and evaluation environment, for a fair and consistent evaluation of approaches. That being said, we are also happy to consider papers [1-4]. Specifically, paper [1] can be viewed as a close variant of our approach. In particular, omitting the predictive loss term in KVAE forms method [1]. The results of this method appear in our ablation study in Table 4 in the submitted paper. Please notice that there was a typo there, and all appearances of $\beta$ should be replaced with $\alpha$ (we have corrected this in the revised manuscript). Further, we are currently running method [2], and we will report the results as soon as we have them. As for method [3], it can be viewed as a VAE only in the time-continuous limit. Nevertheless, we are currently working on migrating the available code to fit into our training and evaluation framework. Finally, we could not find an open-source implementation of method [4], unfortunately. We started to implement this approach from scratch, however, the time constraints of the rebuttal may be too limiting to obtain preliminary results.

Reviewer comment (Experiments):
More discussion in the experiments section is required on the topic of analysing "the behaviour of the system...using dynamical systems theory tools" in order to claim this as an additional feature of KVAE.

Author response:
Following the reviewer's suggestion, we extended the discussion in Section 5.2 on the physics-constrained generation example. The revised text details the particular loss term $\mathcal{L}_\text{eig}$ that we use. Moreover, we analyze the nonlinear pendulum system from a stability viewpoint, where Koopman eigenvalues represent growth and decay modes. The discussion is complemented by Figure 3, which shows the spectra of Koopman operators for unconstrained KVAE vs. constrained KVAE. Our results indicate that constraining the eigenvalue modulus improves generation results and positively affects the operator's spectrum. We also compute the spectra associated with the computed operators in the regular setting and further discuss the results in Appendix F.