HoT-VI: Reparameterizable Variational Inference for Capturing Instance-Level High-Order Correlations

We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer.

Q1: What is the runtime study for the proposed method, and what is the trade-off between the gains from more complex models and their added computational cost?

A1: An empirical runtime study was conducted to measure the cost of increasing the approximation order k. Tables 1-3 below show the average training runtime per epoch on different tasks, demonstrating that the computational cost scales linearly with k.

Table 1: Time series anomaly detection performance on SMAP dataset.

Table 2: Time series forecasting performance on ETTh1 dataset with horizon 24.

Table 3: Constrained clustering performance on MNIST dataset.

The trade-off is between this linear increase in runtime and the corresponding performance gains. While model performance consistently improves as k increases, our experiments show diminishing marginal returns, with higher values of k (such as k beyond 10) yielding progressively smaller gains. Therefore, we focused on k up to 10 in the main paper because this range demonstrates the advantages of our approach while maintaining a practical and efficient balance between performance and computational tractability.

Q2: What are the specific limitations of the method given that its variational posteriors must be Gaussian, particularly concerning its applicability to settings like directional statistics?

A2: The Gaussian posterior is a foundational and highly effective choice, used with great success in many state-of-the-art generative models like VAEs and modern diffusion models. We acknowledge that, while powerful, the Gaussian posterior is not universally applicable and is not directly suited for modeling certain data types, such as discrete data or directional statistics. The core contribution of our work is to introduce a novel and scalable method for structuring correlations, and extending this framework to more flexible distributional families to address these broader settings is a promising direction for our future research.

Q3: How are the constraints specifically incorporated into the model within the constrained clustering experiment?

A3: As detailed in Section 4.2, the constraints are not implemented as hard rules but are integrated as a probabilistic prior over the cluster assignments $\mathbf{c}$. This soft constraint mechanism guides the model to favor clusterings that are consistent with the provided must-link and cannot-link information. Specifically, the key is to define a conditional probability distribution for the cluster assignments $\mathbf{c}$, conditioned on a pairwise prior information matrix $\mathbf{A}$. This distribution $p( \mathbf{c}|\mathbf{A})$, as shown in Line 345, replaces the uninformative prior $p(\mathbf{c})$ used in standard generative clustering models, to assign a higher probability to clusterings that respect the constraints encoded in $\mathbf{A}$ (If there is a must-link between instances $x_i$ and $x_j$, then $[\mathbf{A}]_ {ij}$ is a large positive value; if there is a cannot-link between $x_i$ and $x_j$, then $[\mathbf{A}]_ {ij}$ is a large negative value; if there is no constraint, then $[\mathbf{A}]_ {ij}=0$). This prior is then integrated into the overall generative model $p_{\theta}(\mathbf{X}, \mathbf{Z}, \mathbf{c} | \mathbf{A}) = p_{\theta}(\mathbf{X}|\mathbf{Z})p(\mathbf{Z}|\mathbf{c})p(\mathbf{c}|\mathbf{A})$, guiding the learning process within the variational inference framework towards a posterior that respects the given domain knowledge.

Q4: How does the proposed reparameterization method restrict the expressiveness of the variational family?

A4: Our reparameterization method trades the universal expressiveness of a general multivariate Gaussian for guaranteed validity and computational tractability. It restricts the variational family to a subset of Gaussian distributions whose precision matrices have a k-order connection structure. And the hyperparameter k allows us to explicitly control the trade-off between expressiveness and computational cost.

Q5: What are the key differences between the proposed method and TreeVI, what specific improvements does your approach offer, and what is the comparison in terms of effective runtime?

A5: The key difference between TreeVI and our method is that TreeVI is inherently limited to modeling only first-order correlations, while our method can capture higher-order correlations. This limitation of TreeVI arises from its reliance on an acyclic tree structure to construct its correlation matrix. While this is sufficient for simple pairwise relationships, attempting to model higher-order correlations inherently introduces loops into the underlying correlation structure. The construction that TreeVI uses to build a valid correlation matrix depends on the acyclic property of its backbone and is no longer valid when these loops are present. Moreover, simply modeling higher-order correlation coefficients within the framework of TreeVI does not guarantee the resulting correlation matrix to remain positive definite. So even though TreeVI can capture instance-level correlations, it cannot be easily extended to model higher-order correlations.

