How To Train Your Covariance

We thank you for your review. We would like to clarify below:

1. 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.

2. 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.

3. Although the paper studies a theoretical problem ...
   We have addressed this point by rearranging the content of the paper. The notation table has been moved from the related work section to the introduction. We have moved the negative log-likelihood equation from the related work to the introduction, since we optimize our covariance formulation through the negative log-likelihood. Our introduction already stated our assumption that the target distribution is a multivariate normal distribution, the standard assumption in machine learning. Furthermore, we also specify the standard assumption that we do not know the data generating distribution p(X, Y). Instead, we obtain samples (x, y) from it to form our dataset. We understand that our notation may have been ambiguous, and we have revised the draft to address this.

4. Most figures are missing axis labels ...
   Thank you for pointing this out. We have addressed this in the revised draft.

We would like to conclude the response by reiterating that this work is in line with previous works which propose models and training strategies to learn the covariance in applied settings.

[1] Andrew Stirn, Harm Wessels, Megan Schertzer, Laura Pereira, Neville Sanjana, and David Knowles. Faithful heteroscedastic regression with neural networks. In International Conference on Artificial Intelligence and Statistics, pp. 5593–5613. PMLR, 2023

[2] Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the pitfalls of het- eroscedastic uncertainty estimation with probabilistic neural networks. In International Confer- ence on Learning Representations, 2022

[3] Nitesh B Gundavarapu, Divyansh Srivastava, Rahul Mitra, Abhishek Sharma, and Arjun Jain. Struc- tured aleatoric uncertainty in human pose estimation. In CVPR Workshops 2019.

[4] Ivor JA Simpson, Sara Vicente, and Neill DF Campbell. Learning structured gaussians to approximate deep ensembles. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition 2022

[5] Garoe Dorta, Sara Vicente, Lourdes Agapito, Neill DF Campbell, and Ivor Simpson. Structured uncertainty prediction networks. In Proceedings of the IEEE conference on computer vision and pattern recognition 2018.

Dear Reviewer,

Our work builds upon a wide range of successful parametric variance / covariance estimation methods. This includes the highly cited: Kendall and Gal [NeurIPS ‘17], Dorta et al. [CVPR ‘18], Seitzer et al. [ICLR ‘22], Stirn et al. [AISTATS ‘23]. Many more works use parametric covariance estimation for downstream tasks, e.g.: Lu and Koniusz [CVPR ‘22], Liu et al. [ICRA ‘18], Gundavarapu et al. [CVPRW ‘20], Russel and Reale [T-NNLS ‘21], Simpson et al. [CVPR ‘22] among many others. These references are available in the draft. We follow and describe a similar (and same in some cases) probabilistic framework, and maintain standard machine learning notation which is explicitly described through a table.

We therefore confess that we are perplexed by the evaluation of the draft by the reviewer.

1. In nonparametric statistics literature, it is standard ...
   Neural networks with differentiable activation functions and optimized using negative log-likelihood with Tikhonov regularization are assumed to be Lipschitz and smooth [6]. These assumptions are implied and well known in deep learning research, if needed we can add them in the draft. Additionally, results from the theory of neural networks highlight that neural networks with non-constant, bounded, continuous activation functions, with continuous derivatives up to order K belong to Sobolev spaces of order K [7]. Further, the use of neural networks to model the variance / covariance has already been explored in previous work on parametric covariance estimation, hence we do not emphasize this in the draft. The biggest advantage of neural networks is that it allows us to learn complex representations from varying kinds of inputs (such as images) to predict the covariance.

