Fix-A-Step: Effective Semi-supervised Learning from Uncurated Unlabeled Sets

A key roadblock to applying supervised learning to real applications is the need to assemble a large-enough labeled dataset for the intended task. Modern deep learning pipelines are especially data-hungry. In many cases, the acquisition of a large dataset of unlabeled features is rather affordable. However, providing reliable labels for each example is cost-prohibitive, requiring expensive, time-consuming work, often from human experts. This tradeoff is especially apt in our motivating application: classifying medical images. Images are collected in the course of routine care and easily available by querying a hospital’s electronic records. However, labeling these images well often requires clinical staff with years of training to spend minutes per image.

If only a tiny labeled set is available but access to a huge unlabeled set of images is possible, one promising approach is semi-supervised learning (SSL) (Zhu, 2005; van Engelen and Hoos, 2020). Recent years have seen amazing progress on standard benchmarks such as recognizing digits from photos of address numbers on houses (SVHN; Netzer et al. (2011)). With only 100 labeled examples per digit class, a supervised neural net’s error rate is roughly 12%. Using a large unlabeled set, the FixMatch SSL method (Sohn et al., 2020) delivers error below 2.5%, while even more recent work has pushed below 2% (Xu et al., 2021; Han et al., 2020).

Unfortunately, common benchmarks like SVHN may be too optimistic. In real tasks, unlabeled sets will be collected automatically at scale for convenience, and thus may differ from the labeled set in terms of represented classes or class frequencies, among other differences. To maintain the efficiency that motivates SSL, the effort required to curate the unlabeled set – that is, apply labels and assess differences – remains prohibitive. To be effective in practice, SSL must reliably improve accuracy despite uncurated data.

Off-the-shelf SSL using mismatched unlabeled sets often predicts worse than ignoring the unlabeled set altogether (Oliver et al., 2018). Many recent methods do try to be robust to uncurated unlabeled data (see Tab. 1). Broadly, most follow the same intuitive direction: learn to identify examples in the unlabeled set that are out-of-distribution (OOD), and then remove or downweight these. However, we find existing work delivers at best modest gains in accuracy.

Larger, more reliable gains are needed to enable SSL in real uncurated applications. This study makes 3 contributions toward this goal backed by reproducible experiments¹¹1Code: https://github.com/tufts-ml/fix-a-step.

First, we challenge the dominant paradigm that handles uncurated unlabeled sets by filtering out OOD examples. Our experiments suggest that even perfect OOD filtering (which is unrealistic in practice) does not perform well. Instead, we argue for a new paradigm: OOD images from uncurated unlabeled sets are potentially helpful.

Second, following this paradigm we introduce a new training procedure called Fix-A-Step that delivers accuracy gains from uncurated unlabeled sets. When applied to repair several deep SSL methods across a full range of labeled-unlabeled mismatch levels, our Fix-A-Step improves predictions better than alternative methods while being substantially simpler too.

Finally, we offer a new SSL benchmark using real uncurated medical images that can assess cross-hospital generalization. For too long, SSL methods advances have been evaluated only on the same repurposed generic object recognition datasets. Using three inter-operable open-access datasets – TMED (Huang et al., 2021b), CAMUS (Leclerc et al., 2019), and Unity (Howard et al., 2021) – we pursue a clinically-relevant problem: recognizing the view type of an echocardiogram image of the heart. Future methods for learning from limited data can follow our reproducible protocol. We hope this benchmark enables authentic SSL applications in medicine and ultimately improves the efficiency and quality of care for patients with heart disease.

We pursue semi-supervised learning (SSL) (Zhu, 2005; van Engelen and Hoos, 2020) for the specific problem of image classification with deep neural networks (Oliver et al., 2018). We can observe images $x$ (represented as $D$-dimensional vectors) as well as corresponding labels $y\in \{1,2,\dots C\}$ for $C$ classes of interest. The goal is to train predictors from both a labeled dataset $D^L$ of feature-labeled pairs $x,y$ and an unlabeled dataset $D^U$ containing feature vectors $x$ only.

