Towards Self-Supervised Covariance Estimation in Deep Heteroscedastic Regression

Dear Reviewer ZXre,

We thank you for your review and your suggestions which we have incorporated in the draft. We hope to address your concerns below:

W1 - Problem formulation not clear.: We accept your suggestion and devote a paragraph explaining the problem formulation and notation in the paragraph Preliminaries, Section 2. We also switch our notation to maintain consistency where possible with Seitzer et al. Please let us know if any ambiguities persist.

W2 - Typos and minor mistakes in the derivation: We thank you for going through our derivation! We have cross checked Eq. (2) using [b, c]. While possible, we think it is unlikely that the equation or the derivation is incorrect since it tallies with experimental evaluation.

W3 - Appendix: We have organised the appendix, and would like to thank you for spotting the errors!

Q1 - Normalizing Flows: We do not see a direct connection between the ideas proposed in this work and normalizing flows. If we recall correctly, normalizing flows are a class of generative models which construct a series of invertible transformations to map from a known distribution to an arbitrary distribution. While it may be possible to create applications of Lemma 1 and Theorem 1 in normalizing flows, this is not the scope of our work.

Q2 - Forward KL Divergence: We define the forward KL Divergence $D_{KL}(p || q)$ in Section 3.1, specifically Eq. 2. Here, $p$ corresponds to the target (or true) distribution and $q$ corresponds to the predicted distribution. This form of the KL Divergence also gives rise to other optimization measures such as the negative log-likelihood ($-\mathbb{E}_p \log q$).

Q3 - Consistent Notation: Thank you for this feedback. We have incorporated this in our draft!

[b] Zhang et al. "On the properties of kullback-leibler divergence between multivariate gaussian distributions.", NeurIPS 2024

[c] Joram Soch. Kullback-leibler divergence for the multivariate normal distribution, May 2020. https://statproofbook.github.io/P/mvn-kl.html.

Dear Reviewer 5LBD,

We thank you for your evaluation of the paper! We have noted your suggestion and have made changes in the draft accordingly. We hope to address your concerns below:

W1: We now describe the overall problem setting in Preliminaries, Section 2 of the paper. We separately discuss the formulation of Claim/Lemma 1 in the paragraph preceding it. We have cross-checked and provide an explicit derivation of Claim/Lemma 1 in the appendix. Please let us know if there still exists any ambiguity in the derivation!

W2: We have updated the abstract to specify that our analysis is restricted to the multivariate normal distribution.

W3: Appendix/Table 3 has been updated with the mean and standard deviations for our results on the UCI Regression datasets. The shaded regions in Figure 11 correspond to the standard devation.

W4 and Q4: We merge W4 and Q4 since we believe that they are related. We also believe that there may be some confusion over the scope of this experiment, which we clarify.

The TIC parameterization ([a]) predicts the covariance matrix as a function of the gradient and curvature of the mean estimator $\Sigma(x) = k_1(x) \nabla\widehat{\mu}(x) + k_2(x) \nabla^2\widehat{\mu}(x) + k_3(x)$. Here, $k_1(x), k_2(x), k_3(x)$ are learnable parameters. The authors of [a] show that the TIC parameterization for the covariance when trained using the negative log-likelihood (which is referred to as NLL:TIC) significantly improved upon the standard parameterization.

The experiment on heteroscedastic human pose estimation was introduced in [a] with NLL:TIC as the best performing method. Our preliminary experiments using the 2-Wasserstein failed to match the results of NLL:TIC, which we believe is because our pseudo-labels, which are computed based on a low-dimensional representation of the input images, may not necessarily be accurate. However, we identified an alternate paradigm: training the framework initially using the 2-Wasserstein bound and then switching to NLL:TIC helped retain the advantages across both, the 2-Wasserstein bound and NLL:TIC. Therefore, we highlight this specific experiment not necessarily as one where we are computationally efficient, but one which opens up possibilities of using hybrid approaches to further improve the state-of-the-art (NLL:TIC in this case).

While it is tempting to get pseudo-labels directly using the input images, this is computionally prohibitive: a single colour image with a height and width of 256 when flattened would correspond to a vector of size 196,608. Moreover, the space of images is sparsely populated, making it difficult to extract meaningful patterns. Finally, we would expect a well trained model to be invariant to changes in the image such as background variations. Directly using the images forgoes the invariances learnt by the model.

[a] Shukla et al. TIC-TAC: A Framework For Improved Covariance Estimation in Deep Heteroscedastic Regression, ICML 2024

W5: We have updated Figures 2 and 9 by setting the vertical axis to be log-scale.

