SimPro: A Simple Probabilistic Framework
Towards Realistic Long-Tailed Semi-Supervised Learning

We begin by outlining the formulation for the realistic long-tailed semi-supervised learning (ReaLTSSL) problem, laying the groundwork for our approach. The setup involves a labeled dataset $\mathcal{D}_{l}\!=\!\{(x_{i},y_{i})\}_{i=1}^{N}$ and an unlabeled dataset $\mathcal{D}_{u}\!=\!\{x_{i}\}_{i=1}^{M}$, where $x_{i}\!\in\!\mathbb{R}^{d}$ represents the $i$-th data sample and $y_{i}\!\in\!\{0,1\}^{K}$ is the corresponding one-hot label, with $K$ denoting the number of classes. The objective of ReaLTSSL is to train a classifier $F_{\bm{\theta}}:\mathbb{R}^{d}\!\mapsto\!\{0,1\}^{K}$, parameterized by $\bm{\theta}$.

Further, we consider five diverse distributions for the unlabeled data (Fig. 1), reflecting various real-world situations.

In SSL, pseudo-labeling is a key technique for leveraging unlabeled data. This involves creating pseudo-labels for the unlabeled data using the model and then training the model with both the pseudo-labeled and ground-truth data. This aligns with the Expectation-Maximization (EM) algorithm (Dempster et al., 1977), where the E-step generates pseudo-labels, and the M-step updates the parameters using the pseudo-labels, maximizing the likelihood function.

Our method builds on a popular algorithm FixMatch (Sohn et al., 2020), which integrates consistency regularization in the standard SSL setting. Pseudo-labels are created via weakly-augmented unlabeled data and applied to train strongly-augmented samples based on a confidence threshold. The loss for unlabeled data is

where $q_{\omega}$ and $q_{\Omega}$ represent the prediction logits for weakly and strongly augmented samples, respectively, $\mathcal{H}$ denotes the cross-entropy loss, and $t$ is the confidence threshold.

In typical SSL scenarios, the assumption of identical distributions for labeled, unlabeled, and test data often prevails. However, long-tailed learning tasks usually involve imbalanced training sets and balanced test sets, leading to discrepancies in the prior distribution of $P(y)$ between training and testing data. Some studies (Ren et al., 2020; Menon et al., 2021; Hong et al., 2021) tackle this via Bayesian inference, introducing a prior distribution over class labels:

where $\phi_{y}$ denotes the class frequency in the training or test set, $\bm{\pi}$ is the class distribution parameter and $\bm{\theta}$ is the parameter of $P(x|y)/P(x)$. Here we omit the parameter of $P(x)$ for simplicity. The detailed mathematical derivation is provided in App. A.

While supervised learning allows for a known distribution parameter $\bm{\pi}$, enabling a direct application to model $P(y)$ and explicit decoupling from $\bm{\theta}$, ReaLTSSL poses a greater challenge as the prior $\bm{\pi}$ for unlabeled data is unknown. This necessitates innovative approaches to adapt to the imbalanced data while maintaining model efficacy.

In the realistic semi-supervised learning (ReaLTSSL) context, the conventional assumption of independent and identically distributed (i.i.d.) labeled and unlabeled data is no longer valid. Moreover, the marginal (class) distribution $P(y)$ of the unlabeled data may be inconsistent with that of the labeled data and remains unknown, which challenges the traditional SSL frameworks.

To overcome this, we introduce SimPro, an elegant and effective probabilistic framework adapted for the unique ReaLTSSL setting. Illustrated in Fig. 2, SimPro distinctively decouples $\bm{\pi}$ and $\bm{\theta}$, unlike traditional SSL methods (Sohn et al., 2020). In the E-step, we generate pseudo-labels using the parameters $\bm{\pi}$ and $\bm{\theta}$ obtained from the previous M-step. The M-step then models the conditional distribution $P(x|y)$ using network parameters $\bm{\theta}$, which are optimized through gradient descent. Simultaneously, we derive a closed-form solution for the class distribution $P(y)$, represented by $\bm{\pi}$.

It is worth noting that the treatment of the $\bm{\pi}$ in our framework is not heuristic. It is firmly rooted in probabilistic modeling and the principles of the EM algorithm, providing theoretical soundness as substantiated in Props. 1 and 2.

In addressing the ReaLTSSL challenge, we adopt an Expectation-Maximization (EM) approach, underpinned by a robust probabilistic model. The model is governed by the fundamental principles of conditional probability, as shown in:

Here, we do not explicitly parameterize $P({x})$, as per the independence of parameters through conditional parameterization (Koller & Friedman, 2009). Thus, when ${x}$ is not a condition, the parameters of the relevant notions omit the parameters of $P({x})$, such as $P({x})$, $P({x}|y)$, $P({x},y)$, etc. According to Sec. 3.1, this may lead to a potential misunderstanding, as the equation $P({x})=\sum_{y}P({x}|y;\bm{\theta})P(y;\bm{\pi})$ seems to suggest that $P({x})$ is parameterized by $\bm{\theta}$ and $\bm{\pi}$, which is not the case. The detailed mathematical derivation is provided in App. A.

