TIC-TAC: A Framework for Improved Covariance Estimation
in Deep Heteroscedastic Regression

For a continuously distributed random variable $X$, the probability of exactly observing $p(X\!=\!x)$ is zero. Instead, the standard approach (for example Sec. 2.4 in (Evans & Rosenthal, 2004)) is to observe $X$ in the neighborhood of $x$: $X\in[x-\delta,x+\delta]$. The definition of this neighborhood is not rigid, allowing for a heuristic interpretation. For instance, we can represent this neighborhood stochastically: $X=x+\epsilon$. Here, $x$ is the observation, and $\epsilon$ is a random variable, which we set to be a zero-mean isotropic Gaussian distribution $p(\bm{\epsilon})=\mathcal{N}(0,\sigma^{2}_{\epsilon}({\bm{x}}){\bm{I}}_{m})$ for future analysis.

The advantage of this heuristic is that it allows us to represent $\hat{y}=f_{\theta}({\bm{x}}+\epsilon)$ stochastically. While the variance of $\epsilon$ is unknown (we will later show that it is learnt), we assume heteroscedasticity, which allows us to represent neighborhoods of varying spatial extents for each ${\bm{x}}$. We therefore model $\mathrm{Cov}(f_{\theta}(x+\epsilon))$, and continue by taking the second order Taylor polynomial of $f_{\theta}({\bm{x}}+\epsilon)$.

The second order Taylor polynomial introduces the notion of gradient and curvature in modeling the covariance, and quantifies the rate at which a function can change within a small neighborhood around ${\bm{x}}$. We have

Here, ${\bm{x}}\in\mathbb{R}^{m}$ is the input, $f_{\theta}({\bm{x}})\in\mathbb{R}^{n}$ represents the multivariate prediction, $\epsilon\in\mathbb{R}^{m}$ represents the neighborhood of ${\bm{x}}$, ${\bm{J}}({\bm{x}})\in\mathbb{R}^{n\times m}$ corresponds to the Jacobian matrix, and ${\bm{\mathsfit{H}}}({\bm{x}})\in\mathbb{R}^{n\times m\times m}$ represents the Hessian tensor. We note that all the individual terms in Eq. 2 are $n$-dimensional.

The covariance of Eq. 2, $\mathrm{Cov}f_{\theta}({\bm{x}}+\epsilon)$, with respect to the random variable $\epsilon$ is

We obtain this since $f_{\theta}({\bm{x}})$ is a constant with respect to $\epsilon$. Below, we evaluate the three terms individually.

We defined the covariance through $\epsilon$, the neighborhood random variable which allows us to capture the gradient and curvature of $f_{\theta}(x+\epsilon)$. However, the target $y$ could have stochasticity that does not depend upon the neighborhood. We take as an example the function $y=c+\mathcal{N}(0,\Sigma(x))$. Here, the stochasticity of $y$ is independent of the neighborhood $\epsilon$. To address scenarios such as these, we introduce a new random variable $\varepsilon\sim\mathcal{N}(0,\Sigma(x))$, which is heteroscedastic and is independent of $f_{\theta}(x+\epsilon)$. Indeed, $\varepsilon$ does not depend upon the gradient or curvature of $f_{\theta}$. Subsequently, we can write $\hat{y}=f_{\theta}(x+\epsilon)+\varepsilon$. Since $\varepsilon$ is independent of $f_{\theta}(x+\epsilon)$, we can write the covariance as the sum of $\textrm{Cov}f_{\theta}(x+\epsilon)$ and $\Sigma(x)$. This is possible because the sum of two independent Gaussians results in a Gaussian with the means and covariances summed. We therefore model the covariance through the gradient and curvature as well as account for the inherent stochasticity of the samples.

We learn the covariance of $\varepsilon$ through $k_{3}({\bm{x}})\in\mathbb{R}^{n\times n}$, a learnable positive definite matrix which is optimized via the covariance estimator $g_{\Theta}(x)$. The final expression for TIC is

The covariance estimator $g_{\Theta}(x)$ predicts $k_{1}({\bm{x}}),k_{2}({\bm{x}})$ and $k_{3}({\bm{x}})$, where $k_{1}({\bm{x}}),k_{2}({\bm{x}})$ are positive scalars. We enforce $k_{3}(x)$ to be positive definite by predicting an unconstrained matrix and multiplying it with its transpose, similarly to previous work. The covariance estimator is jointly optimized with the mean and is trained to minimize the negative log-likelihood by substituting Eq. 7 into Eq. 1.

At first, the use of the Hessian in TIC resembles its use in optimization (Gilmer et al., 2022; Kingma & Ba, 2015). The Cramer-Rao bound (Ly et al., 2017) links the variance of a parametric estimator with its inverse Fisher information. However, the Fisher information computes the Hessian with respect to the parameters, measuring its sensitivity over all samples. By contrast, the Hessian in our formulation is computed with respect to the input, allowing us to model heteroscedasticity.

