Variational Imbalanced Regression: Fair Uncertainty Quantification via Probabilistic Smoothing

Assuming an imbalanced dataset in continuous space $\{{\bf x}_{i},y_{i}\}^{N}_{i=1}$ where $N$ is the total number of data points, ${\bf x}_{i}\in\mathbb{R}^{d}$ is the input, and $y_{i}\in\mathcal{Y}\subset\mathbb{R}$ is the corresponding label from a continuous label space $\mathcal{Y}$. In practice, $\mathcal{Y}$ is partitioned into B equal-interval bins $[y^{(0)},y^{(1)}),[y^{(1)},y^{(2)}),...,[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 ${\bf z}_{i}$, and use $(\tilde{{\bf z}}_{i}^{\mu},\tilde{{\bf z}}_{i}^{\Sigma})=q_{\phi}({\bf z}|{\bf x}_{i})$ to denote the probabilistic representations for input ${\bf x}_{i}$ generated by a probabilistic encoder parameterized by $\phi$; furthermore, we denote $\bar{{\bf z}}$ as the mean of representation ${\bf z}_{i}$ in each bin, i.e., $\bar{{\bf z}}=\tfrac{1}{N_{b}}\sum^{N_{b}}_{i=1}{\bf 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}|{\bf z})$.

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}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})$ and the predictor $p_{\theta}(y_{i}|{\bf z}_{i})$.

Typical VAE [25] lower-bounds input ${\bf x}_{i}$’s marginal likelihood; in contrast, VIR lower-bounds the marginal likelihood of input ${\bf x}_{i}$ and labels $y_{i}$:

$\displaystyle\log p_{\theta}({\bf x}_{i},y_{i})=\mathcal{D}_{\mathcal{KL}}\big{(}q_{\phi}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})||p_{\theta}({\bf z}_{i}|{\bf x}_{i},y_{i})\big{)}+\mathcal{L}(\theta,\phi;{\bf x}_{i},y_{i}).$

Note that our variational distribution $q_{\phi}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})$ (1) does not condition on labels $y_{i}$, since the task is to predict $y_{i}$ and (2) conditions on all (neighboring) inputs $\{{\bf x}_{i}\}^{N}_{i=1}$ rather than just ${\bf x}_{i}$. The second term $\mathcal{L}(\theta,\phi;{\bf x}_{i},y_{i})$ is VIR’s evidence lower bound (ELBO), which is defined as:

$\displaystyle\mathcal{L}(\theta,\phi;{\bf x}_{i},y_{i})=\underbrace{\mathbb{E}_{q}\big{[}\log p_{\theta}({\bf x}_{i}|{\bf z}_{i})\big{]}}_{{\mathcal{L}}_{i}^{{\mathcal{D}}}}+\underbrace{\mathbb{E}_{q}\big{[}\log p_{\theta}(y_{i}|{\bf z}_{i})\big{]}}_{{\mathcal{L}}_{i}^{{\mathcal{P}}}}-\underbrace{\mathcal{D}_{\mathcal{KL}}(q_{\phi}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})||p_{\theta}({\bf z}_{i}))}_{{\mathcal{L}}_{i}^{{\mathcal{K}}{\mathcal{L}}}}.$

(1)

