Paper: How To Train Your Covariance (ICLR-2024)

Reviewer comment (Theory):
Most equations in the paper remain at the level of heuristic. While it is okay to have a heuristic explanation of the proposed method, the way it is presented in the current draft is unnecessarily confusing. For example, how do we even define the limit in Eq. (1)? I think the paper should try to minimize the use of non-rigorous math.

Author response:
**Most equations in the paper remain at the level of heuristic... how do we even define the limit in Eq. (1) ...**\
The limit in Eq. 1 is a standard definition for continuous distributions, and we have added references for the same. The only heuristic is in treating the limit as a random variable with zero mean, the variance of which is learnt through neural network optimization. We take a principled approach in deriving our closed-form expression for the covariance, and fail to understand the reviewer's concern.

Reviewer comment (Theory):
To my knowledge, conditional mean/covariance estimation is impossible without further assumption (e.g. regularity of $\mu_{y | x}, \Sigma_{y | x}$). I wonder how the authors avoided this problem.

Author response:
**To my knowledge, conditional mean/covariance estimation is impossible without further assumption.**\
Conditional mean/covariance estimation has been studied in multiple previous works, e.g.: [1-5]. The common approach in these works is to assume that the target distribution $p(Y | X = x)$  conditioned on the input $x$ (can also be an image) follows a normal distribution $\mathcal{N}(\mu_{Y|X}, \Sigma_{Y|X})$. The mean and covariance of the target are predicted using two neural networks which takes $x$ as input and are trained using the negative log-likelihood. Our work follows  this paradigm, and proposes a new way to model the covariance. Therefore, it is not impossible to assume conditional mean/covariance estimation within our problem statement.

Reviewer comment (Presentation):
Most figures are missing axis labels. For example, Figure 1 is placed in the very beginning of paper without explaining what the curves are.

Author response:
**Most figures are missing axis labels ...**\
Thank you for pointing this out. We have addressed this in the **revised draft.**

Reviewer comment (Evaluation):
The computational complexity of the approach could be an issue as the paper tackles a practical estimation problem.

Author response:
Indeed, we note in our limitations section that the primary bottleneck of our approach is in computing the Hessian. While determining the computational complexity of the Hessian for a generalized network is non-trivial, we draw your attention to [1]. Computing the Hessian for a function that maps an m-dimensional x to an n-dimensional y has a complexity of $O(nm^3)$.
In practice, covariance estimation can be performed using a smaller, proxy model in place of a large model (which could be retained for mean estimation). The reduced parameter count would decrease the computational requirements of computing the Hessian.



[1] Yao, Zhewei, et al. "Pyhessian: Neural networks through the lens of the hessian." 2020 IEEE international conference on big data (Big data). IEEE, 2020.

Reviewer comment (Experiments):
The paper does not compare to other approaches to log-likelihood such as Lotfi, S., Izmailov, P., Benton, G., Goldblum, M., & Wilson, A. G. (2022, June). Bayesian model selection, the marginal likelihood, and generalization. In International Conference on Machine Learning (pp. 14223-14247). PMLR.

Author response:
Thank you for pointing us to this paper, which we have included in our related work. However, we believe that not comparing to Lotfi et al. should not be considered a weakness since the two works are fundamentally different. Conditional Marginal Likelihood is a metric to evaluate generalization. The metric does so by creating two non-overlapping subsets of the dataset: $D_1$ and $D_2$, training a model on $D_1$, and quantifying its performance on $D_2$. By contrast, CMAE is a metric with a different objective. Given a sample $(x, y)$ with its prediction $(\hat{y}, \textrm{Cov}(\hat{y}))$, the metric evaluates the covariance by measuring the improvement in $\hat{y}$ when $y$ is partially observed. Hence, the two metrics are fundamentally different.

Reviewer comment (Writing):
In Section 2.1, \sigma_{\Theta} has not been defined. Is it a scalar that is assuming a diagonal covariance matrix? How do you go from Cov(\hat{y}) to \sigma_{\Theta}?