Q1: Yes! The original 2-Wasserstein distance (Eq. 3) has one matrix square root operation that requires gradient and cannot be circumvented. This requires us to use eigendecomposition of the matrix, which may be unstable through backpropagation. The 2-Wasserstein bound does not have a matrix square root operation that requires gradient if we predict the square root of the matrix directly. In addition, the 2-Wasserstein bound is simpler to analyze and more intuitive.

Q2: The memory consumption refers to the additional memory needed by each of the methods to compute the forward and backward pass during training.

Q3: The instabilties unfortunately occur at random and are hard to tabulate. Therefore, we provide a reference in the paper to the documentation which highlights that this is a known issue due to numerical instability. However, we notice that some datasets such as superconductivity almost always face this instability. Since NLL:TIC also takes a large amounts of time and memory on this dataset, we do not report it in the results.

Q5: We thank you for pointing out this ambiguity. We refer you to appendix/Fig. 7: The residual can be treated as a vector which approximately points along the line segment joining the predicted mean and the target mean. However, if the residual is large, we observe that it influences the predicted covariance significantly. Consequently, the inverse covariance is aligned orthogonal to the residual vector. Since the gradient of the mean estimator is directly proportional to the inverse, the gradient as a whole is desensitised to move towards the target. Please let us know if additional clarification is required.

Q1: The major difference is in the philoshophy; since the negative log-likelihood cannot supervise the variance, Skafte et al. constructs the neighborhood on the fly to encode the variance in the batch. We avoid this by developing tools to directly supervise the variance, and pre-compute the covariance matrix to train the network for each sample directly. This also leads to better allocation of the mini-batch capacity, since the samples in the batch need not be restricted to similar neighborhoods, and therefore the mini-batch can optimize on diverse regions of the input space at the same time.

Algorthmically, in Skafte et al., the neighborhood of a sample is computed based on the euclidean distance, and a threshold 'd', which is a hyperparameter. The threshold 'd' may require hyperparameter search. We make a small change: we compute neighborhoods based on the Mahalanobis distance, with number of neighbors 'k'. We use the Mahalobis distance to account for the spread and scale of the samples. We set the number of neighbors 'k' to a constant value (10 * dimensionality of output) and do not perform hyperparameter search. This is based on a heuristic that the number of samples to compute the covariance scales linearly with the dimensionality. Skafte et al. is also sensitive to the number of neighbors 'n' chosen to construct the mini-batch. Finally, in Skafte et al., the samples in the batch are no longer i.i.d, a common assumption in many optimization frameworks.

Q2: While both the covariance pseudo labels and kernel methods rely on measures of similarity for the neighborhoods, we believe that they are fundamentally different. Kernel methods are used for non-linear transformations to solve challenging problems in high dimensional spaces. In contrast, we use the neighborhoods purely to compute pseudo labels for the covariance.

Q3: While we do not impose any specific restriction on the architecture, we recommend following guidelines set by Nix and Weigend, Stirn et al. and Sluijterman et al. to decouple the mean and covariance estimators. While the two estimators can share the backbone (such as in our Human Pose experiments using Vision Transformers), the two should not strongly regularize each other.

Q4: The biggest challenge is a mechanism to supervise the learning of the pseudo-labels. Such a mechanism is inherently infeasible with the negative log-likelihood. Therefore, it is non-trivial for methods relying on NLL to use pseudo labels.

Q5: Since the mean and covariance estimators are decoupled, we do not employ warm-up. Therefore, the mean and covariance are learnt from the start. "Is any information for the mean relevant in the presence of pseudo labels?" We are not quite sure we understand this. We obtain pseudo labels only for the covariance, the 2-Wasserstein distance trains the mean estimator using the mean square error. Therefore the mean is also required along with the pseudo labels for the covariance.

Q6: The pseudo labels for the covariance are computed based on the sample's label $y$. Therefore, if the mean estimator has not converged (underfit), the mean estimate is incoherent with the covariance estimate. If the mean estimator has converged, the predicted covariance will be coherent with the predicted mean.

Q7: If the mean estimator has underfit, the predicted covariance is incoherent with the mean. However, this is unlikely in practice unless the mean estimator is strongly regularized. If it has overfit, the situation is ambiguous since it is difficult to quantify the degree of overfitting. However, we believe that the covariance predictions by themselves can still be useful. Our reasoning is based on the way the pseudo labels are being computed, which is independent of the mean estimate, and leverages patterns within the dataset. The algorithm employs the Mahalanobis distance and softmax which are differentiable, implying continuity in the pseudo labels when the labels are continuous. Therefore, the covariance estimator can learn these patterns without relying on the mean estimator, making it useful as a standalone predictor.

Dear Reviewer VkG5, We thank you for your comments and questions! We hope to address your concerns below:

