SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 2: The barycenter results between noise and a leopard image. It is very blurry without the use of  $\mathbf{Q}$ . However, when we set  $\mathbf{Q} = \mathbf{P}^{\epsilon}$ , the resulting barycenters clearly display the shape of the leopard, but fail to transition smoothly from noise to leopard image as  $\lambda$  changes. To address this, we set  $\mathbf{Q}$  as shown in Eq. 10, which blurs the image while transitioning with the increase of  $\lambda$ .
which aims to get the optimal solution of the non-regularized OT. As discussed in Sinkhorn of RE-OT, we have the optimal solution form:
$$
\begin{array}{l} \mathbf {Q} ^ {(t)} = \operatorname {d i a g} \left(\mathbf {u} ^ {(t - 1)}\right) \mathbf {Q} ^ {(t - 1)} \odot e ^ {- C / \epsilon} \operatorname {d i a g} \left(\mathbf {v} ^ {(t - 1)}\right) \\ = \operatorname {d i a g} \left(\mathbf {u} ^ {(t - 1)} \odot \dots \odot \mathbf {u} ^ {(0)}\right) \mathbf {Q} ^ {(0)} \odot e ^ {- \frac {(t + 1) C}{\epsilon}} \operatorname {d i a g} \left(\mathbf {v} ^ {(0)} \odot \dots \mathbf {v} ^ {(t - 1)}\right). \tag {13} \\ \end{array}
$$
When we set  $\mathbf{Q}^{(0)} = \mathbf{1}_{n\times m}$ , the optimization is equivalent to applying Sinkhorn's algorithm iteratively with a kernel  $e^{-\frac{(t + 1)C}{\epsilon}}$ , i.e., with a decaying regularization coefficient  $\frac{\epsilon}{t + 1}$ . This observation further emphasizes the importance of selecting an appropriate  $\mathbf{Q}$  for RE-OT.
4 RE-OT for Long-tailed Recognition
We now discuss the application of inverse RE-OT to address the long-tailed classification task (including the contrastive learning setting). Unlike previous works that study this problem using conditional probability, our approach aims to learn features by matching samples with their labels using Inverse OT. Specifically, given the mini-batch pair set  $\{(\mathbf{x}_i,\mathbf{y}_i)\}_{i = 1}^n$  where  $\mathbf{x}_i$  is training sample and  $\mathbf{y}_i$  is the  $m$ -dimensional one-hot label of  $\mathbf{x}_i$ . We can set the cost  $\mathbf{C}_{ij}^{\theta}$  between sample  $\mathbf{x}_i$  and another one-hot label given the neural network  $f_{\theta}$ . Without loss of generality, we can set  $\mathbf{C}_{ij}^{\theta} = c - l_{ij}$  where  $l_{ij} = (f_{\theta}(\mathbf{x}_i))_j$  is the  $j$ -th component of logits
Figure 3: RE-OT for unbalanced classification.
$f_{\theta}(\mathbf{x}_i)$  for sample  $\mathbf{x}_i$ . Then we can learn the matching between features and labels with inverse OT via a bi-level optimization:
$$
\min  _ {\theta} K L (\widetilde {\mathbf {P}} | \mathbf {P} ^ {\theta}) \quad s. t. \quad \mathbf {P} ^ {\theta} = \underset {\mathbf {P} \in U} {\arg \min } <   \mathbf {C} ^ {\theta}, \mathbf {P} > - \epsilon H _ {\mathbf {Q}} (\mathbf {P}), \tag {14}
$$
where  $\widetilde{\mathbf{P}}$  is the supervision given the label  $\{\mathbf{y}_i\}$ . For example, we can set  $\widetilde{\mathbf{P}}_{ij} = \mathbf{y}_{ij}$  when  $\mathbf{y}_{ij}$  is the  $j$ -th component of  $\mathbf{y}_i$ . Here  $U$  refers to the constraints for the coupling  $\mathbf{P}$ . We can set  $U = U(\mathbf{a},\mathbf{b})$  with the full matching constraints or its relaxation e.g.  $U = U(\mathbf{a}) = \{\mathbf{P}|\mathbf{P}\mathbf{1}_m = \mathbf{a}\}$ .
For our designed smoothing guidance  $\mathbf{Q}$ , a simple idea is to set it with the long-tailed ratio i.e. the training sample ratio of class labels which are readily available from the training set, which equals to
Table 1: Previous works in the view of inverse RE-OT by setting different ground truth  $\hat{\mathbf{P}}_{ij}$ , cost  $\mathbf{C}_{ij}$ , prior  $\mathbf{Q}_{ij}$  and coefficient  $\epsilon$  under the constraints  $U(\mathbf{a})$ .
Methods	Formulation
Ground Truth \(\tilde{\mathbf{P}}_{ij}\)	Cost \(\mathbf{C}_{ij}\)	Prior \(\mathbf{Q}_{ij}\)	penalty coefficient \(\epsilon\)
Vanilla Softmax	\(\tilde{\mathbf{P}}_{ij} = y_{ij}\)	\(\mathbf{C}_{ij} = c - l_{ij}\)	\(Q_{ij} = 1\)	\(\epsilon = \tau\)
Focal Loss [35]	\(\tilde{\mathbf{P}}_{ij} = y_{ij} * (1 - \mathbf{P}_{ij})^{\gamma}\)	\(\mathbf{C}_{ij} = c - l_{ij}\)	\(Q_{ij} = 1\)	\(\epsilon = \tau\)
A-softmax [37]	\(\tilde{\mathbf{P}}_{ij} = y_{ij}\)	\(\mathbf{C}_{ij} = c - ||l_i||\psi(\theta_{\mathbf{y}_{i,i}})\)	\(Q_{ij} = 1\)	\(\epsilon = \tau\)
Triplet [48]	\(\tilde{\mathbf{P}}_{ij} = y_{ij}\)	\(\mathbf{C}_{ij} = ||f(x_i) - f(x_j)||^2\)	\(Q_{ij} = 1\)	\(\epsilon \rightarrow 0^+\)
Class-Balanced Loss [8]	\(\tilde{\mathbf{P}}_{ij} = (1 - \beta) / (1 - \beta_{y_{i,:}}^n)\)	\(\mathbf{C}_{ij} = c - l_{ij}\)	\(Q_{ij} = 1\)	\(\epsilon = \tau\)
LDAM Loss [3]	\(\tilde{\mathbf{P}}_{ij} = y_{ij}\)	\(\mathbf{C}_{ij} = c - l_{ij}\)	\(Q_{ij} = e^{-1\{i=y\} \cdot C/n_j^{\gamma}}\)	\(\epsilon = \tau\)
Balanced-Softmax [45]	\(\tilde{\mathbf{P}}_{ij} = y_{ij}\)	\(\mathbf{C}_{ij} = c - l_{ij}\)	\(Q_{ij} = n_j\)	\(\epsilon = \tau\)
Our Setting	\(\tilde{\mathbf{P}}_{ij} = y_{ij}\)	\(\mathbf{C}_{ij} = c - l_{ij}\)	\(Q = (1 - \lambda(t)) \text{Uniform} + \lambda(t)\mathbf{r}\)	\(\epsilon = \tau\)
the Balanced Softmax method proposed in [45]. In this paper, motivated by Eq. 10 having a good efficiency in OT and other works [60] which adopt two stages for learning, we let it vary over epochs:
$$
\mathbf {Q} = (1 - \lambda (t)) \operatorname {U n i f o r m} + \lambda (t) \mathbf {r}, \tag {15}
$$
where  $t$  is the training epoch and  $\mathbf{r}$  is the prior probability matrix e.g. the balanced ratio in [45], teacher-based prior predictions or the sample-class-wise setting discussed in Appendix E.  $\lambda(t)$  is an epoch varying weight. When  $t$  is small,  $\lambda(t)$  is close to 0, and  $\mathbf{Q}$  is more likely to be a uniform distribution, which gives the model less prior information. As  $t$  approaches the final training epoch number  $T$ ,  $\lambda(t) \rightarrow 1$  and  $\mathbf{Q} \rightarrow \mathbf{r}$ , which means the model is given the full prior information. This gradual process of introducing prior information can be helpful for training the inverse OT. We provide a specific setting for  $\lambda(t)$  and  $\mathbf{r}$  in Appendix E. During the inference process, the problem reduces to the inner optimization in Eq. 6, which takes the form of traditional unsupervised OT. We use this optimization to obtain the matching between the sample features and labels, which corresponds to classification for the testing data. We do not adopt the long-tailed setting of  $\mathbf{Q}$  during inference because the testing data is assumed to be uniform.
4.1 Generalization to a Family of Losses in Form of Softmax-based Cross-Entropy
We find our new formulation with inverse OT is a more general form if we set  $U = U(\mathbf{a})$  where  $\mathbf{a}_i = 1 / n$  and many Softmax-based cross-entropy losses can be our special case by varying the form of  $\widetilde{\mathbf{P}}$ ,  $\mathbf{C}^\theta$  and  $\mathbf{Q}$ . Specifically, with  $U = U(\mathbf{a})$ , Eq.14 equals to
$$
\min  _ {\theta} L = \sum_ {i, j} \widetilde {\mathbf {P}} _ {i j} \log \frac {\mathbf {Q} _ {i j} e ^ {- \mathbf {C} _ {i j} ^ {\theta} / \epsilon}}{n \sum_ {k = 1} ^ {m} \mathbf {Q} _ {i k} e ^ {- \mathbf {C} _ {i k} ^ {\theta} / \epsilon}}. \tag {16}
$$
The proof is given in Appendix F. We can find the form in Eq. 16 is similar with the Softmax cross-entropy loss. In particular, if we set  $\widetilde{\mathbf{P}}$ ,  $\mathbf{C}^{\theta}$ ,  $\mathbf{Q}$  and  $\epsilon$  in different forms as shown in Tab. 1, many classification loss can be the special cases of our formulation, as discussed in detail as follows.
Special cases by varying  $\mathbf{C}^{\theta}$ . For a matching problem, it is quite important to define a distance among samples. We first set  $\mathbf{C}_{ij}^{\theta} = c - l_{ij}$  where  $c$  is a large enough value. In this case, Eq. 16 reduces to a normal softmax-based cross-entropy loss when  $\mathbf{Q}$  is in the uniformly distributed. Alternatively, we can set  $\mathbf{C}^{\theta}$  differently to suit specific tasks. For example, in learning in the hypersphere space, the Vanilla Softmax will be converted to A-softmax [37] to enforce angular margin constraints for improved face recognition performance.
Deriving more losses by varying  $\mathbf{Q}$ . By varying the setting of  $\mathbf{Q}$ , we can derive many Softmax-based losses as special cases of the RE-OT-based loss. For example, the Balanced Softmax can be obtained by setting  $\mathbf{Q}_{ij} = n_j$ , where  $n_j$  is the frequency of label  $j$  in the training set. This choice of  $\mathbf{Q}$  encourages the model to assign equal weights to each class during training, which helps to mitigate the impact of class imbalance.
Discussion on  $\widetilde{\mathbf{P}}$ . Re-weighting the example is a common way to improve the Long-tailed learning, including the empirical Re-weighting methods (e.g. focal loss [35]) and automatic Re-weighting ways (e.g. L2RW [46]). In the view of IOT, the supervision of the coupling  $\widetilde{\mathbf{P}}$  is self-adapting and may depend on  $\mathbf{P}^{\theta}$ , which may be helpful in IOT application fields such as deep graph matching.
Figure 4: t-SNE for training and testing sample features and class barycenters. The class barycenters are calculated by Eq. 18 with features of training data. It can be observed that in the training data, the class barycenters are mainly on the feature centers of corresponding category samples, which proves the effectiveness of the barycentric projection map defined in Eq. 18. At the same time, we also find that barycenters calculated by features of training samples may not necessarily be centered in the test data. Additionally, the head class occupies a significantly larger space than the tail class in the 2D feature space, which may be one of the reasons for the poor accuracy of the tail class.
Triplet loss for  $\epsilon \rightarrow 0^{+}$ . Though Eq. 16 has a very different form compared to triplet loss, the equivalence can also be established when the temperature  $\epsilon \rightarrow 0^{+}$ . Note that the idea of the relation between triplet loss and our loss is motivated by the work [55] which gives the proof of the equivalence between triplet loss and InfoNCE. We give the proof in Appendix G.
Representation Learning with our Formulation Our matching perspective can be applied not only to unbalanced classification but also to unbalanced representation learning, in which we only need to modify the original one-hot label to the augmentation's features. In Sec. 5, we conducted experiments on long-tailed contrastive learning and verified that by selecting an appropriate  $\mathbf{Q}$ , we can improve the effectiveness of representation learning on long-tailed datasets.
4.2 Bridging Regularized OT to Classification
We have shown the classification can be viewed as a matching problem and if we formulate the matching with optimal transport, we can get the equivalence between Softmax (as well as its variants) and RE-OT under the setting of  $U(\mathbf{a})$ . We can observe that there are some interconnected concepts between OT and classification, which can shed lights on a deeper understanding of OT and in turn the classification problem with different forms. One notable example is the entropy regularization coefficient  $\epsilon$  in OT and the temperature coefficient  $\tau$  in Softmax. There have been numerous works [24, 61] exploring the role of the temperature hyperparameter in Softmax. We can understand it easier from the view of OT, where  $\epsilon$  serves as a relaxation hyperparameter for couplings yielding smoother predictions. In fact, there are other concepts, such as barycentric projection in OT [13], which may help us understand more about classification. We give its definition here:
Definition 1 (Barycentric Projection). Consider the setting of Eq. 6 in which we use entropic regularization to approximate OT between discrete measures. One can define the so-called barycentric projection map with the coupling  $\mathbf{P}$
$$
T: x _ {i} \in X \rightarrow \frac {1}{\mathbf {a} _ {i}} \sum_ {j} \mathbf {P} _ {i j} y _ {j} \in Y, \tag {17}
$$
where  $\{x_{i}\} \subset X$  is the set of locations corresponding to simplex a and  $\{y_j\}$  is the set of locations in the  $Y$  space.  
where  $\{x_{i}\} \subset X$  is the set of locations corresponding to simplex b.
From the classification view, exactly  $T(x_{i})$  is probability confidence of samples  $x_{i}$  which can be understand as the feature in the one-hot label space. Besides, motivated by this barycentric projection in OT, we can exactly define the mapping for the feature of one-hot label in hidden feature space:
$$
T ^ {\prime}: y _ {j} \in Y \rightarrow \frac {1}{\sum_ {i} \mathbf {P} _ {i j}} \sum_ {i} \mathbf {P} _ {i j} \mathbf {f} _ {i} \in \mathcal {F}, \tag {18}
$$
where  $\mathbf{f}_i$  is the feature of  $x_i$  extracting from the neural networks and  $\mathcal{F}$  is the feature space. Eq. 18 maps the (one-hot) label to the feature space, which means no need to learn the centroid features
as done in oneline clustering methods [4, 32] but calculate it statically with Eq. 18 if the coupling is known. As shown in Fig. 4, we first train the models using vanilla softmax, focal loss, and our loss on the CIFar10-LT dataset. We consider the logits as the sample features and save the logits from the three head and three tailed classes, along with the corresponding predicted couplings, as matrices. Then, we calculate the barycenters of the labels using Eq. 18, which effectively computes the weighted average of the features. We concat the logits and barycenters as a new matrix and use it to calculate t-SNE results for training data, as shown in the first row of Fig. 4. For testing data, we concat the testing logits with the barycenters calculated from the training data and use t-SNE dimensional reduction, as shown in the second row in Fig. 4. By comparing the first row and the second row of Fig. 4, we can observe the differences in the positions of the barycenters, which are actually caused by shift in the feature distributions.
Further Discussions between OT and classification. Note that many theoretical aspects and variations of OT can be incorporated into the field of classification. Here are a few advantages unique to the OT perspective that can be achieved but are not possible in the traditional Bayesian view: 1) Variable constraints instead of  $\{\mathbf{P}:\mathbf{P1} = \mathbf{a}\}$ , such as (modified) Optimal Partial Transport (OPT) with  $\{\mathbf{P}:\mathbf{P1}\leq \mathbf{a},\mathbf{1}^{\top}\mathbf{P1} = s\}$  where  $s$  is the number requiring to predict. This allows for distributions that are not just one-vs-all and enables rejection of classification if the model determines that a sample cannot be classified into any of the candidate labels; 2) Generalization of softmax. The current softmax formulation is essentially based on Entropic Regularization of OT. However, OT regularization goes beyond just the entropic regularization and includes other regularizations such as L2, Tsallis entropies, or divergence-based ones. This opens up possibilities for generalizing softmax; 3) Classification cross-representation models. From the OT perspective, classification between the features of samples and labels, based on Gromov-Wasserstein Distance, can be performed in different feature spaces. For example, the sample features may be in 100 dimensions while the label features are in 20 dimensions. This can be helpful for preserving privacy as Gromov-Wasserstein Distance [44] only requires similarity between samples and labels. Due to the various variants and theories within OT, considering classification as OT naturally allows us to incorporate existing OT knowledge into the field of classification, which is the ongoing work we are conducting.
We evaluate our proposed approach for both image and molecule datasets. As the primary focus of this paper is the impact of loss functions on model learning, we do not adopt additional specialized techniques to improve performance, such as transfer learning or resampling. This is to control variables and thus all the baselines used in this study represent only the loss functions proposed by them. Additional experimental settings and details can be found in Appendix H.
Experiments on long-tailed image classification. The long-tailed image classification datasets analyzed in this study include CIFAR10-LT, CIFAR100-LT [30], and Imagenet-LT [38]. To evaluate the performance of the models, we use the corresponding balanced test dataset and report top-1 accuracy. Specifically, for CIFAR-10, we report accuracy on two sets of classes: Many-shot (more than 100 images) and Few-shot (less than 100 images). For CIFAR-100 and Imagenet, we report accuracy on three sets: Many-shot (more than 100 images), Medium-shot  $(20\sim 100$  images), and Few-shot (less than 20 images). The experiments for unbalanced image classification are all conducted with an imbalanced factor of less than 200, which is defined as the ratio of the number of training instances in the largest class to the smallest [45]. The results for long-tailed classification are presented in Table 3. Note that in our approach, the teacher-based method involves setting the prior matrix  $\mathbf{r}$  in Eq. 15 using a vanilla softmax trained teacher model. On the other hand, the ratio-based method refers to setting  $\mathbf{r}$  as the long-tailed ratio. From a comprehensive perspective, at least one of our approaches achieves the best average accuracy across the entire dataset.
Experiments on unbalanced molecule classification. We further examine the application of our proposed method in the context of molecular representation learning task, which holds significant implications for healthcare [59] and drug discovery [17, 58]. Specifically, we perform the experiments on two unbalanced molecule classification datasets, i.e. OGBG-MOLBBBP and OGBG-MOLBACE, from the Open Graph Benchmark (OGB) [21]. We use the default train/val/test split with ratio 8:1:1. Given the presence of class imbalance within the two datasets, we have chosen to adhere to the precedent set by prior studies, also employing the ROC-AUC as the metric for performance evaluation. As the experimental results are easily influenced by initialization, we conducted five repetitions of
Table 3: Top-1 accuracy (%) for long-tailed image classification with 200 imbalanced factor.
Method	CIFAR10-LT	CIFAR100-LT	ImageNet-LT
Many	Few	All	Many	Medium	Few	All	Many	Medium	Few	All
Vanilla Softmax	77.4	68.9	74.9	75.8	48.2	11.0	42.0	57.3	26.2	3.1	35.0
Focal Loss [35]	79.6	58.4	73.3	76.1	46.9	11.1	41.7	57.3	27.6	4.4	35.9
LDAM [3]	80.5	65.2	75.9	75.7	50.6	11.5	42.9	57.3	27.6	4.4	35.9
LogitAdjust [39]	80.0	35.3	66.6	75.7	39.2	4.1	36.5	54.2	14.0	0.4	27.6
CB-CE [8]	76.6	70.7	74.8	53.2	48.8	13.3	36.3	35.3	32.1	21.2	31.9
CB-FC [8]	76.6	70.7	74.8	53.2	48.8	13.3	36.3	35.3	32.1	21.2	31.9
Balanced Softmax [45]	82.2	71.6	79.0	70.3	50.4	26.5	47.0	52.5	38.6	17.8	41.1
Ours(teacher-based)	84.4	63.7	78.2	69.7	55.6	24.9	47.9	41.5	37.6	31.1	38.2
Ours(ratio-based)	81.5	74.6	79.4	70.4	53.0	26.6	47.9	53.5	39.0	17.4	41.6
the experiment and reported the mean accuracy and its standard deviation in the Tab. 2, which show the effectiveness of our method.
Experiments on Long-Tailed Instance Segmentation. We do experiments on long-tailed instance segmentation with LVIS v1.0 datasets [18], as one of the most challenging datasets in vision with a much higher imbalance factor compared to the rest. We use the official splits and evaluation is conducted on validation
Table 2: Test on instance segmentation and molecule classification.
Method	LVIS (instance seg.)	OGBG-MOLBBBP	OGBG-MOLBACE
APm	APr	APc	APf	ROC-AUC	ROC-AUC
Vanilla	20.62	9.66	19.11	27.11	69.21 ± 0.34	79.40 ± 1.20
Focal Loss [35]	19.69	8.75	17.74	26.68	68.15 ± 1.29	80.65 ± 1.50
LDAM [3]	15.53	7.14	13.23	20.89	65.99 ± 0.94	79.20 ± 2.00
LogitAdjust [39]	22.41	11.70	21.18	28.48	69.05 ± 1.59	81.25 ± 0.33
CB-CE [8]	20.38	9.58	18.72	26.98	69.19 ± 1.18	79.39 ± 1.16
CB-FC [8]	21.41	15.67	20.34	25.25	69.51 ± 1.18	80.24 ± 1.45
Balanced Softmax	22.60	12.88	21.20	28.44	68.13 ± 0.87	80.26 ± 2.28
Ours(ratio-based)	22.64	13.15	21.26	28.34	70.48 ± 0.73	82.48 ± 1.59
set. The setting mainly follows the experiments in [45], whose details are given in Appendix H.  $\mathrm{AP}_m$ ,  $\mathrm{AP}_r$ ,  $\mathrm{AP}_c$ ,  $\mathrm{AP}_f$  are reported in Tab. 2.  $\mathrm{AP}_m$  denotes Average Precision of masks,  $\mathrm{AP}_r$ ,  $\mathrm{AP}_c$ ,  $\mathrm{AP}_f$  denote Average Precision of masks on rare classes, common classes and frequent classes. As shown in Tab. 2, we can find that our method achieves the best  $\mathrm{AP}_m$  on the entire set and perform competitively for other evaluations.
More Experiments. We also perform experiments on Long-Tailed contrastive representation learning on CIFAR10-LT, CIFAR100-LT with 100 imbalanced factor. We compare our methods with the methods in [5, 27]. The settings and results are given in Appendix H.5. Besides, to further demonstrate the usefulness of viewing classification as OT, a new inference method is proposed for testing. We replaced the softmax (i.e. constraints within  $U(\mathbf{a})$ ) with Sinkhorn algorithm (i.e. constraints within  $U(\mathbf{a},\mathbf{b})$ ), where  $\mathbf{b}$  represents the assumed ratio in the testing data (e.g., long-tailed, uniform, or reverse long-tailed distribution). Details are given in Appendix H.6.
We have provided a matching perspective to provide new understanding to classification. Under our matching-based framework, we show that the inherent challenge of long-tailed classification lies in the inconsistency of the prior information used for training and testing under the matching formulation. We then develop inverse relative entropic OT approach to revisit the classification problem, and especially for the long-tailed case, we develop a simple yet effective technique based on our theoretical insights. Experimental results verify the effectiveness of our methods.
Broader impacts. Our work establishes a theoretical framework for understanding and solving (self-)supervised problems using OT theory. This not only broadens the range of applications for OT in generative models, such as WGAN and its variants, but also provides a new perspective for redefining, understanding, and solving (self-)supervised problems. Furthermore, the theory and concepts of OT can facilitate the discovery of interesting relationships and insights for supervised problems, such as the connection between open set recognition and unbalanced OT. Limitations. It assumes known prior label distribution testing data, which is often unknown in real-world scenarios.