2. Eq. (2) in the revision is still a mathematically nonsense expression...
   Let us assume we have $N$ predictions $(x, \hat{y})^{(1)} \ldots (x, \hat{y})^{(N)}$. However, the probability of exactly observing $p(X = x)$ is zero for continuous variables. Instead, the standard approach (Section 2.4 in [8]) is to observe over the set $X \in lim\_{\varepsilon \rightarrow 0} [x, x + \varepsilon]$. This implies that for continuous variables, we do not observe a specific but a range of values around $X = x$. We can incorporate this definition in the covariance too. Since $\hat{Y}$ is a deterministic transformation of $X$ through a parametric network, $\hat{Y} = f\_{\theta}(X)$, we have: $Cov(\hat{Y} | X = x) = Cov( f\_{\theta}(X) | X \in lim\_{\varepsilon \rightarrow 0} [x, x + \epsilon])$.
   Finally, the only heuristic in our derivation is in treating this neighborhood $\lim\_{\varepsilon \rightarrow 0} [x, x + \varepsilon]$ as $x + \epsilon$, where $\epsilon$ is an m-dimensional random variable with a zero-mean isotropic Gaussian distribution $p(\epsilon) = \mathcal{N}(0, \sigma^2\_{\epsilon}(x) I_m)$. Introducing $\epsilon$ imposes a distribution on the neighborhood which is centred at $x$. This also allows us to represent $\hat{y} = f\_{\theta}(x + \epsilon)$ stochastically. While the variance of this random variable is unknown (we later show that it can be learnt), we assume heteroscedasticity which allows us to represent neighborhoods of varying spatial extents for each $x$. We elaborate on expression 2 in subsection 3.1 of the revised draft.
   We take a principled approach in deriving our closed-form expression for the covariance, which we believe is not reflected in the current review.

3. Despite the changes made, many of the notations remain unclear and confusing ...
   We already follow standard machine learning conventions for $y, \hat{y}, f\_{\theta}$ and described them in Table 1. Further, our introduction also stated that $p(Y | X)$ is normally distributed, and parameterized by two networks $f\_{\theta}$ and $g\_{\Theta}$, which are trained through the negative log-likelihood. We emphasized that our work assumes the same probabilistic setting as prior literature on parametric covariance estimation, and the only change we propose is in the way we model the covariance through a novel closed form expression. Without specific examples we are not able to understand what in the notation is confusing.

Nevertheless, we thank you for your review.

[6] Du, S., Lee, J., Li, H., Wang, L. and Zhai, X., 2019, May. Gradient descent finds global minima of deep neural networks. In International conference on machine learning (pp. 1675-1685). PMLR.

[7] Czarnecki, W.M., Osindero, S., Jaderberg, M., Swirszcz, G. and Pascanu, R., 2017. Sobolev training for neural networks. Advances in neural information processing systems, 30.

[8] Evans, Michael J., and Jeffrey S. Rosenthal. Probability and statistics: The science of uncertainty. Macmillan, 2004.

We thank you for your review. We address your questions below:

1. In Section 2.1, $\sigma_{\Theta}$ has not been defined ...
   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.
2. The sentence after Equation (6)
   We have completed the sentence.
3. What is the theoretical explanation motivating the use of the thirs matrix term $k_3(x)$ ...
   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.
4. Conditional Marginal Likelihood comparison.
   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.
5. Computational complexity ...
   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.

We thank you for your review. We understand that the concerns majorly revolve around presentation and formatting. We have addressed the figures, tables and capitalisation in the revised draft. We also answer your specific questions:

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.
2. ...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.
3. 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.
4. 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.
5. Another point worth noticing ...
   While we understand your concern, we believe that this point should not be considered as a weakness. The reason we propose a new metric is because of the lack of any measure to directly assess the covariance. Moreover, CMAE in our opinion elegantly combines the notion of L1/L2 error with correlations through the fundamentals of conditioning a distribution. Additionally, the metric is independent of the training objective for all methods, and hence we believe it is a fair comparison. We hope that this paper sparks discussion in the community not only on covariance estimation but more importantly on covariance evaluation.
6. 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.

Dear Reviewer r64o,
We thank you for your time and efforts.
Since the discussion period ends soon, please let us know if we could resolve your questions.
We believe that our revised draft addresses your concerns.
We welcome you to read our general response as well.
If satisfied, we would be grateful if you would consider accepting the manuscript.
Once again, thank you for your time and efforts!

We thank you for your review. We address your questions below:

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.
2. 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.
3. 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.