Uncertainty-aware self-training with expectation maximization basis transformation

TL;DR

This work addresses the over-confident pseudo-label problem in self-training by introducing an uncertainty-aware framework that combines Expectation-Maximization-based label smoothing with a Basis Extraction Net to derive a probabilistic, low-rank basis from data. The approach alternates between EM-driven pseudo-label generation and uncertainty-aware retraining, using an ATT-based basis extractor and an orthogonality regularizer to produce robust latent representations. By modeling pseudo-labels as distributions with variance and incorporating uncertainty into both sample selection and retraining losses, the method achieves improved performance on image classification and semantic segmentation benchmarks, outperforming state-of-the-art confidence-aware self-training methods by 1–3 percentage points in several settings. The proposed framework offers a principled way to encode and leverage data- and model-level uncertainty in semi-supervised learning, with potential applicability to a wide range of self-training tasks.

Abstract

Self-training is a powerful approach to deep learning. The key process is to find a pseudo-label for modeling. However, previous self-training algorithms suffer from the over-confidence issue brought by the hard labels, even some confidence-related regularizers cannot comprehensively catch the uncertainty. Therefore, we propose a new self-training framework to combine uncertainty information of both model and dataset. Specifically, we propose to use Expectation-Maximization (EM) to smooth the labels and comprehensively estimate the uncertainty information. We further design a basis extraction network to estimate the initial basis from the dataset. The obtained basis with uncertainty can be filtered based on uncertainty information. It can then be transformed into the real hard label to iteratively update the model and basis in the retraining process. Experiments on image classification and semantic segmentation show the advantages of our methods among confidence-aware self-training algorithms with 1-3 percentage improvement on different datasets.

Figures (3)

• Figure 1: Uncertainty-aware representations. In the right part of this figure, dashed curves represent the basis distributions while the blue curve represent the uncertainty-aware representation and uncertainty-aware labels of the data. The expectation of the labels could be used as the final label and the variance could be used to evaluate the uncertainty.
• Figure 2: One self training round. Pseudo-label generation (a) use EM algorithm to update the Gaussian basis and the classifier, then it generates some pseudo-labels with uncertainty information while the classifier is also trained in this stage. Then in model retraining stage (b), an uncertainty-aware training strategy is used to update the whole model (CNN and classifier).
• Figure 3: Whole training process for basis initialization net. Concretely, we train the model like classical machine learning training process and add a small module (attention block) to extract the processed weights which then become the initialized basis of EM algorithm.