We have designed a training procedure for deep SSL classifiers that we call Fix-A-Step, which is short for Fix via Augmentation and Step direction modification. Fix-A-Step can be applied to any SSL method that minimizes an SSL objective matching Eq. (1) via gradient descent. The primary goal of Fix-A-Step is to make any SSL classifier far more robust to uncurated unlabeled data.

The overall Fix-A-Step training algorithm is provided in Alg. 1. Two phases of Fix-A-Step happen when visiting each minibatch during gradient descent. First, in the augmentation phase (lines 3-4), our insight is the unlabeled set might be helpful for creating useful augmentations, even when uncurated, by injecting realistic diversity (for a motivating experiment, see App. A.2). Inspired by MixMatch (Berthelot et al., 2019), we transform each labeled pair $(x,y)$ using another pair $(x^{\prime },y^{\prime })$ drawn either from the labeled set or the unlabeled set (where only $x^{\prime }$ is known, so use soft pseudo-label predictions for $y^{\prime }$, see Alg. C.1). Given $x,y$ and $x^{\prime },y^{\prime }$, we build a new labeled pair $\tilde{x},\tilde{y}$ via MixUp (Zhang et al., 2017) interpolation (see Alg. C.2), then use that pair to compute the labeled loss.

Second, in the step direction phase (lines 5-10), we prioritize the labeled loss in parameter updates, only using the unlabeled loss if it improves labeled-set performance. At each batch, we compute two gradient vectors, one for each term in the loss: Let $g^L={\nabla }_w{\ell }^L$ and let $g^U={\nabla }_w{\ell }^U$. The update for weights $w$ using step size $ϵ$ is then

In the top case, we do the standard steepest descent update that minimizes the two-term SSL objective in Eq. (1). In the bottom case, we perform an alternative update, using only the labeled-term gradient. This two-case construction tries to ensure that SSL learning does not harm labeled set performance. Formally, we can show that each case of the gradient updates in Eq. (2) adjusts weights $w$ in a descent directions for the labeled set loss at the current minibatch.

Definition 1: Descent direction of loss $\ell$. For any loss function $\ell$ parameterized by weight vector $w\in R^D$, a vector $v\in R^D$ is a descent direction of $\ell$ at $w$ if it satisfies $v^T{\nabla }_w\ell <0$ (Boyd and Vandenberghe, 2004).

Lemma 1: The update in Eq. (2) steps in a descent direction of the labeled loss ${\ell }^L$ at the current minibatch. We prove for each of the two cases in Eq. (2). Top case: Here by assumption the inner product ${\sum }_d{g}_d^L{g}_d^U$ is positive. This implies that $v=-(g^L+\lambda g^U)$ is a descent direction, because $\lambda >0$ and thus

Bottom case: $-g^L$ is a descent direction for ${\ell }^L$ by definition.

While Lemma 1 provides a justification for our approach, we cannot formally guarantee the labeled-set loss will not increase after each step, for the same reasons that stochastic gradient descent (SGD) does not always decrease its loss after each minibatch update. First, a descent direction of a small minibatch may not be a descent direction of the entire dataset. Second, even though the direction of the step points locally downhill, the length of the step matters; if the step size $ϵ>0$ is too large, the loss may increase. Nevertheless, with proper step size tuning SGD has been wildly successful despite following minibatch-specific descent directions without formal guarantees of non-increasing loss. Thus far, we find our approach also successful in practice.

Geometric Intuition. Recall that two vectors $g^L$ and $g^U$ have positive inner product if and only if the angle between the vectors is below 90 degrees (less than perpendicular). Thus, we can motivate the alternative lower case in Eq. (2) geometrically. We usually care most about classifier accuracy, so when the two gradients point in opposing directions (bottom case), we might be better off ignoring the unlabeled and updating parameters using only the labeled loss.

Inspiration from multi-task learning. Our design of the step direction modification in Eq. (2) is motivated by previously suggested step direction modifications for multi-task learning in the presence of a “main” task and an “auxiliary” task (Du et al., 2020). This is part of a broader thread of work on step modification for multi-task (Yu et al., 2020b) and continual learning (Lopez-Paz and Ranzato, 2017). To our knowledge, such ideas have not yet been suggested or validated for SSL, or been combined with augmentation.

