Towards Self-Supervised Covariance Estimation in Deep Heteroscedastic Regression
Authors: Megh Shukla
Aziz Shameem
Mathieu Salzmann
Alexandre Alahi
Venue: ICLR 2025 Poster (2025)
Keywords: deep regression
heteroscedastic
uncertainty
2-Wasserstein
KL-Divergence
Negative Log-Likelihood
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We study the KL Divergence and 2-Wasserstein distance for self-supervised covariance estimation in deep regression

Abstract: Deep heteroscedastic regression models the mean and covariance of the target distribution through neural networks. The challenge arises from heteroscedasticity, which implies that the covariance is sample dependent and is often unknown. Consequently, recent methods learn the covariance through unsupervised frameworks, which unfortunately yield a trade-off between computational complexity and accuracy. While this trade-off could be alleviated through supervision, obtaining labels for the covariance is non-trivial.
Here, we study self-supervised covariance estimation in deep heteroscedastic regression. We address two questions: (1) How should we supervise the covariance assuming ground truth is available? (2) How can we obtain pseudo labels in the absence of the ground-truth? We address (1) by analysing two popular measures: the KL Divergence and the 2-Wasserstein distance. Subsequently, we derive an upper bound on the 2-Wasserstein distance between normal distributions with non-commutative covariances that is stable to optimize. We address (2) through a simple neighborhood based heuristic algorithm which results in surprisingly effective pseudo labels for the covariance. Our experiments over a wide range of synthetic and real datasets demonstrate that the proposed 2-Wasserstein bound coupled with pseudo label annotations results in a computationally cheaper yet accurate deep heteroscedastic regression.

Published as a conference paper at ICLR 2025
TOWARDS SELF-SUPERVISED COVARIANCE ESTIMATION IN DEEP HETEROSCEDASTIC REGRESSION
Megh Shukla EPFI.
Aziz Shameem ITT Bombay
Mathieu Salzmann SDSC and EPFI.
Alexandre Alahi EPFL
(eeg-rogreesionn.iithh.
ABSTRACT
Deep heteroscedastic regression models the mean and covariance of the target distribution through neural networks. The challenge arises from heteroscedassicity, which implies that the covariance is sample dependent and is often unknown. Con- sequently, recent methods learn the covariance through unsupervised frameworks, which unfortunately yield a trade-off between computational complexity and ac- curacy. While this trade-off could be alleviated through supervision, obtaining labels for the covariance is non-trivial. Here, we study self-supervised covariance estimation in deep heteroscedastic regression. We address two questions: (1) How should we supervise the covariance assuming ground truth is available? (2) How can we obtain pseudo- laeuls-labels the absence of the ground-truth? We address (1) by analysing two popular measures: the KI. Divergence and the 2- 2-sasserstein distance. Subsequently, we derive an upper bound on the Wassersttei distance between normal distributions with non-commutative covariances that is stable optimize. We address (2) through a simple neighborhood based heuristic algorithm which results in surprisingly effective pseudo-labed labels for the covariance. Our ex- periments over a wide range of synthetic and real datasets demonstrate that the proposed 2- W-ssseesttein bound coupled with pseudo-label annotations results in computationally cheaper yet accurate deep heteroscedastic regression.
INTRODUCTION
Deep heteroscedastic regression leverages neural networks powerful feature extractors to regress the mean and covariance of the target distribution. The target distribution typically used for downstream tasks such as uncertainty estimation, correlation analysis, sampling, and in bayesian frameworks. The key challenge in deep heteroscedastic regression lies in estimating heteroscedasticity. which implies that the variance of the target is input dependent and variable. Moreover, unlike the mean, the covariance lacks direct supervision and needs to be inferred.
The standard approach without the ground-truth covariance relies on optimizing the negative log- likelihood to jointly learn the mean and covariance (Dorta et al. 2018). However, Skafte et al. (2019); Seitzer al. (2022) show that the absence of supervision, incorrect variance predictions lead to sub-optimal convergence. Subsequently, a flurry of recent works addresses this by either modifying the negative likelihol (Skafte al., 2019 Seitzer et al., [2022] SStrn et al.) [223] [mmer et al..
2023) or improving the covariance through better parameterization (Shukla et al. 2024 However, the methods introduce a trade-off between computational complexity and accuracy. Moreover, the thematic message underlying these works is that estimating heteroscedasticity is challenging when annotations for the covariance are not available. Therefore, would having annotations for the covariance alleviate this trade-off? To answer this, explore the use of self-supervision for covariance estimation in deep heteroscedastic regression. We study two questions:
(Q1) How supervise the learning of the covariance assuming annotations are available? Since the negative log-likelihood is not formulated for supervising the covariance, we analyse the KI. Divergence and the 2- Wasserstein distance to supervise the learning of the mean and covariance.
Work done by Aziz Shameem during his internship at EPFI.. Contact: fi rstnnmm. 1 lastnare@epf)).cm

---

