SlowGuess/NeurIPS_2023_OCR huggingface dataset:
wzy@gatech.edu
Existing regression models tend to fall short in both accuracy and uncertainty estimation when the label distribution is imbalanced. In this paper, we propose a probabilistic deep learning model, dubbed variational imbalanced regression (VIR), which not only performs well in imbalanced regression but naturally produces reasonable uncertainty estimation as a byproduct. Different from typical variational autoencoders assuming I.I.D. representations (a data point's representation is not directly affected by other data points), our VIR borrows data with similar regression labels to compute the latent representation's variational distribution; furthermore, different from deterministic regression models producing point estimates, VIR predicts the entire normal-inverse-gamma distributions and modulates the associated conjugate distributions to impose probabilistic reweighting on the imbalanced data, thereby providing better uncertainty estimation. Experiments in several real-world datasets show that our VIR can outperform state-of-the-art imbalanced regression models in terms of both accuracy and uncertainty estimation. Code will soon be available at https://github.com/Wang-ML-Lab/variational-imbalanced-regression.
Deep regression models are currently the state of the art in making predictions in a continuous label space and have a wide range of successful applications in computer vision [50], natural language processing [22], healthcare [43, 45], recommender systems [16, 42], etc. However, these models fail however when the label distribution in training data is imbalanced. For example, in visual age estimation [30], where a model infers the age of a person given her visual appearance, models are typically trained on imbalanced datasets with overwhelmingly more images of younger adults, leading to poor regression accuracy for images of children or elderly people [48, 49]. Such unreliability in imbalanced regression settings motivates the need for both improving performance for the minority in the presence of imbalanced data and, more importantly, providing reasonable uncertainty estimation to inform practitioners on how reliable the predictions are (especially for the minority where accuracy is lower).
Existing methods for deep imbalanced regression (DIR) only focus on improving the accuracy of deep regression models by smoothing the label distribution and reweighting data with different labels [48, 49]. On the other hand, methods that provide uncertainty estimation for deep regression models operates under the balance-data assumption and therefore do not work well in the imbalanced setting [1, 8, 29].
To simultaneously cover these two desiderata, we propose a probabilistic deep imbalanced regression model, dubbed variational imbalanced regression (VIR). Different from typical variational autoencoders assuming I.I.D. representations (a data point's representation is not directly affected by other data points), our VIR assumes Neighboring and Identically Distributed (N.I.D.) and borrows
data with similar regression labels to compute the latent representation's variational distribution. Specifically, VIR first encodes a data point into a probabilistic representation and then mix it with neighboring representations (i.e., representations from data with similar regression labels) to produce its final probabilistic representation; VIR is therefore particularly useful for minority data as it can borrow probabilistic representations from data with similar labels (and naturally weigh them using our probabilistic model) to counteract data sparsity. Furthermore, different from deterministic regression models producing point estimates, VIR predicts the entire Normal Inverse Gamma (NIG) distributions and modulates the associated conjugate distributions by the importance weight computed from the smoothed label distribution to impose probabilistic reweighting on the imbalanced data. This allows the negative log likelihood to naturally put more focus on the minority data, thereby balancing the accuracy for data with different regression labels. Our VIR framework is compatible with any deep regression models and can be trained end to end.
- We identify the problem of probabilistic deep imbalanced regression as well as two desiderata, balanced accuracy and uncertainty estimation, for the problem.
- We propose VIR to simultaneously cover these two desiderata and achieve state-of-the-art performance compared to existing methods.
- As a byproduct, we also provide strong baselines for benchmarking high-quality uncertainty estimation and promising prediction performance on imbalanced datasets.
Variational autoencoder. Variational autoencoder (VAE) [25] is an unsupervised learning model that aims to infer probabilistic representations from data. However, as shown in Figure 1, VAE typically assumes I.I.D. representations, where a data point's representation is not directly affected by other data points. In contrast, our VIR borrows data with similar regression labels to compute the latent representation's variational distribution.
Imbalanced Regression. Imbalanced regression is under-explored in the machine learning community. Most existing methods for imbalanced regression are direct extensions of the SMOTE algorithm [9], a commonly
Figure 1: Comparing inference networks of typical VAE [25] and our VIR. In VAE (left), a data point's latent representation (i.e.,  $\mathbf{z}$ ) is affected only by itself, while in VIR (right), neighbors participate to modulate the final representation.
used algorithm for imbalanced classification, where data from the minority classes is over-sampled. These algorithms usually synthesize augmented data for the minority regression labels by either interpolating both inputs and labels [40] or adding Gaussian noise [5, 6] (more discussion on augmentation-based methods in the Appendix).
Such algorithms fail to measure the distance in continuous label space and fall short in handling high-dimensional data (e.g., images and text). Recently, DIR [49] addresses these issues by applying kernel density estimation to smooth and reweight data on the continuous label distribution, achieving state-of-the-art performance. However, DIR only focuses on improving the accuracy, especially for the data with minority labels, and therefore does not provide uncertainty estimation, which is crucial to assess the predictions' reliability. [32] focuses on re-balancing the mean squared error (MSE) loss for imbalanced regression, and [13] introduces ranking similarity for improving deep imbalanced regression. In contrast, our VIR provides a principled probabilistic approach to simultaneously achieve these two desiderata, not only improving upon DIR in terms of performance but also producing reasonable uncertainty estimation as a much-needed byproduct to assess model reliability. There is also related work on imbalanced classification [11], which is related to our work but focusing on classification rather than regression.
Uncertainty Estimation in Regression. There has been renewed interest in uncertainty estimation in the context of deep regression models [1, 12, 18, 23, 26, 29, 36-38, 51]. Most existing methods directly predict the variance of the output distribution as the estimated uncertainty [1, 23, 52], rely on post-hoc confidence interval calibration [26, 37, 51], or using Bayesian neural networks [44, 46, 47]; there are also training-free approaches, such as Infer Noise and Infer Dropout [29], which produce multiple predictions from different perturbed neurons and compute their variance as uncertainty estimation. Closest to our work is Deep Evidential Regression (DER) [1], which attempts to estimate both aleatoric and epistemic uncertainty [20, 23] on regression tasks by training the neural networks to directly infer the parameters of the evidential distribution, thereby producing uncertainty measures. DER [1] is designed for the data-rich regime and therefore fails to reasonably estimate the uncertainty if the data is imbalanced; for data with minority labels, DER [1] tends produce unstable distribution parameters, leading to poor uncertainty estimation (as shown in Sec. 5). In contrast, our proposed VIR explicitly handles data imbalance in the continuous label space to avoid such instability; VIR does so by modulating both the representations and the output conjugate distribution parameters according to the imbalanced label distribution, allowing training/inference to proceed as if the data is balance and leading to better performance as well as uncertainty estimation (as shown in Sec. 5).
In this section we introduce the notation and problem setting, provide an overview of our VIR, and then describe details on each of VIR's key components.
3.1 Notation and Problem Setting
Assuming an imbalanced dataset in continuous space  $\{\mathbf{x}_i, y_i\}_{i=1}^N$  where  $N$  is the total number of data points,  $\mathbf{x}_i \in \mathbb{R}^d$  is the input, and  $y_i \in \mathcal{V} \subset \mathbb{R}$  is the corresponding label from a continuous label space  $\mathcal{V}$ . In practice,  $\mathcal{V}$  is partitioned into B equal-interval bins  $[y^{(0)}, y^{(1)}], [y^{(1)}, y^{(2)}], \ldots, [y^{(B-1)}, y^{(B)})$ , with slight notation overload. To directly compare with baselines, we use the same grouping index for target value  $b \in \mathcal{B}$  as in [49].
We denote representations as  $\mathbf{z}_i$ , and use  $(\widetilde{\mathbf{z}}_i^\mu, \widetilde{\mathbf{z}}_i^\Sigma) = q_\phi(\mathbf{z}|\mathbf{x}_i)$  to denote the probabilistic representations for input  $\mathbf{x}_i$  generated by a probabilistic encoder parameterized by  $\phi$ ; furthermore, we denote  $\bar{\mathbf{z}}$  as the mean of representation  $\mathbf{z}_i$  in each bin, i.e.,  $\bar{\mathbf{z}} = \frac{1}{N_b}\sum_{i=1}^{N_b}\mathbf{z}_i$  for a bin with  $N_b$  data points. Similarly, we use  $(\widehat{y}_i, \widehat{s}_i)$  to denote the mean and variance of the predictive distribution generated by a probabilistic predictor  $p_\theta(y_i|\mathbf{z})$ .
Figure 2: Overview of our VIR method. Left: The inference model infers the latent representations given input  $\mathbf{x}$ 's in the neighborhood. Right: The generative model reconstructs the input and predicts the label distribution (including the associated uncertainty) given the latent representation.
In order to achieve both desiderata in probabilistic deep imbalanced regression (i.e., performance improvement and uncertainty estimation),
our proposed variational imbalanced regression (VIR) operates on both the encoder  $q_{\phi}(\mathbf{z}_i|\{\mathbf{x}_i\}_{i=1}^N)$  and the predictor  $p_{\theta}(y_i|\mathbf{z}_i)$ .
Typical VAE [25] lower-bounds input  $\mathbf{x}_i$ 's marginal likelihood; in contrast, VIR lower-bounds the marginal likelihood of input  $\mathbf{x}_i$  and labels  $y_i$ :
$$
\log p _ {\theta} \left(\mathbf {x} _ {i}, y _ {i}\right) = \mathcal {D} _ {\mathcal {K L}} \left(q _ {\phi} \left(\mathbf {z} _ {i} \mid \{\mathbf {x} _ {i} \} _ {i = 1} ^ {N}\right) | | p _ {\theta} \left(\mathbf {z} _ {i} \mid \mathbf {x} _ {i}, y _ {i}\right)\right) + \mathcal {L} (\theta , \phi ; \mathbf {x} _ {i}, y _ {i}).
$$
Note that our variational distribution  $q_{\phi}(\mathbf{z}_i|\{\mathbf{x}_i\}_{i=1}^N)$  (1) does not condition on labels  $y_i$ , since the task is to predict  $y_i$  and (2) conditions on all (neighboring) inputs  $\{\mathbf{x}_i\}_{i=1}^N$  rather than just  $\mathbf{x}_i$ . The second term  $\mathcal{L}(\theta, \phi; \mathbf{x}_i, y_i)$  is VIR's evidence lower bound (ELBO), which is defined as:
$$
\mathcal {L} (\theta , \phi ; \mathbf {x} _ {i}, y _ {i}) = \underbrace {\mathbb {E} _ {q} \left[ \log p _ {\theta} \left(\mathbf {x} _ {i} \mid \mathbf {z} _ {i}\right) \right]} _ {\mathcal {L} _ {i} ^ {\mathcal {D}}} + \underbrace {\mathbb {E} _ {q} \left[ \log p _ {\theta} \left(y _ {i} \mid \mathbf {z} _ {i}\right) \right]} _ {\mathcal {L} _ {i} ^ {\mathcal {P}}} - \underbrace {\mathcal {D} _ {\mathcal {K L}} \left(q _ {\phi} \left(\mathbf {z} _ {i} \mid \{\mathbf {x} _ {i} \} _ {i = 1} ^ {N}\right) \mid p _ {\theta} \left(\mathbf {z} _ {i}\right)\right)} _ {\mathcal {L} _ {i} ^ {\mathcal {K L}}}. \tag {1}
$$
where the  $p_{\theta}(\mathbf{z}_i)$  is the standard Gaussian prior  $\mathcal{N}(\mathbf{0},\mathbf{I})$ , following typical VAE [25], and the expectation is taken over  $q_{\phi}(\mathbf{z}_i|\{\mathbf{x}_i\}_{i = 1}^N)$ , which infers  $\mathbf{z}_i$  by borrowing data with similar regression labels to produce the balanced probabilistic representations, which is beneficial especially for the minority (see Sec. 3.3 for details).
Different from typical regression models which produce only point estimates for  $y_{i}$ , our VIR's predictor,  $p_{\theta}(y_i|\mathbf{z}_i)$ , directly produces the parameters of the entire NIG distribution for  $y_{i}$  and further imposes probabilistic reweighting on the imbalanced data, thereby producing balanced predictive distributions (more details in Sec. 3.4).
3.3 Constructing  $q(\mathbf{z}_i|\{\mathbf{x}_i\}_{i=1}^N)$
To cover both desiderata, one needs to (1) produce balanced representations to improve performance for the data with minority labels and (2) produce probabilistic representations to naturally obtain reasonable uncertainty estimation for each model prediction. To learn such balanced probabilistic representations, we construct the encoder of our VIR (i.e.,  $q_{\phi}(\mathbf{z}_i|\{\mathbf{x}_i\}_{i=1}^N)$ ) by (1) first encoding a data point into a probabilistic representation, (2) computing probabilistic statistics from neighboring representations (i.e., representations from data with similar regression labels), and (3) producing the final representations via probabilistic whitening and recoloring using the obtained statistics.
Intuition on Using Probabilistic Representation. DIR uses deterministic representations, with one vector as the final representation for each data point. In contrast, our VIR uses probabilistic representations, with one vector as the mean of the representation and another vector as the variance of the representation. Such dual representation is more robust to noise and therefore leads to better prediction performance. Therefore, We first encode each data point into a probabilistic representation. Note that this is in contrast to existing work [49] that uses deterministic representations. We assume that each encoding  $\mathbf{z}_i$  is a Gaussian distribution with parameters  $\{\mathbf{z}_i^{\mu},\mathbf{z}_i^{\Sigma}\}$ , which are generated from the last layer in the deep neural network.
From I.I.D. to Neighboring and Identically Distributed (N.I.D.). Typical VAE [25] is an unsupervised learning model that aims to learn a variational representation from latent space to reconstruct the original inputs under the I.I.D. assumption; that is, in VAE, the latent value (i.e.,  $\mathbf{z}_i$ ) is generated from its own input  $\mathbf{x}_i$ . This I.I.D. assumption works well for data with majority labels, but significantly harms performance for data with minority labels. To address this problem, we replace the I.I.D. assumption with the N.I.D. assumption; specifically, VIR's variational latent representations still follow Gaussian distributions (i.e.,  $\mathcal{N}(\mathbf{z}_i^\mu, \mathbf{z}_i^\Sigma)$ , but these distributions will be first calibrated using data with neighboring labels. For a data point  $(\mathbf{x}_i, y_i)$  where  $y_i$  is in the  $b$ 'th bin, i.e.,  $y_i \in [y^{(b-1)}, y^{(b)})$ , we compute  $q(\mathbf{z}_i|\{\mathbf{x}_i\}_{i=1}^N) \triangleq \mathcal{N}(\mathbf{z}_i; \widetilde{\mathbf{z}}_i^\mu, \widetilde{\mathbf{z}}_i^\Sigma)$  with the following four steps.
(1) Mean and Covariance of Initial  $\mathbf{z}_i$ :
$$
\mathbf {z} _ {i} ^ {\mu}, \mathbf {z} _ {i} ^ {\Sigma} = \mathcal {I} (\mathbf {x} _ {i}),
$$
(2) Statistics of Bin  $b$ 's Statistics:
$$
\boldsymbol {\mu} _ {b} ^ {\mu}, \boldsymbol {\mu} _ {b} ^ {\Sigma}, \boldsymbol {\Sigma} _ {b} ^ {\mu}, \boldsymbol {\Sigma} _ {b} ^ {\Sigma} = \mathcal {A} \left(\left\{\mathbf {z} _ {i} ^ {\mu}, \mathbf {z} _ {i} ^ {\Sigma} \right\} _ {i = 1} ^ {N}\right),
$$
(3) Smoothed Statistics of Bin  $b$ 's Statistics:
$$
\widetilde {\boldsymbol {\mu}} _ {b} ^ {\mu}, \widetilde {\boldsymbol {\mu}} _ {b} ^ {\Sigma}, \widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\mu}, \widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\Sigma} = \mathcal {S} \left(\left\{\boldsymbol {\mu} _ {b} ^ {\mu}, \boldsymbol {\mu} _ {b} ^ {\Sigma}, \boldsymbol {\Sigma} _ {b} ^ {\mu}, \boldsymbol {\Sigma} _ {b} ^ {\Sigma} \right\} _ {b = 1} ^ {B}\right),
$$
(4) Mean and Covariance of Final  $\mathbf{z}_i$ :
$$
\widetilde {\mathbf {z}} _ {i} ^ {\mu}, \widetilde {\mathbf {z}} _ {i} ^ {\Sigma} = \mathcal {F} (\mathbf {z} _ {i} ^ {\mu}, \mathbf {z} _ {i} ^ {\Sigma}, \boldsymbol {\mu} _ {b} ^ {\mu}, \boldsymbol {\mu} _ {b} ^ {\Sigma}, \boldsymbol {\Sigma} _ {b} ^ {\mu}, \boldsymbol {\Sigma} _ {b} ^ {\Sigma}, \widetilde {\boldsymbol {\mu}} _ {b} ^ {\mu}, \widetilde {\boldsymbol {\mu}} _ {b} ^ {\Sigma}, \widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\mu}, \widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\Sigma}),
$$
where the details of functions  $\mathcal{I}(\cdot),\mathcal{A}(\cdot),\mathcal{S}(\cdot)$  , and  $\mathcal{F}(\cdot)$  are described below.
(1) Function  $\mathcal{I}(\cdot)$ : From Deterministic to Probabilistic Statistics. Different from deterministic statistics in [49], our VIR's encoder uses probabilistic statistics, i.e., statistics of statistics. Specifically, VIR treats  $\mathbf{z}_i$  as a distribution with the mean and covariance  $(\mathbf{z}_i^\mu, \mathbf{z}_i^{\Sigma}) = \mathcal{I}(\mathbf{x}_i)$  rather than a deterministic vector.
As a result, all the deterministic statistics for bin  $b$ ,  $\pmb{\mu}_b$ ,  $\pmb{\Sigma}_b$ ,  $\widetilde{\pmb{\mu}}_b$ , and  $\widetilde{\pmb{\Sigma}}_b$  are replaced by distributions with the means and covariances,  $(\pmb{\mu}_b^\mu, \pmb{\mu}_b^\Sigma)$ ,  $(\pmb{\Sigma}_b^\mu, \pmb{\Sigma}_b^\Sigma)$ ,  $(\widetilde{\pmb{\mu}}_b^\mu, \widetilde{\pmb{\mu}}_b^\Sigma)$ , and  $(\widetilde{\pmb{\Sigma}}_b^\mu, \widetilde{\pmb{\Sigma}}_b^\Sigma)$ , respectively (more details in the following three paragraphs on  $\mathcal{A}(\cdot)$ ,  $\mathcal{S}(\cdot)$ , and  $\mathcal{F}(\cdot)$ ).
(2) Function  $\mathcal{A}(\cdot)$ : Statistics of the Current Bin  $b$ 's Statistics. In VIR, the deterministic overall mean for bin  $b$  (with  $N_{b}$  data points),  $\pmb{\mu}_{b} = \bar{\mathbf{z}} = \frac{1}{N_{b}}\sum_{i=1}^{N_{b}}\mathbf{z}_{i}$ , becomes the probabilistic overall mean, i.e., a distribution of  $\pmb{\mu}_{b}$  with the mean  $\pmb{\mu}_{b}^{\mu}$  and covariance  $\pmb{\mu}_{b}^{\Sigma}$  (assuming diagonal covariance) as
$$
\boldsymbol {\mu} _ {b} ^ {\mu} \triangleq \mathbb {E} [ \bar {\mathbf {z}} ] = \frac {1}{N _ {b}} \sum_ {i = 1} ^ {N _ {b}} \mathbb {E} [ \mathbf {z} _ {i} ] = \frac {1}{N _ {b}} \sum_ {i = 1} ^ {N _ {b}} \mathbf {z} _ {i} ^ {\mu},
$$
$$
\boldsymbol {\mu} _ {b} ^ {\Sigma} \triangleq \mathbb {V} [ \bar {\mathbf {z}} ] = \frac {1}{N _ {b} ^ {2}} \sum_ {i = 1} ^ {N _ {b}} \mathbb {V} [ \mathbf {z} _ {i} ] = \frac {1}{N _ {b} ^ {2}} \sum_ {i = 1} ^ {N _ {b}} \mathbf {z} _ {i} ^ {\Sigma}.
$$
Similarly, the deterministic overall covariance for bin  $b$ ,  $\Sigma_{b} = \frac{1}{N_{b}}\sum_{i = 1}^{N_{b}}(\mathbf{z}_{i} - \bar{\mathbf{z}})^{2}$ , becomes the probabilistic overall covariance, i.e., a matrix-variate distribution [15] with the mean:
$$
\pmb {\Sigma} _ {b} ^ {\mu} \triangleq \mathbb {E} [ \pmb {\Sigma} _ {b} ] = \frac {1}{N _ {b}} \sum_ {i = 1} ^ {N _ {b}} \mathbb {E} [ (\pmb {z} _ {i} - \bar {\pmb {z}}) ^ {2} ] = \frac {1}{N _ {b}} \sum_ {i = 1} ^ {N _ {b}} \left[ \pmb {z} _ {i} ^ {\Sigma} + (\pmb {z} _ {i} ^ {\mu}) ^ {2} - \left([ \pmb {\mu} _ {b} ^ {\Sigma} ] _ {i} + ([ \pmb {\mu} _ {b} ^ {\mu} ] _ {i}) ^ {2}\right) \right],
$$
since  $\mathbb{E}[\bar{\mathbf{z}}] = \pmb{\mu}_b^\mu$  and  $\mathbb{V}[\bar{\mathbf{z}}] = \pmb{\mu}_b^{\Sigma}$ . Note that the covariance of  $\pmb{\Sigma}_{b}$ , i.e.,  $\pmb{\Sigma}_{b}^{\Sigma} \triangleq \mathbb{V}[\pmb{\Sigma}_{b}]$ , involves computing the fourth-order moments, which is computationally prohibitive. Therefore in practice, we directly set  $\pmb{\Sigma}_{b}^{\Sigma}$  to zero for simplicity; empirically we observe that such simplified treatment already achieves promising performance improvement upon the state of the art. More discussions on the idea of the hierarchical structure of the statistics of statistics for smoothing are in the Appendix.
(3) Function  $S(\cdot)$ : Neighboring Data and Smoothed Statistics. Next, we can borrow data from neighboring label bins  $b'$  to compute the smoothed statistics of the current bin  $b$  by applying a symmetric kernel  $k(\cdot, \cdot)$  (e.g., Gaussian, Laplacian, and Triangular kernels). Specifically, the probabilistic smoothed mean and covariance are (assuming diagonal covariance):
$$
\widetilde {\boldsymbol {\mu}} _ {b} ^ {\mu} = \sum_ {b ^ {\prime} \in \mathcal {B}} k (y _ {b}, y _ {b ^ {\prime}}) \boldsymbol {\mu} _ {b ^ {\prime}} ^ {\mu}, \quad \widetilde {\boldsymbol {\mu}} _ {b} ^ {\Sigma} = \sum_ {b ^ {\prime} \in \mathcal {B}} k ^ {2} (y _ {b}, y _ {b ^ {\prime}}) \boldsymbol {\mu} _ {b ^ {\prime}} ^ {\Sigma}, \quad \widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\mu} = \sum_ {b ^ {\prime} \in \mathcal {B}} k (y _ {b}, y _ {b ^ {\prime}}) \boldsymbol {\Sigma} _ {b ^ {\prime}}.
$$
(4) Function  $\mathcal{F}(\cdot)$ : Probabilistic Whitening and Recoloring. We develop a probabilistic version of the whitening and re-coloring procedure in [39, 49]. Specifically, we produce the final probabilistic representation  $\{\widetilde{\mathbf{z}}_i^\mu, \widetilde{\mathbf{z}}_i^{\Sigma}\}$  for each data point as:
$$
\widetilde {\mathbf {z}} _ {i} ^ {\mu} = \left(\mathbf {z} _ {i} ^ {\mu} - \boldsymbol {\mu} _ {b} ^ {\mu}\right) \cdot \sqrt {\frac {\widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\mu}}{\boldsymbol {\Sigma} _ {b} ^ {\mu}}} + \widetilde {\boldsymbol {\mu}} _ {b} ^ {\mu}, \quad \widetilde {\mathbf {z}} _ {i} ^ {\Sigma} = \left(\mathbf {z} _ {i} ^ {\Sigma} + \boldsymbol {\mu} _ {b} ^ {\Sigma}\right) \cdot \sqrt {\frac {\widetilde {\boldsymbol {\Sigma}} _ {b} ^ {\mu}}{\boldsymbol {\Sigma} _ {b} ^ {\mu}}} + \widetilde {\boldsymbol {\mu}} _ {b} ^ {\Sigma}. \tag {2}
$$
During training, we keep updating the probabilistic overall statistics,  $\{\pmb{\mu}_b^\mu, \pmb{\mu}_b^\Sigma, \pmb{\Sigma}_b^\mu\}$ , and the probabilistic smoothed statistics,  $\{\widetilde{\pmb{\mu}}_b^\mu, \widetilde{\pmb{\mu}}_b^\Sigma, \widetilde{\pmb{\Sigma}}_b^\mu\}$ , across different epochs. The probabilistic representation  $\{\widetilde{\mathbf{z}}_i^\mu, \widetilde{\mathbf{z}}_i^\Sigma\}$  are then re-parameterized [25] into the final representation  $\mathbf{z}_i$ , and passed into the final layer (discussed in Sec. 3.4) to generate the prediction and uncertainty estimation. Note that the computation of statistics from multiple  $\mathbf{x}$ 's is only needed during training. During testing, VIR directly uses these statistics and therefore does not need to re-compute them.
3.4 Constructing  $p(y_{i}|\mathbf{z}_{i})$
Our VIR's predictor  $p(y_{i}|\mathbf{z}_{i}) \triangleq \mathcal{N}(y_{i};\widehat{y}_{i},\widehat{s}_{i})$  predicts both the mean and variance for  $y_{i}$  by first predicting the NIG distribution and then marginalizing out the latent variables. It is motivated by the following observations on label distribution smoothing (LDS) in [49] and deep evidential regression (DER) in [1], as well as intuitions on effective counts in conjugate distributions.