where the $p_{\theta}({\bf z}_{i})$ is the standard Gaussian prior ${\mathcal{N}}({\bf 0},{\bf I})$, following typical VAE [25], and the expectation is taken over $q_{\phi}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})$, which infers ${\bf 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}|{\bf 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).

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}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})$) 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 ${\bf z}_{i}$ is a Gaussian distribution with parameters $\{{\bf z}_{i}^{\mu},{\bf 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., ${\bf z}_{i}$) is generated from its own input ${\bf 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}({\bf z}_{i}^{\mu},{\bf z}_{i}^{\Sigma}$), but these distributions will be first calibrated using data with neighboring labels. For a data point $({\bf 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({\bf z}_{i}|\{{\bf x}_{i}\}_{i=1}^{N})\triangleq{\mathcal{N}}({\bf z}_{i};\tilde{{\bf z}}_{i}^{\mu},\tilde{{\bf z}}_{i}^{\Sigma})$ with the following four steps.

	(1) Mean and Covariance of Initial ${\bf z}_{i}$:	$\displaystyle{\bf z}_{i}^{\mu},{\bf z}_{i}^{\Sigma}=\mathcal{I}({\bf x}_{i}),$
	(2) Statistics of Bin $b$’s Statistics:	$\displaystyle\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\Sigma}=\mathcal{A}(\{{\bf z}_{i}^{\mu},{\bf z}_{i}^{\Sigma}\}_{i=1}^{N}),$
	(3) Smoothed Statistics of Bin $b$’s Statistics:	$\displaystyle\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\Sigma},\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\Sigma}=\mathcal{S}(\{\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\Sigma}\}_{b=1}^{B}),$
	(4) Mean and Covariance of Final ${\bf z}_{i}$:	$\displaystyle\tilde{{\bf z}}_{i}^{\mu},\tilde{{\bf z}}_{i}^{\Sigma}=\mathcal{F}({\bf z}_{i}^{\mu},{\bf z}_{i}^{\Sigma},\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\Sigma},\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\Sigma},\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{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 ${\bf z}_{i}$ as a distribution with the mean and covariance $({\bf z}_{i}^{\mu},{\bf z}_{i}^{\Sigma})={\mathcal{I}}({\bf x}_{i})$ rather than a deterministic vector.

As a result, all the deterministic statistics for bin $b$, $\mbox{\boldmath$\mu$\unboldmath}_{b}$, $\mbox{\boldmath$\Sigma$\unboldmath}_{b}$, $\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}$, and $\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}$ are replaced by distributions with the means and covariances, $(\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma})$, $(\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\Sigma})$, $(\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\Sigma})$, and $(\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{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), $\mbox{\boldmath$\mu$\unboldmath}_{b}=\bar{{\bf z}}=\tfrac{1}{N_{b}}\sum^{N_{b}}_{i=1}{\bf z}_{i}$, becomes the probabilistic overall mean, i.e., a distribution of $\mbox{\boldmath$\mu$\unboldmath}_{b}$ with the mean $\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu}$ and covariance $\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma}$ (assuming diagonal covariance) as follows:

	$\displaystyle\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu}$	$\displaystyle\triangleq\mathbb{E}[\bar{{\bf z}}]=\tfrac{1}{N_{b}}\sum\nolimits^{N_{b}}_{i=1}\mathbb{E}[{\bf z}_{i}]=\tfrac{1}{N_{b}}\sum\nolimits^{N_{b}}_{i=1}{\bf z}_{i}^{\mu},$
	$\displaystyle\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma}$	$\displaystyle\triangleq\mathbb{V}[\bar{{\bf z}}]=\tfrac{1}{N_{b}^{2}}\sum\nolimits^{N_{b}}_{i=1}\mathbb{V}[{\bf z}_{i}]=\tfrac{1}{N_{b}^{2}}\sum\nolimits^{N_{b}}_{i=1}{\bf z}_{i}^{\Sigma}.$

Similarly, the deterministic overall covariance for bin $b$, $\mbox{\boldmath$\Sigma$\unboldmath}_{b}=\tfrac{1}{N_{b}}\sum^{N_{b}}_{i=1}({\bf z}_{i}-\bar{{\bf z}})^{2}$, becomes the probabilistic overall covariance, i.e., a matrix-variate distribution [15] with the mean:

$\displaystyle\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu}$

