Doubly Robust Self-Training

The proposed algorithm is based on a simple modification of the standard loss for self-training. Assume that we are given a set of unlabeled samples, $\mathcal{D}_{1}=\{X_{1},X_{2},\cdots,X_{m}\}$, drawn from a fixed distribution $\mathbb{P}_{X}$, a set of labeled samples $\mathcal{D}_{2}=\{(X_{m+1},Y_{m+1}),(X_{m+2},Y_{m+2}),\cdots,(X_{m+n},Y_{m+n})\}$, drawn from some joint distribution $\mathbb{P}_{X}\times\mathbb{P}_{Y|X}$, and a teacher model $\hat{f}$. Let $\ell_{\theta}(x,y)$ be a pre-specified loss function that characterizes the prediction error of the estimator with parameter $\theta$ on the given sample $(X,Y)$. Traditional self-training aims at minimizing the combined loss for both labeled and unlabeled samples, where the pseudo-labels for unlabeled samples are generated using $\hat{f}$:

Note that this can also be viewed as first using $\hat{f}$ to predict all the data, and then replacing the originally labeled points with the known labels:

Our proposed doubly robust loss instead replaces the coefficient $1/(m+n)$ with $1/n$ in the last two terms:

This seemingly minor change has a major beneficial effect—the estimator becomes consistent and doubly robust.

The result shows that the true parameter $\theta^{\star}$ is also a local minimum of the doubly robust loss, but not a local minimum of the original self-training loss. We flesh out this comparison for the special example of mean estimation in Section 2.1, and present empirical results on image and driving datasets in Section 3.

Missing-data inference and causal inference. The general problem of causal inference can be formulated as a missing-data inference problem as follows. For each unit in an experiment, at most one of the potential outcomes—the one corresponding to the treatment to which the unit is exposed—is observed, and the other potential outcomes are viewed as missing (Holland, 1986, Ding and Li, 2018). Two of the standard methods for solving this problem are data imputation Rubin (1979) and propensity score weighting Rosenbaum and Rubin (1983). A doubly robust causal inference estimator combines the virtues of these two methods. The estimator is referred to as “doubly robust” due to the following property: if the model for imputation is correctly specified then it is consistent no matter whether the propensity score model is correctly specified; on the other hand, if the model propensity score model is correctly specified, then it is consistent no matter whether the model for imputation is correctly specified (Scharfstein et al., 1999, Bang and Robins, 2005, Birhanu et al., 2011, Ding and Li, 2018).

We note in passing that double machine learning is another methodology that is inspired by the doubly robust paradigm in causal inference (Semenova et al., 2017, Chernozhukov et al., 2018a, b, Foster and Syrgkanis, 2019). The problem in double machine learning is related to the classic semiparametric problem of inference for a low-dimensional parameter in the presence of high-dimensional nuisance parameters, which is different goal than the predictive goal characterizing semi-supervised learning.

The recent work of prediction-powered inference (Angelopoulos et al., 2023) focuses on confidence estimation when there are both unlabeled data, labeled data, along with a teacher model. Their focus is the inferential problem of obtaining a confidence set, while ours is the doubly robust property of a point estimator. Since they focus on confidence estimation, an important, strong, yet biased baseline point-estimate algorithm that directly combines the ground-truth labels and pseudo-labels is not considered in their case. In our paper, we show with both theory and experiments that the proposed doubly-robust estimator achieves better performance than the naive combination of ground-truth labels and pseudo-labels.

Self-training. Self-training is a popular semi-supervised learning paradigm in which machine-generated pseudo-labels are used for training with unlabeled data (Lee, 2013, Berthelot et al., 2019b, a, Sohn et al., 2020a, Zhao and Zhu, 2023). To generate these pseudo-labels, a teacher model is pre-trained on a set of labeled data, and its predictions on the unlabeled data are extracted as pseudo-labels. Previous work seeks to address the noisy quality of pseudo-labels in various ways. MixMatch (Berthelot et al., 2019b) ensembles pseudo-labels across several augmented views of the input data. ReMixMatch (Berthelot et al., 2019a) extends this by weakly augmenting the teacher inputs and strongly augmenting the student inputs. FixMatch (Sohn et al., 2020a) uses confidence thresholding to select only high-quality pseudo-labels for student training.