We focus on estimating the parameters $\bm{\theta}$ and $\bm{\pi}$, pivotal for learning a discriminative model. Consequently, we concentrate on those terms dependent on $\bm{\theta}$ and $\bm{\pi}$, sidelining those independent of these parameters.

The complete data log-likelihood is thus articulated as:

where $\bm{\pi}\!=\!\{\bm{\pi}_{l},\bm{\pi}_{u}\}$ signifies the class distributions for labeled and unlabeled data, respectively, with $N$ and $M$ representing the number of labeled/unlabeled samples.

By Eq. 4, the expected complete data log-likelihood $\mathcal{Q}$ function is derived from the preceding iteration’s parameters, $\bm{\theta}^{\prime}$ and $\bm{\pi}^{\prime}$:

The E-step involves generating soft pseudo-labels $P(y|x_{j};\bm{\theta}^{\prime},\bm{\pi}^{\prime})$ under the current $\bm{\theta}^{\prime}$ and $\bm{\pi}^{\prime}$. These soft pseudo-labels are specifically defined by Eq. 10, which is detailed in Prop. 2. In the subsequent M-step, these pseudo-labels are used alongside the one-hot labels of the labeled data to compute the cross-entropy loss.

The M-step focuses on optimizing the expected complete data log-likelihood $\mathcal{Q}$-function concerning the parameters $\bm{\theta}$ and $\bm{\pi}$.

Building upon Prop. 2, it is crucial to acknowledge that the parameter vector $\bm{\phi}$ is vital for learning Bayes classifiers. Consequently, to delve deeper into the theoretical foundations, we evaluate the impact of $\bm{\phi}$ on the model’s performance. In line with the principles of online decision theory, we establish a regret bound for the decision error rate on the test set, denoted as $P(e;\bm{\phi})$. Our analysis is simplified by concentrating on a binary classification scenario, where the labels $y$ belong to $\{-1,+1\}$.

Prop. 3 illustrates that the regret bound is primarily governed by the first-order term of the estimation deviation. Additionally, it is inversely proportional to the ground truth $\bm{\phi}^{*}$, highlighting the learning challenges associated with imbalanced training data from a regret-bound perspective.

In SimPro, the E-step primarily involves generating pseudo-labels using parameters $\bm{\theta}$ and $\bm{\pi}$. Consequently, in the M-step, we first focus on optimizing the network parameter $\bm{\theta}$ guided by Eq. 8 via Stochastic Gradient Descent (SGD). Building on the FixMatch algorithm (Sohn et al., 2020), the overall training objective is formulated as:

where $\mathcal{L}_{l}$ and $\mathcal{L}_{u}$ represent the losses on labeled and unlabeled data, respectively. The hyper-parameter $\alpha$ acts as a scaling factor, the specifics of which are elucidated later.

For $\mathcal{L}_{l}$, we modify it (originally the standard cross-entropy loss) following Eq. 8:

where $\tau$ is a hyper-parameter for enhancing adaptability to long-tail distributions (Menon et al., 2021). $B$ is batch size.

For $\mathcal{L}_{u}$, we implement the standard SSL format (Eq. 1) and adapt it for ReaLTSSL:

where $\mu$ controls the number of unlabeled samples, and $t$ is a confidence threshold. The pseudo-label of weak augmentation $\omega$ from the Bayes classifier (Eq. 10) is denoted by

and the actual prediction $p_{y}$ is obtained using strong augmentation $\Omega$ and calibrated with $\bm{\phi}$ as shown in Eq. 8:

Moreover, in practical situations, the size of the unlabeled dataset $M$ is generally larger than that of the labeled dataset $N$. To ensure a balanced sample size in each iteration, we usually set $\mu\!=\!M/N$ in Eq. 15. In specific scenarios, we further adjust the balance factor in Eq. 13 to $\alpha\!=\!\mu\cdot N/M\!<\!1$. This methodology effectively mitigates overfitting to the labeled data (see Tabs. 6 and 7).

As discussed in Sec. 3.2, the parameter $\bm{\pi}$ of marginal distribution $P(y)$, has a closed-form solution in the M-step. Therefore, unlike $\bm{\theta}$, which requires SGD optimization, $\bm{\pi}$ (Eq. 7) and $\bm{\phi}$ (Eq. 9) are computed and updated via a moving average during training.

Based on the previous analysis, Our SimPro framework can be summarized as an extended EM algorithm, which includes:

$\bm{\cdot}$ E-step (Eq. 10): Generating pseudo-labels using model parameters $\bm{\theta}$ and estimated distribution parameters $\bm{\pi}$;

$\bm{\cdot}$ M-step: Optimizing network parameters $\bm{\theta}$ via SGD using Eq. 8 (actually using Eq. 13), and solving distribution parameters $\bm{\pi}$ and hyper-parameters $\bm{\phi}$ by Eq. 7 and Eq. 9.