We incorporate TIC within the negative log-likelihood formulation and do not employ covariance specific regularization. Similar to previous works (Kendall & Gal, 2017; Skafte et al., 2019; Seitzer et al., 2022; Stirn et al., 2023; Immer et al., 2023), our method is an approximation without theoretical guarantees. This approximation results from a heuristic interpretation of the neigborhood, as well as the use of the second order taylor polynomial. However, our experimental evaluations show that TIC provides accurate covariance estimates and works well in practice.

Limitations. The computational complexity in TIC arises from computing the Hessian. While determining this for a general network architecture is non-trivial, computing the Hessian for a function that maps $x\in\mathbb{R}^{m}$ to $y\in\mathbb{R}^{n}$ has a complexity of $O(nm^{3})$ (Yao et al., 2020), which is large. There are multiple possible ways to mitigate this in practice. The simplest approach would be to use parallelization, which we provide in our code. The second approach would be to use a smaller, proxy model in place of a large model (which could be retained for mean estimation). This smaller model could be trained through a student-teacher setup using techniques from knowledge distillation (Gou et al., 2021). The reduced parameter count would decrease the computational requirements of the Hessian. An interesting direction for future research would be to find useful approximations of the Hessian with respect to the input, similarly to research in optimization which approximates the Hessian with respect to the parameters.

Formally, given an $n$-dimensional target prediction $\hat{{\bm{y}}}$, the ground truth ${\bm{y}}$, and the predicted covariance $\mathrm{Cov}(\hat{Y}|X\!\!=\!\!{\bm{x}})$, we define the TAC error as $\sum_{i}|y_{i}-\tilde{y}_{i}|/n$, where $\tilde{y}_{i}$ is the updated mean obtained after conditioning $\mathcal{N}(\tilde{y}_{i},\mathrm{Cov}(\hat{Y}|X)\,\,|\,\,{\bm{y}}_{j\neq i},{\bm{x}})$. For each prediction $\hat{y}_{i}$, we obtain a revised estimate $\tilde{y}_{i}$ by conditioning it over the ground truth of the remaining variables ${\bm{y}}_{i\neq j}$. We measure the absolute error of this revised estimate against the ground truth of the unobserved variable and repeat for all $i$. An accurate estimate of $\mathrm{Cov}(\hat{Y}|X\!\!=\!\!{\bm{x}})$ will decrease the error whereas an incorrect estimate will cause an increase. We describe this in Algorithm 2.

This evaluation bears resemblance with leave-one-out, where we observe one $\tilde{y}_{i}$ given other observations ${\bm{y}}_{j\neq i}$. While leave-one-out can be generalized to leave-k-out, we do not observe any change in the evaluation trend. A method having lower leave-one-out also has a lower leave-k-out error. Moreover, leave-k-out requires taking $\binom{n}{k}$ combinations, which is significantly higher than taking $n$ combinations in leave-one-out. This motivates the use of the leave-one-out.

We highlight that this metric is agnostic of downstream tasks involving covariance estimation. We also note that TAC and the log-likelihood are complementary: while log-likelihood is a measure of optimization, TAC is a measure of accuracy of the learnt correlations. Hence, we use TAC as an additional metric for all multivariate experiments.

Univariate. We repeat the experiments of Seitzer et al. (2022) with a major revision. First, we introduce heteroscedasticity and substantially increase the variance of the samples. Second, we simulate different sinusoidals having constant and varying amplitudes. We draw 50,000 samples and train a fully-connected network with batch normalization (Ioffe & Szegedy, 2015) for 100 epochs. Our results are shown in Figure 3. We observe that in the absence of direct supervision, the negative log-likelihood incorrectly overestimates the variance since it does not represent the randomness of the predicted mean. Furthermore, both $\beta$-NLL and Faithful are susceptible to incorrect variance predictions because the methods regularize the variance, which compromises on variance fits. While Natural Laplace fits the constant amplitude sinusoidal, the method results in unstable optimization for sinusoidals of varying amplitude.

Multivariate. We propose a new experiment for multivariate analysis to study heteroscedastic covariance. We let $X,Y$ be jointly distributed and sample ${\bm{x}}$ from this distribution. Subsequently, we sample ${\bm{y}}$ conditioned on ${\bm{x}}$. To simulate heteroscedasticity, we draw samples from $Z$, a zero mean random variable whose covariance $\Sigma_{Z}=\text{diag}(\sqrt{|{\bm{x}}|})$ depends on ${\bm{x}}$. Since $Y$ and $Z$ are independent given $X$, their sum also satisfies the normal distribution $Q|X\sim\mathcal{N}(\mu_{Y|X},\Sigma_{Y|X}+\Sigma_{Z|X})$. Therefore, the goal of this experiment is to model the mean and the heteroscedastic covariance of $Q|X$ given samples $({\bm{x}},{\bm{q}})$. The schematic for our experimental design is shown in Fig. 4.