Our approach, in contrast, is generalized to explicitly model k-order correlations. For $k=1$, our model is conceptually equivalent to TreeVI, but for $k>1$, it becomes substantially more expressive. As shown in our constrained clustering results (refer to Table 3 above), the runtime of our method is approximately k times that of TreeVI. But even with a small k (e.g., k=3), our method achieves significantly better performance across all metrics compared to TreeVI, demonstrating the benefit of modeling higher-order correlations.

Q6: More detailed information about the actual sizes of the datasets used in the experiments (e.g., the number of observations for SMAP)?

A6: The actual sizes of the datasets used in the constrained clustering task can be found in Appendix D.1. Detailed information about the other datasets is shown as follows and will be included in the appendix of our final version.

Q7: Clarify extending the method to dimension-level correlations.

A7: While we focused on instance-level correlations in the paper for clarity, the method can be extended to model dimension-level correlations by introducing a structured covariance to the sampling procedure's noise term.

In the current framework, each dimension is sampled independently, meaning the noise vector $\epsilon$ in the re-parameterization step is drawn from a standard Gaussian with a diagonal covariance matrix $\mathcal{N}(0, I)$. To incorporate dimension-level correlations, we can replace this with a low-rank Gaussian distribution for the noise. By sampling $\epsilon$ from a distribution with a low-rank covariance, we can efficiently induce correlations across the dimensions of the generated latent variables $\mathbf{z}$. This provides a practical and computationally efficient path to a more expressive posterior that captures both instance-level and dimension-level dependencies within a unified framework.

We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer.

Q1: Are the chosen baselines strong or merely easy to beat? How were they configured and tuned during experimentation?

A1: Our chosen baselines represent strong and state-of-the-art methods in their respective domains, selected to ensure a fair evaluation of our model. For the time series modeling tasks, we compared our model against both leading generative methods and prominent forecasting models. For constrained clustering, we benchmarked against state-of-the-art algorithms.

To ensure a fair comparison, we configured and tuned all methods to achieve their best possible performance. Hyperparameters of baseline methods are set consistent with our method as far as possible.

Q2: What is the specific likelihood or loss function used in the Variational Autoencoder (VAE) component? And how many posterior samples were utilized during both the training and evaluation phases of the model?

A2: The specific likelihood function $p( \mathbf{X} |\mathbf{Z})$ used in our VAE is task-dependent. For the time series tasks, the likelihood is a Gaussian distribution with parameters produced by a GRU-based decoder. For constrained clustering, the likelihood is a normal distribution conditioned on a cluster-dependent latent variable. Across all applications, the loss function is the negative of the ELBO. For all experiments, we utilized a single sample from the posterior distribution during both training and evaluation.

Q3: Why are there no error bars or uncertainties reported for the evaluation metrics, and how robust are the presented results without them?

A3: We would like to clarify that we did report error bars in Table 3 across 10 independent runs. For Tables 1 and 2, we omitted the error bars in the main paper due to space constraints, and we will include them in the appendix in our revised manuscript. To address your concern directly, we present the F1-Scores and ELBOs for the time series anomaly detection task with their corresponding standard deviations from 10 runs below:

Q4: Were experiments conducted with approximation orders exceeding 10? If so, did performance continue to improve at higher orders?

A4: Yes, we have conducted additional experiments with orders k exceeding 10, specifically for k=50 and k=100, to explore the performance limits of our model, as shown in Tables 1-3 below.

Table 1: Time series anomaly detection performance on SMAP dataset.

Table 2: Time series forecasting performance on ETTh1 dataset with horizon 24.

Table 3: Constrained clustering performance on MNIST dataset.

The results show that performance does continue to improve as the order k increases beyond 10, which further validates our central claim that capturing higher-order dependencies is beneficial.

However, we observed that the marginal performance gains diminish as k becomes excessively large, and concurrently the computational runtime increases linearly with the order k. We focused on k up to 10 in the main paper because this range demonstrates the advantages of our approach while maintaining a practical and efficient balance between performance and computational tractability.

Q5: What is the effect of systematically increasing the approximation order from a mean-field approach all the way to a fully flexible Gaussian (potentially on a smaller dataset)? How does the proposed method compare to other variational families, such as low-rank Gaussians and normalizing flows?

A5: Increasing the approximation order k consistently improves model performance by capturing higher-order correlations, as seen from Tables 1-3. To capture all correlations, k would need to equal $N-1$, making our model equivalent to a fully flexible Gaussian. However, the computational cost scales linearly with k, with the runtime being approximately k times that of a mean-field method. This makes a fully flexible Gaussian computationally intractable for instance-level correlation modeling, even on smaller datasets.