For further clarity, the pseudo-code of SimPro is provided in Alg. 1. The modifications we made to the core training code, in comparison to FixMatch, are highlighted in bold. In the SimPro, we incorporate just a single additional line of code in the M-step to compute the closed-form solution of $\pi$ (Prop. 1). Furthermore, only four lines of code need to be modified to construct a Bayes classifier (Prop. 2) and to balance the loss between labeled and unlabeled data (denoted as $\alpha$). These minor yet crucial adjustments demonstrate that our SimPro framework is not only grounded in rigorous theoretical derivation but is also straightforward to implement in practice, exemplifying both simplicity and elegance.

The results are presented in Tab. 1 (for CIFAR10-LT), Tab. 2 (for CIFAR100-LT and STL10-LT), and Tab. 3 (for ImageNet-127/1k). It can be concluded that our method consistently outperforms the competitors across all distributions of unlabeled data and achieves SOTA performance. Notably, SimPro exhibits significant performance improvements on our two newly introduced distributions of unlabeled data: middle and head-tail. This substantiates the robust generalization capabilities of SimPro across various distributions that could potentially appear in real-world scenarios.

It is worth noting that compared to CIFAR10/100-LT, STL10-LT is a more challenging dataset that mirrors the real-world data distribution scenarios: an unknown distribution for the unlabeled data. The results in Tab. 2 demonstrate the significant improvements of SimPro over baseline methods.

Moreover, we also conduct experiments on ImageNet-127, whose test dataset is imbalanced and consistent with the labeled data and unlabeled data. However, this is not suitable as a benchmark for long-tail learning, as biased classifiers tend to perform well in such scenarios, which is precisely what we aim to avoid. Therefore, we resample it to achieve a uniform test distribution ($\gamma_{t}=1$). The results highlight that our SimPro achieves substantial performance enhancements when evaluated against this balanced test dataset. Beyond this, we further conduct experiments on ImageNet-1k to validate the performance of our method across a broader range of classes. The results in Tab. 3 demonstrate that our SimPro achieves state-of-the-art performance on ImageNet-1k.

To investigate the impact of the anchor distributions in ACR and CPE (Ma et al., 2024), we also incorporate them into our approach, referred to as SimPro^⋆. Instead of calculating the distribution distance and adjusting the consistency regularization as in ACR or employing multiple anchor experts as in CPE, our usage of these anchors is more straightforward: after training for five epochs, we calculate the distance between our estimated distribution $P(y|\bm{\pi})$ and the three anchor distributions (i.e. consistent, uniform, and reversed). This calculation helps us predict the actual distribution and construct the Bayes classifier. Then we fix the marginal distribution parameters $\bm{\pi}$ in the remainder of the training.

The results in Tab. 4 indicate that (1) the usage of anchor distributions can significantly enhance the performance of SimPro^⋆, consistently outperforming the original ACR and CPE; (2) our estimation for $\bm{\pi}$ is accurate (Fig. 5 further validates the accurate estimation of $\bm{\pi}$); (3) when the pre-defined anchors fail to cover the evaluated distributions (middle and head-tail), SimPro outperforms ACR by a large margin; (4) even compared to the original ACR, SimPro exhibits enhanced performance across all scenarios except the consistent distribution. This demonstrates the superior robustness and generalization ability of our method. We believe these advantages guarantee a better adaptable nature for SimPro in real applications and are more valuable than the accuracy improvements when using the anchor distributions.

Fig. 3 reports the performance of SimPro under more imbalance ratios of unlabeled data. The results indicate that our method consistently outperforms ACR across all imbalance ratios, further substantiating the robustness of our method.

We conduct a series of ablation studies to validate the effectiveness of different components in SimPro.

As outlined in App. B, we discover that reducing the threshold value improves performance for the CIFAR100-LT dataset. The rationale behind adjusting the confidence threshold is based on the observation that an increase in the number of classes typically results in a corresponding decrease in the confidence of the predictions. Consequently, it becomes necessary to lower the threshold to accommodate this change in confidence levels. Sensitivity analysis regarding the threshold value is presented in Fig. 4. It is consistently observed across different settings that the optimal performance is achieved when the threshold is set at $t=0.2$ and $t=0.95$ for the CIFAR100-LT and CIFAR10-LT, respectively. Moreover, to compare with ACR, we also conduct a sensitivity analysis of the confidence threshold $t$ for ACR. The results in Fig. 6 of App. C demonstrate that lowering the threshold does not improve performance for ACR.

Our study includes a visualization of the estimated distribution quality in Fig. 5. The vertical axis quantifies the Kullback-Leibler (KL) divergence, which measures the deviation of the estimated distribution from the ground truth. The results indicate a significant improvement in estimation accuracy after only a few training epochs. This empirical evidence validates the effectiveness of the theoretically derived estimation method of distribution, as outlined in Prop. 1.