Self-training has been applied in both 2D computer vision problems (Liu et al., 2021a, Jeong et al., 2019, Tang et al., 2021, Sohn et al., 2020b, Zhou et al., 2022) and 3D problems (Park et al., 2022, Wang et al., 2021, Li et al., 2023, Liu et al., 2023) object detection. STAC (Sohn et al., 2020b) enforces consistency between strongly augmented versions of confidence-filtered pseudo-labels. Unbiased teacher (Liu et al., 2021a) updates the teacher during training with an exponential moving average (EMA) of the student network weights. Dense Pseudo-Label (Zhou et al., 2022) replaces box pseudo-labels with the raw output features of the detector to allow the student to learn richer context. In the 3D domain, 3DIoUMatch (Wang et al., 2021) thresholds pseudo-labels using a model-predicted Intersection-over-Union (IoU). DetMatch (Park et al., 2022) performs detection in both the 2D and 3D domains and filters pseudo-labels based on 2D-3D correspondence. HSSDA (Liu et al., 2023) extends strong augmentation during training with a patch-based point cloud shuffling augmentation. Offboard3D (Qi et al., 2021) utilizes multiple frames of temporal context to improve pseudo-label quality.

There has been a limited amount of theoretical analysis of these methods, focusing on semi-supervised methods for mean estimation and linear regression (Zhang et al., 2019, Azriel et al., 2022). Our analysis bridges the gap between these analyses and the doubly robust estimators in the causal inference literature.

As a concrete example, in the case of one-dimensional mean estimation we take $\ell_{\theta}(X,Y)=(\theta-Y)^{2}$. Our target is to find some $\theta^{\star}$ that satisfies

One can see that $\theta^{\star}=\mathbb{E}[Y]$. In this case, the loss for training only on labeled data becomes

Moreover, the optimal parameter is $\hat{\theta}_{\mathsf{TL}}=\frac{1}{n}\sum_{i=m+1}^{m+n}Y_{i}$, which is a simple empirical average over all observed $Y$’s.

For a given pre-existing predictor $\hat{f}$, the loss for self-training becomes

It is straightforward to see that the minimizer of the loss is the unweighted average between the unlabeled predictors $\hat{f}(X_{i})$’s and the labeled $Y_{i}$’s:

In the case of $m\gg n$, the mean estimator is almost the same as the average of all the predicted values on the unlabeled dataset, which can be far from $\theta^{\star}$ when the predictor $\hat{f}$ is inaccurate.

On the other hand, for the proposed doubly robust estimator, we have

Note that the loss is still convex, and we have

This recovers the estimator in prediction-powered inference (Angelopoulos et al., 2023). Assume that $\hat{f}$ is independent of the labeled data. We can calculate the mean-squared error of the three estimators as follows.

The proof is deferred to Appendix F. The proposition illustrates the double-robustness of $\hat{\theta}_{\mathsf{DR}}$—no matter how poor the estimator $\hat{f}(X)$ is, the rate is always upper bounded by $\frac{4}{n}(\mathsf{Var}[Y]+\mathsf{Var}[\hat{f}(X)])$. On the other hand, when $\hat{f}(X)$ is an accurate estimator of $Y$ (i.e., $\mathsf{Var}[{\hat{f}(X)-Y}]$ is small), the rate can be improved to $\frac{2}{m+n}\mathsf{Var}[Y]$. In contrast, the self-training loss always has a non-vanishing term, $\frac{2m^{2}}{(m+n)^{2}}\mathbb{E}[(\hat{f}(X)-Y)]^{2}$, when $m\gg n$, unless the predictor $\hat{f}$ is accurate.

On the other hand, when $\hat{f}(x)=\hat{\beta}_{(-1)}^{\top}x+\hat{\beta}_{1}$ is a linear predictor trained on the labeled data with $\hat{\beta}=\operatorname*{arg\,min}_{\beta=[\beta_{1},\beta_{(-1)}]}\frac{1}{n}\sum_{i=m+1}^{m+n}(\beta_{(-1)}^{\top}X_{i}+\beta_{1}-Y_{i})^{2}$, our estimator reduces to the semi-supervised mean estimator in Zhang et al. (2019). Let $\tilde{X}=[1,X]$. In this case, we also know that the self-training reduces to training only on labeled data, since $\hat{\theta}_{\mathsf{TL}}$ is also the minimizer of the self-training loss. We have the following result that reveals the superiority of the doubly robust estimator compared to the other two options.

In the general case, we show that the doubly robust loss function continues to exhibit desirable properties. In particular, as $n,m$ goes to infinity, the global minimum of the original loss is also a critical point of the new doubly robust loss, no matter how inaccurate the predictor $\hat{f}$.

Let $\theta^{\star}$ be the minimizer of $\mathbb{E}_{\mathbb{P}_{X,Y}}[\ell_{\theta}(X,Y)]$. Let $\hat{f}$ be a pre-existing model that does not depend on the datasets $\mathcal{D}_{1},\mathcal{D}_{2}$. We also make the following regularity assumptions.

