SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kaiwen Zha $^{1,*}$  Peng Cao $^{1,*}$  Jeany Son $^{2}$  Yuzhe Yang $^{1}$  Dina Katabi $^{1}$
$^{1}$ MIT CSAIL  $^{2}$ GIST
Deep regression models typically learn in an end-to-end fashion without explicitly emphasizing a regression-aware representation. Consequently, the learned representations exhibit fragmentation and fail to capture the continuous nature of sample orders, inducing suboptimal results across a wide range of regression tasks. To fill the gap, we propose Rank-N-Contrast (RNC), a framework that learns continuous representations for regression by contrasting samples against each other based on their rankings in the target space. We demonstrate, theoretically and empirically, that RNC guarantees the desired order of learned representations in accordance with the target orders, enjoying not only better performance but also significantly improved robustness, efficiency, and generalization. Extensive experiments using five real-world regression datasets that span computer vision, human-computer interaction, and healthcare verify that RNC achieves state-of-the-art performance, highlighting its intriguing properties including better data efficiency, robustness to spurious targets and data corruptions, and generalization to distribution shifts. Code is available at: https://github.com/kaienzha/Rank-N-Contrast.
Figure 1: Learned representations of different methods on a real-world temperature regression task [7] (details in Sec. 5). Existing general regression learning  $(L_{1})$  or representation learning (SupCon) schemes fail to recognize the underlying continuous information in data. In contrast, RNC learns continuous representations that capture the intrinsic sample orders w.r.t. the regression targets.
Regression problems are pervasive and fundamental in the real world, spanning various tasks and domains including estimating age from human appearance [36], predicting health scores via human physiological signals [11], and detecting gaze directions using webcam images [52]. Due to the
*The two authors contributed equally and order was determined by a random coin flip. Correspondence to Kaiwen Zha <kzha@mit.edu>, Peng Cao <pengcao@mit.edu>.
continuity in regression targets, the most widely adopted approach for training regression models is to directly predict the target value and employ a distance-based loss function, such as  $L_{1}$  or  $L_{2}$  distance, between the prediction and the ground-truth target [51, 52, 38]. Methods that tackle regression tasks using classification models trained with cross-entropy loss have also been studied [36, 33, 40].
However, previous methods focus on imposing constraints on the final predictions in an end-to-end fashion, but do not explicitly emphasize the representations learned by the model. Unfortunately, these representations are often fragmented and incapable of capturing the continuous relationships that underlie regression tasks. Fig. 1(a) highlights the representations learned by the  $L_{1}$  loss on SkyFinder [7], a regression dataset for predicting weather temperature from webcam outdoor images captured at different locations (details in Sec. 5). Rather than exhibiting the continuous ground-truth temperatures, the learned representations are grouped by different webcams in a fragmented manner. Such unordered and fragmented representation is suboptimal for the regression task and can even hamper performance by including irrelevant information, such as the capturing webcam.
Furthermore, despite the great success of representation learning schemes on solving discrete classification or segmentation tasks (e.g., contrastive learning [4, 20] and supervised contrastive learning (SupCon) [25]), less attention has been paid to designing algorithms that capture the intrinsic continuity in data for regression. Interestingly, we highlight that existing representation learning methods inevitably overlook the continuous nature in data: Fig. 1(b) shows the representation learned by SupCon on the SkyFinder dataset, where it again fails to capture the underlying continuous order between the samples, resulting in a suboptimal representation for regression tasks.
To fill the gap, we present Rank-N-Contrast (RNC), a novel framework for generic regression learning. RNC first learns a regression-aware representation that orders the distances in the embedding space based on the target values, and then leverages it to predict the continuous targets. To achieve this, we propose the Rank-N-Contrast loss  $(\mathcal{L}_{\mathrm{RNC}})$ , which ranks the samples in a batch according to their labels and then contrasts them against each other based on their relative rankings. Theoretically, we prove that optimizing  $\mathcal{L}_{\mathrm{RNC}}$  results in features that are ordered according to the continuous labels, leading to improved performance in downstream regression tasks. As confirmed in Fig. 1(c), RNC learns continuous representations that capture the intrinsic ordered relationships between samples. Notably, our framework is orthogonal to existing regression methods, allowing for the use of any regression method to map the learned representation to the final prediction values.
To support practical evaluations, we benchmark RNC against state-of-the-art (SOTA) regression and representation learning schemes on five real-world regression datasets that span computer vision, human-computer interaction, and healthcare. Rigorous experiments verify the superior performance, robustness, and efficiency of RNC on learning continuous targets. Our contributions are as follows:
- We identify the limitation of current regression and representation learning methods for continuous targets, and uncover intrinsic properties of learning regression-aware representations.
- We design RNC, a simple & effective method that learns continuous representations for regression.
- We conduct extensive experiments on five diverse regression datasets in vision, human-computer interaction, and healthcare, verifying the superior performance of RNC against SOTA schemes.
- Further analyses reveal intriguing properties of RNC on its data efficiency, robustness to spurious targets & data corruptions, and better generalization to unseen targets.
Regression Learning. Deep learning has achieved great success in addressing regression tasks [36, 51, 52, 38, 11]. In regression learning, the final predictions of the model are typically trained in an end-to-end manner to be close to the targets. Standard regression losses include the  $L_{1}$  loss, the mean squared error (MSE) loss, and the Huber loss [22]. Past work has also proposed several variants. One branch of work [36, 13, 14, 35] divides the regression range into small bins, converting the problem into a classification task. Another line of work [12, 2, 40] casts regression as an ordinal classification problem [33] using ordered thresholds and employing multiple binary classifiers. Recently, a line of work proposes to regularize the embedding space for regression, ranging from modeling feature space uncertainty [28], encouraging higher-entropy feature spaces [50], to regularizing features for imbalanced regression [44, 17]. In contrast to existing works, we provide a regression-aware representation learning approach that emphasizes the continuity in the features space w.r.t. the targets, which enjoys better performance while being compatible to prior regression schemes. In addition,
Figure 2: Illustration of  $\mathcal{L}_{\mathrm{RNC}}$  in the context of positive and negative pairs. (a) An example batch of input data and their labels. (b) Two example positive pairs and corresponding negative pair(s) when the anchor is the 20-year-old man (shown in gray shading). When the anchor forms a positive pair with a 30-year-old man, their label distance is 10, hence the corresponding negative samples are the 0-year-old baby and the 70-year-old man, whose label distances to the anchor are larger than 10. When the 0-year-old baby creates a positive pair with the anchor, only the 70-year-old man has a larger label distance to the anchor, thus serving as a negative sample.
C-Mixup [47] adapts the original mixup [49] by adjusting the sampling probability of the mixed pairs according to the target similarities. It is also worth noting that our method is orthogonal and complementary to data augmentation algorithms for regression learning, such as C-Mixup [47].
Representation Learning. Representation learning is crucial in machine learning, often studied in the context of classification. Recently, contrastive learning has emerged as a popular technique for self-supervised representation learning [4, 20, 5, 8, 3]. The supervised version of contrastive learning, SupCon [25], has been shown to outperform the conventional cross-entropy loss on multiple discrete classification tasks, including image recognition [25], noisy labels [26], long-tailed classification [23, 27], and out-of-domain detection [48]. A few recent papers propose to adapt SupCon to tackle ordered labels in specific downstream applications, including gaze estimation [42], medical imaging [9, 10], and neural behavior analysis [37]. Different from prior works that directly adapt SupCon, we incorporate the intrinsic property of label continuity for designing representation learning scheme tailored for regression, which offers a simple and principled approach for generic regression tasks.
3 Our Approach: Rank-N-Contrast (RNC)
Problem Setup. Given a regression task, we aim to train a neural network model composed of a feature encoder  $f(\cdot): X \to \mathbb{R}^{d_e}$  and a predictor  $g(\cdot): \mathbb{R}^{d_e} \to \mathbb{R}^{d_t}$  to predict the target  $y \in \mathbb{R}^{d_t}$  based on the input data  $x \in X$ .
For a positive integer  $I$ , let  $[I]$  denote the set  $\{1,2,\dots,I\}$ . Given a randomly sampled batch of  $N$  input and label pairs  $\{(\pmb{x}_n,\pmb{y}_n)\}_{n\in [N]}$ , we apply standard data augmentations to obtain a two-view batch  $\{(\tilde{\pmb{x}}_{\ell},\tilde{\pmb{y}}_{\ell})\}_{\ell \in [2N]}$ , where  $\tilde{\pmb{x}}_{2n} = t(\pmb{x}_n)$  and  $\tilde{\pmb{x}}_{2n - 1} = t^{\prime}(\pmb{x}_n)$ , with  $t$  and  $t^\prime$  being independently sampled augmentation operations, and  $\tilde{\pmb{y}}_{2n} = \tilde{\pmb{y}}_{2n - 1} = \pmb{y}_n$ ,  $\forall n\in [N]$ . The augmented batch is then fed into the encoder  $f$  to obtain the feature embedding for each augmented input data, i.e.,  $\pmb{v}_l = f(\tilde{\pmb{x}}_l)\in \mathbb{R}^{d_e},\forall l\in [2N]$ . The representation learning phase is then performed over the feature embeddings. To harness the acquired representation for regression, we freeze the encoder  $f(\cdot)$  and train the predictor  $g(\cdot)$  on top of it using a regression loss (e.g.,  $L_{1}$  loss).
In this context, a natural question arises: How to design a regression-aware representation learning scheme tailored for continuous and ordered samples?
The Rank-N-Contrast Loss. In order to align distances in the embedding space ordered by distances in their labels, we propose Rank-N-Contrast loss  $(\mathcal{L}_{\mathrm{RNC}})$ , which first ranks the samples according to their target distances, and then contrasts them against each other based on their relative rankings.
Following [16, 45], given an anchor  $\pmb{v}_i$ , we model the likelihood of any other  $\pmb{v}_j$  to increase exponentially with respect to their similarity in the representation space [15]. Inspired by the listwise ranking methods [43, 6], we introduce  $S_{i,j} \coloneqq \{ \pmb{v}_k \mid k \neq i, d(\tilde{\pmb{y}}_i, \tilde{\pmb{y}}_k) \geq d(\tilde{\pmb{y}}_i, \tilde{\pmb{y}}_j) \}$  to denote the set of samples that are of higher ranks than  $\pmb{v}_j$  in terms of label distance w.r.t.  $\pmb{v}_i$ , where  $d(\cdot, \cdot)$  is the distance measure between two labels (e.g.,  $L_1$  distance). Then the normalized likelihood of  $\pmb{v}_j$  given  $\pmb{v}_i$  and  $S_{i,j}$  can be written as
$$
\mathbb {P} \left(\boldsymbol {v} _ {j} \mid \boldsymbol {v} _ {i}, \mathcal {S} _ {i, j}\right) = \frac {\exp \left(\operatorname {s i m} \left(\boldsymbol {v} _ {i} , \boldsymbol {v} _ {j}\right) / \tau\right)}{\sum_ {\boldsymbol {v} _ {k} \in \mathcal {S} _ {i , j}} \exp \left(\operatorname {s i m} \left(\boldsymbol {v} _ {i} , \boldsymbol {v} _ {k}\right) / \tau\right)}, \tag {1}
$$
where  $\mathrm{sim}(\cdot, \cdot)$  is the similarity measure between two feature embeddings (e.g., negative  $L_{2}$  norm) and  $\tau$  denotes the temperature parameter. Note that the denominator is a sum over the set of samples that possess higher ranks than  $\pmb{v}_{j}$ ; Maximizing  $\mathbb{P}(\pmb{v}_j|\pmb{v}_i,\mathcal{S}_{i,j})$  effectively increases the probability that  $\pmb{v}_{j}$  outperforms the other samples in the set and emerges at the top rank within  $\mathcal{S}_{i,j}$ .
As a result, we define the per-sample RNC loss as the average negative log-likelihood over all other samples in a given batch:
$$
l _ {\mathrm {R N C}} ^ {(i)} = \frac {1}{2 N - 1} \sum_ {j = 1, j \neq i} ^ {2 N} - \log \frac {\exp \left(\operatorname {s i m} \left(\boldsymbol {v} _ {i} , \boldsymbol {v} _ {j}\right) / \tau\right)}{\sum_ {\boldsymbol {v} _ {k} \in \mathcal {S} _ {i , j}} \exp \left(\operatorname {s i m} \left(\boldsymbol {v} _ {i} , \boldsymbol {v} _ {k}\right) / \tau\right)}. \tag {2}
$$
Intuitively, for an anchor sample  $i$ , any other sample  $j$  in the batch is contrasted with it, enforcing the feature similarity between  $i$  and  $j$  to be larger than that of  $i$  and any other sample  $k$  in the batch, if the label distance between  $i$  and  $k$  is larger than that of  $i$  and  $j$ . Minimizing  $l_{\mathrm{RNC}}^{(i)}$  will align the orders of feature embeddings with their corresponding orders in the label space w.r.t. anchor  $i$ .
$\mathcal{L}_{\mathrm{RNC}}$  is then enumerating over all  $2N$  samples as anchors to enforce the entire feature embeddings ordered according to their orders in the label space:
$$
\mathcal {L} _ {\mathrm {R N C}} = \frac {1}{2 N} \sum_ {i = 1} ^ {2 N} l _ {\mathrm {R N C}} ^ {(i)} = \frac {1}{2 N} \sum_ {i = 1} ^ {2 N} \frac {1}{2 N - 1} \sum_ {j = 1, j \neq i} ^ {2 N} - \log \frac {\exp (\sin (\boldsymbol {v} _ {i} , \boldsymbol {v} _ {j}) / \tau)}{\sum_ {\boldsymbol {v} _ {k} \in S _ {i , j}} \exp (\sin (\boldsymbol {v} _ {i} , \boldsymbol {v} _ {k}) / \tau)}. \tag {3}
$$
Interpretation. To exploit the inherent continuity underlying the labels,  $\mathcal{L}_{\mathrm{RNC}}$  ranks samples in a batch with respect to their label distances to the anchor. When contrasting the anchor to the sample in the batch that is closest in the label space, it enforces their similarity to be larger than all other samples in the batch. Similarly, when contrasting the anchor to the second closest sample in the batch, it enforces their similarity to be larger than only those samples that have a rank of three or higher in terms of distance to the anchor. This process is repeated for higher-rank samples (i.e., the third closest, fourth closest, etc.) and for all anchors in a batch.
Feature Ordinality. We further examine the impact of RNC on the ordinality of learned features. Fig. 3 visualizes the feature similarity matrices obtained from 2,000 randomly sampled data points in a real-world temperature regression task [7] for models trained using the vanilla  $L_{1}$  loss and RNC. For clarity, the data points are sorted based on their ground-truth labels, with the expectation that the matrix values decrease progressively from the diagonal to the periphery. Notably, our method, RNC, exhibits a more discernible pattern compared to the  $L_{1}$  loss. Furthermore, we calculate two quantitative metrics, the Spearman's rank correlation coefficient [41] and the Kendall rank correlation coefficient [24], between the label similarities and the feature similarities for both methods. The results in Table 1 confirm that the feature similarities learned by our method have a significantly higher correlation with label similarities than those by the  $L_{1}$  loss.
Table 1: Correlation between feature & label similarities.
Spearman's ρ†	Kendall's τ†
L1	0.822	0.664
RNC	0.971	0.870
Figure 3: Feature similarity matrices sorted by labels.
Connections to Contrastive Learning. The loss can also be explained in the context of positive and negative pairs in contrastive learning. Contrastive learning and SupCon are designed for classification tasks, where positive pairs consist of samples that belong to the same class or the same input image, while all other samples are considered as negatives. In regression, however, there are no distinct classes but rather continuous labels [44, 46]. Thus in  $\mathcal{L}_{\mathrm{RNC}}$ , any two samples can be considered as a positive or negative pair depending on the context. For a given anchor sample  $i$ , any other sample  $j$  in the same batch can be used to construct a positive pair with the corresponding negative samples set to all samples in the batch whose labels differ from  $i$ 's label by more than the label of  $j$ . Fig. 2(a) shows an example batch, and Fig. 2(b) shows two positive pairs and their corresponding negative pair(s).
In this section, we theoretically prove that optimizing  $\mathcal{L}_{\mathrm{RNC}}$  results in an ordered feature embedding that corresponds to the ordering of the labels. All proofs are in Appendix A.
Notations. Let  $s_{i,j} \coloneqq \sin(\pmb{v}_i, \pmb{v}_j) / \tau$  and  $d_{i,j} \coloneqq d(\tilde{\pmb{y}}_i, \tilde{\pmb{y}}_j), \forall i, j \in [2N]$ . Let  $D_{i,1} < D_{i,2} < \dots < D_{i,M_i}$  be the sorted label distances starting from the  $i$ -th sample (i.e., sort  $(\{d_{i,j} | j \in [2N] \setminus \{i\}\})$ ),
$\forall i \in [2N]$ . Let  $n_{i,m} := |\{j \mid d_{i,j} = D_{i,m}, j \in [2N] \setminus \{i\}\}|$  be the number of samples whose distance from the  $i$ -th sample equals  $D_{i,m}, \forall i \in [2N], m \in [M_i]$ .
First, to formalize the concept of ordered feature embedding according to the order in the label space, we introduce a property termed as  $\delta$ -ordered for a set of feature embeddings  $\{\pmb{v}_l\}_{l \in [2N]}$ .
Definition 1 ( $\delta$ -ordered feature embeddings). For any  $0 < \delta < 1$ , the feature embeddings  $\{\pmb{v}_l\}_{l \in [2N]}$  are  $\delta$ -ordered if  $\forall i \in [2N], j, k \in [2N] \setminus \{i\}$ ,
$$
\left\{ \begin{array}{l} s _ {i, j} > s _ {i, k} + \frac {1}{\delta} \text {i f} d _ {i, j} <   d _ {i, k} \\ | s _ {i, j} - s _ {i, k} | <   \delta \text {i f} d _ {i, j} = d _ {i, k} \\ s _ {i, j} <   s _ {i, k} - \frac {1}{\delta} \text {i f} d _ {i, j} > d _ {i, k} \end{array} . \right.
$$
Definition 1 implies that a set of feature embeddings that is  $\delta$ -ordered satisfies the following properties: 1 For any  $j$  and  $k$  such that  $d_{i,j} = d_{i,k}$ , the difference between  $s_{i,j}$  and  $s_{i,k}$  is no more than  $\delta$ ; and 2 For any  $j$  and  $k$  such that  $d_{i,j} < d_{i,k}$ , the value of  $s_{i,j}$  exceeds that of  $s_{i,k}$  by at least  $\frac{1}{\delta}$ . Note that  $\frac{1}{\delta} > \delta$ , which indicates that the feature similarity gap between samples with different label distances to the anchor is always larger than that between samples with equal label distances to the anchor.
Next, we demonstrate that the batch of feature embeddings will be  $\delta$ -ordered as the optimization of  $\mathcal{L}_{\mathrm{RNC}}$  approaches its lower bound. In order to prove this, it is necessary to derive a tight lower bound for  $\mathcal{L}_{\mathrm{RNC}}$ . Let  $L^{\star} := \frac{1}{2N(2N - 1)}\sum_{i = 1}^{2N}\sum_{m = 1}^{M_i}n_{i,m}\log n_{i,m}$ , we have:
Theorem 1 (Lower bound of  $\mathcal{L}_{\mathrm{RNC}}$ ).  $L^{\star}$  is a lower bound of  $\mathcal{L}_{\mathrm{RNC}}$ , i.e.,  $\mathcal{L}_{\mathrm{RNC}} > L^{\star}$ .
Theorem 2 (Lower bound tightness). For any  $\epsilon >0$ , there exists a set of feature embeddings such that  $\mathcal{L}_{\mathrm{RNC}} < L^{\star} + \epsilon$ .
The above theorems verify the lower bound of  $\mathcal{L}_{\mathrm{RNC}}$  as well as its tightness. Given that  $\mathcal{L}_{\mathrm{RNC}}$  can approach its lower bound  $L^{\star}$  arbitrarily, we demonstrate that the feature embeddings will be  $\delta$ -ordered when  $\mathcal{L}_{\mathrm{RNC}}$  is sufficiently close to  $L^{\star}$  for any  $0 < \delta < 1$ .
Theorem 3 (Main theorem). For any  $0 < \delta < 1$ , there exist  $\epsilon > 0$ , such that if  $\mathcal{L}_{\mathrm{RNC}} < L^{\star} + \epsilon$ , then the feature embeddings are  $\delta$ -ordered.
From Batch to Entire Feature Space. Now we have examined the property of a batch of feature embeddings optimized using  $\mathcal{L}_{\mathrm{RNC}}$ . However, what will be the final outcome for the entire feature space when  $\mathcal{L}_{\mathrm{RNC}}$  is optimized? In fact, if any batch of feature embeddings is optimized to achieve a low enough loss such that it is  $\delta$ -ordered, the entire feature embedding will also be  $\delta$ -ordered. This is because any triplet  $(i,j,k)$  in the entire feature embeddings is certain to appear in some batch, thus their feature embeddings  $(v_{i},v_{j},v_{k})$  will satisfy the condition in Definition 1. Nevertheless, to achieve  $\delta$ -ordered features for the entire feature embeddings, do we need to optimize all batches to achieve a sufficiently low loss? The answer is no. Optimizing every batch is not only unnecessary, but also practically infeasible. In fact, one should consider the training process as a cohesive whole, which is effectively optimizing the expectation of the loss over all possible random batches. Then, the Markov's inequality [18] guarantees that when the expectation of the loss is optimized to be sufficiently low, the loss on any batch will be low enough with a high probability.
Connections to Final Performance. Suppose we have a  $\delta$ -ordered feature embedding, how can it help to boost the final performance of a regression task? In Appendix B, we present an analysis based on Rademacher Complexity [39] to prove that a  $\delta$ -ordered feature embedding results in a better generalization bound. To put it intuitively, fitting an ordered feature embedding reduces the complexity of the regressor, which enables better generalization ability from training to testing, and ultimately leads to the final performance gain. Relatedly, we note that the enhanced generalization ability is further empirically verified in Sec. 5. Specifically, if not constrained, the learned feature embeddings could capture spurious or easy-to-learn features that are not generalizable to the real continuous targets. Moreover, such property also leads to better robustness to data corruptions, better resilience to reduced training data, and better generalization to unseen targets.