Since other variational families like low-rank Gaussians and normalizing flows are not designed for instance-level modeling, we apply our method to model dimension-level correlations within the latent space of a VAE to provide a fair comparison. We configured a VAE with 30 latent dimensions on MNIST and report the ELBO on the test set below.

The results show that our method bridges the gap between mean-field approximations and fully expressive models. Even with a small order (k=3), our method outperforms the mean-field approach and becomes competitive with the low-rank Gaussian. As k increases to 10, our method surpasses both the low-rank Gaussian and the normalizing flow in this setup, and approaches the performance of the full-rank Gaussian, demonstrating that it offers a competitive and highly flexible alternative to other variational families.

Q6: What is the trade-off in computational complexity and runtime when the approximation order is increased from mean-field to a full Gaussian? At what point does it become more computationally efficient to simply use a full Gaussian rather than the structured approximation?

A6: The choice of the approximation order k is a direct trade-off between model expressiveness and computational cost. As k increases, our model captures higher-order correlations, leading to consistently better performance, as shown in Tables 1-3 above. However, the computational cost also scales linearly with k, with the runtime being approximately k times that of a mean-field model. Because of this trade-off between diminishing performance gains and linearly increasing cost, we selected a moderate k (up to 10) in our main experiments to strike a balance between performance and efficiency.

When k is set to $N-1$, where $N$ is the number of instances, our method becomes equivalent to a full Gaussian. However, a full Gaussian would require evaluating $\mathcal{O}(N^2)$ times of neural networks $f_\phi(\cdot)$ per epoch, which is computationally intractable for instance-level correlation modeling, especially considering that the number of instances can be very large in practice.

We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer.

Q1: How difficult or involved is it to construct more detailed correlation structures within the proposed framework?

A1: Constructing more detailed correlation structures within our framework is straightforward by design. We can increase detail in two main ways: (1) by increasing the approximation order k to capture higher-order correlations, (2) by generalizing the dependency backbone from a simple chain to a tree structure to model more flexible correlations among data instances.

Q2: What are the specific computational overheads of the proposed method? How do these overheads scale as the approximation order k increases?

A2: The specific computational overheads involve three parts: (i) the overhead of neural network evaluations for re-parameterizing correlation coefficients, (ii) the overhead of sampling from the variational posterior, and (iii) entropy calculation.

In our method, to define a k-order posterior over $N$ latent variables, we need to estimate exactly $(N – 1)$ first-order, $(N – 2)$ second-order, and so on, up to $(N – k)$ k-order correlations, yielding a total of $(N-1) + (N-2) + \cdots + (N-k) = k(2N-k-1)/2$ correlation coefficients. Each coefficient is produced by a single forward pass through our re-parameterization network $f_{\phi}$, so this cost amount to the cost of evaluating $\mathcal{O}(kN)$ times of the neural network $f_{\phi}$.

The sampling and entropy calculations involve operations like inversion and determinant computation on $k \times k$ sub-matrices, incurring a cost of $\mathcal{O}(k^3)$ FLOPs. Considering that $k$ is typically much smaller than $N$ and the complexity of evaluating neural networks, the cost of these operations is negligible compared to the cost of neural network evaluations.

Therefore, the total cost of the variational inference approximately amounts to the cost of evaluating $\mathcal{O}(kN)$ times of the neural network $f_{\phi}$ per epoch, which scales linearly with the approximation order k.

Q3: What are the practical limits for k? Is k=10 the reasonable maximum, or are much larger values like k=50, 100, 250 also feasible?

A3: Yes, larger values of k beyond 10 are also feasible for our method, and we have conducted experiments with approximation orders up to k=50 and k=100, as shown in Tables 1-3 below.

Table 1: Time series anomaly detection performance on SMAP dataset.

Table 2: Time series forecasting performance on ETTh1 dataset with horizon 24.

Table 3: Constrained clustering performance on MNIST dataset.

The results show that performance continues to improve beyond k=10, but with diminishing marginal returns. Since the runtime scales linearly with k, the significant increase in computational cost for these very high orders may not be justified by the smaller performance improvements. Therefore, we set k to moderate values (up to 10) in our main experiments, as it provides substantial performance gains over lower-order methods while maintaining computational efficiency. Ultimately, the optimal choice for k depends on the specific application and the available computational budget.

Thank you for your constructive feedback. We confirm that we will incorporate the clarifications from our previous response into the revised version of our paper to improve clarity and completeness, considering that one additional page is allowed.