$\displaystyle\triangleq{\mathbb{E}}[\mbox{\boldmath$\Sigma$\unboldmath}_{b}]=\tfrac{1}{N_{b}}\sum\nolimits_{i=1}^{N_{b}}{\mathbb{E}}[({\bf z}_{i}-\bar{{\bf z}})^{2}]=\tfrac{1}{N_{b}}\sum\nolimits_{i=1}^{N_{b}}\Big{[}{\bf z}_{i}^{\Sigma}+({\bf z}_{i}^{\mu})^{2}-\Big{(}[\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma}]_{i}+([\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu}]_{i})^{2}\Big{)}\Big{]},$

since $\mathbb{E}[\bar{{\bf z}}]=\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu}$ and $\mathbb{V}[\bar{{\bf z}}]=\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma}$. Note that the covariance of $\mbox{\boldmath$\Sigma$\unboldmath}_{b}$, i.e., $\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\Sigma}\triangleq{\mathbb{V}}[\mbox{\boldmath$\Sigma$\unboldmath}_{b}]$, involves computing the fourth-order moments, which is computationally prohibitive. Therefore in practice, we directly set $\mbox{\boldmath$\Sigma$\unboldmath}_{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 ${\mathcal{S}}(\cdot)$: Neighboring Data and Smoothed Statistics. Next, we can borrow data from neighboring label bins $b^{\prime}$ 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):

$\displaystyle\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\mu}=\sum\nolimits_{b^{{}^{\prime}}\in\mathcal{B}}k(y_{b},y_{b^{\prime}})\mbox{\boldmath$\mu$\unboldmath}_{b^{\prime}}^{\mu},\ \ \ \tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\Sigma}=\sum\nolimits_{b^{{}^{\prime}}\in\mathcal{B}}k^{2}(y_{b},y_{b^{\prime}})\mbox{\boldmath$\mu$\unboldmath}_{b^{\prime}}^{\Sigma},\ \ \ \tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu}=\sum\nolimits_{b^{{}^{\prime}}\in\mathcal{B}}k(y_{b},y_{b^{\prime}})\mbox{\boldmath$\Sigma$\unboldmath}_{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 $\{\tilde{{\bf z}}_{i}^{\mu},\tilde{{\bf z}}_{i}^{\Sigma}\}$ for each data point as:

$\displaystyle\tilde{{\bf z}}_{i}^{\mu}$

$\displaystyle=({\bf z}_{i}^{\mu}-\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu})\cdot\sqrt{\tfrac{\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu}}{\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu}}}+\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\mu},\ \ \ \ \ \ \tilde{{\bf z}}_{i}^{\Sigma}=({\bf z}_{i}^{\Sigma}+\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma})\cdot\sqrt{\tfrac{\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu}}{\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu}}}+\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\Sigma}.$

(2)