Published as a conference paper at ICLR 2025
Our analysis highlights the utility of the KL Divergence as a regularizer of the covariance. However, we observe that despite supervision, the KL divergence underperforms when compared to the 2- Wasssesstein distance, it shares susceptibility to residuals similar to that of the negative log-likelihood. Next, we study the 2- 2-Wasserstein distance between normal distributions with non- commutative matrices. We specifically note optimization challenge (PyTorch, (024) due to the eigendecomposition involved. Consequently, we extend the formulation for the 2- Wasserstein distance between commutative covariance matrices the general case of con-commutative matrices, eliminating the need for eigendecomposition. This makes the optimization process stable.
(Q2) How can we obtain pseudo-labels for the covariance when annotations are not available? In the absence of priors, we propose a neighborhood based heuristic algorithm to generate pseudo- labels for the covariance. Specifically, for a given sample, the pseudo-label corresponds to the covariance over the targets of all the samples in the neighborhood of the specified sample. The contri- bution of the neighboring samples are weighed by their Mahalanobis distance to the specified sample. We show that this simple strategy provides effective self-supervision for covariance estimation.
We perform extensive experiments across a wide range of synthetic and real world settings. Our results show that the proposed self-supervised framework is (I) computationally cheaper and (2) retains accuracy with respect to the oftte-of-the-arr Finally, our experiments on human pose also highlight the possibility of using self- self-super and unsupervised objectives to further improve optimization. Our code is available at our project page.
2 DEEP HETEROSCEDASTHR REGRESSION
Heteroscedastic regression is the task of modeling the mean and the variance of the target distribution. In contrast to homoscedasticity, heteroscedastic models allow the variance to vary as function of the input. Deep heteroscedastic regression provides an advantage over non-parametric methods like the Gaussian Processes (Le et al. 2005) because it can model complex features from inputs such as images. This attribute has made it widely applicable in fields like active learning (Houlsby et al.
2011 Gal et al., 2017). uncertainty estimation (Gal & Ghahramani 2016 Kendall & Gal 2017]
Lakshminarayanan al., [2177] Russell & Reale (022)) image reconstruction (Dorta al. 2018).
human pose (Gundavarapu et al.
2019
Nakka & Salzmann
2023 Tekin et al. 2011)) and other
vision-based tasks (Lu & Koniusz
2022 Simpson al. 20222 Liu et al. 2018 Bertoni et al., 2019).
Preliminaries. Our goal is to learn the target distribution P(Y[X) for different X, where X E RT
the input and Y E R" is the target variable. While PXY is unknown, it is assumed to be normally distributed: PXY = ((u((X)
(X)). Ex))),, Our estimate of the target is P(] = (By (X), (X)), where the mean poy (X) fe(X) and covariance Ex(X) = ge(X) are parameterized by neural networks. The standard approach in the literature to learn the target distribution is through minimizing the negative log-likelihood, -Ep(x.y P() Specifically, the mean and covariance networks are trained minimize (Nix & Weigend 1994)
LNL (f, θ) = Ep(x)) log [Ex(X] X
+ (Y - -xx( (x) (Y -xx) (I)
Challenges. The lack of supervision for the covariance results in optimization challenge which is formalized in Skafte et al. (2019); Seitzer et al. (2022). The works observed that an underestimated variance can increase the effective learning rate and disrupt optimization (Skafte et al. 2019) whereas an overestimated variance can decrease the effective learning rate and stop optimization (Seitzer el al. 2022). A number of recent approaches modify the negative log-liilidood reduce the effect of the predicted covariance during optimization. 3INLL (Seitzer et al. (002) scales the negative log- likelihood (Eq. 1 by the predicted variance resulting the objective: Le LL-NLL NIL (ā()) 23 x LALL- However, since J-NLL does not originate from valid distribution, the optimized values do not estimate the true variance. [Stirn et al. 2023) decouples the estimation of the mean and variance by scaling the gradient of the mean estimator with the covariance, thereby eliminating the effect of the predicted covariance on the mean. This leads to confficting assumptions: the mean estimator assumes that the multivariate residual is uncorrelated while the covariance estimator is expected to identify correlations. [mmer et al. (2023) proposed the use of natural parameterization of the univariate normal distribution: nn 22 and n2 2 for regression. While principled, the method
2

---

Published as a conference paper at ICLR 2025
assumes univariate targets. TIC-TAC (Shukla et al. 2024 in contrast to previous works, retains the negative log- oog-liielihood and formulates the predicted covariance through the gradient and curvature of the mean. However, the improvement in accuracy comes the expense of increased computational requirements. A parallel line of works studies the impact of training dynamics in deep heteroscedastic regression. [Wong- Toi et al. (2024) provide theoretical study linking the training of heteroscedastic regression to phase transitions, however the experimental evaluation is limited to univariate outputs. Sluijterman al. (024)) experimentally show that decoupling the mean and variance networks can lead to improved performance, similar to Stirn et al. (2023). However, while the authors suggested a warm-up schedule, we observed that this may not necessarily improve performance.
An overview of the related works reveals shared theme: estimating heteroscedasticity is difficult without annotations. Further, existing works trade-off accuracy for lower computational requirements. This trade-off could be mitigated with supervision; however, acquiring labels for the covariance is challenging, which restricts further analysis. Moreover, the negative log-likelihood is not formulated to supervise the covariance, requiring new approach supervision. Therefore, investigate two key aspects of the problem: (I) How can we supervise the covariance when negative log-like log-likelihood not specifically designed for this task? and (2) How generate labels- for the covariance in the absence of ground truth?
3 ANALYSIS
We analyze the KL Divergence and the 2- Wasserstten distance, two widely used metries for comparing and optimizing distributions. Our analysis focuses on multivariate normal distributions, which follows a common assumption in machine learning that the residuals are normally distributed. We support our analysis by studying Problem 1, which lets us visualize the convergence process of various methods using bivariate normal distributions. We then seek to answer which metric is better suited for deep heteroscedastic regression.
Problem 1. (Bivariate Normal Distribution) We consider the task of learn- ing a bivariate normal distribution. We initialize the target and predicted distributions to have different means. While the predicted distribution is initialized with identity covariance matrix, the covariance for the target distribution is initialized randomly and such that it exhibits a high degree of correlation (> 0.5). Given samples y R22 from the target distribution, the goal is to compare different methods in optimizing the predicted distribution match the target one.
Ground Truth Prediction
3.1 KL DIVERGENCE
The KL Divergence between two continuous distributions p. is D Ep logí p(x)/q(x))
For two multivariate normal distributions (Zhang et al., 2024 Soch 2020) we have
DK (1) = Tr(2;''2 2pp) + (My Hp) T 2i (My pp) - k + In
det Eg det Eip
(2)
Since the KL Divergence is asymmetric, we often let p and & be the target and predicted distributions respectively. This definition also gives rise to popular alternatives such as the negative log- likelihood. While the KL Divergence has been well explored (Goodfellow et al. 2016 Arjovsky et al. (017)) from statistical viewpoint, we ask: how can the KL Divergence be formulated for deep heteroscedastic regression? Answering this is necessary since the KI. Divergence is defined in terms of the means and covariances, which are unknown for the target distribution.
Formulation. A logical approach would be to replace each label with a distribution. Specifically, for sample F; y, from the dataset, the pseudo target distribution be set to /N Exprio (X)). However, this approach requires calibrating the KL Divergence since the optimal value for the covariance is not the same the prior. We demonstrate this through simple setting in Lemma]]
3

---

Published as a conference paper at ICLR 2025
Methed: Negative Loglikelihoad
Method: Divergence
Metric: Epech
Metric: 2-W vis Epoch
Ground truth
Ground Truth
Prediction
Prediction
Sample
Sample
Initislization
Initalization
In addition to feature granularity (Seitzer et al. 2022), subpar convergence may occur due to the sensitvity of the negative log-likelihood and the KL- Di-Divergence to residuals in the mini- While we show that the KL-Divergence can act as a regularizer over the learnt covariance, the gradients for both the methods are dominated by the residual term, slowing down convergence.
Figure 1: Sub-optimal convergence due residuals (Section: 3.7
Lemma 1 (Calibration). Let == Vi) be set of samples drawn from the unknown larget PX Nuy
for fora given æ. We write each label yi as distribution N(g,, Eiprion (X)). Then, the optimal solution using the K1. Divergence for the predicted covariance over the S is Ey (X) Ey (X) + Exprion
(X). Consequently, if the target covariance is known and set as the prior, we have (x) x (X).
refer the reader to the appendix, section A.2) for the proof.
Discussion. The optimal solution for Σy (X) ) is twice the target covariance. This can be addressed by simple calibration of the KL Divergence (Eq. 2):
1 2
Tr Exprion)
(ūy - y) y) T (By V) k + In
det Ey det >y
2
In the general scenario where the true covariance is known priori, the optimal solution the average of the prior and the sample covariance. This introduces a notion of regularization on the predicted covariance. Finally, the predicted covariance also robust outliers.
Impact of residuals. In general, the solution in Lemma iis reached only when the mean estimator has converged the true mean and when we observe multiple targets y, for the same observation æ. This may not hold true in practical settings because: (1) samples in batch are i.i.d, implying that the same observation
unlikely to be repeated, and (2) the mean estimator may not have converged.
In practice, for each sample in the batch, we take noisy gradient step towards Ey (X) Σ >(prion) (X) (X) (py
3 (appendix/eq 13). If the residual term (X) - y) large, the gradient step due to the residual dominates over the prior covariance, moving us closer to Ey Ex(X) Ro (X) y) (āx y)² This residual matrix can be interpreted
'covariance" matrix aligned along the line segment joining and (Xy((X (Eq. 3; OLS estimate of the slope, [oohh (222 ). However, this residual matrix desensitizes the mean estimator to move along y and juy (X). slowing down optimization. This is because the 'variance' induced by the residual matrix is large along and ğo (X), and the gradient of the mean estimator is proportional the inverse of the covariance (appendix/eq (T). This is pictorially depicted in appendix/Fig. 7
We study this phenomenon through Problem in Fig. [ We observe that after a few iterations, the predicted covariance is aligned along the means of the target and predicted distribution. This observation a result the residual term appearing in the optimal solution. Moreover, while the KL Divergence incorporates our prior knowledge of the covariance, the prior term is dwarfed in magnitude when compared to the residual term. We also note increased optimization instability at higher learning rates (uppendix/fig. 8 While the KL Divergence leverages the prior covariance regularizer, shares drawbacks pertaining to the residual with the negative log-likelibee liglliiohhodd motivating analysis of the 2- Wasserstein distance.
4

---

Published as a conference paper at ICLR 2025
Negative Loglieehhood
Pattitu
KL Divergence
2-Wasserstein stance
2 2-Waasserteein (Bound)
Groune
Fattrnl
nattrul
KL
241
24W
24Y
24 (Boun
(Boundd
W (ouudd
1.0
380
21.5
20.5
Figure 2: Visualizing convergence in bivariate regression (Section: 3.2 We observe that the -L-Divergence and likelihood based methods: vanilla negative log-like bog-likelihood and Faithful (Stir al. 2023) result in unstable convergence due the sensitivity of the methods the residuals. In comparison, the 2- 2-Wasserstein based methods are more stable and accurate. This observation can also be replicated when the predicted mean is initialized the same location the mean, shown in appendix/Fig ( (b) (Note: metrics NLL/ / KL and 2-W are plotted in log-scale)
3.2 2- WASSERSTEIN DISTANCE
The Wasserstein distance is a metric for quantifying the distance between two probability distributions. It defines the minimum "cost" required to morph distribution into another. The 2- Wasserstein distance measures the cost in proportion to the squared Euclidean distance. It is widely used in optimal transport theory and generative modeling (Arjovsky al. 2017 Li al. 2024), both the shape and spread of distributions while penalizing long-distance transport more heavily, Let (14 E,,, Ng(22, E2) be two multivariate normal distributions. The 2- Wasserstein distance between them is given by
captures
| #2 1 Tr + 2i
2, /3/////11
(3)
This formulation, however, requires computing the root of a matrix, which typically involves eigen- decomposition. Unfortunately, the eigendecoompoition in popular deep learning frameworks can potentially lead to unstable gradients (PyTorch 2024). If En and Eg are commutative (implying Σ Σ En
E,,,,, then the 2- Wasserstein distance reduced to (M = |||| -#2 |2 | ||||||/// 1/22 This formulation allows to directly predict the square roots, thereby avoiding the eigendecom- position. However, for two covariance matrices to be commutative, they need to share the same eigenbasis, implying that the matrices differ only in the variance of the individual random variables. Fortunately, Theorem Tallows us expand this formulation to non-commutative covariance matrices by linking upper bound the true 2- Wassersttii distance.
Theorem (2-Wasserstein bound for non-commutative covariances). Let N1 Ex), (2 Σ2) be two multivariate normal distributions, where Σ, and E. are non-comee commuuattee matrices. Then, the 2- 2-Wasserstein distance between the two distributions has an upper bound of
(2(N1, & ||и #2l13 + +|=| 1
where [[()]| represents the Frobenius norm of a matrix.
We refer the reader to the appendix, section A.1 for the proof.
Significance. Deriving this bound allows us to extend the simplification for the Wassertstein dis- tance between commutative covariance matrices to the more general scenario of non-commutative matrices. The simplification allows us to directly supervise the covariance without the use of eigende- composition, making optimization inherently more stable.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1: Covariance Pseudo-Label
Input: a: E IR:: Given observation
1/b 1.24
More important
Input: X: All observations; Y: All targets
Output: (X)): Covariance pseudo-label
11 2.95
1/2 3.05
// Mahalanobis distance
Σ = Cov(2))
// d(x, X, ). shape = Asamples du( = (x - x) -) x)*
3/5 5.1
3
1/8
5.1
// k Nearest neighbours
dist, idx = bottom-k (d, Σ),
V/ 3.1
2.9
'robabillsttcc' interpret softmax (dist)
// yy for nearest neighbours
3/7
1.24
important
y y_nbr= fidx]
// 'Expected' mean, covariance
Given 2o. its pseudo-label is the variance in the targets corresponding to samples which are the nearest neighbors of øo- Samples closer to Ey given more importance than samples further away.
Figure Pseudo-Label (Section 3.3
By (dist mumn) (dim=C)
Σy (X) dist (y_nbr - 7 -
// Pseudo-label for given æ.y
return (x))
We return to Problem and visually compare the 2- W-sssesstein distance with variants of the log- likelihood (vanilla negative log-likelihood, Faithful (Stirn et al. ) and KL-Divergence. Fig. 2
we observe that in comparison to the likelihood based methods, the 2-Wasserstein is significantly
more stable since the covariance does not depend on the residual and the convergence of the mean estimator. We also study the impact of warm-up (Sluijterman et al., 2024), where the mean estimator is allowed to converge before training the covariance estimator: To do so, we directly initialize the mean of the predicted distribution to the target mean, and the covariance to identity. However, our results in appendix/Fig 9 due residuals. In contrast, the learning of the covariance in 2-Wasserstein does not depend upon the residual, leading to stable convergence.
(b) (appendix) show that these methods are still susceptible to instability
3.3 GENERATING PSRUDO-LABELS FOR THE COVARIANCE
In the absence of labels for the covariance, existing approaches rely on the residual of the mean estimator as a signal to optimize the covariance. However, optimizing in this manner trades-off accuracy with computational complexity. While having labels would allow to directly optimize the covariance estimator, obtaining annotations for the covariance is non-trivial. Therefore, we take step in this direction and explore the possibility of self-supervision for the covariance. To this end, we propose simple heuristic, which when combined with the 2- 2-ssssrsteein distance, is surprisingly effective supervising the covariance.
Intuition. The neighborhood of a sample has been widely used in uncertainty quantification
(Van Amersfoort et al. 2020 Skafte et al., 2019) and kernel methods (Hofmann et al., 2008).
The key idea is to infer properties of a sample based on its neighborhood. TIC (Shukla et al. 2024
learns the covariance through a learnt neighbboroood of the input, Cov(Ý + +c). We extend upon this idea to obtain pseudo-labels for the covariance. Specifically, we use two concepts:
1. The target y has high (co-)variance if it exhibits large variations in small vicinity of æ.
2. The closer 2, to æ; the likelier it is that yj a potential label for ac.
We quantify these concepts through the use of (a) the Mahalanobis distance, to measure the degree of closeness between samples æ, and zy; and (b) probabilistic interpretation of this distance to weight different targets yj as being a potential label for x;,
The Mahalanobis distance between two points u, with respect to covariance matrix M is
du v; Σ) V(u (u-v)TE v).
(4)

---

Published as a conference paper at ICLR 2025
Ground
MLL
Bets-NLL
Wmmallll MWl
Faithiial
1. 1.
m
Figure 4: We sample from the ground truth sinusoidal y |x| sin (2mæ) with ((x) = |a:] and train networks using different objectives. The 2- 2-Wasserstein distance trained using pseudo-labels is able to converge to the accurate mean and variance faster since it does not depend upon residuals or convergence of the mean estimator learn the variance.
Unlike the Euclidean distance, the Mahalanobis distance accounts for the spread of the samples. This is crucial since not only does the distance scale according to the alignment of u, v w.t.t the covariance, but it also scales based the spatial extent of the samples.
Therefore, we define Σ = Cov(X) to quantify the alignment and scale of all the samples X. We compute the pairwise Mahalanobis distance between the given sample and all other samples, choosing the top-k nearest neighbors and their associated distances. This also includes the given sample itself. Next, compute the softmax over these distances, giving them probabilistic interpretation: the closer the sample, the higher the likeliness of it being the true mean. The pseudo-label covariance uses this probabilistic interpretation compute the expected mean and covariance over the neighboring targets. A concise description of these steps is available in Algorithm
Self-Supervised Training. We use the pseudo-labels Σy (x) in conjunction with the input x and target y train the mean and covariance estimators simultaneously using the 2- Wasseerstein bound.
EXPERIMENTS
How effective is self-supervision in deep heteroscedastic regression? We study this question through series of synthetic and real world datasets for regression. We the same setup Shukla
et al. (2024) which experiments on univariate sinusoidals, synthetic multivariate data, UCI Machine
Learning repository (Markelle Kelly) and 2D human pose estimation (Andriluka al. 2014 Johnson
& Everingham, 2010 2011). We provide a detailed description of the experimental setup and implementation details the appendix (B. For all our experiments, we set the nearest neighbors hyperparameter in the pseudo- algorithm to ten times the dimensionality of the target. We compare our approach with popular and state-of-the-ar methods such as the vanilla negative log- likelihood, 3-NIL (Seitzer et al. 2022), Faithful heteroscedastic regression (Stirn et al.) (023). Empirical-Baan Bayes (Immer al.) (023) and the Taylor Induced Covariance parameterization (NL.E:TIC) (Shukla et al. 2024). In addition to the mean square error and the negative log-likelihood, we use the KL-Divergence and the 2- Wasserstein distance as metrics when the ground truth covariance is known. We also use the Task Agnostic Correlations (TAC) metric introduced in (Shukla et al. 2024) to evaluate the covariance through its learnt correlations. Finally, we also report the additional memory consumed and the time required for each method to run for different experiments.
4.1 SYNTHETIC DATA
Univariate. We use samples from different varying amplitude sinusoidals to compare the methods. In Fig. we observe that faithful heteroscedastic regression overestimates the covariance because of the lack of synergy between the mean and variance estimator. While the mean estimator assumes homoscedastic unit variance, the variance estimator models heteroscedasticity, Although the TIC formulation stabilizes convergence, unfortunately, convergence itself is slow. The KI. Divergence
7

---

Published as a conference paper at ICLR 2025
25.0
382
22.5
Fattha
202
21.
TC
422
TKC
215
15.0
Wantrrten (iound))
382
152
N.Wase (teuse)
32.5
12.0
107
21423
Suapa
Stapa
300
NLL
Exenta
70000
NJLL
692
KLD Darrggnncc
50000
Co Uudheraano
BL.Chesgence
Hewow den
20000
203
23
60 100 20 60 100
22
40
20
Figure 5: (Multivariatee: Metrics.) We simulate multivariate data with heteroscedastic covariance of increasing dimensionality (top row: 8, bottom row: 24). We observe that modeling heteroscedasticity challenging without annotations, with some popular approaches diverging away from the true distribution. Our results highlight the potential of self-supervision for improved convergence.
Table (Multivariate: Computational Costs). While TIC is able accurately model the covariance in comparison to other likelihood based approaches, it has a significantly increased computational cost. The Wasserstein (bound) has a significantly lower cost without sacrificing accuracy.
(a) Compute time (in milliseconds)
Dimensions
4
8
12
16
20
24
28
32
Beta-NLL, Diagonal
2.88
3.15
2.17
2.06
1.83
1.74
2.00
2.04
Faithful, NIJL
4.56
4.74
3.94
3.69
3.76
3.66
4.08
4.85
NLL: TIC
56.60
56.81
59.28
95.54
197.58
448.58
943.79
1961.08
KL-Divergence
4.79
5.06
4.05
4.05
4.10
3.94
4.81
5.24
2. Wasserstein
5.10
5,43
4.47
4.38
4.31
4.14
4,88
5.20
2- 2-Wasserstein (Bound)
4.59
4.79
3.72
3.73
3.64
3.56
3.91
4.72
(b) Compute memory (in megabytes)
Dimensions
4
12
16
20
24
28
32
NLL: TIC
11.68
120.84
523.22
1543.77
3625.51
7333.84
13313.31
22398.00
2-Wasserstein (Bound)
3.45
17.10
29.51
54.51
111.73
201.51
339.55
46 other methods
and vanilla negative log-likelihood suffer from large residuals which prevents further optimisation. In comparison, the 2- 2-Wasserstein distance combined with the lseudo-labels is able to converge faster while being accurate. An additional study comparing the methods on different variations of the sinusoidal is presented in the appendix (Fig: 10
Multivariate. Unlike in our real world experiments, synthetic datasets allow us to define the ground truth covariance to evaluate different approaches. We use the same setup as previous work, which defines the multivariate target NN Σy x + Σ2ix) as a function of the input with heteroscedastic variance. X and Y are jointly distributed following the normal distribution, with Z being variable conditionally independent of and a function of X. In addition to evaluating different methods through optimization metrics, also compare them through their computational requirements (memory and time). We vary the dimensionality our targets ranging from 4 to 32, and report results in Fig. (appendix), and Table T We observe that while TIC facilitated improved
8

---

Published as a conference paper at ICLR 2025
Table 2: UCI Regression. The 2-Wasserstein distance using pseudo-labels for supervision accurately estimates the mean and covariance while having low compute requirements. In contrast, the negative log- legelike and KL. Divergence sub-optimally converge due to large residual errors. While methods such Faithful encourage convergence by using the mean squared error, the mean and covariance inconsistently model the residual leading to sub-optimal covariance estimates. While TIC accurately models the covariance, it has high computational costs and lags in mean estimation.
(a) Mean Square Error (MSE)
Method
Abalone
Air
Appliances
Concrcte
Hlectrical
Energy
Gas
Naval
Parkinson
Power
Red Wine
White Wine
53.49
10.98
KL- XL-Divegence
1.90
14,20
90.90
1.84
15.57
4.20
10.16
12. 39
59.39
9.97
7.26
8.17
Beta-NLL
0.35
1.58
3.69
2,02
3,62
1,87
1.50
0.72
8.11
3.06
2.15
3.43
NLL: Diagonal
1.32
8,90
37.91
4.28
6.58
3.99
5.73
9.60
27.35
6.52
5.75
6.01
Fuithfal
0.20
0.06
NLL: TIC
0.21
0.82
4.45
0.96
0.91
0.61
0.67
1.36
8.89
0.66
0.97
0.92
2-W (Bound)
0.16
0.34
0.20
0.72
0.90
0.41
0.45
0.07
0.30
0.61
0.71
0.79
(h) Negative Log-Likelihood (NLL)
Method
Ahalone
Appliances
Concrete
Flestrical
Fnergy
Gas
Naval
Parkinson
Power
Red Wine
White Wine
345.99
29.85
41.03
262.59
49.37
Xi--Divereence
18.27
83.18
413.96
38.23
73.87
28.50
34.38
49. 24
257.36
45.49
61.96
48.66
Beta-NUL
9.80
29.38
60.30
20.45
35.44
20.15
20.26
20.81
59.98
27.64
34.05
30.95
NE Diagonal
369.67
46.8
65.73
36.09
47.06
238.71
77.98
Foiihful
11.86
33.31
65.15
17.42
34.73
19:41
22.47
27.70
57.04
24.08
24.34
26.01
NLL TIC
4.71
16.46
30:41
11.36
14.97
12.06
9.96
14.99
42.52
9.31
14.66
12.33
2-W (Bound)
6.32
13.58
22.72
8.96
15.57
8.85
10.49
11.44
21.48
11,31
11.65
12.12
(c) Compute time (in milliseconds)
Method
Ahalone
Air
Appliances
Concrete
Flectrical
Fnergy
Gas
Naval
Parkinsnn
Power
Red Wine
White Wine
5.85
5.53
Beta-NLL
4.71
4.58
4.98
4.35
5.20
4.41
4.62
5.8)
6.70
4.81
6,83
4.41
Faithful
5.28
5.35
5.62
5.73
6.02
5.03
5.03
7.02
6.47
5.81
7.40
5:05
NLL: Diagonal
4.50
4.19
1.81
1.67
5.13
4.32
4.38
5.81
4.98
4.61
5.77
4,36
NLL: TIC
68.55
49.37
47.57
36.04
59.25
49.74
KL KL--Divegence
5.30
6.65
6.08
5.47
8.09
5.16
5.14
6.85
9.15
5.36
6.93
5.23
2-W (Bound)
4.51
5.83
5.17
4.51
5.36
4.50
4.38
5.23
7.16
4.51
5.28
4.48
(d) Compute memory (in megabytes)
Method
Ahalone
Applionces
Concrete
Electrical
Energy
Naval
Parkinson
Power
Red Wine
White Wine
108.59
637.74
41.57
2-W (Bound)
66 other methods
3.10
9.02
25.30
1.09
7.75
4.63
5.56
9.02
23.05
5.56
5.56
5.56
covariance estimation, this resulted in slower convergence of the mean estimator. This trend evident as the dimensionality increases. Moreover, TIC requires significantly more computational resources. In contrast, the 2- Wasserstein bound is significantly cheaper to compute while maintaining accurate convergence of both, the mean and covariance estimator.
4.2
REAL DATASETS
UCI Regression. We evaluate mean and covariance predictions by performing the same study as Shukla et al.] (2024) on regression datasets from the UCI Machine Learning repository. We standardize each dataset to have zero mean and unit variance. We randomly choose 25% of the features as observations and the remaining 75% as the targets, adding considerable heteroscedasticity in the data. We conduct five trials and report the mean, highlighting top- performing methods which are statistically indistinguishable. We evaluate different methods not only using performance on various metrics (Table appendix/Table but also through computational costs. While TIC outperforms other likelihood based baselines significantly in our TAC and NLL evaluation, this comes at the cost of significantly higher computational requirements. Although compute efficient methods such as Faithful leverage the mean squared error to accurately converge to the mean, it does not accurately converge on the optimal covariance. In contrast, the 2- Wassersteii bound accurately converges in both, the mean and covariance without additional computation overhead. Moreover, the vanilla 2- Wasserstein formulation exhibited significant training instabilities on certain datasets such as superconductivity, motivating the use of the proposed bound which is stable to train.

---

Published as a conference paper at ICLR 2025
250
275
140
250
Waasamree Soordd
130
200
120
TIC
TIC
37.5
2
110
2-wecsette
Wiserror (Bound]
175
Divessetta n
Wesseror (Roand)
15-2
350
100
125
1.3
100
5pach
Fpentss
Figure 6: (Human Pose: Improving state-of-the-ar heteroscedastic pose estimation) We explore hybrid training strategy by combining the 2- 2-Wasserstein bound with the negative log- lieeliikdoood. We train ViTPose for the first 20 epochs using the bound and then switch to negative log- logliikeihood. We observe that the hybrid approach retains best of both the worlds: improved mean and covariance estimates, as measured by the mean square error and the lig-likeldoood. (aooodd(iife (Different learning rates are
explored in Fig. 13
We also study the impact of warmup Sluijterman et al. (0024) in the training process, where we train the mean estimator for half the number of epochs, and jointly train the mean and covariance for the remaining half. We share our results in appendix/Fig. T2. Noticeably, the training diverges due to the effect of residuals coupled with incorrect covariance estimates, effectively nullifyying the use of warm-up.
Human Pose Estimation. We perform experiments on 2D human pose estimation using the same setup as previous work. We use the ViTPose (Xu et al., 2022) architecture as our base model, which popular vision transformer model for human pose estimation. Since we are introducing the covariance in the training process, this also requires us to modify the covariance in response to image and keypoint augmentations. Popular augmentations use affine transformations, which linearly transform the keypoints. Let Y RY represent the transformed keypoints using the matrix R. The new covariance underlying Ý is (X) = R (X) RT.
We highlight this experiment not which motivates the use of hybrid training to further improve the state-of-the-arr Our preliminary experiments using self-supervision fell short of matching (Shukla al.. (2024), which we believe is because the pseudo-labels, computed based on a low-dimensional representation of the input images, may not necessarily be accurate. However, identify alternate paradigm: pre-training using self-supervision and then switching to NLL:TIC helped retain the advantages across both, the 2- Wasserstein bound and NLL:TIC.
where self-supervision is computationally efficient, but one
We show our results in Fig. 6 and Fig. 13 We use the TIC parameterization and train our models using the self-supervision for the first 20 epochs. After this, we switch to the negative log-likelihood which provides more freedom to explore the optimal covariance. Our results show that the hybrid approach outperforms its individual components and retains both: a low mean square error and low negative ligliikeld
CONCLUSION
We study deep heteroscedastic regression, noting the optimization challenges present due to the lack of annotations for the covariance. Therefore, we study methods for self-supervision which requires us to define (I) a framework for supervision, and (2) method to obtain pscudo-labels for the covariance. We critically study the KI.-Divergence, highlighting the need for calibration and noting its susceptibility to residuals. Next, we study the 2- W-sssersteein distance, proposing a bound on the latter that is stable to optimize. Finally, we propose a simple neighborhood based heuristic which is effective in providing pseudo-labels for the covariance. Our experiments show that, unlike existing approaches, the use of the 2- Wasserstein bound and pseudo-labed yields accurate mean and covariance estimation while remaining computationally inexpensive. Our experiments on human pose show the potential for a hybrid approach, where combining the 2- Wasserstein and NLL frameworks enables superior performance compared to using either method alone.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We make our code available on complete with a docker image and documentation for reproducibility. We have taken sufficient care to perform multiple trials and report the mean and standard deviation.
http://rgggessi....g..bbb.b.... github. 1o The code comes
ACKNOWLEDGEMENT
We thank the reviewers for their valuable comments, questions as well as participation in the rebuttal period. We also thank Reyhanch Hosseininejad for her help in preparing the paper and for having insighful discussions. Finally, we are grateful to RCP, EPFL for their support in compute.
This research is funded by the Swiss National Science Foundation (SNSF) through the project Narratives from the Long Tail: Transforming Access to Audiovisual Archives (Grant: CRSIIS 198632). The project description is available on: https: www. futureci nema live/projecct/
REFERENCES
Mykhaylo Andriluka, Leonid Pishchulin, Peter Gehler, and Bernt Schiele. 2d human pose estimation: New benchmark and state the art analysis. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2014.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings the 34th International Conference on Ma- chine Learning, volume of Proceedings of Machine Learning Reseurch, pp. 214-223, PMLR, 06- 11 Aug 2017. URL https:: htt..:/........rrre/////r//rrrra
html
Lorenzo Bertoni, Sven Kreiss, and Alexandre Alahi. Monoloco: Monocular 3d pedestrian localization and uncertainty estimation. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 6861-6871, 2019.
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, pp. 5477-5485, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL httpattt/pppeereeee// net / worumoriioooo PINNIINNIIyyYc
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059.
PMLR, 2016.
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep bayesian active learning with image data. arXiv preprint 1xXx:1703.009111 2017.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Nitesh B Gundavarapu, Divyansh Srivastava, Rahul Mitra, Abhishek Sharma, and Arjun Jain. Struc- tured aleatoric uncertainty human pose estimation. In CVPR Workshops, volume 2, 2019.
Thomas Hofmann, Bernhard Schölkopf, and Alexander Smola. Kernel methods in machine learning. 2008.
Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning. stat, 1050:24, 2011.

---

Published as a conference paper at ICLR 2025
Alexander Immer, Emanuele Palumbo, Alexander Marx, and Julia E Vogt. Effective bayesian heteroscedastic regression with deep neural networks. In Thirty Conference on Neural Information Processing Systems, 2023. URL
A6FquHDenk)
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by re- ducing internal covariate shift. pp. 448-456, 2015. URL http://pp/////..oo/////rereeee/ /
stppers/v37/ooffeet
Johnson and Mark Everingham. Clustered pose and nonlinear appearance models for human pose estimation. In Proceedings of the British Machine Vision Conference, 2010. doi: 10.144.5244/C.24.
Sam Johnson and Mark Everingham. Learning effective human pose estimation from inaccurate annotation. In Proceedings of IEEE Conference on Computer Vision and Pattern Recognition, 2011.
Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? In Advances in neural information processing systems, pp. 5574-5584, 2017.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles, Advances in neural information processing systems, 30, 2017.
Quoc V Le, Alex J Smola, and Stéphane Canu. Heteroscedastic gaussian process regression. In Proceedings of the 22nd international conference on Machine learning, pp. 489-496, 2005.
Cheuk Ting Li, Jingwei Zhang, and Farzan Farnia. On convergence in wasserstein distance and I-divergence minimization problems. In International Conference on Artificial Intelligence and Statistics, pp. 2062-2070. PMLR, 2024.
Jiefeng Li, Tong Chen, Ruiqi Shi, Yujing Lou, Yong-Lu Li, and Cewu Lu. Localization with Advances in Neural Information Processing Systems, 34:27236-2777 272488, 2021a.
sampling argmax. aarglix.-argmax.
Jiefeng Li, Chao Xu, Zhicun Chen, Siyuan Bian, Lixin Yang, and Cewu Lu. Hybrik: hybrid analytical-neuree inverse kinematics solution for 3d human pose and shape estimation. In Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 3383-3393, 2021b.
Katherine Liu, Kyel Ok, William Vega-Brown and Nicholas Roy. Deep inference for covariance estimation: Learning gaussian noise models for state estimation. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 1436-1443, 2018. doi: 10.1109//CRA.201
8461047.
Changsheng and Piotr Koniusz. Few-shot keypoint detection with uncertainty learning for unseen species. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 19416-19426, 2022.
Kolby Nottingham Markelle Kelly, Rachel Longjohn. Uci machine learning repository. URL http://archive. ics www..ed////
Krishna Kanth Nakka and Mathieu Salzmann. Understanding pose and appearance disentanglement in human pose estimation. arXiv preprint arXiv:2309.1166) 2023.
David A Nix and Andreas S Weigend. Estimating the mean and variance of the target probability distribution. Proceedings of 1994 ieee international conference on neural networks (ICNN"9H) 94), volume 1, pp. 55-60. IEFE, 1994.
K. B. Petersen and M. S. Pedersen. The matrix cookbook, nov 2012. URL http://wwZ compute.dtu /twpdd wtt.:////pbbpppp///////// 33744 html Version 20121115.
PyTorch. torch.linalg.eighff. Oct 2024. URL
generated/torch.. 1ina 1.. el high. html
12

---

Published as a conference paper at ICLR 2025
Rebecca L Russell and Christopher Reale. Multivariate uncertainty in deep learning. IEEE Transac- tions on Neural Networks and Learning Systems, 33(12):7937-7944 2021.
Maximilian Seitzer, Arash Tavakoli, Dimitrije Antic, and Georg Martius. On the pitfalls of het- eroscedastic uncertainty estimation with probabílistic neural networks. In International Confer- ence on Learning Representations, 2022. URL bttt:ss///penrevee...o net www.:nee/foummmmmm di
aPOpXinV1T
Megh Shukla, Mathieu Salzmann, and Alexandre Alahi. TIC-TAC: framework for improved covariance estimation in deep heteroscedastic regression. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 45244 45257. PMLR, 21-27 Jul 2024. URL https:/
proceedi ngs m press v235/ Shuk1 a24a. ht.ml
Ivor JA Simpson, Sara Vicente, and Neill DF Campbell. Learning structured gaussians approximate deep ensembles. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 366-374, 2022.
Nicki Skafte, Martin Jørgensen, and Søren Hauberg. Reliable training and estimation of variance networks. Advances in Neural Information Processing Systems, 32, 2019.
Laurens Sluijterman, Eric Cator, and Tom Heskes. Optimal training of mean variance estimation neural networks. Neurocomputiing, pp. 127929, 2024.
Joram Soch. Kullback-leibler diverggne for the multivariate normal distribution, May 2020. URL httpp:///saapppoffoooo....ss....
html
Joram Soch, Relationship between correlation coefficcenn and slope estimate simple linear regres- sion, Oct 2021. URI. /tt.ntttss//tttttroooooooo. wwt.u....u..ooo...oo http:/ooo///soooooo httm
Andrew Stir, 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.
Bugra Tekin, Pablo Marquez-Neila, Mathieu Salzmann, and Pascal Fua. Learning to fuse 2d and 3d image cues for monocular body pose estimation. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Oct 2017.
Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. Uncertainty estimation using a single deep deterministic neural network. In International Conference on Machine Learning, pp. 9690-9700. PMLR, 2020.
Eliot Wong-Toi, Alex James Boyd, Vincent Fortuin, and Stephan Mandt. Understanding pathologies of deep heteroskedastic regression. The 40th Conference on Uncertainty in Artificial Intelligence, 2024. [RRLittts///opeeeeeeee
net, wtt.:e/wwwwe.ffoumtmmmmem ftw.:mmmdnnSSSaarsr.......
Yufei Xu, Jing Zhang, Qiming Zhang, and Dacheng Tao. Vitpose: Simple vision transformer baselines for human pose estimation. Advances in Neural Information Processing Systems, 35:38571-38588 2022.
Yufeng Zhang, Jialu Pan, Li Ken Li, Wanwei Liu, Zhenbang Chen, Xinwang Liu, and Ji Wang. On the properties of kullback-leibler divergence between multivariate gaussian distributions. Advances in Neural Information Processing Systems, 36, 2024.
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PROOF OF THEOREM 1
Proof. We hegin with the definition of the 2- 2-Wasserstein distance (Definition rewriting 21 + 2g by adding (and subtracting) new terms to get
First, focus
E - - - > + u + Σ2 ///2
The advantage of introducing new terms is to write Σ1 Σg as
Σ Σ₂ (E) Σ ) Σ + / > + /
Substituting this the definition of the 2- Wasserstein distance, we get
W2(Ni.N2) = 11 - -22l2 + Tr (x) - -
+ Σ 2
+ -Σ - 2(22)/
(5)
We proceed by noting that the Trace operator is linear; implying Tr(A+ + B) = Tr(A) + (B) allowing us to analyse the terms separately. Next, we note that the Frobenius norm of a matrix related to its trace by: ||A| 3/- = Tr(AA²). Therefore,
Tr [=) /2 -1/ (E) - ///3)7] = ||=|// - IE
(6)
Since the Trace operator is cyclic; implying Tr( Tr(BA), we have
Tr(E)/ -2 + Σグ ||/ =2
Substituting Eqs. 6 and into Eq. we have
W2(Ni, || 22ll?
||
/ || 2Tr
///202//
We note that in the trivial case where Σ₁ and Eg commutative, the trace is reduced to However, what happens in the general case when the covariance matrices are not commutative? We address this through proposition]] [ which shows that
///3 2 Tr(E//2/11
Therefore, substitution the trace terms cancel out, leading familiar expression wrapped in an inequality:
W2(AG, < || uell
|||// 1
(9)
L
Proposition 1. Let A, B be any two positive definite matrices not necessarily commutative. Then, B= 2 /1/22
> Tr(X). Let X = PU be the polar decomposition X where Pss a positive definite matrix and U is a orthogonal matrix since X is a real matrix. On substitution, we have TF(((P))PU)))) by definition is symmetric and is orthogonal, we need to show that Tr [r[P] 2 Tr(PU).
Proof. Let X = A1/2 consequently we need to prove that Tr
2 Tr(PU). Since
Let = QAQT be the eigendecomposition of where Q is the orthonormal basis and A is the diagonal matrix consisting of the eigenvalues X, Since the trace is the sum of eigenvalues of a matrix, we need to show that E, A > Tr(QAQTU). Since the trace is cyclic, we have E, A 2 TU(AQTUQ) = Tr(AK),, where IC = QTUQ is an orthogonal matrix by virtue of being product of orthogonal matrices. Since A is diagonal matrix, Tr(AK) Eg AcKa. Moreover, as IC is an orthogonal matrix, Icx <= 1. As a consequence, Σ. A 2 Σ, ArKis concluding our proof.
14

---

Published as a conference paper at ICLR 2025
Method: Negative Loglikelihoad
Method: Divergence
Metric: Epech
Metric: 2-W vis Epoch
NLL
Invers Cyvveiance
KL
Ground Truth
Ground Truth
Prediction
Prediction
Sample Initialization
Samplc
Intiaaiization
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 aligned orthogonal to the residual vector. Since the gradient of the mean estimator is directly proportional to the inverse (Eq. [T). the gradient as whole is desensitised to move towards the target. In fact, the inverse magnifies gradient updates in the direction orthogonal the residual vector, potentially leading to oscillations.
Figure 7: Residuals and sub-optimal convergence (Section: 3.7
A.2 PROOF OF LEMMA II
Proof. The optimal solution for Ey (X) involves minimizing
Σ (N E Exrrio (X)) ([ Ex(x))).
Using the definition of the KI. Divergence (Eq. and dropping the non-parametric terms, we get
1 N 3 (x){"" + (fx (X) y:) Txp (X)((x()) v.) In 1x]2)((x]]] (10)
Setting the derivative w.n. the predicted mean 7y (X) to O, we get
N
Σ
(X) - yi)= =
(11)
1=1
N N y
Ay (X)
(12)
This is indeed the same solution minimizing the negative log-likelihood, and therefore jay (X)
predicts the correct mean. However, setting the derivative w.r/ the predicted precision (X) gives us three terms of the form (1) Tr(AB), (2) 6 Ab and (3) In [A], where A is the shorthand for precision and B represented different terms. The derivative of the form Tr( (r(AB) w.r. A is Bd (Eq. 100 in
Petterse & Pedersen ) Here, B is the prior term >xpriar)
(X). Since the prior is symmetric,
the derivative of term is My Exprior) (X). . The derivative of the form Dt Ab is 6bT (Eq. 72 in PPteesen
& Pedersen 2012)). Therefore, the derivative of term 2 is (X) - y: (2) (X) ys) Finally,
the derivative of the form In |A| is (A-1)" (Eq. 57 in (Petersen & Pedersen 2012 Therefore, the
derivate of term is (X)). By combining the three terms, the derivative of Eq.
N
1
((prion) (X) + (X) - (i)( - y¹) - [x =
N
1=1
N
Ey(X)=
Expmion (X) +
1
Σ
(Tr - v.) -1)
(13)
N
-=1
Ex(X) x 2 xmfrrorr) (X) + Σy (X)
(14)
Note: In comparison, the optimal value of the covariance using the negative log-likelihood is Ex(X)= y/2 (fr v:)(T(()) Dr
15

---

Published as a conference paper at ICLR 2025
B EXPERIMENT DETAILS
Training. We use separate networks to estimate the mean and covariance, with no overlapping parameters, following the results of Stirn et al. (2023). This is also advocated for by [Slujjeerman al. (2024). The architectures of these network are described in the different subsections. The mean and covariance have the same architecture with the exception of the final layer. We use the pseudo- labels (z) together with the input and target to train the mean and covariance estimators simultaneously using the 2-Wasserstein bound. Hence, instead of having training pairs (æ. y). we have triplets (x, y, (x)). Unless specified, we do not use warm-up in our experiments. Moreover, the bound does not require warm-up since the mean and covariance estimator training is decoupled. All the methods are trained using the Adam W optimizer, which implicitly imposes a weight decay of 0.01 on the parameters. We ensure a fair comparison by randomly initializing all methods with the same mean and covariance estimators, and each method uses its own learning rate scheduler. Additionally, the batching and sample ordering are the exact same across all methods. At training time, we sample a batch which is simultaneously used by all the baselines for optimization.
B.1 SYNTHETIC DATA
Univariate. We draw 50,000 samples for each of the three different sinusoidal distributions: (1) |x]] sin (2Tx) (2) (5 - (x|) sin (2Tx) (3) 5 sin (2πr), all of them with heteroscedastic noise (2)) Φ (2). We train a fully connected feed-forward neural network with batch normalization (Ioffe & Szegedy 2015) and tanh() activation to learn the mean and variance, Specifically, we use four hidden layers with a latent dimension of fifty. Every alternate layer is followed by the batch normalization layers. The full results are shown in Fig. TO
Multivariate. The dimensionality of the input and target, x and q is varied from 4 to 32 in steps of 4, and the mean and standard deviation are reported over ten trials for each dimension. Depending on the dimensionality, between 4000 and 20000 samples are drawn. Similar to our univariate setup, we train fully connected feed-forward neural network with batch normalization but ELU() activation to learn the mean and covariance. Specifically, we use ten hidden layers with a latent dimension that is the dimensionality of the input squared. The idea is that the size of the network increases the dimensionality of the network increases to account for increasing complexity. Every alternate layer is followed by the batch normalization layers. The full results are shown in Figure [1] We report the computational requirements in Table I
B.2 UCI REGRESSION
We follow the experimental setup in Shukla et al. (024).. For each of the twelve datasets, 25% of the features are randomly selected as inputs, with the remaining 75% used as multivariate targets during run-time. Although some of the resulting input-target pairings may yield sub-optimal performance in prediction, this presents a valuable test for the covariance estimator, which needs to identify correlations even in challenging scenarios. Moreover, the random assignment of features guarantees that our experiments are unbiased, as the selection process is not manipulated. All datasets are standardized mean of zero and variance of one. We reuse our neural network architecture from the multivariate experiments. We perform five trials for each dataset and report the mean and standard deviation in 3
B.3 2D HUMAN POSE ESTIMATION
ViTPose (Xu et al. 2022) is a recent state-of-the-an model that adapts vision transformers (Doso vitskiy et al. 2022 for the task of human pose estimation. We use the base version (ViTPose-B) for our task Additionally we use soft-argmax (Li et al. 2021b a) which is applied to reduce the heatmap, initially tensor of shape N x 64 x 64, to a ID vector of length 2.V, where is the number of joints in the human pose. To obtain the input for the covariance estimator, we use residual connections which involves downscaling and upscaling of the I-D features predicted by the backbone network. The output of the downscaling is used to predict the covariance. We perform our experi- ments on the MPI (Andriluka et al.) 2014) and LSP/LSPET (Johnson & Everingham 2010 201T) datasets, with the latter focusing poses related sports. We merge the MPI1 and LSP-L.PET
16

---

Published as a conference paper at ICLR 2025
datasets to increase the sample size. The pose estimator is trained using the Adam optimizer with a 'ReduceL_ROnPlaeeuu learning rate scheduler for 100 epochs, with the learning rate set to 1e-3. Two augmentations, Shift-Scalc+Rotate and horizontal flip, are applied. For details the specific implementation, readers are referred to the code.
c COMPILATION OF RESULTS
Method: Negalive Loglikelihoad
Method: Divergence
Epoch
NLL
NLL
KL
Ground Truth
Ground Truth
Prediction
Sample
Samp
Initialization
Iritialization
(a) We observe that the KL Divergence can act a regularizer over the learnt covariance, thereby stabilizing optimization. However, the covariance for both the methods is dominated by the residual term, slowing down convergence.
Method: Negative Loglikelihood
Method: KL Divergence
Metric: KL vis Epoch
Metric: 2-W v/s Epoch
KL
Ground Truth
Ground Truth
Prediction
Prediction
Sample
Intiialzaaiion
Intialization
(b) At higher learning rates (le- I), both the Negative Log-Likelihood and the KL Divergence oscillate around the true distribution resulting in unstable optimization.
Figure 8: Bivariate Normal Distribution (A) Impact of residuals optimization (Section: 3./) In addition to feature granularity (Seitzer et al. 2022), we show that source for subpar convergence arises from the susceptibility of the negative log-likelihood and KL-Divergence residuals in optimization.
17

---

Published as a conference paper at ICLR 2025
Negative Loglieeihood
Partiuu
KL Divergence
2-Wasserstein stance
2-Waasserteein (Bound)
Groune
Fattrnl
nattrul
KL
24W
24Y
24 (Boun
2-W (Boundd)
N (ouudd)
130
1800
21.0
20.5
(a) We observe that the KL- Diverreenee and likelihood based methods: vanilla negative log log-likelihood and faithful (Strrn et al. 2023) result in unstable convergence. In comparison, the 2- 2-Wassrsstin based methods are much more stable and accurate since they not affected by residuals nor is the covariance affected by the convergence the estimator.
Negathie Inglttethood
Twithld
Ki Diveggnnce
2-4lesserttiin stance
2-4lesserdeen (loand)
Sample
Sample
Sample
4214
NLL
Fultrul
Nattril
KL
2447
2-17
2
2-41 (Round)
(Bound)
(Round)
00.0
13.0
45.0
30.
20.1
(b) we initialize the predicted mean to the true mean, we still observe unstable convergence for the divergence and likelihood based methods due to perturbations caused by residuals (Section: 3.1
Figure Bivariate Normal Distribution (B) Visualizing convergence bivariate regression (Section: 3.2) We perform analysis on two settings (a) the mean and covariance of the predicted distribution are initialized away from the target (b) the mean the predicted distribution is initialized at the mean of the target.
18

---

Published as a conference paper at ICLR 2025
Ground
MLL
Bets-NLL
M MWl
Falthiia
1 14
m
Natural Laplace
NLL
TOC
Diveggnnc
Wabse Steci
Ground Truth
Beto NLL
Falthrull
Natural Laplaco
NLL TIC
Divegenne
Wisse sttein
Figure 10: Univariate. We define the ground truth sinusoidals as (Top) y |ac| sin (2mw) with ( |o| (Middle) y (5 (ær) sin (2Tι) and (zx = |x|| (Bottom) y = sin (2ww) and
( = [x]. Given samples from the ground truth, the networks are trained to learn the underlying distribution using different objectives.
19

---

Published as a conference paper at ICLR 2025
Battfu Mil
HLL
Th
302
FL-C RUChorgenca
HLL: TIC
2-Waro sdaln
Mis Wiissesion Iteurd
10
202
150
100
M8L
602
MLL
502
TIC
&-Oraagene
402
Messee stein (heune)
RLL IIC
52 302
202
102
Stapa
Staza
Staps
175
NLLL
Eaenta
20060
802
702
NLL
KL Drvrrgnn
CL Coyeeerance
KL, KL.Chrsgencs
125-
120
L0660
202
25-
105
60
02
300
20
00
00
1000
22
40
30
000
300
Fatnty
Fulbeul
230
MLC TI
60600
HLL: TIC
BL.Chaagence
220
Hosss (Bard))
602
2-Wasv Iteurd)
40000
350
2060
402
220
202
10000
Stapa
Stupa
NLL
341033
Fabtaul
N.L-TIC
Severgance
220000
Wesoemoe dboued
500
250
200
370-
40000
20030
60
100
10
550
M1
3.00
2000
N-A
420-
200030
RLL
200
200003
400
320-
30030
200
Figure 11: (Multivariate). We simulate multivariate data with increasing dimensionality (top row: 12, bottom row: 32, middle rows: increment of four ). An increase in dimensionality causes the mean estimator converge slow (or diverge) for likelihood and KL-DDvergence. The 2- bound is succesfully able to learn the mean and covariance across all dimensions.
20

---

Published as a conference paper at ICLR 2025
Table 3: UCT Regression. Results Across Different Metries: Mean Square Error (MSE), Task
Agnostic Correlations (TAC), and Negative Log-Likelihood (NLL) with standard deviations.
58.06 = 16.90
16.96
2.799
67.21 # 16.24
1.33
12.33 1,68
1.68
12.12
48.66
30.55
Wine
46,58 + 10.26
+ 24.34
17661
14.66 = 1.24
While Wine
9.40
1.97
+ 0/8
0.12
0.79 0.05
3 1
Red
61.96
77.98
11.65
6.01
0.78
0.92
Wine
4.8% 1.06
1.17 0:29
1.83 0,20
1.1 0:22
0.67 0011
Wixs
4.53
0.64
0.05
0.34
0.04
=:
21.28
11.37
4.6
51.60 = 3.45
9.31 9.31 0.92
11.31 0.39
Power
Rod
8.9
7.26
215
5.75
0.70
0.97
0.7
3 1.10
015
0.33
1.55 0014
49.37
45,49
27.64
1.061:033
0.
0.49 =
Power
9.09 2.67
+
3,06 3: 1:10
0.10
0
0.61 1 0.10
4.32
100
1,66
0.79
+ 73.52
257 36 92.79
238 4 38
42.52 1 7.52
21.48 0.63
6.52
0,61
0,66
4.2 1 1.22
188 016
1,90 0:39
161 0000
0.51 0.03
1500=0.03
262.59
59.88
Parkinson
54.51 = 23.94
8.1 = 8.25
0.20 0.08
6.46
0..00 0.08
4.49
0.50
=
38.18 3:001
49.24 .4.44 - 23.79
2.21
77.47 = 45.1
14.99 2 3.43
8,59
14.07 3.16
=: 0:46
3,53 1: 0.36
0.61 = 0.10
+ 0,03
=
11.4
10.34 5.07
0.5.5
0.05
1044
0.04
1.34
1.54
Navil
10.34
0.72
9.00
0.06
1.36
0.00
1 1.17
012
106
0,25 0004
41.03 + 21
34.55 34.38 + 96
1.43
1: 13 98
9,96 1. 2.73
10.49 2:
0.12 =
Negottive Log-1 Likellhood (NLL)
20 10.26
47.06
8.35
1.50 1. 0.33
0.13
0.37
0.13
3.36
041
2.88
0.37
(MSE)
Gss
10.98
00.8188.
10.16
5,73
0.45
0.67
0.45
(TAC)
1: 69
0.06
0.20
0.13
= 0.10
0.35 0.02
Energy
89 = 8.99
+ 4.81
5.27
= 8:57
12.06 +: 1.44
3.85
0.79
1.78
1.16
0.43
28 50
20.15
36.00
19.41
8.N5
1.20 1.75
0.99
0:18
0.14
0.18
Energy
Mean Square
1,87
3.99
0.41
0,61
0.41
0.78
= 20
0.16
0.07
0.59 0.10
0.02
Electrical
63.15 15.57
73.87 + 23.77
63.73 + 10.47
1 0,80
3
0.98
1.45
121
0.36
Eleetrical
9 3.65
+ 11.22
3,62 H 2.63
0.07
01.07
0.18
Task Agnoctic
1497
15.57
Electrial
5.24 1 1.41
1.31::03 0.32
1: 0.17
2:30 0.07
+ 0.04
001
6,58
0
0,91
0.99
28.50 6.66
23 + 13.1
4.74
13.39
11.36 : 1.50
1.31
1.94
2.30
0.70
0.67
Concreto
4.57 = 1.07
- 0.57
0.08
0.25
0.08
20.55
4688
8.96
Cone
0.65
= 0 14
0.29
0.3
033
= 0.16
0.51 0.02
=
2.02
1.28
0.72
0.36
0.72
2.94
100
1.50
1.48
0,70
2 73:35
413 + 151.56
1.67
1 88.48
1041:1334
13:42
22.72 0LJS
53,49 = 24.10
= 20.79
Applianees
1 93
0:30- 0.03
3,78
0.20 L 0.03
15.71 8:4 4.54
0015
0.33
151 0017
044
245.99
60.30
369.67
3044
3.69
37,91
4.45
105
4.09
9.23
15.55
8.46
16.46 = 0.89
5.23
14700 3.63
0.81
0.03
0.37
0.83
5.31 1:31
0
0.15
1:31 0.17
0.04
Air
2.15
0.34 =
56.98
63.18
80.30
13.58
17.92
1,58
9.0
0.33
0.82
0.34
7.07
088
2.15
0.53
Ahalooe
3.74 1.65
1
0.35 06 0.3
0.16 = 0.03
0.04
0.16 = 0.03
1.26
0:40 = 0.08
0.37 0,08
0.54 = 0.07
0.02
Abalonc
28.55
4.74
1.17
3/6
3.96
4.71 1: 1.19
0.21
3.03
0.29
0.24 =
35.89
18.27
9,80
18.61
6,32
KL Diverreeeee
Diagonal
2-W (Bonnd)
NLL: Diagonal
2W (Band))
KT U...Da Divergence
NLL Dingonal
2AW (Bound)
Meshod
Beta NL.
Faithful
TIC
Bein-NLL
Fantiful
TIC
Method
Beta NL.
Faithful
NL.: TIC
NLL
NI N...:
NLL.
NLJ
NLL
21

---

Published as a conference paper at ICLR 2025
60
NLL
NLL
Faithful
NLL: TIC
TIE
KL-Divergence
KL Diverrencee
-- 2-Wasserstein (Bound)
125
2-Masssarttiin (Bound)
40
75-
50
25
20
Epochs
Epochs
(a) Abalone
(b) Air
260
NLL
700
NLL
NLL: TIC
600
NLL: TIC
Ki-Divergence
KL-Divergence
120
2 Wesserstein (Bound)
500
Wesserstein (Bound)
100
400
300
60-
200
100
10 Epochs
20
25
0
10
15
20
25
Epochs
(c) Electrical
(d) Parkinson
100
NLL
120
NLL
80-
NLL: TIC
100
NLL: TIC
KL- Divergenee
KL- GL-Divergence
7WWascerstein (Bound)
M.ssstrrttiin (Bound)
80
NLL
60
40
20
10
20
25
10
15
20
Epachs
Epochs
(c) Red Wine
(f) White Wine
Figure 12: (Warm-up, UCT) We explore training deep heteroscedastic regression models using warm-up proposed in Sluijterman et al.) (2024). We train only the mean estimator for half the training epochs, and jointly train the mean and covariance estimator for the remaining half. We observe trend across datasets that the training is unstable and momentarily diverges for the negative log- reeliikl and KL-Divergence which are especially sensitive to residuals and incorrect covariance estimates. This trend is similar to our observations for the methods on our bivariate normal distribution experiments (Fig. 9). While we plot for six datasets here, this trend is representative of all datasets.
22

---

Published as a conference paper at ICLR 2025
ML TIC
TCC
MLL: TC
2-Wartensten
enceererr (Bannd)
210esuentein (oound
24lase rasn (Round
100
100
520
420
140
300
120-
100
320
Aparts
150
140
275
22.
MLL
230
130
20.0
225
170
ML IIC
300
Wissersian
100
125
00.5
320
30
10
15
20
25
32
25
Epoc7a
Spocha
Epec
140
5030
TIC Wissersteen (teund) ML
MIL
20.0
MLL: TCC
1.00
Misierser (Boand ML
12.5
2033
100
7.5
20
15
33
35
60
40
5
00 15 20
25
30
35
10
3
10
15
20
25
32
25
Figure 13: Human Pose. (Improving state-of-the-art heteroscedastic pose estimation) (Top row:
learning rate: le-2, Middle row: learning rate: 1c-3, Bottom row: learning rate: Ic-4) We explore hybrid training strategy by combining the 2- 2-Wasserstein bound with the negative log-likelihood. We train ViTPose for the first 20 epochs using the bound, and then switch to negative log-likelihood. We use the TIC parameterization for the covariance which when trained with the negative log- likelihoodd showed state-of-the-art performance in heteroscedastic pose estimation. We observe that using the hybrid approach retains the competitiveness of the 2- Wassersttei bound on the mean square error, and the competitiveness of the negative ligliiillhood the negative log- likelihoddoo metric.
23


Reviews:
[{"review_id": "HG0J0ZxTaz", "reviewer": "Reviewer_VkG5", "summary": "Deep heteroscedastic regression models provide both a mean and covariance estimate for each input (typically assuming a Gaussian likelihood). These models are difficult to fit due to overfitting, which can lead to (co)variance estimates collapsing to zero. While recent work has primarily focused on how to perform stable training of these models through adaptations to the objective or tinkering with the training process this paper suggests using pseudo labels to learn the covariance. The authors analyze fitting the covariance with the KL divergence and the 2-Wasserstein distance while using pseudo labels that come from a local heuristic algorithm. To optimize against the 2-Wasserstein distance they derive an upper bound. On synthetic and real-world datasets this method was found to perform well in comparison to recent baselines.", "questions": "- This method of using local information for the variance pseudo labels seems similar in spirit to the local mini-batching in Skafte 2019. Can you comment on how these differ?\n- Does this relate to kernel methods?\n- What sort of architectures can this be used for? Would the mean and covariance networks share some (all up to the final layer) parameters?\n- Could pseudo labels also be incorporated into other existing methods for fitting heteroscedastic regression models?\n- Took a few reads a little bit confused on how training works. Is there a need for any warmup period for the mean estimate (as with other methods)? Are mean and covariance learned together from the start? Is any information for the mean relevant in the presence of pseudo labels? \n- When training against the negative log likelihood, though fits may be imperfect at least the mean and covariance will be coherent together. Are the situations where the covariance conditional on the mean is nonsensical?\n- Do the usage of pseudo labels still make sense/work well when the mean model does not fit well to the data (overfit/underfit)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Different approach from much of the recent literature in the space\n- Good balance of theoretical and empirical support/motivation for this method\n- Will provide resources to reproduce results", "weaknesses": "- Related works is nearly identical to [Shukla 2024](https://openreview.net/pdf?id=zdNTiTs5gU)\n- Missing some recent literature: [Optimal training of Mean Variance Estimation neural networks](https://www.sciencedirect.com/science/article/pii/S0925231224007008), [Understanding Pathologies of Deep Heteroskedastic Regression](https://openreview.net/pdf?id=n5faLvrsA0)\n- Somewhat unclear how this methodology is implemented\n- See questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7vU7qsdBPS", "reviewer": "Reviewer_6hor", "summary": "The paper deals with heteroscedastic regression, and proposes to derive a signal from the neighborhood of a given observation to supervise the covariance. The authors derive pseudo labels for the covariance by looking at the $k$-NNs to a given observation wrt the Malanobis distance. The distance is interpreted in a probabilistic sense and used to compute the expected mean and covariance over the neighboring targets. To integrate the pseudo labels for the covariance in the loss, the authors propose an upper-bound over the 2-Wasserstein distance to be optimized, where the observed target value for each observation and the pseudo label for the covariance allow supervision of both mean and covariance.", "questions": "- Format of citations is often wrong (missing parenthesis). see lines 152,153. \n- Why in e.g. the UCI benchmark the comparison in terms of NLL is relegated to the Appendix, while the MSE results are in the main text? Especially in heteroscedastic regression NLL can be much indicative of overall performance.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper deals with the relevant problem of reaching accurate and efficient estimation of both the covariance and mean parameters in deep heteroscedastic regression\n- The idea of relying on the neighbourhood of a given observation to obtain a signal to use as supervision for the covariance is sensible and, while this idea was already introduced in [1], the authors propose a more computationally efficient solution. \n- The empirical results show that the method improves over alternative baselines. \n- The paper is well-structured and easy to follow. \n\n[1] Shukla et al. TIC-TAC: A Framework For Improved Covariance Estimation In Deep Heteroscedastic Regression, ICML 2024.", "weaknesses": "- Claim 1, how it is formulated, is confusing to me, and it does not seem like a well-posed setup. In particular, if the true covariance is assumed to be known, it does not appear to make sense to estimate it. Also, if there's a proof, I'd suggest to use \"Lemma\" instead of \"Claim\".\n- The comparison on UCI regression is in my opinion not exhaustive. Why is the method from [4] not reported on the UCI benchmark despite being SOTA? Also, from the values that the authors obtain for e.g. NLL for UCI regression, which are very different from previous works (e.g. [2, 3, 4]) but much more similar to [1], it looks like they did a similar adaptation to [1], where the authors \"adapt the datasets for covariance estimation\". This appears to be the case, also given the authors mention that they rescale the data to have variance of ten. Can the authors report the results on the original UCI regression dataset, as used in a large number of previous works?  This allows for transparent comparisons with a large number of previous methods. \n- The authors advertise the method as computationally efficient, e.g. compared to [1]. However, do the authors grid-search hyperparameters, like weight decay? I think it would help to make this explicit and compare on overall compute time with other state-of-the-art methods like [4] , where grid-search is not needed due to automatic regularization due to Empirical Bayes.\n- Do the authors try different values of $\\beta$ for the $\\beta$-NLL objective? Or which value is chosen?  \n\n[1] Shukla et al. TIC-TAC: A Framework For Improved Covariance Estimation In Deep Heteroscedastic Regression, ICML 2024. [2] Gal et al. Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning, ICML 2015. [3] Seitzer et al. On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks, ICLR 2022. [4] Immer et al. Effective Bayesian Heteroscedastic Regression with Deep Neural Networks, NeurIPS 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WmabSbXGv4", "reviewer": "Reviewer_5LBD", "summary": "This work studies the problem of estimating the covariance of a heteroscedastic random variable, i.e., a random variable $y$ whose covariance is a function of an associated covariate $x$. (in contrast to the usual homoscedastic case, where the variance of $y$ is independent of $x$). Existing approaches tackle this problem as an unsupervised problem that minimizes a KL-like loss, which tend to be unstable or slow. In this work, the authors study how to better supervise this problem, and how to obtain sensible pseudo labels to solve a self-supervised problem instead. After proposing a simple-to-optimize bound based on the 2-Wasserstein distance, and a way of collecting pseudo-labels by assuming x-continuity of the variance, the authors show that their approach is able to model heteroscedastic cases well and significantly quicker than previous methods.", "questions": "- **Q1.** I am struggling to understand why the W2-bound is better. Are you parametrizing $\\Sigma_1^{1/2}$ directly, so you do not need to compute that square root?\n- **Q2.** What is that memory consumption exactly referring to? Those values look quite large.\n- **Q3.** Where are the instabilities mentioned in L482? I couldn't find them in the results.\n- **Q4.** What is the TIC parametrization exactly? \n- **Q5.** The claim in the lines 210-212 is not clear to me from the results. Could you clarify it?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **S1.** Properly modelling heteroscedastic random variables is an important and often overlooked problem.\n- **S2.** The paper is generally well written, and contains enough explanations and intuitions to help the reader.\n- **S3.** The experiments are convincing and show the points the authors were trying to make.\n- **S4.** The arguments and derivations that precede the proposed methods are clear and convincing.", "weaknesses": "- **W1.** The particularities specific to the heteroscedastic setting should be better stressed. For example, I understand that replacing $\\mu_p$ by $y_i$ in the claim 1 is done since the variance $\\Sigma_y$ depends on the sample $i$ only, but this is not clear.\n  - I would also double-check and explicitly write the derivations of Eq. 3, as I am not sure that they are correct.\n- **W2.** The derivations are limited to the case of Gaussian predictions. A common assumption, but not clear from the abstract.\n- **W3.** The experiments lack standard deviations, yet they were repeated 5 times each.\n- **W4.** It is strange that for the last experiment the authors proposed a hybrid approach that combines 3 different methods. \n   - It would be at least nice to see what happens if the TIC parametrization is dropped. \n   - And, since speed and low-memory were a selling point of the proposed method, it would be necessary to see the penalty in time/memory that the TIC parametrization incurs. \n   - It would only be fair to compare with TIC-TAC in this experiment too, as the authors are using half of their contribution.\n   - Finally, I would love to see the proposed method applied here with the pseudo-labels taken from the input images themselves.\n- **W5.** Some figures are hard to read. I would encourage the authors to consider using a log-scale for the y-axis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BONldUmxyK", "reviewer": "Reviewer_ZXre", "summary": "The paper under review is concerned with the problem of Deep Heteroscedastic Regression in the case where the variance is known. In the classical problem, given samples from the joint distribution of $(X,Y)$, one aims to estimate the mean and covariance matrix of the conditional distribution $Y|X$ as a function of $x$ using maximum likelihood estimation. The authors consider the case where a reference distribution is known and consider two distance measures, the KL divergence and the 2-Wasserstein metric between the distribution to be optimized (given by mean and covariance) and a distribution with known covariance. The authors then also consider the problem of how to obtain pseudo-labels in the absence of a ground truth. The paper concludes with extensive experiments, suggesting that using the 2-Wasserstein distance coubpled with preudo-label generation results in a competitive performance.", "questions": "* Could you comment on how your approach relates to the literature on normalizing flows?\n* Could you specify how the KL-divernce minimization problem is formulated? Since the KL divergence is not symmetric, it makes a difference which distribution is used as first argument and which as second. You use the term \"forward KL divergence\" but I don't see it defined.\n* Use consistent notation for the KL divernce (on page 3 you write KL, while on page 4  you write $D_{\\mathrm{KL}}$)\n* In the proof of Claim 1, can you check whether (2) is correct or whether the $\\ln(\\cdot)$ term should be outside of the square brackets? My calculation in the univariate special case indicates that the logarithmic term should not have a factor $1/2$. This also makes the factor $2$ disappear in the conclusion.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is generally well written, and apart from a few minor mistakes indicated below, the mathematical derivations are correct and sound. The graphics are informative and well readable. The results on synthetic and real datasets look convincing, though in the absence of the code (which is due to be published) I can't verify everything.", "weaknesses": "The problem formulation is not clear. The authors could start by 1) specifying the terminology they use. For example, what is the 'target'? What is the data given and what do we want to get from the data? What are the precise assumptions on the distribution? The preliminaries section in Seitzer et al. 2022 is a good example on how to set the stage (incidentally, you could call the mean $\\mu$, which seems more common). There are a few typos and some minor mistakes in the mathematical derivations (see questions below). The appendix on experiments seems to have been written a bit in a hurry (there are some typos, like capitalizing 'univariate').", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)