Author response:
We have clarified this in the related work in the **revised draft**. $\sigma_{\Theta}$ (or Var$_\theta$) is a vector encoding a diagonal covariance matrix. The authors of Beta-NLL specify Cov($\hat{y}$) to be a diagonal matrix.

Reviewer comment (Writing):
The sentence after Equation (6) is incomplete: "We note that both both Cov() and Cov()..."

Author response:
We have completed the sentence.

Reviewer comment (Theory):
What is the theoretical explanation motivating the use of the thirs matrix term k_3(x)? It is said in the paper, that the curvature of the function at x cannot alone explain the stochasticity of the samples which motivate the use of k_3(x). It would then be appropriate to motivate the definition of "spatial covariance" by the "curvature of x" and of "?" coming from k_3(x). Could the authors please elaborate more on this?

Author response:
We have addressed this in the revised draft. The sum of two independent Gaussians is the sum of the means and covariances of the two Gaussians. Therefore, the covariance of a sample can be attributed to the nature of the function underlying the sample as well as independent stochasticity, which is not a function of the input. We attributed $k_1(x)$ and $k_2 (x)$ to the former and $k_3(x)$ to the latter. Given the independence assumption, the final formulation can be written as the sum of all three terms.

Reviewer comment (Novelty):
The following paper proposes "Conditional Marginal likelihood" as an alternative to the likelihood (although for generalisation context): Lotfi, S., Izmailov, P., Benton, G., Goldblum, M., & Wilson, A. G. (2022, June). Bayesian model selection, the marginal likelihood, and generalization. In International Conference on Machine Learning (pp. 14223-14247). PMLR. This relates to your definition of C-MAE. Could you elaborate on the differences between the two metrics and possibly compare them in the current setting?

Author response:
Thank you for pointing us to this paper, which we have included in our related work. However, we believe that not comparing to Lotfi et al. should not be considered a weakness since the two works are fundamentally different. Conditional Marginal Likelihood is a metric to evaluate generalization. The metric does so by creating two non-overlapping subsets of the dataset: $D_1$ and $D_2$, training a model on $D_1$, and quantifying its performance on $D_2$. By contrast, CMAE is a metric with a different objective. Given a sample $(x, y)$ with its prediction $(\hat{y}, \textrm{Cov}(\hat{y}))$, the metric evaluates the covariance by measuring the improvement in $\hat{y}$ when $y$ is partially observed. Hence, the two metrics are fundamentally different.

Reviewer comment (Theory):
The computational complexity of the approach could be an issue as the paper tackles a practical estimation problem. Could the authors please provide an exact analysis of the computational complexity of the method and suggest possible ways of improvement?

Author response:
Indeed, we note in our limitations section that the primary bottleneck of our approach is in computing the Hessian. While determining the computational complexity of the Hessian for a generalized network is non-trivial, we draw your attention to [1]. Computing the Hessian for a function that maps an m-dimensional x to an n-dimensional y has a complexity of $O(nm^3)$.
In practice, covariance estimation can be performed using a smaller, proxy model in place of a large model (which could be retained for mean estimation). The reduced parameter count would decrease the computational requirements of computing the Hessian.



[1] Yao, Zhewei, et al. "Pyhessian: Neural networks through the lens of the hessian." 2020 IEEE international conference on big data (Big data). IEEE, 2020.

Reviewer comment (Writing):
Generally, the paper could be clearer, and its format can be improved. For instance: the abstract (5th line) defines the mean of a Gaussian as $f(x)$, and the covariance as $Cov(f(x))$. So, it is not clear whether $f(x)$ refers to the RV to be modelled or its mean.

Author response:
1. **So, it is not clear whether  refers to the RV to be modelled or its mean.** \
$f(x)$ refers to the mean of the random variable being modelled. We have updated this in the abstract.

Reviewer comment (Reproducibility):
A few times, it is mentioned that the experiments are run over _multiple network architectures_; however, in the experiments, there is no mention of specific architectures used.