First, we will expand our explanation of how the correlation structure is built upon a tree backbone. The initial step is to establish this backbone, for which there are several methods. A straightforward approach is to compute a similarity matrix between all data instances and then use a classic algorithm like a DFS to sample a tree. There are also some other approaches, such as the continuous optimization framework proposed by TreeVI, which directly learns the adjacency matrix $\mathbf{A}$ subject to an acyclicity constraint that guarantees the learned structure to represent a tree. In our paper, we adopt the simpler and more direct approach of using a similarity matrix. Once this tree backbone is established, our k-order structure is imposed upon it by defining correlations over the local (k+1)-vertex cliques on the tree.

Additionally, as you suggested, we will add the detailed discussion on computational overheads to our main text and include the full experimental results for larger k values (up to k=100) in the supplementary material.

We thank the reviewer for the comments and suggestions. Below we respond to each point raised by the reviewer.

Q1: What is the exact cost of inference with $P$?

A1: The exact cost of inference with our proposed k-order precision matrix $P$ involves three parts: (i) the cost of neural network evaluations for re-parameterizing correlation coefficients, (ii) the cost of sampling from the variational posterior, and (iii) entropy calculation.

In our method, to define a posterior with k-order correlation over $N$ latent variables, we need to specify exactly $(N–1)$ first-order, $(N–2)$ second-order, and so on, up to $(N–k)$ k-order correlations, yielding a total of $(N-1)+(N-2)+\cdots+(N-k)=k(2N-k-1)/2$ correlation coefficients. In our method, each coefficient is parameterized by the output of a re-parameterization network $f_{\phi}(\cdot)$. Therefore, to re-parameterize these coefficients, we need to run the network $f_{\phi}(\cdot)$ for $\mathcal{O}(kN)$ times.

The sampling and entropy calculations involve operations like inversion and determinant computation on $k \times k$ sub-matrices, incurring a cost of $\mathcal{O}(k^3)$ FLOPs. Considering that k is typically much smaller than $N$ and the complexity of evaluating neural networks, the cost of these operations is negligible compared to the cost of neural network evaluations.

Therefore, the total cost approximately amounts to the cost of evaluating $\mathcal{O}(kN)$ times of the neural network $f_{\phi}$ per epoch, which is approximately k times of the cost of mean-field amortized VI methods.

Q2: What is the intuition behind k-level correlation modeling? How should $k$ be chosen? What are you losing here? Does k=1 imply a Markov assumption in the posterior?

A2: The choice of k is a trade-off between model expressiveness and computational cost. Generally, as k increases, the model's performance consistently improves, as demonstrated in our experimental results. However, as seen from Table 1 below, the performance gains diminish as the order k (e.g. 50, 100) goes higher. Meanwhile, the computational cost always scales linearly with k according to our analysis in answer A1. To balance the gains and cost, we set k to moderate values (up to 10) in our main experiments.

By setting k to a moderate value (e.g. 10), we can only model correlation up to 10-th order, losing the ability to model higher-order correlations. But as observed from Table 1, the gains become increasely weak as the order goes higher. Therefore, setting k to a moderate value is a reasonable choice.

Yes, when k=1, our variational posterior is reduced to the first-order Markov random field distribution.

Table 1: Time series anomaly detection on SMAP dataset.

Q3: Which baselines actually model instance-level correlation, how do they differ from the proposed method?

A3: Due to the difficulties of modeling instance-level correlation, existing works mostly only consider the modeling of correlation within one instance. In our baselines, only OmniAnomaly and TreeVI explicitly model instance-level correlations in their posterior distributions.

OmniAnomaly imposes a first-order Markov assumption among latent variables and employs an MLP to parameterize the one-step conditional Gaussian, where $(\mu_{z_t},\sigma_{z_t})=\text{MLP}([z_{t-1},e_t])$. This design makes it essentially a recurrent model, just like an RNN learning a transition function for its hidden state. As a result, it is inherently limited to modeling only first-order correlation. More importantly, OmniAnomaly can only be applied to model chain-structured correlation due to its inherent RNN-like structure, making it unable to model correlation with more complicated structure like trees. In contrast, our method not only can model higher-order correlation, but also can be applied to model correlations with tree backbones.

The key difference between TreeVI and our method is that TreeVI is inherently limited to modeling only first-order correlations, while our method can capture higher-order correlations. This limitation of TreeVI arises from its reliance on an acyclic tree structure to construct its correlation matrix. While this is sufficient for simple pairwise relationships, attempting to model higher-order correlations inherently introduces loops into the underlying correlation structure. The construction of TreeVI depends on the acyclic property of its backbone and is no longer valid when these loops exist. Moreover, simply modeling higher-order correlation coefficients within the framework of TreeVI does not guarantee the correlation matrix to be positive definite. So even though TreeVI can capture instance-level correlations, it cannot be easily extended to model higher-order correlations.

