SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhihan Zhou $^{1}$  Jiangchao Yao $^{1,2\dagger}$  Feng Hong $^{1}$  Ya Zhang $^{1,2}$  Bo Han $^{3}$  Yanfeng Wang $^{1,2\dagger}$
$^{1}$ Cooperative Medianet Innovation Center, Shanghai Jiao Tong University
$^{2}$ Shanghai AI Laboratory  $^{3}$ Hong Kong Baptist University
{zhihanzhou, Sunarker, feng.hong, ya_zhang, wangyanfeng}@sjtu.edu.cn
bhanml@comp.hkbu.edu.hk
Self-supervised learning (SSL) as an effective paradigm of representation learning has achieved tremendous success on various curated datasets in diverse scenarios. Nevertheless, when facing the long-tailed distribution in real-world applications, it is still hard for existing methods to capture transferable and robust representation. Conventional SSL methods, pursuing sample-level uniformity, easily leads to representation learning disparity where head classes dominate the feature regime but tail classes passively collapse. To address this problem, we propose a novel Geometric Harmonization (GH) method to encourage category-level uniformity in representation learning, which is more benign to the minority and almost does not hurt the majority under long-tailed distribution. Specially, GH measures the population statistics of the embedding space on top of self-supervised learning, and then infer an fine-grained instance-wise calibration to constrain the space expansion of head classes and avoid the passive collapse of tail classes. Our proposal does not alter the setting of SSL and can be easily integrated into existing methods in a low-cost manner. Extensive results on a range of benchmark datasets show the effectiveness of GH with high tolerance to the distribution skewness. Our code is available at https://github.com/MediaBrain-SJTU/Geometric-Harmonization.
Recent years have witnessed a great success of self-supervised learning to learn generalizable representation [7, 9, 15, 63]. Such rapid advances mainly benefit from the elegant training on the label-free data, which can be collected in a large volume. However, the real-world natural sources usually exhibit the long-tailed distribution [50], and directly learning representation on them can lead to the distortion issue of the embedding space, namely, the majority dominates the feature regime [75] and the minority collapses [47]. Thus, it becomes urgent to pay attention to representation learning disparity, especially as fairness of machine learning draws increasing attention [28, 41, 68, 77].
Different from the flourishing supervised long-tailed learning [29, 46, 68], self-supervised learning under long-tailed distributions is still under-explored, since there is no labels available for the calibration. Existing explorations to overcome this challenge mainly resort to the possible tailed sample discovery and provide the implicit bias to representation learning. For example, BCL [77] leverages the memorization discrepancy of deep neural networks (DNNs) on unknown head classes and tail classes to drive an instance-wise augmentation. SDCLR [28] contrasts the feature encoder and its pruned counterpart to discover hard examples that mostly covers the samples from tail classes, and efficiently enhance the learning preference towards tailed samples. DnC [59] resorts to a divide-and-conquer methodology to mitigate the data-intrinsic heterogeneity and avoid the representation
† The corresponding authors are Jiangchao Yao and Yanfeng Wang.
Figure 1: Comparison of Geometric Harmonization and the plain SSL method on a 2-D imbalanced synthetic dataset. (Left) Visualization of the 2-D synthetic dataset. (Middle) The embedding distribution of each category learnt by the vanilla contrastive learning loss is approximately proportional to the number of samples, leading to the undesired representation learning disparity. (Right) GH mitigates the adverse effect of class imbalance and approaches to the category-level uniformity.
collapse of minority classes. Liu et al. [41] adopts a data-dependent sharpness-aware minimization scheme to build support to tailed samples in the optimization. However, few works hitherto have considered the intrinsic limitation of the widely-adopted contrastive learning loss, and design the corresponding balancing mechanism to promote the representation learning parity.
We rethink the characteristic of the contrastive learning loss, and try to understand "Why the conventional contrastive learning underperforms in self-supervised long-tailed context?" To answer this question, let us consider two types of representation uniformity: (1) Sample-level uniformity. As stated in [62], the contrastive learning targets to distribute the representation of data points uniformly in the embedding space. Then, the feature span of each category is proportional to their corresponding number of samples. (2) Category-level uniformity. This uniformity pursues to split the region equally for different categories without considering their corresponding number of samples [49, 19]. In the class-balanced scenarios, the former uniformity naturally implies the latter uniformity, resulting in the equivalent separability for classification. However, in the long-tailed distributions, they are different: sample-level uniformity leads to the feature regime that is biased towards the head classes considering their dominant sample quantity and sacrifices the tail classes due to the limited sample quantity. By contrast, category-level uniformity means the equal allocation w.r.t. classes, which balances the space of head and tail classes, and is thus more benign to the downstream classification [17, 19, 37]. Unfortunately, there is no support for promoting category-level uniformity in contrastive learning loss, which explains the question arisen at the beginning.
In this study, we propose a novel method, termed as Geometric Harmonization (GH) to combat representation learning disparity in SSL under long-tailed distributions. Specially, GH uses a geometric uniform structure to measure the uniformity of the embedding space in the coarse granularity. Then, a surrogate label allocation is computed to provide a fine-grained instance-wise calibration, which explicitly compresses the greedy representation space expansion of head classes, and constrain the passive representation space collapse of tail classes. The alternation in the conventional loss refers to an extra efficient optimal-transport optimization that dynamically pursues the category-level uniformity. In Figure 1, we give a toy experiment<sup>1</sup> to visualize the distribution of the embedding space without and with GH. In a nutshell, our contributions can be summarized as follows,
1. To our best knowledge, we are the first to investigate the drawback of the contrastive learning loss in self-supervised long-tailed context and point out that the resulted sample-level uniformity is an intrinsic limitation to the representation parity, motivating us to pursue category-level uniformity with more benign downstream generalization (Section 3.2).
2. We develop a novel and efficient Geometric Harmonization (Figure 2) to combat the representation learning disparity in SSL, which dynamically harmonizes the embedding space of SSL to approach the category-level uniformity with the theoretical guarantee.
3. Our method can be easily plugged into existing SSL methods for addressing the data imbalance without much extra cost. Extensive experiments on a range of benchmark datasets demonstrate the consistent improvements in learning robust representation with our GH.
For more details, please refer to Appendix E.3.
Self-Supervised Long-tailed Learning. There are several recent explorations devoted to this direction [40, 28, 59, 41, 77]. BCL [77] leverages the memorization effect of DNNs to automatically drive an instance-wise augmentation, which enhances the learning of tail samples. SDCLR [28] constructs a self-contrast between model and its pruned counterpart to learn more balanced representation. Classic Focal loss [40] leverages the loss statistics to putting more emphasis on the hard examples, which has been applied to self-supervised long-tailed learning [77]. DnC [59] benefits from the parameter isolation of multi-experts during the divide step and the information aggregation during the conquer step to prevent the dominant invasion of majority. Liu et al. [41] proposes to penalize the sharpness surface in a reweighting manner to calibrate class-imbalance learning. Recently, TS [35] employs a dynamic strategy on the temperature factor of contrastive loss, harmonizing instance discrimination and group-wise discrimination. PMSN [2] proposes the power-law distribution prior, replacing the uniform prior, to enhance the quality of learned representations.
Hyperspherical Uniformity. The distribution uniformity has been extensively explored from the physica area, e.g., Thomson problem [58, 53], to machine learning area like some kernel-based extensions, e.g., Riesz s-potential [21, 43] or Gaussian potential [12, 4, 62]. Some recent explorations regarding the features of DNNs [49, 17, 47] discover a terminal training stage when the embedding collapses to the geometric means of the classifier w.r.t. each category. Specially, these optimal class means specify a perfect geometric uniform structure with clear geometric interpretations and generalization guarantees [78, 67, 31]. In this paper, we first extend this intuition into self-supervised learning and leverage the specific structure to combat the representation disparity in SSL.
We denote the dataset  $\mathcal{D}$ , for each data point  $(\boldsymbol{x}, \boldsymbol{y}) \in \mathcal{D}$ , the input  $\boldsymbol{x} \in \mathbb{R}^m$  and the associated label  $\boldsymbol{y} \in \{1, \dots, L\}$ . Let  $N$  denote the number of samples,  $\mathrm{R} = N_{max} / N_{min}$  denote the imbalanced ratio (IR), where  $N_{max}, N_{min}$  is the number of samples in the largest and smallest class, respectively. Let  $n_i$  denote the number of samples in class  $i$ . In SSL, the ground-truth  $\boldsymbol{y}$  can not be accessed and the goal is to transform an image to an embedding via DNNs, i.e.,  $f_\theta: \mathbb{R}^m \to \mathbb{R}^d$ . In the linear probing evaluation, we construct a supervised learning task with balanced datasets. A linear classifier  $g(\cdot)$  is built on top of the frozen  $f_\theta(\cdot)$  to produce the prediction, i.e.,  $g(f_\theta(\boldsymbol{x}))$ .
3.2 Geometric Harmonization
As aforementioned, most existing SSL methods in long-tailed context leverage the contrastive learning loss, which encourages the sample-level uniformity in the embedding space. Considering the intrinsic limitation illustrated in Figure 1, we incorporate the geometric clues from the embedding space to calibrate the current loss, enabling our pursuit of category-level uniformity. In the following, we first introduce a specific geometric structure that is critical to Geometric Harmonization.
Definition 3.1. (Geometric Uniform Structure). Given a set of vertices  $\mathbf{M} \in \mathbb{R}^{d \times K}$ , the geometric uniform structure satisfies the following between-class duality
$$
\mathbf {M} _ {i} ^ {\top} \cdot \mathbf {M} _ {j} = C, \forall i, j \in \{1, 2, \dots , K \}, i \neq j, \tag {1}
$$
where  $\| \mathbf{M}_i\| = 1,\forall i\in \{1,2,\ldots ,K\}$ $K$  is the number of geometric vertices and  $C$  is a constant.
Above structure provides a characteristic that any two vectors in  $\mathbf{M}$  have the same angle, namely, the unit space are equally partitioned by the vectors in  $\mathbf{M}$ . This fortunately follows our expectation about category-level uniformity. Specially, if we use  $\mathbf{M}$  as a constant classifier to involve into training, and have the oracle label  $\pmb{y}$  of the long-tailed data  $(K = L)$  to supervise the below prediction
$$
\boldsymbol {q} = p (\boldsymbol {y} \mid f _ {\theta} (\boldsymbol {x}), \mathbf {M}) = \exp \left(\mathbf {M} _ {y} ^ {\top} \cdot f _ {\theta} (\boldsymbol {x}) / \gamma_ {\mathrm {G H}}\right) / \left(\sum_ {i = 1} ^ {K} \exp \left(\mathbf {M} _ {i} ^ {\top} \cdot f _ {\theta} (\boldsymbol {x}) / \gamma_ {\mathrm {G H}}\right)\right), \tag {2}
$$
then according to the neural collapse theory [49], the representation of all samples will fit the geometric uniform structure of  $\mathbf{M}$  in the limit, namely, approach category-level uniformity. However,
: Geometric Uniform Structure
Surrogate Label Allocation
Figure 2: Overview of Geometric Harmonization  $(w/ InfoNCE)$ . To achieve harmonization with the category-level uniformity, we require ground-truth labels for supervision during training. However, these labels are not available under the self-supervised paradigm. Moreover, estimating surrogate labels directly from the geometric uniform structure is challenging and noisy, especially when the representation is not ideally distributed. To fulfill this gap, we utilize the geometric uniform structure to measure the embedding space, and the captured population statistics are used for an instance-wise calibration by surrogate label allocation, which provides a supervision feedback to counteract the sample-level uniformity. Specially, our method are theoretically grounded to approach category-level uniformity at the loss minimum. The additional model parameters incurred by GH are analytically and empirically demonstrated to be trained in an efficient manner.
the technical challenge is the complete absence of annotation  $\pmb{y}$  in our context, making directly constructing the objective min  $\mathbb{E}[y\log q]$  with Eq. (2) impossible to combat the representation disparity.
Surrogate Label Allocation. To address the problem of unavailable labels, we explore constructing the surrogate geometric labels  $\hat{\pmb{q}}$  to supervise the training of Eq. (2). Concretely, we utilize the recent discriminative clustering idea [1] to acquire such geometric labels, formulated as follows
$$
\min  _ {\hat {\mathbf {Q}} = \left[ \hat {q} _ {1}, \dots , \hat {q} _ {N} \right]} \mathcal {L} _ {\mathrm {G H}} = - \frac {1}{| \mathcal {D} |} \sum_ {\boldsymbol {x} _ {i} \sim \mathcal {D}} \hat {\boldsymbol {q}} _ {i} \log \boldsymbol {q} _ {i}, \quad \text {s . t .} \hat {\mathbf {Q}} \cdot \mathbb {1} _ {N} = N \cdot \boldsymbol {\pi}, \hat {\mathbf {Q}} ^ {\top} \cdot \mathbb {1} _ {K} = \mathbb {1} _ {N}, \tag {3}
$$
where  $\hat{\pmb{q}}_i$  refers to the soft assignment constrained in a  $K$ -dimensional probability simplex and  $\pi \in \mathbb{R}_+^K$  refers to the distribution constraint. As we are dealing with long-tailed data, we propose to use the geometric uniform structure  $\mathbf{M}$  to automatically compute the population statistic of the embedding space as  $\pi$ . Finally, Eq. (3) can be analytically solved by Sinkhorn-Knopp algorithm [14] (refer to Appendix D for Algorithm 1). Note that, the above idea builds upon a conjecture: the constructed surrogate geometric labels are mutually correlated with the oracle labels so that they have the similar implication to approach category-level uniformity. We will empirically verify the rationality of such an assumption via normalized mutual information in Section 4.4.
Overall Objective. Eq. (3) can be easily integrated into previous self-supervised long-tailed learning methods for geometric harmonization, e.g., SDCLR [28] and BCL [77]. For simplicity, we take their conventional InfoNCE loss [48] as an example and write the overall objective as follows
$$
\min  _ {\theta , \hat {\mathbf {Q}}} \mathcal {L} = \mathcal {L} _ {\text {I n f o N C E}} + w _ {\mathrm {G H}} \mathcal {L} _ {\mathrm {G H}}, \tag {4}
$$
where  $w_{\mathrm{GH}}$  represents the weight of the geometric harmonization loss. Optimizing Eq. (4) refers to a bi-level optimization style: in the outer-loop, optimize  $\min_{\hat{\mathbf{Q}}} \mathcal{L}_{\mathrm{GH}}$  with fixing  $\theta$  to compute the surrogate geometric labels; in the inner-loop, optimize  $\min_{\theta} \mathcal{L}_{\mathrm{InfoNCE}} + \mathcal{L}_{\mathrm{GH}}$  with fixing  $\hat{\mathbf{Q}}$  to learn the representation model. The additional cost compared to the vanilla method primarily arises from from the outer-loop, which will be discussed in Section 3.4 and verified in Section 4.4.
Compared with previous explorations [1, 7], the uniqueness of GH lies in the following three aspects: (1) Geometric Uniform Structure. The pioneering works mainly resort to a learnable classifier to perform clustering, which can easily be distorted in the long-tailed scenarios [17]. Built on
Table 1: Linear probing of vanilla discriminative clustering methods and variants of GH on CIFAR-100-LT.
IR	SimCLR	+SeLA	+SwAV	+GH	w/o GUS	w/o π
100	50.7	50.5	52.2	54.0	53.3	53.1
50	52.2	52.0	53.0	55.4	54.6	54.4
10	55.7	56.0	56.1	57.4	56.7	56.4
the geometric uniform classifier, our method is capable to provide high-quality clustering results with clear geometric interpretations. (2) Flexible Class Prior. The class prior  $\pi$  is assumed to be uniform among the previous attempts. When moving to the long-tailed case, this assumption will strengthen the undesired sample-level uniformity. In contrast, our methods can potentially cope with any distribution with the automatic surrogate label allocation. (3) Theoretical Guarantee. GH is theoretically grounded to achieve the category-level uniformity in the long-tailed scenarios (refer to Section 3.3), which has never been studied in previous methods. To gain more insights into our method, we further compare GH with discriminative clustering methods (SeLA [1], SwAV [7]) and investigate the impact of various components in GH. From the results in Table 1, we can see that GH consistently outperforms the vanilla discriminative clustering baselines in the linear probing evaluation. Notably, we observe that both GUS and the class prior  $\pi$  play a critical role in our GH.
3.3 Theoretical Understanding
Here, we reveal the theoretical analysis of GH on promoting the representation learning to achieve category-level uniformity instead of sample-level uniformity. Let us begin with a deteriorated case of sample-level uniformity under the extreme imbalance, i.e., minority collapse [17].
Lemma 3.2. (Minority collapse) Assume the samples follow the uniform distribution  $n_1 = n_2 = \dots = n_{L_H} = n_H$ ,  $n_{L_H + 1} = n_{L_H + 2} = \dots = n_L = n_T$  in head and tail classes respectively. Assume  $d \geq L$  and  $n_H / n_T \to +\infty$ , the lower bound (Lemma C.1) of  $\mathcal{L}_{\mathrm{InfoNCE}}$  achieves the minimum when the class means of the tail classes collapse to an identical vector:
$$
\lim  \boldsymbol {\mu} _ {i} - \boldsymbol {\mu} _ {j} = \mathbf {0} _ {L}, \forall L _ {H} \leq i \leq j \leq L,
$$
where  $\pmb{\mu}_l = \frac{1}{n_l}\sum_{i=1}^{n_l}f_\theta(\pmb{x}_{l,i})$  denotes the class means and  $\pmb{x}_{l,i}$  is the  $i$ -th data point with label  $l$ .
This phenomenon indicates all representations of the minority will collapse completely to one point without considering the category discrepancy, which aligns with our observation regarding the passive collapse of tailed samples in Figure 1. To further theoretically analyze GH, we first quantitatively define category-level uniformity in the following, and then theoretically claim that with the geometric uniform structure (Definition. 3.1) and the perfect aligned allocation (Eq. (3)), we can achieve the loss minimum at the stage of realizing category-level uniformity.
Definition 3.3. (Categorical-level Uniformity) We define categorical-level uniformity on the embedding space  $w.r.t$  the geometric uniform structure  $\mathbf{M}$  when it satisfies
$$
\boldsymbol {\mu} _ {k} ^ {*} = \mathbf {M} _ {k}, \forall k = 1, 2, \dots , K,
$$
where  $\pmb{\mu}_k^* = \frac{1}{n_k}\sum_{i=1}^{n_k}f_\theta^*(\pmb{x}_{k,i})$  represents the class mean for samples assigned with the surrogate geometric label  $k$  in the embedding space.
Theorem 3.4. (Optimal state for  $\mathcal{L}$ ) Given Eq. (4) under the proper optimization strategy, when it arrives at the category-level uniformity (Definition 3.3) defined on the geometric uniform structure M (Definition 3.1), we will achieve the minimum of the overall loss  $\mathcal{L}^*$  as
$$
\mathcal {L} ^ {*} = - 2 \sum_ {l = 1} ^ {L} \pi_ {l} ^ {\mathcal {Y}} \log (1 / (1 + (K - 1) \exp (C - 1))) + \log (J / L), \tag {5}
$$
where  $J$  denotes the size of the collection of the negative samples and  $\pi^y$  refers to the marginal distribution of the latent ground-truth labels  $\pmb{y}$ .
This guarantees the desirable solution with the minimal intra-class covariance and the maximal interclass covariance under the geometric uniform structure [49, 17], which benefits the downstream generalization. Notably, no matter the data distribution is balanced or not, our method can persistently maintain the theoretical merits on calibrating the class means to achieve category-level uniformity. We also empirically demonstrate the comparable performance with GH on the balanced datasets in Section 4.4, as in this case category-level uniformity is equivalent to sample-level uniformity.
3.4 Implementation and Complexity analysis
In Algorithm 2 of Appendix D, we give the complete implementation of our method. One point that needs to be clarified is that we learn the label allocation  $\hat{q}$  in the mini-batch manner. In addition, the geometric prediction  $q$  and the adjusted  $\hat{q}$  are computed at the beginning of every epoch as the population-level statistic will not change much in a few mini-batches. Besides, we maintain a momentum update mechanism to track the prediction of each sample to stabilize the training, i.e.,  $q^{m} \gets \beta q^{m} + (1 - \beta)q$ . When combined with the joint-embedding loss, we naturally adopt a cross-supervision mechanism  $\min \mathbb{E}[\hat{q}^+ \log q]$  for the reconciliation with contrastive baselines. The proposed method is illustrated in Figure 2 for visualization.
For complexity, assume that the standard optimization of deep neural networks requires forward and backward step in each mini-batch update with the time complexity as  $\mathcal{O}(B\Lambda)$ , where  $B$  is the minibatch size and  $\Lambda$  is the parameter size. At the parameter level, we add an geometric uniform structure with the complexity as  $\mathcal{O}(BKd)$ , where  $K$  is the number of geometric labels and  $d$  is the embedding dimension. For Sinkhorn-Knopp algorithm, it only refers to a simple matrix-vector multiplication as shown in Algorithm 1, whose complexity is  $\mathcal{O}(E_s(B + K + BK))$  with the iteration step  $E_{s}$ . The complexity incurred in the momentum update is  $\mathcal{O}(BK)$ . Since  $K,d$  and  $E_{s}$  are significantly smaller than the model parameter  $\Lambda$  of a million scale, the computational overhead involved in GH is negligible compared to  $\mathcal{O}(B\Lambda)$ . The additional storage for a mini-batch of samples is the matrix  $\mathbf{Q}^m\in \mathbb{R}^{K\times B}$ , which is also negligible to the total memory usage. To the end, GH incurs only a small computation or memory cost and thus can be plugged to previous methods in a low-cost manner. The empirical comparison about the computational cost is summarized in Table 17.
Baselines. We mainly choose five baseline methods, including (1) plain contrastive learning: SimCLR [9], (2) hard example mining: Focal [40], (3) asymmetric network pruning: SDCLR [28], (4) multi-expert ensembling: DnC [59], (5) memorization-guided augmentation: BCL [77]. Empirical comparisons with more baseline methods can be referred to Appendix F.3.