Comparison to OOD filtering. Using the Fix-A-Step update, we know that each (stochastic) gradient step could beneficially reduce the labeled set loss, given a small-enough step size $ϵ$. We argue our step direction modification is safer than simply learning to downweight individual terms in the unlabeled loss. Without care, the latter could still take steps in a problematic non-descent direction of the labeled loss.

Simplicity compared to related work. We emphasize a key overall advantage of Fix-A-Step is extreme simplicity. Beyond the modest cost of MixMatch-like augmentation, we compute exactly the same losses and gradients as any standard deep SSL solving Eq. (1) would. Each possible weight parameter update is also simple and straightforward. Determining which update to perform depends purely on computing the inner product between two vectors $g^L$ and $g^U$, which is efficient on modern hardware and adds negligible runtime cost. Table 1 suggests our approach is favorable to other open-set or safe SSL approaches in this regard. There is no added complexity from extra backward passes through the classifier, extra neural networks that must be trained for OOD discrimination, no need for curriculum learning, and no expensive bi-level optimization problem to solve.

Our open source code uses PyTorch Paszke et al. (2019) and allows reproducing each experiment (see App. E for details). Following Oliver et al. (2018), for all methods we use the same ResNet-28-2 architecture Zagoruyko and Komodakis (2016), apply standard augmentation (random crops, flips) on the labeled set, and regularize via weight decay.

SSL Baselines. We compared to 6 SSL methods (Pi-Model, Mean-Teacher, Pseudo-label, VAT, MixMatch, and FixMatch) as well as the baseline that minimizes labeled loss ${\ell }^L$ on the labeled set (“labeled only”). We also compare to 3 methods intended for open-set/safe SSL: UASD, DS3L, and MTCF. When possible, we use our own adapted implementations of baselines, to ensure that architectures, training, and hyperparameters are comparable and reproducible. Throughout, if a result is copied from another publication, we mark that via an asterisk ($\ast$).

Training. Following choices in original implementations, each method is trained using either Adam with fixed learning rate or SGD with a cosine-annealing schedule for learning rate (Sohn et al., 2020). We found cosine-annealing and a slow linear ramp-up schedule for the unlabeled-loss-weight $\lambda$ particularly helpful for several baselines (see App. E). Each training run used one NVIDIA A100 GPU.

In pursuit of realistic evaluation, we consider a reproducible, clinically-relevant SSL task that we call Heart2Heart. The key question is this: can we generalize classifiers of ultrasound images of the heart from one hospital system to heart images from different hospitals in other countries with no overlap in staff or ownership. For training, we use the Tufts Medical Echocardiogram Dataset 2 (TMED-2) (Huang et al., 2022, 2021b). TMED2 contains a small labeled set of echocardiogram studies and a larger uncurated unlabeled set. Thanks to common device standards, these images are interoperable with two other open-access datasets of “echo” images: Unity from 17 hospitals in the UK (Howard et al., 2021) and CAMUS from a hospital in France (Leclerc et al., 2019). We emphasize that all datasets are deidentified and open to any researcher for non-commercial purposes. We will release code so this Heart2Heart evaluation of SSL can be readily reproduced for future SSL methods.

Classification task: View type of 2D TTE image. Trans-thoracic echocardiography (TTE) is a gold-standard way to non-invasively capture the heart’s 3-dimensional anatomy for measurement and diagnosis. A human sonographer wields a handheld transducer over the patient’s chest at different angles in order to provide clear views of each facet of the heart. A routine TTE scan of a patient, called a study, produces many images (median=68, 10-90th percentile range=27-97 in TMED-2), each showing a canonical 2D view of the heart. No view type annotation is recorded with any image. In later analysis, clinicians must manually search over all recorded images to find a desired view type. Automated interpretation of echocardiograms must also be able to pick out specific view types before any useful measurements or diagnosis can be made, making view classification a prediction task with immediate potential clinical impact (Madani et al., 2018a; Huang et al., 2021b).