Q4: At what level of k should the proposed method recover the same correlation structure and performance as baselines? In Table 1, I would have expected the 1 order method to perform as well as LSTM-VAE.

A4: At k=1, our method is reduced to TreeVI, with both methods modeling first-order instance-level correlation in the same way. But for other baselines, although some of them also attempt to model first-order instance-level correlation, they are fundamentally different from us in terms of correlation modelling methods and the ability of generalizing to more complicated correlation structures.

As for the performance gap between LSTM-VAE and our method, the reason is explained below. Although LSTM-VAE also attempts to capture instance-level correlation, it is achieved via the deterministic dependency among the hidden states in RNN. Given the hidden states' values, the latent variables $z_i$ of instances at different time steps are independent. By contrast, our method not only captures the deterministic dependency, but also models the correlation among the stochastic variations of latent variables $z_i$ corresponding to different instances. That is, given the hidden states' values, in our method, the latent variables $z_i$ are still correlated. Thus, our method is able to capture more accurate and complicated correlation among data. That explains why our method performs better than LSTM-VAE even under the case of k=1.

Q5: Why is it necessary to directly parameterize the correlation coefficients $R$? Why not parameterize the Cholesky decomposition of the covariance matrix to ensure positive definiteness by construction?

A5: We first want to emphasize that this work aims to capture instance-level correlation, under which the dimension of correlation matrix is $N\times N$, with $N$ denoting the number of instances. Parameterizing the Cholesky factor $\mathbf{L}$ of a correlation matrix $\mathbf{R} = \mathbf{L} \mathbf{L}^\top$ does ensure its positive definiteness. However, we have to re-parameterize as many as $N(N+1)/2$ elements, i.e., all elements from the lower-triangular positions of $\mathbf{L}$ need to be re-parameterized. That means we need to evaluate the neural network $f_\phi(\cdot,\cdot)$ by $\mathcal{O}(N^2)$ times for every epoch, which is computationally unacceptable, especially considering that $N$ could be very large in practice.

To reduce the cost, a common practice is to restrict $\mathbf{L}$ to be low-rank, which reduces the number of parameters to be re-parameterized from $\mathcal{O}(N^2)$ to $\mathcal{O}(kN)$ for a low rank $k$. However, each element in $\mathbf{L}$ does not explicitly correspond to the correlation between any two specific data instances. Therefore, to parameterize each element in $\mathbf{L}$, we have to feed all $N$ data instances into the neural network $f_\phi$. Obviously, this will result in the network to be huge, and thereby incurring unacceptable computation and storage burden.

In contrast, although our method also need to re-parameterize $\mathcal{O}(kN)$ parameters, in our method, the parameter is the correlation coefficient. For each correlation coefficient, such as $\gamma_{i,i+t}^c$, it directly represents the correlation between a specific pair of instances $(x_i,x_{i+t})$. Therefore, we only need to feed the pair of instances into the neural network $f_\phi$, making the network to be small.

Q6: What kind of neural network is $f_{\phi}$?

A6: The function $f_{\phi}$ is implemented as a standard MLP. Inputs $x_i$ and $x_{i+t}$ are concatenated across the feature dimension to form a single input, which is then processed by the MLP with a tanh activation function to constrain the output to the range of (-1,1).

Q7: Any relation to inducing points used in sparse GPs, given that both aim to approximate large matrices?

A7: While both our method and the inducing points method aim for scalability, the inducing points framework cannot handle instance-level correlation modeling effectively. The inducing points method works by a low-rank approximation based on a small set of $M$ representative points. Since these points must summarize the properties of the entire dataset, the network used to learn their optimal locations needs to see all $N$ data instances to derive this global summary. This results in a network with an $\mathcal{O}(N)$ input size, which is computationally intractable for large datasets. Therefore, low-rank methods like inducing points are not usable in our setting. In contrast, in our method, the re-parameterization neural network used to compute each correlation coefficient only requires a pair of instances as input.

Q8: Computational cost of the sampling procedure laid out in line 121?

A8: The cost is composed of two parts: $\mathcal{O}(kN)$ times of neural network evaluations for re-parameterizing correlation coefficients, and $\mathcal{O}(k^3)$ FLOPs for matrix inversion. For more details, please refer to A1.