SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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:
$$
\mathcal {L} _ {i} ^ {D E R} = \frac {1}{2} \log (\frac {\pi}{\nu}) + (\alpha + \frac {1}{2}) \log ((y _ {i} - \gamma) ^ {2} \nu + \Omega) - \alpha \log (\Omega) + \log (\frac {\Gamma (\alpha)}{\Gamma (\alpha + \frac {1}{2})}), \tag {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  $(\mathbf{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:
$$
\mu \sim \mathcal {N} (\gamma , \Sigma / \nu), \quad \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:
$$
\gamma_ {n} = \frac {\gamma_ {0} \nu_ {0} + n \Psi}{\nu_ {n}}, \quad \nu_ {n} = \nu_ {0} + n, \quad \alpha_ {n} = \alpha_ {0} + \frac {n}{2}, \quad \beta_ {n} = \beta_ {0} + \frac {1}{2} (\gamma_ {0} ^ {2} \nu_ {0}) + \Phi , \qquad (4)
$$
where  $\Psi = \bar{u}$  and  $\Phi = \frac{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 as 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}|\mathbf{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  $\{\mathbf{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  $(\mathbf{x}_i, y_i)$ :
$$
n _ {i}, \Psi_ {i}, \Phi_ {i} = \mathcal {G} (\mathbf {z} _ {i}), \quad \mathbf {z} _ {i} \sim q (\mathbf {z} _ {i} | \{\mathbf {x} _ {i} \} _ {i = 1} ^ {N}) = \mathcal {N} (\mathbf {z} _ {i}; \widetilde {\mathbf {z}} _ {i} ^ {\mu}, \widetilde {\mathbf {z}} _ {i} ^ {\Sigma})
$$
We then apply the importance weights  $\left(\sum_{b' \in \mathcal{B}} k(y_b, y_{b'}) p(y_{b'})\right)^{-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  $(\mathbf{x}_i, y_i)$ :
$$
\gamma_ {i} ^ {*} = \frac {\gamma_ {0} \nu_ {0} + \left(\sum_ {b ^ {\prime} \in \mathcal {B}} k (y _ {b} , y _ {b ^ {\prime}}) p (y _ {b ^ {\prime}})\right) ^ {- 1 / 2} \cdot n _ {i} \Psi_ {i}}{\nu_ {n} ^ {*}}, \quad \nu_ {i} ^ {*} = \nu_ {0} + \left(\sum_ {b ^ {\prime} \in \mathcal {B}} k (y _ {b}, y _ {b ^ {\prime}}) p (y _ {b ^ {\prime}})\right) ^ {- 1 / 2} \cdot n _ {i},
$$
$$
\alpha_ {i} ^ {*} = \alpha_ {0} + \left(\sum_ {b ^ {\prime} \in \mathcal {B}} k \left(y _ {b}, y _ {b ^ {\prime}}\right) p \left(y _ {b ^ {\prime}}\right)\right) ^ {- 1 / 2} \cdot \frac {n _ {i}}{2}, \quad \beta_ {i} ^ {*} = \beta_ {0} + \frac {1}{2} \left(\gamma_ {0} ^ {2} \nu_ {0}\right) + \Phi_ {i}.
$$
Based on the NIG posterior distribution, we can then compute final prediction and uncertainty estimation as
$$
\widehat {y} _ {i} = \gamma_ {i} ^ {*}, \quad \widehat {s} _ {i} = \frac {\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:
$$
\mathcal {L} _ {i} ^ {\mathcal {P}} = \mathbb {E} _ {q _ {\phi} (\mathbf {z} _ {i} | \{\mathbf {x} _ {i} \} _ {i = 1} ^ {N})} \left[ \frac {1}{2} \log \left(\frac {\pi}{\nu_ {i} ^ {*}}\right) + \left(\alpha_ {i} ^ {*} + \frac {1}{2}\right) \log \left(\left(y _ {i} - \gamma_ {i} ^ {*}\right) ^ {2} \nu_ {n} ^ {*} + \Omega\right) - \alpha_ {i} ^ {*} \log \left(\omega_ {i} ^ {*}\right) + \log \left(\frac {\Gamma \left(\alpha_ {i} ^ {*}\right)}{\Gamma \left(\alpha_ {i} ^ {*} + \frac {1}{2}\right)}\right) \right], \tag {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:
$$
\mathcal {L} _ {i} ^ {\mathcal {R}} = (\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(\mathbf{y}_i|\mathbf{z}_i)$ .
Table 1: Accuracy on AgeDB-DIR.
Metrics	MAE ↓	GM ↓
Shot	all	many	medium	few	all	many	medium	few
VANILLA [49]	7.77	6.62	9.55	13.67	5.05	4.23	7.01	10.75
VAE [25]	7.63	6.58	9.21	13.45	4.86	4.11	6.61	10.24
DEEP ENS. [27]	7.73	6.62	9.37	13.90	4.87	4.37	6.50	11.35
INFER NOISE [29]	8.53	7.62	9.73	13.82	5.57	4.95	6.58	10.86
SMOTER [40]	8.16	7.39	8.65	12.28	5.21	4.65	5.69	8.49
SMOGN [5]	8.26	7.64	9.01	12.09	5.36	4.9	6.19	8.44
SQINV [49]	7.81	7.16	8.80	11.2	4.99	4.57	5.73	7.77
DER [1]	8.09	7.31	8.99	12.66	5.19	4.59	6.43	10.49
LDS [49]	7.67	6.98	8.86	10.89	4.85	4.39	5.8	7.45
FDS [49]	7.69	7.10	8.86	9.98	4.83	4.41	5.97	6.29
LDS + FDS [49]	7.55	7.01	8.24	10.79	4.72	4.36	5.45	6.79
RANKSIM [13]	7.02	6.49	7.84	9.68	4.53	4.13	5.37	6.89
LDS + FDS + DER [1]	8.18	7.44	9.52	11.45	5.30	4.75	6.74	7.68
VIR (OURS)	6.99	6.39	7.47	9.51	4.41	4.07	5.05	6.23
OURS VS. VANILLA	+0.78	+0.23	+2.08	+4.16	+0.64	+0.16	+1.96	+4.52
OURS VS. INFER NOISE	+1.54	+1.23	+2.26	+4.31	+1.16	+0.88	+1.53	+4.63
OURS VS. DER	+1.10	+0.92	+1.52	+3.15	+0.78	+0.52	+1.38	+4.26
OURS VS. LDS + FDS	+0.56	+0.62	+0.77	+1.28	+0.31	+0.29	+0.40	+0.56
OURS VS. RANKSIM	+0.03	+0.10	+0.37	+0.17	+0.12	+0.06	+0.32	+0.66
Table 2: Accuracy on IW-DIR.
Metrics	MAE ↓	GM ↓
Shot	All	Many	Medium	Few	All	Many	Medium	Few
VANILLA [49]	8.06	7.23	15.12	26.33	4.57	4.17	10.59	20.46
VAE [25]	8.04	7.20	15.05	26.30	4.57	4.22	10.56	20.72
DEEP ENS. [27]	8.08	7.31	15.09	26.47	4.59	4.26	10.61	21.13
INFER NOISE [29]	8.11	7.36	15.23	26.29	4.68	4.33	10.65	20.31
SMOTER [40]	8.14	7.42	14.15	25.28	4.64	4.30	9.05	19.46
SMOGN [5]	8.03	7.30	14.02	25.93	4.63	4.30	8.74	20.12
SQINV [49]	7.87	7.24	12.44	22.76	4.47	4.22	7.25	15.10
DER [1]	7.85	7.18	13.35	24.12	4.47	4.18	8.18	15.18
LDS [49]	7.83	7.31	12.43	22.51	4.42	4.19	7.00	13.94
FDS [49]	7.83	7.23	12.60	22.37	4.42	4.20	6.93	13.48
LDS + FDS [49]	7.78	7.20	12.61	22.19	4.37	4.12	7.39	12.61
RANKSIM [13]	7.50	6.93	12.09	21.68	4.19	3.97	6.65	13.28
LDS + FDS + DER [1]	7.24	6.64	11.87	23.44	3.93	3.69	6.64	16.00
VIR (OURS)	7.19	6.56	11.81	20.96	3.85	3.63	6.51	12.23
OURS VS. VANILLA	+0.87	+0.67	+3.31	+5.37	+0.72	+0.54	+4.08	+8.23
OURS VS. INFER NOISE	+0.92	+0.80	+3.42	+5.33	+0.83	+0.70	+4.14	+8.08
OURS VS. DER	+0.66	+0.62	+1.54	+3.16	+0.62	+0.55	+1.67	+2.95
OURS VS. LDS + FDS	+0.59	+0.64	+0.8	+1.23	+0.52	+0.49	+0.88	+0.38
OURS VS. RANKSIM	+0.31	+0.37	+0.28	+0.72	+0.34	+0.34	+0.14	+1.05
3.5 Final Objective Function
Putting together Sec. 3.3 and Sec. 3.4, our final objective function (to minimize) for VIR is:
$$
\mathcal {L} ^ {\mathcal {V I R}} = \sum_ {i = 1} ^ {N} \mathcal {L} _ {i} ^ {\mathcal {V I R}}, \quad \mathcal {L} _ {i} ^ {\mathcal {V I R}} = \lambda \mathcal {L} _ {i} ^ {\mathcal {R}} - \mathcal {L} (\theta , \phi ; \mathbf {x} _ {i}, y _ {i}) = \lambda \mathcal {L} _ {i} ^ {\mathcal {R}} - \mathcal {L} _ {i} ^ {\mathcal {P}} - \mathcal {L} _ {i} ^ {\mathcal {D}} + \mathcal {L} _ {i} ^ {\mathcal {K L}},
$$
where  $\mathcal{L}(\theta, \phi; \mathbf{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.
Notation. As mentioned in Sec. 3.1, we partitioned  $\{Y_j\}_{j=1}^N$  into  $|\mathcal{B}|$  equal-interval bins (denote the set of bins as  $\mathcal{B}$ ), and  $\{Y_j\}_{j=1}^N$  are sampled from the label space  $\mathcal{V}$ . In addition, We use the binary set  $\{P_i\}_{i=1}^{|\mathcal{B}|}$  to represent the label distribution (frequency) for each bin  $i$ , i.e.,  $P_i \triangleq \mathbb{P}(Y_j \in \mathcal{B}_i)$ . We also use the binary set  $\{O_i\}_{i=1}^N$  to represent whether the data point  $(\mathbf{x}_j, y_j)$  is observed (i.e.,  $O_j \sim \mathbb{P}(O_j = 1) \propto P_{B(Y_j)}$ , and  $\mathbb{E}_O[O_j] \sim P_{B(Y_j)}$ ), where  $B(Y_j)$  represents the bin which  $(x_j, y_j)$  belongs to. For each bin  $i \in \mathcal{B}$ , we denote the associated set of data points as
$$
\mathcal {U} _ {i} \triangleq \{j: Y _ {j} = i \}.
$$
When the imbalanced dataset is partially observed, we denote the observation set as:
$$
\mathcal {S} _ {i} \triangleq \left\{j: O _ {j} = 1 \& B \left(Y _ {j}\right) = i \right\}.
$$
Definition 4.1 (Expectation over Observability  $\mathbb{E}_O$ ). We define the expectation over the observability variable  $O$  as  $\mathbb{E}_O[\cdot ]\equiv \mathbb{E}_{O_j\sim \mathbb{P}(O_j = 1)}[\cdot ]$ .
Definition 4.2 (True Risk). Based on the previous definitions, the true risk is defined as:
$$
R (\widehat {Y}) = \frac {1}{| \mathcal {B} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \frac {1}{| \mathcal {U} _ {i} |} \sum_ {j \in \mathcal {U} _ {i}} \delta_ {j} (Y, \widehat {Y}),
$$
where  $\delta_j(Y, \widehat{Y})$  refers to some loss function (e.g. MAE, MSE). In this paper, we assume the loss is upper bounded by  $\Delta$ , i.e.,  $0 \leq \delta_j(Y, \widehat{Y}) \leq \Delta$ .
Below we define the Naive Estimator.
Definition 4.3 (Naive Estimator). Given the observation set, the Naive Estimator is defined as:
$$
\widehat {R} _ {\mathrm {N A I V E}} (\widehat {Y}) = \frac {1}{\sum_ {i = 1} ^ {| \mathcal {B} |} | \mathcal {S} _ {i} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \sum_ {j \in \mathcal {S} _ {i}} \delta_ {j} (Y, \widehat {Y}).
$$
It is easy to verify that the expectation of this naive estimator is not equal to the true risk, as  $\mathbb{E}_O[\widehat{R}_{\mathrm{NAIVE}}(\widehat{Y})] \neq R(\widehat{Y})$ .
Considering an imbalanced dataset as a subset of observations from a balanced one, we contrast it with the Inverse Propensity Score (IPS) estimator [34].
Definition 4.4 (Inverse Propensity Score Estimator). The inverse propensity score (IPS) estimator (an unbiased estimator) is defined as
$$
\widehat {R} _ {\mathrm {I P S}} (\widehat {Y} | P) = \frac {1}{| \mathcal {B} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \frac {1}{| \mathcal {U} _ {i} |} \sum_ {j \in \mathcal {S} _ {i}} \frac {\delta_ {j} (Y , \widehat {Y})}{P _ {i}}.
$$
The IPS estimator is an unbiased estimator, as we can verify by taking the expectation value over the observation set:
$$
\begin{array}{l} \mathbb {E} _ {O} [ \widehat {R} _ {\mathrm {I P S}} (\widehat {Y} | P) ] = \frac {1}{| \mathcal {B} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \frac {1}{| \mathcal {U} _ {i} |} \sum_ {j \in \mathcal {U} _ {i}} \frac {\delta_ {j} (Y , \widehat {Y})}{P _ {i}} \cdot \mathbb {E} _ {O} [ O _ {j} ] \\ = \frac {1}{| \mathcal {B} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \frac {1}{| \mathcal {U} _ {i} |} \sum_ {j \in \mathcal {U} _ {i}} \delta_ {j} (Y, \widehat {Y}) = R (\widehat {Y}). \\ \end{array}
$$
Finally, we define our VIR/DIR estimator below.
Definition 4.5 (VIR Estimator). The VIR estimator, denoted by  $\widehat{R}_{\mathrm{VIR}}(\widehat{Y}|\widetilde{P})$ , is defined as:
$$
\widehat {R} _ {\mathrm {V I R}} (\widehat {Y} | \widetilde {P}) = \frac {1}{| \mathcal {B} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \frac {1}{| \mathcal {U} _ {i} |} \sum_ {j \in \mathcal {S} _ {i}} \frac {\delta_ {j} (Y , \widehat {Y})}{\widetilde {P} _ {i}}, \tag {6}
$$
where  $\{\widetilde{P}_i\}_{i=1}^{|\mathcal{B}|}$  represents the smoothed label distribution used in our VIR's objective function. It is important to note that our VIR estimator is biased.
For multiple predictions, we select the "best" estimator according to the following definition.
Definition 4.6 (Empirical Risk Minimizer). For a given hypothesis space  $\mathcal{H}$  of predictions  $\widehat{Y}$ , the Empirical Risk Minimization (ERM) identifies the prediction  $\widehat{Y} \in \mathcal{H}$  as
$$
\widehat {Y} ^ {\text {E R M}} = \operatorname {a r g m i n} _ {\widehat {Y} \in \mathcal {H}} \left\{\widehat {R} _ {\text {V I R}} (\widehat {Y} | \widetilde {P}) \right\}
$$
With all the aforementioned definitions, we can derive the generalization bound for the VIR estimator.
Theorem 4.1 (Generalization Bound of VIR). In imbalanced regression with bins  $\mathcal{B}$ , for any finite hypothesis space of predictions  $\mathcal{H} = \{\widehat{Y}_1, \ldots, \widehat{Y}_{\mathcal{H}}\}$ , the transductive prediction error of the empirical risk minimizer  $\widehat{Y}^{ERM}$  using the VIR estimator with estimated propensities  $\widetilde{P}(P_i > 0)$  and given training observations  $O$  from  $\mathcal{V}$  with independent Bernoulli propensities  $P$ , is bounded by:
$$
R \left(\widehat {Y} ^ {E R M}\right) \leq \widehat {R} _ {\mathrm {V I R}} \left(\widehat {Y} ^ {E R M} \mid \widetilde {P}\right) + \underbrace {\frac {\Delta}{| \mathcal {B} |} \sum_ {i = 1} ^ {| \mathcal {B} |} \left| 1 - \frac {P _ {i}}{\widetilde {P} _ {i}} \right|} _ {\text {B i a s T e r m}} + \underbrace {\frac {\Delta}{| \mathcal {B} |} \sqrt {\frac {\log (2 | \mathcal {H} | / \eta)}{2}} \sqrt {\sum_ {i = 1} ^ {| \mathcal {B} |} \frac {1}{\widetilde {P} _ {i} ^ {2}}}} _ {\text {V a r i a n c e T e r m}}. \tag {7}
$$
Remark. The naive estimator (i.e., Definition 4.3) has large bias and large variance. If one directly uses the original label distribution in the training objective (i.e., Definition 4.4), i.e.,  $\widetilde{P}_i = P_i$ , the "bias" term will be 0. However, the "variance" term will be extremely large for minority data because  $\widetilde{P}_i$  is very close to 0. In contrast, under VIR's N.I.D.,  $\widetilde{P}_i$  used in the training objective function will be smoothed. Therefore, the minority data's label density  $\widetilde{P}_i$  will be smoothed out by its neighbors and becomes larger (compared to the original  $P_i$ ), leading to smaller "variance" in the generalization error bound. Note that  $\widetilde{P}_i \neq P_i$ , VIR (with N.I.D.) essentially increases bias, but significantly reduces its variance in the imbalanced setting, thereby leading to a lower generalization error.
Datasets. We evaluate our methods in terms of prediction accuracy and uncertainty estimation on four imbalanced datasets<sup>1</sup>, AgeDB-DIR [30], IMDB-WIKI-DIR [33], STS-B-DIR [7], and NYUD2-DIR [35]. Due to page limit, results for NYUD2-DIR [35] are moved to the Appendix. We follow the preprocessing procedures in DIR [49]. Details for each datasets are in the Appendix, and please refer to [49] for details on label density distributions and levels of imbalance.
Baselines. We use ResNet-50 [17] (for AgeDB-DIR and IMDB-WIKI-DIR) and BiLSTM [19] (for STS-B-DIR) as our backbone networks, and more details for baseline are in the Appendix. We describe the baselines below.
- Vanilla: We use the term VANILLA to denote a plain model without adding any approaches.
- Synthetic-Sample-Based Methods: Various existing imbalanced regression methods are also included as baselines; these include Deep Ensemble [27], Infer Noise [29], SMOTER [40], and SMOGN [5].
- Cost-Sensitive Reweighting: As shown in DIR [49], the square-root weighting variant (SQINV) baseline (i.e.,  $\left(\sum_{b' \in \mathcal{B}} k(y_b, y_{b'}) p(y_{b'})\right)^{-1/2}$ ) always outperforms Vanilla. Therefore, for fair comparison, all our experiments (for both baselines and VIR) use SQINV weighting.
Evaluation Metrics. We follow the evaluation metrics in [49] to evaluate the accuracy of our proposed methods; these include Mean Absolute Error (MAE), Mean Squared Error (MSE). Furthermore, for AgeDB-DIR and IMDB-WIKI-DIR, we use Geometric Mean (GM) to evaluate the accuracy; for STS-B-DIR, we use Pearson correlation and Spearman correlation. We use typical evaluation metrics for uncertainty estimation in regression problems to evaluate our produced uncertainty estimation; these include Negative Log Likelihood (NLL), Area Under Sparsification Error (AUSE). Eqn. 5 shows the formula for NLL, and more details regarding to AUSE can be found in [21]. We also include calibrated uncertainty results for VIR in the Appendix.
Evaluation Process. Following [28, 49], for a data sample  $x_{i}$  with its label  $y_{i}$  which falls into the target bins  $b_{i}$ , we divide the label space into three disjoint subsets: many-shot region  $\{b_{i} \in \mathcal{B} \mid y_{i} \in b_{i} \& |y_{i}| > 100\}$ , medium-shot region  $\{b_{i} \in \mathcal{B} \mid y_{i} \in b_{i} \& 20 \leq |y_{i}| \leq 100\}$ , and few-shot region  $\{b_{i} \in \mathcal{B} \mid y_{i} \in b_{i} \& |y_{i}| < 20\}$ , where  $|\cdot|$  denotes the cardinality of the set. We report results on the overall test set and these subsets with the accuracy metrics and uncertainty metrics discussed above.
Implementation Details. We conducted five separate trials for our method using different random seeds. The error bars and other implementation details are included in the Appendix.
Table 3: Accuracy on STS-B-DIR.
Metrics	MSE ↓	Pearson ↑
Shot	All	Many	Medium	Few	All	Many	Medium	Few
VANILLA [49]	0.974	0.851	1.520	0.984	0.742	0.720	0.627	0.752
VAE [25]	0.968	0.833	1.511	1.102	0.751	0.724	0.621	0.749
DEEP ENS. [27]	0.972	0.846	1.496	1.032	0.746	0.723	0.619	0.750
INFER NOISE [29]	0.954	0.980	1.408	0.967	0.747	0.711	0.631	0.756
SMOTER [40]	1.046	0.924	1.542	1.154	0.726	0.693	0.653	0.706
SMOGN [5]	0.990	0.896	1.327	1.175	0.732	0.704	0.655	0.692
INV [49]	1.005	0.894	1.482	1.046	0.728	0.703	0.625	0.732
DER [1]	1.001	0.912	1.368	1.055	0.732	0.711	0.646	0.742
LDS [49]	0.914	0.819	1.319	0.955	0.756	0.734	0.638	0.762
FDS [49]	0.927	0.851	1.225	1.012	0.750	0.724	0.667	0.742
LDS + FDS [49]	0.907	0.802	1.363	0.942	0.760	0.740	0.652	0.766
RANKSIM [13]	0.903	0.908	0.911	0.804	0.758	0.706	0.690	0.827
LDS + FDS + DER [1]	1.007	0.880	1.535	1.086	0.729	0.714	0.635	0.731
VIR (OURS)	0.892	0.795	0.899	0.781	0.776	0.752	0.696	0.845
OURS VS. VANILLA	+0.082	+0.056	+0.621	+0.203	+0.034	+0.032	+0.069	+0.093
OURS VS. INFER NOISE	+0.062	+0.185	+0.509	+0.186	+0.029	+0.041	+0.065	+0.089
OURS VS. DER	+0.109	+0.117	+0.469	+0.274	+0.044	+0.041	+0.050	+0.103
OURS VS. LDS + FDS	+0.015	+0.007	+0.464	+0.161	+0.016	+0.012	+0.044	+0.079
OURS VS. RANKSIM	+0.011	+0.113	+0.012	+0.023	+0.018	+0.046	+0.006	+0.018
5.1 Imbalanced Regression Accuracy
We report the accuracy of different methods in Table 1, 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
<sup>1</sup>Among the five datasets proposed in [49], only four of them are publicly available.
Table 4: Uncertainty estimation on AgeDB-DIR.
Metrics	NLL ↓	AUSE ↓
Shot	All	Many	Med	Few	All	Many	Med	Few
DEEP ENS. [27]	5.311	4.031	6.726	8.523	0.541	0.626	0.466	0.483
INFER NOISE [29]	4.616	4.413	4.866	5.842	0.465	0.458	0.457	0.496
DER [1]	3.918	3.741	3.919	4.432	0.523	0.464	0.449	0.486
LDS + FDS + DER [1]	3.787	3.689	3.912	4.234	0.451	0.460	0.399	0.565
VIR (OURS)	3.703	3.598	3.805	4.196	0.434	0.456	0.324	0.414
OURS vs. DER	+0.215	+0.143	+0.114	+0.236	+0.089	+0.008	+0.125	+0.072
OURS vs. LDS + FDS + DER	+0.084	+0.091	+0.107	+0.038	+0.017	+0.004	+0.075	+0.151
Table 5: Uncertainty estimation on IW-DIR.
Metrics	NLL ↓	AUSE ↓
Shot	All	Many	Medium	Few	All	Many	Medium	Few
DEEP ENS. [27]	5.219	4.102	7.123	8.852	0.846	0.862	0.745	0.718
INFER NOISE [29]	4.231	4.078	5.326	8.292	0.732	0.728	0.561	0.478
DER [1]	3.850	3.699	4.997	6.638	0.813	0.802	0.650	0.541
LDS + FDS + DER [1]	3.683	3.602	4.391	5.697	0.784	0.670	0.459	0.483
VIR (OURS)	3.651	3.579	4.296	5.518	0.634	0.649	0.434	0.379
OURS vs. DER	+0.199	+0.120	+0.701	+1.120	+0.179	+0.153	+0.216	+0.162
OURS vs. LDS + FDS + DER	+0.032	+0.023	+0.095	+0.179	+0.150	+0.021	+0.025	+0.104
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.
5.2 Imbalanced Regression Uncertainty Estimation
Different from DIR [49] which only focuses on accuracy, we create a new benchmark for uncertainty estimation in imbalanced regression. Table 4, 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 1, 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  $(\mathrm{LDS} + \mathrm{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.