W1: Since we addressed a direct limitation of Shukla et al., we were influenced by their related works. We have updated the related works (Section 2) to include discussion of the work of Sluijterman et al. and Wong-Toi et al.

W2: We thank you for bringing these works to our notice, and have referenced them in the paper. Sluijterman et al. studies recommendations from Nix et al. which advocates for (1) separate networks for the mean and (co)variance, and (2) warm-up to train heteroscedastic models. We also separate the mean and (co)variance networks following a similar observation made by Stirn et al. We study warm-up using our toy example of bivariate normal distributions in Appendix/Fig. 9 (b). Specifically, we initialize the mean of our predicted distribution at the same location as the target distribution, recreating the convergence of the mean estimator. However, we note that convergence is still unstable since residuals can randomly destabilize the mean and covariance estimators. We discuss this in detail in Section 3. We also explore warm-up using the UCI Regression datasets (Section 4.2 and Figure 12) Our findings indicate that deep heteroscedastic regression is susceptible to residuals, which cannot be remedied by using warm-up.

We enjoyed reading Wong-Toi et al. which studies the training dynamics of deep heteroscedastic regression through phase transitions and have added it in our related works. We believe that the works address parallel problem statements. While Wong-Toi et al. restrict their theoretical analysis to the negative log-likelihood, we study the KL Divergence and the 2-Wasserstein distance. In the process, we prove Lemma 1 and Theorem 1 which we believe are sufficiently general enough to be of interest across communities. Moreover, the tools we introduce help us find accurate yet computationally cheap estimates of the covariance, addressing a direct limitation of prior work. Finally, the nature of our results apply to multivariate outputs, whereas Wong-Toi et al. experiment on univariate outputs Lines 73 and 74 explictly define univariate targets in their code (link).

W3: We describe our implementation in Appendix B. Specifically, we use separate networks to predict and regularize the mean and the covariance. While in general we do not use warm-up, we conducted experiments on the UCI regression with warm-up. This involved training the mean estimator for half the number of total training epochs, and jointly training the mean and covariance estimator for the other half. We report the results in appendix/Fig. 12. We would be happy to answer any specific queries regarding the implementation.

Dear Reviewer 6hor,

We thank you for your review! We have incorporated your feeback in the revised draft, and we hope to address your concerns below:

W1: Let us assume we have an input $x$ with the corresponding label $y$. While we know the label, the neural network needs to be trained to predict the label. We draw a similar analogy with Lemma 1; if we knew the ground truth covariance, how do we train the network to predict it? Lemma 1 addresses a specific setting for training the covariance estimator using the KL Divergence. The goal of this lemma is to present necessary tools needed to supervise the covariance estimator using our pseudo labels for the covariance. We acknowledge that the wording may have been ambiguous, and have updated it. In particular, we believe that Lemma 1, along with Theorem 1, is sufficiently general to be of interest across communities.

W2: We have re-run our experiments on the UCI Regression dataset with zero-mean-unit-variance standardization, and report the results in Table 2 and appendix/Table 3, the latter containing all three metrics with mean and standard deviation. We note that our observations and learnings still remain the same.

While [4] is indeed state-of-the-art, the method assumes that the output is univariate, as stated in the introduction section of [4]. This is also seen in the evaluation code used by multiple recent works, including [4]. For instance, Lines 56 to 64 specify the target to be univariate (link). Another recent work using UCI Regression makes the same assumption of the targets being univariate Lines 73 and 74 explictly define univariate targets in their code (link).

Our work addresses limitations in estimating the covariance, and proposes an accurate yet compute efficient way of modelling the heteroscedasticity through the full covariance matrix. As a result, we follow the experimental protocol of previous works [1] that specifically focus on multivariate predictions. Therefore we believe we should not be penalized for the same. Furthermore, we strongly believe that the experiments are well aligned to evaluate the claim of computationally efficient yet accurate covariance estimation in deep heteroscedastic regression.

W3: We do not grid-search over the space of hyperparameters, for instance we implement the default weight decay of 0.01 through the AdamW optimizer. Moreover, we use the exact same configuration to ensure fair comparison when training all methods. We compare the compute time with other state-of-the-art methods in Table 1 and 2. The analysis in [4] for regularization is done under the assumption that the target is univariate, thereby making direct comparisons infeasible.

W4: We choose $\beta$=0.5 which is the value recommended by the authors of $\beta$-NLL in the conclusion section of the work.

Q1: We have corrected our citation format, thank you for bringing it to our notice!

Q2: We have moved our NLL results to the main paper. The goal was not to relegate NLL to the appendix. Since we address a direct limitation of [1], we wanted to show improved performance on their proposed metric, TAC, and highlight that the method of [1] has room for improvement in the mean square error evaluation.