Given this assumption, we denote $\Sigma_{\theta}^{Y-\hat{f}}=\mathsf{Cov}[\nabla_{\theta}\ell_{\theta}(X,\hat{f}(X))-\nabla_{\theta}\ell_{\theta}(X,Y)]$ and let $\Sigma_{\theta}^{\hat{f}}=\mathsf{Cov}[\nabla_{\theta}\ell_{\theta}(X,\hat{f}(X))]$, $\Sigma_{\theta}^{Y}=\mathsf{Cov}[\nabla_{\theta}\ell_{\theta}(X,Y)]$.

The proof is deferred to Appendix G. From the example of mean estimation we know that one can design instances such that $\|\nabla_{\theta}\mathcal{L}^{\mathsf{SL}}_{\mathcal{D}_{1},\mathcal{D}_{2}}(\theta^{\star})\|_{2}\geq C$ for some positive constant $C$.

When the loss $\nabla_{\theta}\mathcal{L}^{\mathsf{DR}}_{\mathcal{D}_{1},\mathcal{D}_{2}}$ is convex, the global minimum of $\nabla_{\theta}\mathcal{L}^{\mathsf{DR}}_{\mathcal{D}_{1},\mathcal{D}_{2}}$ converges to $\theta^{\star}$ as both $m,n$ go to infinity. When the loss $\nabla_{\theta}\mathcal{L}^{\mathsf{DR}}_{\mathcal{D}_{1},\mathcal{D}_{2}}$ is strongly convex, it also implies that $\|\hat{\theta}-\theta^{\star}\|_{2}$ converges to zero as both $m,n$ go to infinity, where $\hat{\theta}$ is the minimizer of $\nabla_{\theta}\mathcal{L}^{\mathsf{DR}}_{\mathcal{D}_{1},\mathcal{D}_{2}}$.

When $\hat{f}$ is a perfect predictor with $\hat{f}(X)\equiv Y$ (and $Y|X=x$ is deterministic), one has $\mathcal{L}^{\mathsf{DR}}_{\mathcal{D}_{1},\mathcal{D}_{2}}(\theta^{\star})=\frac{1}{m+n}\sum_{i=1}^{m+n}\ell_{\theta}(X_{i},Y_{i})$. The effective sample size is $m+n$ instead of $n$ in $\mathcal{L}^{\mathsf{SL}}_{\mathcal{D}_{1},\mathcal{D}_{2}}(\theta)$.

When $\hat{f}$ is also trained from the labeled data, one may apply data splitting to achieve the same guarantee up to a constant factor. We provide further discussion in Appendix E.

We also consider the case in which the marginal distributions of the covariates for the labeled and unlabeled datasets are different. Assume in particular that we are given a set of unlabeled samples, $\mathcal{D}_{1}=\{X_{1},X_{2},\cdots,X_{m}\}$, drawn from a fixed distribution $\mathbb{P}_{X}$, a set of labeled samples, $\mathcal{D}_{2}=\{(X_{m+1},Y_{m+1}),(X_{m+2},Y_{m+2}),$ $\cdots,(X_{m+n},Y_{m+n})\}$, drawn from some joint distribution $\mathbb{Q}_{X}\times\mathbb{P}_{Y|X}$, and a pre-trained model $\hat{f}$. In the case when the labeled samples do not follow the same distribution as the unlabeled samples, we need to introduce an importance weight $\pi(x)$. This yields the following doubly robust estimator:

Note that we not only introduce the importance weight $\pi$, but we also change the first term from the average of all the $m+n$ samples to the average of $n$ samples.

The proof is deferred to Appendix H. The proposition implies that as long as either $\pi$ or $\hat{f}$ is accurate, the expectation of the loss is the same as that of the target loss. When the distributions for the unlabeled and labeled samples match each other, this reduces to the case in the previous sections. In this case, taking $\pi(x)=1$ guarantees that the expectation of the doubly robust loss is always the same as that of the target loss.

Datasets and settings. We evaluate our doubly robust self-training method on the ImageNet100 dataset, which contains a random subset of 100 classes from ImageNet-1k (Russakovsky et al., 2015), with 120K training images (approximately 1,200 samples per class) and 5,000 validation images (50 samples per class). To further test the effectiveness of our algorithm in a low-data scenario, we create a dataset that we refer to as mini-ImageNet100 by randomly sampling 100 images per class from ImageNet100. Two models were evaluated: (1) DaViT-T (Ding et al., 2022), a popular vision transformer architecture with state-of-the-art performance on ImageNet, and (2) ResNet50 (He et al., 2016), a classic convolutional network to verify the generality of our algorithm.

Baselines. To provide a comparative evaluation of doubly robust self-training, we establish three baselines: (1) ‘Labeled Only’ for training on labeled data only (partial training set) with a loss $\mathcal{L}^{\mathsf{TL}}$, (2) ‘Pseudo Only’ for training with pseudo labels generated for all training samples, and (3) ‘Labeled + Pseudo’ for a mixture of pseudo-labels and labeled data, with the loss $\mathcal{L}^{\mathsf{SL}}$. See the Appendix for further implementation details and ablations.