We vary the dimensionality of ${\bm{x}}$ and ${\bm{q}}$ from 4 to 20 in steps of 2, and report the mean and standard deviation over ten trials for each dimension. Depending on the dimensionality, we draw from 4000 up to 20000 samples and report our results using TAC (Fig. 5) and the log-likelihood (Table 1). We observe two trends in Fig. 5: First, as the dimensionality of the samples increases, the gap between TIC and the other methods widens. This is because, with increasing dimensionality, the number of free parameters to estimate in the covariance matrix grows quadratically. An increase in parameters typically requires a non-linear growth in the number of samples for robust fitting. As a result, the difficulty of mapping the input to a positive definite matrix increases with dimensionality. Second, we note that TIC allows for better convergence in comparison to the naive parameterization of the covariance in NLL.

We perform our analysis on twelve multivariate UCI regression (Dua & Graff, 2017) datasets, which have been used in previous work on negative log-likelihood (Stirn et al., 2023; Seitzer et al., 2022). Nevertheless, our goal of studying covariance estimation in deep heteroscedastic regression requires us to use a different pre-processing, as many of the datasets have univariate or low-dimensional targets.

Specifically, for each dataset we randomly allocate 25% of the features as input and the remaining 75% features as multivariate targets at run-time. Some combinations of input variables may fare poorly at predicting the target variables. However, this is an interesting challenge for the covariance estimator, which needs to learn the underlying correlations even in unfavorable circumstances. Moreover, random splitting also allows our experiments to remain unbiased as we do not control the split of variables at any instant. For all datasets, we standardize our variables with zero mean and a variance of ten (which yields better convergence for all methods). We perform 10 trials for each dataset and report the TAC error and likelihood in Table 2(b).

While TIC outperforms the baselines on a majority of the datasets, we particularly focus on the Naval dataset which highlights a limitation of the TIC parameterization. We observe that TIC may not be suitable if all samples have a low degree of variance. A low degree of variance (as indicated by the likelihood) results in accurate mean fits, which implies that small gradients are being backpropagated, and in turn affecting the TIC parameterization. However, we argue that datasets with a small degree of variance may not benefit from heteroscedastic modelling.

We introduce experiments on human pose estimation since the human pose is an organised collection of points and is naturally suited for correlation analysis (Shukla et al., 2022). Moreover, popular human pose architectures (Newell et al., 2016; Sun et al., 2019; Kreiss et al., 2019, 2021; Xu et al., 2022) are either convolutional or transformer based, presenting a viable challenge to modelling the Taylor Induced Covariance. This is because TIC assumes vector inputs ${\bm{x}}\in\mathbb{R}^{m}$ and predictions $\hat{{\bm{y}}}\in\mathbb{R}^{n}$, whereas popular architectures rely on input images ${\bm{\mathsfit{X}}}\in\mathbb{R}^{C\times H\times W}$ and output heatmaps $\hat{{\bm{\mathsfit{Y}}}}\in\mathbb{R}^{\#\text{joints}\times 64\times 64}$.

We therefore perform experiments on two architectures: the Stacked Hourglass (Newell et al., 2016) and ViTPose (Xu et al., 2022). The Stacked Hourglass is a popular method which extends the convolutional U-Net (Ronneberger et al., 2015) architecture to predict heatmaps for human pose estimation. ViTPose is a recent state-of-the-art architecture which extends vision transformers (Dosovitskiy et al., 2021) to the task of human pose estimation.

For both architectures, we use soft-argmax (Li et al., 2021b, a) to reduce the heatmap of shape $\hat{{\bm{\mathsfit{Y}}}}\in\mathbb{R}^{\#\text{joints}\times 64\times 64}$ to a vector of shape $\mathbb{R}^{\#\text{joints}*2}$. Next, we recursively call the hourglass module until we obtain a one-dimensional vector encoding (Shukla, 2022) for the image, which serves as the input to the covariance estimator. For ViTPose, we obtain vector embeddings from a simple residual connection involving a one-dimensional downscaling and upscaling of the features predicted by the backbone network.

We use popular single person datasets: MPII (Andriluka et al., 2014) and LSP/LSPET (Johnson & Everingham, 2010, 2011), with the latter emphasizing on poses involving sports. We perform our analysis by merging the MPII and LSP-LSPET datasets to increase the number of samples. We train the pose estimator using the Adam optimizer with a ’ReduceLROnPlateau’ learning rate scheduler for 100 epochs with the learning rate set to 1e-3. We use two augmentations: Shift+Scale+Rotate and horizontal flip. We refer the reader to the code for implementation details.

In addition to the likelihood, we continue to use TAC as our metric since for single person estimation, the scale of the person is fixed. Hence, TAC is highly correlated with PCKh/PCK, the preferred metric for multi-person multi-scale pose estimation. We perform five trials and report our results for the ViTPose backbone in Table 3. We report results on the Stacked Hourglass backbone in the appendix. Our experiments show that TIC outperforms all baselines, especially on challenging joints.