During training, we keep updating the probabilistic overall statistics, $\{\mbox{\boldmath$\mu$\unboldmath}_{b}^{\mu},\mbox{\boldmath$\mu$\unboldmath}_{b}^{\Sigma},\mbox{\boldmath$\Sigma$\unboldmath}_{b}^{\mu}\}$, and the probabilistic smoothed statistics, $\{\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\mu},\tilde{\mbox{\boldmath$\mu$\unboldmath}}_{b}^{\Sigma},\tilde{\mbox{\boldmath$\Sigma$\unboldmath}}_{b}^{\mu}\}$, across different epochs. The probabilistic representation $\{\tilde{{\bf z}}_{i}^{\mu},\tilde{{\bf z}}_{i}^{\Sigma}\}$ are then re-parameterized [25] into the final representation ${\bf 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 ${\bf x}$’s is only needed during training. During testing, VIR directly uses these statistics and therefore does not need to re-compute them.

Our VIR’s predictor $p(y_{i}|{\bf 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 evidental regression (DER) in [1], as well as intuitions on effective counts in conjugate distributions.

LDS’s Limitations in Our Probabilistic Imbalanced Regression Setting. The motivation of LDS [49] is that the empirical label distribution can not reflect the real label distribution in an imbalanced dataset with a continuous label space; consequently, reweighting methods for imbalanced regression fail due to these inaccurate label densities. By applying a smoothing kernel on the empirical label distribution, LDS tries to recover the effective label distribution, with which reweighting methods can obtain ‘better’ weights to improve imbalanced regression. However, in our probabilistic imbalanced regression, one needs to consider both (1) prediction accuracy for the data with minority labels and (2) uncertainty estimation for each model. Unfortunately, LDS 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.

DER’s Limitations in Our Probabilistic Imbalanced Regression Setting. In DER [1], the predicted labels with their corresponding uncertainties are represented by the approximate posterior parameters $(\gamma,\nu,\alpha,\beta)$ of the NIG distribution $NIG(\gamma,\nu,\alpha,\beta)$. A DER model is trained via minimizing the negative log-likelihood (NLL) of a Student-t distribution:

$\displaystyle\mathcal{L}_{i}^{DER}=\tfrac{1}{2}\log(\tfrac{\pi}{\nu})+(\alpha+\tfrac{1}{2})\log((y_{i}-\gamma)^{2}\nu+\Omega)-\alpha\log(\Omega)+\log(\tfrac{\Gamma(\alpha)}{\Gamma(\alpha+\tfrac{1}{2})}),$

(3)

where $\Omega=2\beta(1+\nu)$. It is therefore nontrivial to properly incorporate a reweighting mechanism into the NLL. One straightforward approach is to directly reweight $\mathcal{L}_{i}^{DER}$ for different data points $({\bf x}_{i},y_{i})$. However, this contradicts the formulation of NIG and often leads to poor performance, as we verify in Sec. 5.

Intuition of Pseudo-Counts for VIR. To properly incorporate different reweighting methods, our VIR relies on the intuition of pseudo-counts (pseudo-observations) in conjugate distributions [4]. Assuming Gaussian likelihood, the conjugate distributions would be an NIG distribution [4], i.e., $(\mu,\Sigma)\sim NIG(\gamma,\nu,\alpha,\beta)$, which means:

$\displaystyle\mu\sim\mathcal{N}(\gamma,\Sigma/\nu),\ \ \ \Sigma\sim\Gamma^{-1}(\alpha,\beta),$

where $\Gamma^{-1}(\alpha,\beta)$ is an inverse gamma distribution. With an NIG prior distribution $NIG(\gamma_{0},\nu_{0},\alpha_{0},\beta_{0})$, the posterior distribution of the NIG after observing $n$ real data points $\{u_{i}\}_{i=1}^{n}$ are:

$\displaystyle\gamma_{n}=\tfrac{\gamma_{0}\nu_{0}+n\Psi}{\nu_{n}},\quad\nu_{n}=\nu_{0}+n,\quad\alpha_{n}=\alpha_{0}+\tfrac{n}{2},\quad\beta_{n}=\beta_{0}+\tfrac{1}{2}(\gamma_{0}^{2}\nu_{0})+\Phi,$

(4)

where $\Psi=\bar{u}$ and $\Phi=\tfrac{1}{2}(\sum_{i}u_{i}^{2}-\gamma_{n}^{2}\nu_{n})$. Here $\nu_{0}$ and $\alpha_{0}$ can be interpreted as virtual observations, i.e., pseudo-counts or pseudo-observations that contribute to the posterior distribution. Overall, the mean of posterior distribution above can be interpreted as an estimation from $(2\alpha_{0}+n)$ observations, with $2\alpha_{0}$ virtual observations and $n$ real observations. Similarly, the variance can be interpreted an estimation from $(\nu_{0}+n)$ observations. This intuition is crucial in developing our VIR’s predictor.

From Pseudo-Counts to Balanced Predictive Distributions. Based on the intuition above, we construct our predictor (i.e., $p(y_{i}|{\bf z}_{i})$) by (1) generating the parameters in the posterior distribution of NIG, (2) computing re-weighted parameters by imposing the importance weights obtained from LDS, and (3) producing the final prediction with corresponding uncertainty estimation.

Based on Eqn. 4, we feed the final representation $\{{\bf z}_{i}\}_{i=1}^{N}$ generated from the Sec. 3.3 (Eqn. 2) into a linear layer to output the intermediate parameters $n_{i},\Psi_{i},\Phi_{i}$ for data point $({\bf x}_{i},y_{i})$:

$\displaystyle n_{i},\Psi_{i},\Phi_{i}={\mathcal{G}}({\bf z}_{i}),\quad{\bf z}_{i}\sim q({\bf z}_{i}|\{{\bf x}_{i}\}_{i=1}^{N})={\mathcal{N}}({\bf z}_{i};\tilde{{\bf z}}_{i}^{\mu},\tilde{{\bf z}}_{i}^{\Sigma})$

We then apply the importance weights $\big{(}\sum_{b^{\prime}\in\mathcal{B}}k(y_{b},y_{b^{\prime}})p(y_{b^{\prime}})\big{)}^{-1/2}$ calculated from the smoothed label distribution to the pseudo-count $n_{i}$ to produce the re-weighted parameters of posterior distribution of NIG, where $p(y)$ denotes the marginal distribution of $y$. Along with the pre-defined prior parameters $(\gamma_{0},\nu_{0},\alpha_{0},\beta_{0})$, we are able to compute the parameters of posterior distribution $NIG(\gamma_{i},\nu_{i},\alpha_{i},\beta_{i})$ for $({\bf x}_{i},y_{i})$:

	$\displaystyle\gamma_{i}^{*}$	$\displaystyle=\tfrac{\gamma_{0}\nu_{0}+\big{(}\sum_{b^{\prime}\in\mathcal{B}}k(y_{b},y_{b^{\prime}})p(y_{b^{\prime}})\big{)}^{-1/2}\cdot n_{i}\Psi_{i}}{\nu_{n}^{*}},\ \ \ \ \ \nu_{i}^{*}=\nu_{0}+\big{(}\sum_{b^{\prime}\in\mathcal{B}}k(y_{b},y_{b^{\prime}})p(y_{b^{\prime}})\big{)}^{-1/2}\cdot n_{i},$
	$\displaystyle\alpha_{i}^{*}$	$\displaystyle=\alpha_{0}+\big{(}\sum_{b^{\prime}\in\mathcal{B}}k(y_{b},y_{b^{\prime}})p(y_{b^{\prime}})\big{)}^{-1/2}\cdot\tfrac{n_{i}}{2},\ \ \ \ \ \beta_{i}^{*}=\beta_{0}+\tfrac{1}{2}(\gamma_{0}^{2}\nu_{0})+\Phi_{i}.$

Based on the NIG posterior distribution, we can then compute final prediction and uncertainty estimation as

$\displaystyle\widehat{y}_{i}=\gamma_{i}^{*},\ \ \ \widehat{s}_{i}=\tfrac{\beta_{i}^{*}}{\nu_{i}^{*}(\alpha_{i}^{*}-1)}.$

We use an objective function similar to Eqn. 3, but with different definitions of $(\gamma,\nu,\alpha,\beta)$, to optimize our VIR model:

$\displaystyle\mathcal{L}_{i}^{\mathcal{P}}=\mathbb{E}_{q_{\phi}({\bf z}_{i}|\{{\bf x}_{i}\}^{N}_{i=1})}\big{[}\tfrac{1}{2}\log(\tfrac{\pi}{\nu_{i}^{*}})+(\alpha_{i}^{*}+\tfrac{1}{2})\log((y_{i}-\gamma_{i}^{*})^{2}\nu_{n}^{*}+\Omega)-\alpha_{i}^{*}\log(\omega_{i}^{*})+\log(\tfrac{\Gamma(\alpha_{i}^{*})}{\Gamma(\alpha_{i}^{*}+\tfrac{1}{2})})\big{]},$

(5)

where $\omega_{i}^{*}=2\beta_{i}^{*}(1+\nu_{i}^{*})$. Note that $\mathcal{L}_{i}^{\mathcal{P}}$ is part of the ELBO in Eqn. 1. Similar to [1], we use an additional regularization term to achieve better accuracy :

$\displaystyle\mathcal{L}_{i}^{\mathcal{R}}$

$\displaystyle=(\nu+2\alpha)\cdot|y_{i}-\widehat{y}_{i}|.$

$\mathcal{L}_{i}^{\mathcal{P}}$ and $\mathcal{L}_{i}^{\mathcal{R}}$ together constitute the objective function for learning the predictor $p({\bf y}_{i}|{\bf z}_{i})$.

Putting together Sec. 3.3 and Sec. 3.4, our final objective function (to minimize) for VIR is:

$\displaystyle\mathcal{L}^{\mathcal{VIR}}=\sum\nolimits_{i=1}^{N}\mathcal{L}_{i}^{\mathcal{VIR}},\quad\mathcal{L}_{i}^{\mathcal{VIR}}=\lambda\mathcal{L}_{i}^{\mathcal{R}}-\mathcal{L}(\theta,\phi;{\bf x}_{i},y_{i})=\lambda\mathcal{L}_{i}^{\mathcal{R}}-\mathcal{L}_{i}^{\mathcal{P}}-\mathcal{L}_{i}^{\mathcal{D}}+\mathcal{L}_{i}^{\mathcal{KL}},$

where $\mathcal{L}(\theta,\phi;{\bf x}_{i},y_{i})=\mathcal{L}_{i}^{\mathcal{P}}+\mathcal{L}_{i}^{\mathcal{D}}-\mathcal{L}_{i}^{\mathcal{KL}}$ is the ELBO in Eqn. 1. $\lambda$ adjusts the importance of the additional regularizer and the ELBO, and thus lead to a better result both on accuracy and uncertainty estimation.

We report the accuracy of different methods in Table 2, Table 2, and Table 3 for AgeDB-DIR, IMDB-WIKI-DIR and STS-B-DIR, respectively. In all the tables, we can observe that our VIR consistently outperforms all baselines in all metrics.

As shown in the last four rows of all three tables, our proposed VIR compares favorably against strong baselines including DIR variants [49] and DER [1], Infer Noise [29], and RankSim [13], especially on the imbalanced data samples (i.e., in the few-shot columns). Notably, VIR improves upon the state-of-the-art method RankSim by $9.6\%$ and $7.9\%$ on AgeDB-DIR and IMDB-WIKI-DIR, respectively, in terms of few-shot GM. This verifies the effectiveness of our methods in terms of overall performance. More accuracy results on different metrics are included in the Appendix. Besides the main results, we also include ablation studies for VIR in the Appendix, showing the effectiveness of VIR’s encoder and predictor.

Different from DIR [49] which only focuses on accuracy, we create a new benchmark for uncertainty estimation in imbalanced regression. Table 5, Table 5, and Table 6 show the results on uncertainty estimation for three datasets AgeDB-DIR, IMDB-WIKI-DIR, and STS-B-DIR, respectively. Note that most baselines from Table 2, Table 2, and Table 3 are deterministic methods (as opposed to probabilistic methods like ours) and cannot provide uncertainty estimation; therefore they are not applicable here. To show the superiority of our VIR model, we create a strongest baseline by concatenating the DIR variants (LDS + FDS) with the DER [1].

Results show that our VIR consistently outperforms all baselines across different metrics, especially in the few-shot metrics. Note that our proposed methods mainly focus on the imbalanced setting, and therefore naturally places more emphasis on the few-shot metrics. Notably, on AgeDB-DIR, IMDB-WIKI-DIR, and STS-B-DIR, our VIR improves upon the strongest baselines, by $14.2\%\sim 17.1\%$ in terms of few-shot AUSE.