Results on ImageNet100. We first conduct experiments on ImageNet100 by training the model for 20 epochs using different fractions of labeled data from 1% to 100%. From the results shown in Fig. 1, we observe that: (1) Our model outperforms all baseline methods on both two networks by large margins. For example, we achieve 5.5% and 5.3% gains (Top-1 Acc) on DaViT over the ‘Labeled + Pseudo’ method for 20% and 80% labeled data, respectively. (2) The ‘Labeled + Pseudo’ method consistently beats the ‘Labeled Only’ baseline. (3) While ‘Pseudo Only’ works for smaller fractions of the labeled data (less than 30%) on DaViT, it is inferior to ‘Labeled Only’ on ResNet50.

Results on mini-ImageNet100. We also perform comparisons on mini-ImageNet100 to demonstrate the performance when the total data volume is limited. From the results in Table 1, we see our model generally outperforms all baselines. As the dataset size decreases and the number of training epochs increases, the gain of our algorithm becomes smaller. This is expected, as (1) the models are not adequately trained and thus have noise issues, and (2) there are an insufficient number of ground truth labels to compute the last term of our loss function. In extreme cases, there is only one labeled sample (1%) per class.

Doubly robust object detection. Given a visual representation of a scene, 3D object detection aims to generate a set of 3D bounding box predictions $\{b_{i}\}_{i\in[m+n]}$ and a set of corresponding class predictions $\{c_{i}\}_{i\in[m+n]}$. Thus, each single ground-truth annotation $Y_{i}\in Y$ is a set $Y_{i}=(b_{i},c_{i})$ containing a box and a class. During training, the object detector is supervised with a sum of the box regression loss $\mathcal{L}_{loc}$ and the classification loss $\mathcal{L}_{cls}$, i.e. $\mathcal{L}_{obj}=\mathcal{L}_{loc}+\mathcal{L}_{cls}$.

In the self-training protocol for object detection, pseudo-labels for a given scene $X_{i}$ are selected from the labeler predictions $f(X_{i})$ based on some user-defined criteria (typically the model’s detection confidence). Unlike in standard classification or regression, $Y_{i}$ will contain a differing number of labels depending on the number of objects in the scene. Furthermore, the number of extracted pseudo-labels $f(X_{i})$ will generally not be equal to the number of scene ground-truth labels $Y_{i}$ due to false positive/negative detections. Therefore it makes sense to express the doubly robust loss function in terms of the individual box labels as opposed to the scene-level labels. We define the doubly robust object detection loss as follows:

where $M$ is the total number of pseudo-label boxes from the unlabeled split, $N$ is the total number of labeled boxes, $X^{\prime}_{i}$ is the scene with pseudo-label boxes from the labeled split, and $N_{ps}$ is the total number of pseudo-label boxes from the labeled split. We note that the last two terms now contain summations over a differing number of boxes, a consequence of the discrepancy between the number of manually labeled boxes and pseudo-labeled boxes. Both components of the object detection loss (localization/classification) adopt this form of doubly robust loss.

Dataset and setting. To evaluate doubly robust self-training in the autonomous driving setting, we perform experiments on the large-scale 3D detection dataset nuScenes (Caesar et al., 2020). The nuScenes dataset is comprised of 1000 scenes (700 training, 150 validation and 150 test) with each frame containing sensor information from RGB camera, LiDAR, and radar scans. Box annotations are comprised of 10 classes, with the class instance distribution following a long-tailed distribution, allowing us to investigate our self-training approach for both common and rare classes. The main 3D detection metrics for nuScenes are mean Average Precision (mAP) and the nuScenes Detection Score (NDS), a dataset-specific metric consisting of a weighted average of mAP and five other true-positive metrics. For the sake of simplicity, we train object detection models using only LiDAR sensor information.

Results. After semi-supervised training, we evaluate our student model performance on the nuScenes val set. We compare three settings: training the student model with only the available labeled data (i.e., equivalent to teacher training), training the student model on the combination of labeled/teacher-labeled data using the naive self-training loss, and training the student model on the combination of labeled/teacher-labeled data using our proposed doubly robust loss. We report results for training with 1/24, 1/16, and 1/4 of the total labels in Table 2. We find that the doubly robust loss improves both mAP and NDS over using only labeled data and the naive baseline in the lower label regime, whereas performance is slightly degraded when more labels are available. Furthermore, we also show a per-class performance breakdown in Table 3. We find that the doubly robust loss consistently improves performance for both common (car, pedestrian) and rare classes. Notably, the doubly robust loss is even able to improve upon the teacher in classes for which pseudo-label training decreases performance when using the naive training (e.g., barriers and traffic cones).