Author response:
1. **...multiple network architectures...** \
The architecture used depends on the task being solved. For our synthetic experiments and UCI Regression, we use fully connected networks to learn $y$ in $R^m$ from $x$ in $R^n$. We specifically choose the task of human pose estimation to showcase the use of convolutional neural networks, and regress heatmaps of shape $R^{j \times 64 \times 64}$ from images of shape $R^{256 \times 256}$. Collectively, our experiments address different shapes and types of input-target pairs. We have rephrased this sentence in the draft.

Reviewer comment (Theory):
In the line of clarity, the paper is motivated by the pitfalls of maximum likelihood (ML). However, the proposal in the paper results in a specific parametrisation of the covariance, which ties the structure of the covariance and the mean (eq 8 shows how the covariance contains the Jacobian of the mean function). Therefore, the proposal is not _another training strategy_ but rather a covariance parametrisation. As a matter of fact, after eq 8 the authors state that their parametrisation is used alongside ML.

Author response:
1. **Also, in line of clarity ...** \
We agree with your observations. The proposal is indeed a way to model the covariance which is learnt via the negative log-likelihood. To avoid ambiguity, we have renamed the paper from "_How To Train Your Covariance_" to "_How To Model Your Covariance_". We also rephrased the text to address this ambiguity wherever applicable in the draft.

Reviewer comment (Experiments):
Benchmarks are unclear: The experiments compare the proposed parametrisation against **NLL** (though the proposed method also uses NLL as far as I understand), **Diagonal**, which I assume also uses NLL, and other methods. It is thus confusing if the paper compares approaches to training covariances or models for covariances.

Author response:
1. **Benchmarks are unclear.** \
Indeed, we compare our method against other approaches to model and train the covariance. While NLL, Diagonal and our proposed model for the covariance use NLL as a training objective, Beta-NLL and Faithful Heteroscedastic Regression have their own training objectives, which are variants of NLL. All these models/training strategies give a covariance prediction with which we compare our proposed method.

Reviewer comment (Reproducibility):
There should be given more details about the choice of $k_1,k_2,k_3$, the networks, and the learning objectives.

Author response:
1. **There should be more details about $k_1, k_2, k_3$ ...** \
We have addressed this towards the end of our methodology section.  The covariance estimator $g_{\Theta}(x)$ predicts $k_1(x), k_2(x)$ and $k_3(x)$, where $k_1(x), k_2(x)$ are positive scalars. We enforce $k_3(x)$ to be positive semi-definite by predicting an unconstrained matrix and multiplying it with its transpose, similar to previous work.

Reviewer comment (Experiments):
The C-MAE operates like a leave one out fashion. In general for multi-variate Gaussian, we can do any leave-k-out. Would it provide more info if the consider any k greater than 1? Or mainly, the authors may want to illustrate the specific choice of leave one out here.

Author response:
1. **leave-one-out vs leave-k-out** \
We have updated the CMAE section in the **revised draft** to address this. While _leave-one-out_ can be generalized to _leave-k-out_, we do not observe any change in the evaluation trend. A method having a lower _leave-one-out_ error also has a lower _leave-k-out_ one. Moreover, _leave-k-out_ requires taking nCk combinations, i.e., significantly more than the $n$ combinations required in _leave-one-out_. This motivates the use of the _leave-one-out_ strategy.

Reviewer comment (Experiments):
What are the specific choice considerations for k1, k2 and k3? Would any regularity conditions further help with supervision?

Author response:
1. **What are the specific choices for $k_1, k_2, k_3$?** \
We have addressed this towards the end of our methodology in our revised draft.  The covariance estimator $g_{\Theta}(x)$ predicts $k_1(x), k_2(x)$ and $k_3(x)$, where $k_1(x), k_2(x)$ are positive scalars. We enforce $k_3(x)$ to be positive semi-definite by predicting an unconstrained matrix and multiplying it with its transpose, similar to previous work. \
1. **Would any regularity conditions further help with supervision?** \
We believe that having priors would allow for faster and better convergence of the covariance. To maintain parity with previous work, we do not assume priors. Moreover, we work within the framework of unsupervised heteroscedastic covariance estimation, which means that we do not have supervisory signal for the covariance.