SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sarah M. Hooper
Mayee F. Chen
Khaled Saab
Kush Bhatia
Curtis Langlotz
Radiology and Biomedical Data Science
Christopher Ré
Image classification and segmentation are common applications of deep learning to radiology. While many tasks can be framed using either classification or segmentation, classification has historically been cheaper to label and more widely used. However, recent work has drastically reduced the cost of training segmentation networks. In light of this recent work, we reexamine the choice of training classification vs. segmentation models. First, we use an information theoretic approach to analyze why segmentation vs. classification models may achieve different performances on the same dataset and task. We then implement methods for using segmentation models to classify medical images, which we call segmentation-for-classification, and compare these methods against traditional classification on three retrospective datasets  $(n = 2,018 - 19,237)$ . We use our analysis and experiments to summarize the benefits of using segmentation-for-classification, including: improved sample efficiency, enabling improved performance with fewer labeled images (up to an order of magnitude lower), on low-prevalence classes, and on certain rare subgroups (up to  $161.1\%$  improved recall); improved robustness to spurious correlations (up to  $44.8\%$  improved robust AUROC); and improved model interpretability, evaluation, and error analysis.
Classification and segmentation are two popular applications of deep learning to radiology [1]. Classification produces coarse-grained, image-level predictions, while segmentation produces fine-grained, pixel-level maps. While many tasks can be framed using either classification or segmentation (Figure 1), classification is often the default framing. For one, classification has been cheaper to label, requiring only one label per image compared to segmentation's tedious, pixel-wise labels. Second, training classification networks to produce image-level outputs mirrors radiologists' standard workflows, where radiologists summarize findings at the image-level and only routinely perform segmentation in a few sub-disciplines that rely on quantitative measurements over the segmentation masks. These factors have led to an abundance of medical image classification models [2, 3]. In contrast, because segmentation has been so cumbersome to label and is performed relatively infrequently by radiologists, segmentation networks are often only trained when the downstream application requires it (e.g., for computing the volumes or diameters of structures in an image).
Recent work in label-efficient training enables us to reexamine this paradigm. Self-supervised learning, in-context learning, weakly-supervised learning, and semi-supervised learning can substantially reduce labeling burden [4, 5, 6, 7, 8, 9, 10]. Additionally, more public datasets and broad-use pre
Classification Segmentation
Does this chest x-ray contain a pneumothorax?
Figure 1: Illustration of three medical image analysis problems that can be framed using either classification or segmentation.
Classification
Is this lesion melanoma?
Classification
Does this CT contain a vertebral fracture?
trained networks are coming online [11, 12, 13, 14, 15]. Together, these new methods, datasets, and pretrained backbones are enabling users to develop networks for new segmentation tasks with less and less annotation burden—a trend that will likely continue.
In light of this progress, we reexamine the convention around when to employ segmentation networks. Specifically, we investigate using segmentation networks for medical image classification, exploring if segmentation's fine-grained annotation and output can lead to benefits compared to classification's coarse annotation and output. Our contributions include:
- Exposition. We analyze why classification and segmentation networks may perform differently on the same datasets. We show that segmentation leads to more separable and robust embedding spaces, guiding what benefits we should expect to see (Section 2).
- Best practices. We implement multiple methods to obtain classification labels from segmentation networks, which we call segmentation-for-classification (Section 3). We empirically confirm our analysis (Section 4.1) and compare methods across datasets, tasks, and training conditions to build up best practices for training segmentation-for-classification networks (Section 4.2). We show segmentation-for-classification can improve aggregate performance compared to traditional classification models by up to  $16.2\%$  (from 0.74 AUROC to 0.86 AUROC). Finally, we explore semi-supervised segmentation-for-classification, showing existing label-efficient training methods for segmentation can directly benefit segmentation-for-classification.
- Trade-offs. We pull together our analysis and experiments to provide consolidated and expanded evidence of the benefits of using segmentation-for-classification, including: improved sample efficiency, enabling improved performance on small datasets, on low-prevalence classes, and on certain rare subgroups (up to  $161.1\%$  improved recall); improved robustness to spurious correlations (up to  $44.8\%$  improved robust AUROC); and improved model interpretability, evaluation, and error analysis (Section 5).
In summary, conventional wisdom driving the choice between classification and segmentation was formed when labeling burden was a major bottleneck to building neural networks. With recent progress in label-efficient training, we should reconsider conventions. Ultimately, we show that the choice of training a classification vs. segmentation network not only impacts how the model can be used but changes the properties of the model itself, and we show that leveraging segmentation models can lead to higher quality classifiers in common settings. These results indicate segmentation may be a more natural way to train neural networks to interpret and classify images, even though it's not the conventional way humans classify images.
Related work. Our study is motivated by past works that employ segmentation networks in classification problems and report various benefits of doing so [16, 17, 18, 19, 20, 21]. A full discussion of related work is included in Appendix A1.
2 Analysis
In this section we develop an understanding of why segmentation and classification networks may perform differently given the same dataset and overarching task (i.e., to find a class-of-interest). We cover task specification and notation in Section 2.1 then build up intuition for the benefits of segmentation in Section 2.2.
Figure 2: Visualization of the positive and negative class distributions that classification vs. segmentation networks aim to discriminate. In classification, supervision is applied at the image-level; in segmentation, supervision is applied at the pixel-level.
Task specification. We first specify the classification tasks we consider replacing with segmentation. We study classification tasks where the following properties are true: the class-of-interest can be localized in the foreground and does not occupy the entire image; the class-of-interest appears if and only if the classification label is positive; the background should not inform the classification label. Example tasks that meet these specifications are shown in Figure 1. Examples of tasks that don't meet these specifications include classifying the imaging modality or radiographic bone age, where global features of an image inform the class label and there is no clear segmentation target.
Notation. We consider a dataset of  $N$  images, class labels, and segmentation masks  $\mathcal{D} = \{(X^i,y^i,S^i)\}_{i = 1}^N$ . Each image  $X^{i}$  is made up of  $L$  pixels,  $\{X_1^i,X_2^i,\ldots ,X_L^i\}$ . For simplicity of notation, we assume binary classification, so the image-level classification label is  $y^{i}\in \{0,1\}$  where 1 is the positive label. The segmentation mask  $S^i$  contains  $L$  pixel-level labels  $\{S_1^i,S_2^i,\dots ,S_L^i\}$  indicating which pixels contain the class-of-interest. The image-level class label and pixel-level segmentation mask are related by  $y^{i} = \mathbb{1}\left\{\sum_{k = 1}^{L}S_{k}^{i}\geq 1\right\}$ . Put simply, this means if the class label is positive, there is at least one pixel labeled as positive in the segmentation mask and vice versa. We train a segmentation network  $f(\cdot)$  to produce a predicted segmentation mask  $f(\pmb {X}^i) = \hat{\pmb{S}}^i$  by minimizing the cross entropy loss  $CE(\hat{S}^i,S^i)$ . For notation's sake, we will drop the sample index  $i$ , thus referring to the  $i^{th}$  image  $X^i$  simply as  $\pmb{X}$  and the  $k^{th}$  pixel  $X_{k}^{i}$  simply as  $X_{k}$ .
Data generating process. To study segmentation vs. classification, we take an information theoretic approach and use the following simplified data generating process. First, the pixel-level segmentation labels  $\{S_1,\ldots ,S_L\}$  are drawn from a joint distribution  $P_{S}$ . Then, each pixel  $X_{k}$  is drawn from the distribution  $P_{X}(\cdot |S_{k})$ . As above, the image-level classification label  $y$  can be defined in terms of the segmentation mask,  $y = \mathbb{1}\left\{\sum_{k = 1}^{L}S_{k}\geq 1\right\}$ .
2.2 Analyzing segmentation vs. classification
To understand why segmentation and classification may perform differently on the same dataset, we consider the distributions that classification vs. segmentation networks are trained to discriminate. As an example, we'll use chest x-ray pneumothorax classification. In classification, supervision is applied at the image-level—the network seeks to discriminate images with a pneumothorax vs. images without a pneumothorax (Figure 2, left). In segmentation, supervision is applied at the pixel-level—the network seeks to identify pixels that contain a pneumothorax (Figure 2, right). Looking at these distributions, differences emerge.
Sample complexity. First, we see that segmentation has more annotated input-output pairs, as each pixel is annotated. Moreover, we see that classification's positive and negative classes share more features than segmentation's—implying segmentation may have greater class separability. We capture this notion of separability using the Kullback-Leibler (KL) divergence and show the following:
Proposition 1 It holds that
$$
D _ {K L} \left(\Pr \left(X _ {k} \mid S _ {k} = 0\right) \mid \mid \Pr \left(X _ {k} \mid S _ {k} = 1\right)\right) \geq \frac {\Pr (y = 1)}{\Pr \left(S _ {k} = 1\right)} D _ {K L} \left(\Pr \left(X _ {k} \mid y = 0\right) \mid \mid \Pr \left(X _ {k} \mid y = 1\right)\right). \tag {1}
$$
Classification workflow
Segmentation-for-classification workflow
Figure 3: Diagram showing how a chest x-ray disease classification problem can be framed using either standard classification or segmentation-for-classification. In the traditional classification workflow, the input image is processed by a classification network, which outputs a classification label vector. In the segmentation-for-classification workflow, an input image is processed by a segmentation network, which outputs binary masks showing where each abnormality is found. These masks are then converted into a classification label vector via a summarizing function,  $g(\cdot)$ .
In other words, the KL divergence between data given pixel-level labels (as is done in segmentation) is greater than or equal to the scaled KL divergence given image-level labels (as is done in classification). We note that the scaling factor  $\frac{\operatorname*{Pr}(y = 1)}{\operatorname*{Pr}(S_k = 1)}$  is always greater than or equal to one. We provide the proof for Proposition 1 in Appendix A3.1, confirming our intuition that segmentation supervision results in more separable data distributions. Because of the greater quantity of input-output pairs and the greater KL divergence, we expect segmentation's discrimination function to be easier to learn, requiring fewer images or a simpler function class [22, 23]. Since we study the setting in which the segmentation mask directly indicates the classification label ( $y = \mathbb{1}\left\{\sum_{k=1}^{L} S_k \geq 1\right\}$ ), we similarly expect segmentation-for-classification to have higher performance than standard classification in the limited data regime. Finally, from Proposition 1, we expect segmentation-for-classification to impart more benefit for tasks with small targets; for large targets, the terms in Proposition 1 approach one another.
Robustness to spurious correlations. The second observation we make from Figure 2 is that, because background features appear in classification's positive class, it is possible for the classification network to rely on background features to identify positive samples if those background features spuriously correlate with the class label. Such failure modes have been observed in prior work. For example, classifiers relying on metal L/R tokens to identify pneumonia [24]; classifiers looking for chest drains to classify pneumothorax [25]; and classifiers looking to image edges (which can change with patient positioning and radiographic projection) to classify COVID-19 [26]. Segmentation should be more robust to such background features, since background features are less correlated with the target class. This notion has been formalized in past work [21], which we recall here.
Specifically, for spurious features that do not overlap with the segmentation target, the mutual information between positive features and spurious features decreases as we transition from supervising with image-level labels to pixel-level labels. Since the spurious feature is less informative of the pixel-level labels than the classification label, we expect segmentation to be more robust to the presence or absence of spuriously correlated background features.
To use segmentation for classification tasks, we need to determine how to use segmentation networks to obtain classification labels. We'll do this by defining a summarizing function,  $g(\cdot)$ , which takes as input the image and segmentation network and outputs a class label  $g(\mathbf{X}, f(\cdot)) = \hat{y}$ . The segmentation-for-classification workflow is shown in Figure 3. In this section, we describe the summarizing functions that we evaluate. We note these methods can be used for both binary class problems and multiclass problems, as long as the classes of interest follow the task specification set above. We provide multiclass results in Appendix A5.6.
3.1 Rule-based summarizing functions
We first consider a simple, rule-based summarizing function. This method is intuitive: if there is a positive region in the predicted segmentation mask, we infer a positive classification label.
CANDID
Figure 4: Visualization of the datasets used in this manuscript: the synthetic dataset, where the task is to classify if a navy blue circle is preset; CANDID, a pneumothorax classification task; ISIC, a lesion classification task; and SPINE, a cervical fracture classification task. For each dataset, we show a positive and negative image and their corresponding segmentation masks.
ISIC
SPINE
Specifically, at inference, we binarize the probabilistic mask output from the segmentation network using threshold  $t$ . If the binarized mask contains over threshold  $\tau$  total positive pixels, we return a positive classification label. To compute the class probability, we average the pixel-wise probabilities of all pixels in the class. Algorithms for this method are provided in Appendix A2.1.
3.2 Trained summarizing functions
Next we consider summarizing functions trained to transform the segmentation information into classification labels. The intuition behind these methods is that trained functions may learn more nuanced differences between true positive and true negative predicted segmentation masks. We train these summarizing functions using the same splits used to train the segmentation network, and freeze the segmentation networks before training the summarizing functions.
We first consider trained summarizing functions that operate on the predicted segmentation masks. We consider three architectures for the summarizing function, each with increasing complexity: a fully connected layer, a global average pooling layer followed by a fully connected layer, and existing image classification architectures such as SqueezeNet [27] and ResNet50 [28].
We also consider summarizing functions that operate on segmentation network embeddings, similar to traditional transfer learning. We consider both shallow and deep embeddings and two different summarizing function architectures: a simple architecture consisting of pooling and fully connected layers, and a more complex classification head proposed in past work [20].
We fully describe these summarizing functions and training processes in Appendix A2.2.
In this section, we first confirm the takeaways from our analysis in Section 2 on a synthetic dataset (Section 4.1). Then, we evaluate the summarizing functions we described in Section 3 on medical imaging datasets (Section 4.2). We briefly describe the datasets and training below; we provide full details of the synthetic dataset and training in Appendix A4 and the medical datasets and training in Appendix A5.
Datasets. The synthetic dataset consists of many colored shapes on a gray background; the task is to determine if a navy blue circle is present. We also evaluate three medical imaging datasets: CANDID, in which we aim to classify pneumothorax in chest x-rays (n=19,237) [29]; ISIC, in which we aim to classify melanoma from skin lesion photographs (n=2,750) [30]; and SPINE, in which we aim to classify cervical spine fractures in CT scans (n=2,018, RSNA 2022 Cervical Spine Fracture Detection Challenge). These datasets are visualized in Figure 4. The SPINE dataset is an example of where the growing number of off-the-shelf networks enables segmentation-for-classification: we use a pretrained CT segmentation network to generate segmentation targets for this task, as detailed in Appendix A5.
Training. To establish baseline classification performance, we train a standard supervised classification network. To establish baseline segmentation+classification performance, we train a multitask learning network where the classification and segmentation networks are trained concurrently and
Figure 5: Results from experiments on the synthetic datasets. Each panel shows classification vs. segmentation-for-classification performance, as measured by the balanced accuracy, on the y-axis as a characteristic of the training dataset is changed on the x-axis.
share the same network backbone using the Y-Net structure [31]. To obtain segmentation-for-classification performance, we implement each summarizing function as described in Section 3. We train classification, multitask, and segmentation networks by minimizing the cross entropy loss between the predicted and ground truth labels using the Adam optimizer. We provide additional training details and information on hyperparameter tuning in Appendix A4.2 and A5.2, and we provide ablation experiments controlling for segmentation vs. classification model capacity in Appendix A5.3.
4.1 Experiments on synthetic data
We first perform experiments on synthetic data to confirm the takeaways from our analysis. By using a synthetic dataset, we have precise control over dataset and task characteristics, allowing us to isolate and adjust these variables to assess model performance. We provide highlights of our synthetic experiments below; a full description of the experimental procedure on synthetic data can be found in Appendix A4. In these synthetic experiments, we use the rule-based summarizing function.
The first takeaway from Section 2 is that segmentation-for-classification should require fewer images to achieve high performance, particularly for tasks with small targets. We examine this with our synthetic dataset in the following experiments.
- We sweep the number of training images from 1,000 to 100,000 and plot classification and segmentation-for-classification balanced accuracy in Figure 5a, where we see segmentation-for-classification has improved performance in the limited data regime.
- Similar to learning with small datasets, we also expect segmentation to perform better with high or low class prevalence, which we see in Figure 5b and Figure 5c.
- Segmentation's ability to learn with fewer examples should also extend to rare subtypes—we expect higher performance on class subtypes that appear infrequently during training. To evaluate this, we train on small, medium, and large navy blue circles and evaluate network performance on the small circles as they are made increasingly rare during training (Figure 5d). We see segmentation-for-classification achieve higher performance on this data slice with few training examples.
- We vary the object size in the train and test dataset and find segmentation-for-classification provides more boost with smaller targets (Figure 5e).
The second takeaway from our analysis is that segmentation-for-classification should be more robust to background features that spuriously correlate with the class of interest. We evaluate this with two synthetic experiments.
- We first evaluate an obvious case of spurious correlation: when a background object correlates with the class label. We make a large, pink square increasingly correlated with
Figure 6: Comparison of different segmentation-for-classification methods on the CANDID, ISIC, and SPINE datasets in both the limited and abundant training data settings. As baselines, we include classification and multitask network performance. FC: fully connected.
the target (a small navy blue circle) in our synthetic dataset—then remove this correlation in the test dataset—and show that segmentation is robust to this correlation (Figure 5f).
- Spurious background correlations can also include less obvious features, like where the target object occurs; in other words, we expect segmentation to be more robust to target location. We explore this by changing where the target navy blue circle appears during training, but allowing the navy blue circle to appear anywhere in the test dataset. As the target is increasingly restricted to a particular location during training, classification is less able to generalize to finding the target in other regions of the image (Figure 5g).
4.2 Experiments on medical imaging data
Given the promising results on synthetic data, we next evaluate performance on the medical datasets and evaluate segmentation-for-classification training methods. In Section 4.2.1 we compare the performance of different summarizing functions, classification networks, and multitask networks. In Section 4.2.2 we evaluate semi-supervised segmentation-for-classification.
4.2.1 Comparing summarizing functions, classification, and multitask training
We evaluate each summarizing function and model on the three medical datasets in two regimes: limited labeled data, which consists of  $10\%$  of randomly sampled training data, and all labeled training data. We report mean AUROC on the test set in Figure 6. We observe that the simple, rule-based segmentation-for-classification method achieves higher mean AUROC than both the traditional classification and multitask networks for all tasks and amounts of training data. Specifically, we observe a  $10.0\%$ ,  $6.9\%$ , and  $2.8\%$  improvement in mean test AUROC in the limited training data regime for CANDID, ISIC, and SPINE respectively over traditional classification, and a  $2.6\%$ ,  $0.6\%$ , and  $6.1\%$  performance improvement in the all training data setting.