TMED-2 provides set of labeled images of four specific view types, known as PLAX, PSAX, A2C, and A4C, gathered from certified annotators. Reliably identifying these views would be particularly useful for key valve disease diagnostic tasks (Huang et al., 2021b). TMED-2 also contains an uncurated, truly unlabeled set of 353,500 images from routine TTEs from 5486 patient-studies. This is an authentic open-set SSL task, because at least 9 canonical view types frequently appear in routine TTEs (Mitchell et al., 2019).

Protocol. We train SSL methods on images from 56 labeled studies in TMED-2’s training set as well as all unlabeled studies (353,500 images). We report balanced accuracy on each split’s test set of 120 studies ($\sim$2104 view-labeled images per split). We then assess generalization of these U.S.-hospital-trained classifiers to images from different European hospitals. We report balanced accuracy on 7231 available PLAX, A2C, and A4C images from the Unity dataset, as well as all 2000 images (A2C and A4C views only) in the CAMUS dataset.

Results on Heart2Heart. Fig. 3 shows classifier performance on held-out data from all 3 datasets. TMED-2 evaluations (first panel) show that our Fix-A-Step procedure yields gains across all tested SSL methods (Pi-Model, VAT, FixMatch). Only with Fix-A-Step do the methods convincingly outperform the labeled-only baseline.

External evaluation on Unity and CAMUS (Fig. 3 panels 2-3) show that these gains transfer to new hospitals. Each tested SSL method performs better with Fix-A-Step than standard training. Across splits we see larger performance variation on Unity and CAMUS than on TMED-2, which highlights the difficulty of generalizing across hospitals as well as importance of external validation. All methods perform worse on CAMUS than other datasets; see App. B for further investigations. Overall, this Heart2Heart benchmark task shows the promise of Fix-A-Step to deliver gains from unlabeled data that generalize better than alternatives.

In summary, this paper makes three contributions to deep SSL image classification. First, we argue that uncurated or OOD data in the unlabeled set can be quite helpful, and should not merely be filtered out; experiments in Fig. 1 show that even with perfect OOD filtering most SSL methods deliver underwhelming accuracy gains. Additional evidence for the helpfulness of OOD unlabeled data is provided in App. A.2. Second, building on insights from multi-task learning (Du et al., 2020), our Fix-A-Step SSL training procedure uses gradient step modifications that prioritize the labeled-set loss, leading to effective SSL that is substantially simpler than alternatives (no new loss terms or extra neural networks). Finally, we hope our new Heart2Heart benchmark for SSL evaluation inspires robust studies of clinical model transportability across global populations.

Limitations. Our experiments on artificial unlabeled sets focused exclusively on mismatch in the labels. We did not explore shifts in the features $x$, though some modest shifts may be present in TMED2 data due to its uncurated nature. Our work’s exclusive focus is image classification. Fix-A-Step could apply to other types of data whose feature vectors have fixed dimension (required by MixUp). More work is needed on alternative data types or multi-modal SSL.

Our clinical applications have shown only a proof-of-concept for generalization of view classifiers across hospitals. Far more work is needed to rigorously assess generalizability and translate to improve clinical workflow and care.

Outlook. Fix-A-Step is a promising new first-line approach to SSL that can unlock the promise of uncurated unlabeled sets. We hope future work explores both augmentation and step direction modification ideas further, while maintaining our focus on simplicity and reproducibility.

Ethics statement. Semi-supervised image classification has many positive applications. Indeed, work on SSL is often specifically motivated by the promise of improved efficiency in environments where labels are expensive and time-consuming as is the case in medical imaging (Huang et al., 2021b; Madani et al., 2018b). However, care must be taken to ensure that automated methods actually benefit patients and do not widen current disparities (Celi et al., 2022). Our present Heart2Heart evaluations are an important step beyond single-center evaluations though do not reflect the true geographic and racial diversity of many patient populations. While all images used in our Heart2Heart task are completely de-identified and come from public open-access datasets (and thus did not require ethics review), we stress the responsibility we carry as researchers to protect the best interests of the individuals who contributed data.