SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Comparing summarizing functions, we find that summarizing functions which operate over segmentation network outputs generally outperform those that operate over embeddings. Among summarizing functions that operate over segmentation outputs, three perform similarly well: thresholding the segmentation output, learning a fully connected layer, and training a full classification network on top of the segmentation output. The threshold-based summarizing function has additional advantages: it does not require training, thus saving training effort and avoiding potential overfitting to training data; it is fast to compute at inference; and the results are intuitive—when there is a segmentation mask, there is an associated positive class label. For these reasons, we recommend simple rule-based summarizing functions for converting a segmentation mask into a classification label. We compare the error modes of these methods in Section 5.
Figure 7: Semi-supervised segmentation-for-classification. As the x-axis increases, all models are trained with additional segmentation labels; the semi-supervised model is also trained with unlabeled data, and the boosted semi-supervised model is also trained with data that only has classification labels. Each line plots the average AUROC of the test set with shaded  $95\%$  confidence intervals.
4.2.2 Semi-supervised segmentation-for-classification
As discussed in the Introduction, a historical obstacle to widespread application of segmentation is the difficulty of curating labeled training datasets. However, recent work has led to more public datasets, pretrained networks, and methods to reduce the labeling burden of training segmentation models. We posit that this progress in segmentation will also benefit segmentation-for-classification. In the previous section, we showed that an off-the-shelf spine segmentation model could be used to generate segmentation labels for training the SPINE segmentation-for-classification models. In this section, we verify that a semi-supervised segmentation method can improve performance of segmentation-for-classification models. We consider two common settings:
- Setting 1, additional unlabeled data. Because unlabeled data is often abundant, practitioners may have a subset of their dataset labeled with segmentation masks and a larger set of unlabeled images.
- Setting 2, additional classification data. Other practitioners may have a subset of their dataset labeled with segmentation masks and a larger set of images with classification labels.
For Setting 1, we apply a previously-published semi-supervised segmentation method which takes a small set of labeled data and a large set of unlabeled data as input and trains a segmentation model using a combination of data augmentation, consistency regularization, and pseudo labeling [32].
For Setting 2, we update the previously-proposed methods by using the classification labels to improve the pseudo segmentation labels, which we call "boosted" semi-supervision. The original semi-supervised approach uses pseudo labeling, in which a small amount of labeled segmentation data is used to train an initial segmentation network; the initial network predicts pseudo segmentation masks for all unlabeled data; and a final network is trained on the pseudo segmentation masks. In the "boosted" algorithm, we use the available classification labels to modify the pseudo segmentation masks. Specifically, for images that have negative classification labels, we zero out the pseudo segmentation masks: because we know the classification label is negative, we know that there should not be any positive region in the segmentation mask. For images that have a positive classification label, we check to see if the pseudo segmentation mask has a positive segmentation region; if it does not have a positive segmentation region, we know that the pseudo segmentation mask is wrong and we exclude the sample from training of the final segmentation network.
These semi-supervised approaches learn from the unlabeled data to train a better segmentation model, which we hypothesize will also improve segmentation-for-classification performance. In these experiments, we use the threshold-based summarizing function and train models using a balanced training dataset. We report model performance for varying amounts of labeled training data in Figure 7. We report additional performance metrics and the number of labels for each method at the most limited and abundant data settings in Table A8.
From these results, we observe that all of the segmentation-for-classification models trained with  $10\%$  of the training data labeled exceed the performance of the classification model trained with  $100\%$  of the training data labeled. These results confirm one of the takeaways from our analysis, which is that segmentation-for-classification can achieve higher performance with fewer images. Additionally, we see that the semi-supervised training further improves segmentation-for-classification performance,
confirming that semi-supervised segmentation methods can directly be used for segmentation-for-classification. Finally, we see that the boosted semi-supervised approach indeed improves performance beyond standard semi-supervision in the limited data regime.
5 Benefits and drawbacks of segmentation-for-classification
In this section, we tie together our analysis, synthetics, and experiments on real data to summarize the benefits and drawbacks of using segmentation-for-classification.
Segmentation can improve aggregate performance, particularly with small datasets. Using segmentation-for-classification can lead to higher-performing models compared to traditional classification. Our analysis in Section 2 shows this benefit is due in part to greater divergence between segmentation's positive and negative classes and higher quantity of annotation (we evaluate the contributions of each of these factors with an additional experiment in Appendix A4.3). We observe this takeaway empirically in Figure 6, showing segmentation-for-classification improves mean performance on all datasets and training settings. Our analysis suggests greater performance differences occur in the limited data regime, which can include:
- Small datasets. We expect greater boosts from segmentation-for-classification with limited dataset sizes. We observe this with synthetics in Figure 5a and on two of the three medical datasets in Figure 6. Small datasets often occur when developing models on custom data or when training models for rare diseases.
- Low class prevalence. Segmentation can also achieve higher performance than classification when the class balance is very low or high (shown with synthetic datasets in Figure 5b,c), since segmentation is able to better learn from the few examples of the low-prevalence class.
- Rare subtypes. We expect segmentation-for-classification models to have improved ability to identify rare subtypes, which appear infrequently in the training dataset. We observe this on the synthetics (Figure 5d) and see this property revealed in the medical datasets as well: on small lesions that are positive for melanoma, which are less commonly seen during training, traditional classification achieves a recall of 0.18 and segmentation-for-classification achieves a recall of 0.47.
Finally, the analysis in Section 2 suggests segmentation provides greater benefit for tasks with small targets. We observe this with synthetics in Figure 5e. On the real data experiments, models trained with the full datasets reflect this property as well: the SPINE dataset has the lowest average target-to-background size ratio (with the target taking up an average  $0.6\%$  of the image) and achieves the greatest performance boost with segmentation-for-classification, while the ISIC dataset has the highest target-to-background size ratio (with the target taking up an average  $26.3\%$  of the image) and sees the lowest performance boost.
Segmentation reduces susceptibility to spurious correlations. A worrying failure mode of classification networks is that instead of learning to identify the pathology of interest, the model instead looks for easier-to-identify features that spuriously correlate with the true target—such as the examples given in Section 2. Segmentation-for-classification should be more robust to background features spuriously correlated with the target task.
Empirically, we observe this in our synthetic data (Figure 5f) and evaluate a spurious correlation on the medical datasets here. Specifically, in the CANDID dataset, we evaluate performance under a known spurious correlation condition: most pneumothoraces in chest x-ray datasets co-occur with a chest tube, which are used to treat pneumothorax [25]. Thus, it is important that the pneumothorax classification algorithms are able to find pneumothoraces without chest tubes, as those are the patients who have not yet been identified and treated. To evaluate each model's robustness to this spurious correlation, we evaluate model performance on patients who do not exhibit the spurious correlation (i.e., patients with pneumothorax but no chest tube or patients without pneumothorax but do have a chest tube) and observe that segmentation-for-classification achieves an AUROC of 0.84 while the classifier's AUROC drops to 0.58.
We also expect segmentation-for-classification to be more robust to the target location, which is another type of spurious correlation. We observe this in the synthetics (Figure 5g) and on the medical data: recall of the CANDID classifier drops up to  $22.5\%$  conditioned on the lung region in which the pneumothorax occurred (dropping from 0.92 recall in the right lower hemithorax to 0.71 recall on the
left upper hemithorax) while segmentation-for-classification recall drops at most  $7.5\%$  (from 0.94 in the left lower hemithorax to 0.87 in the right upper hemithorax).
Segmentation facilitates human assessment. Segmentation-for-classification inherently delivers location information. While there are methods to probe classification models for which areas of an image contribute to the model's prediction (e.g., saliency maps [33], class activation maps [34, 35]) these methods are often unreliable; shortcomings are further discussed in prior work [36, 37]. By providing location information as well as classification information, segmentation-for-classification models enable the user to more effectively adjudicate model findings. This location information is more pertinent for some applications where the abnormality may be hard to find (e.g., pneumothorax, fracture) versus use cases where the visual features used by the model are obvious (e.g., skin lesion).
Segmentation enables more precise model evaluation. We can perform more detailed model evaluation and error analysis using segmentation-for-classification's location and size information. For example, we can evaluate models with a stricter definition of recall: instead of identifying a true positive as when the image-level label is positive and correct, we can also require that the correct region of the image is identified. When we compute the stricter recall metric for models trained with all labeled data, we see reductions in recall: CANDID drops from 0.87 to 0.85 and SPINE drops from 0.74 to 0.55. We can't similarly assess the classification algorithms with this stricter definition of recall because of the lack of reliable methods to acquire accurate location information from classifiers. However, precise model evaluation is important for low-failure-tolerance applications such as medicine, and segmentation-for-classification provides an obvious means of doing so. Finally, we can perform additional error analysis using segmentation-for-classification's size and location information, such as assessing the average abnormality size and region.
Segmentation has a higher per-image labeling cost. The burden of generating segmentation training labels is a drawback of training segmentation networks and can change for different tasks (e.g., 3D images have a higher labeling burden). In comparison, classification is cheap to generate labels for—particularly for cases where classification labels can be pulled from electronic health records [38]. In this work, we've explored two ways of reducing segmentation's labeling burden while improving classification performance: using an off-the-shelf segmentation model and using semi-supervised training methods. As new tools emerge and make it easier to generate segmentation training labels [13, 14, 15], we expect the labeling burden gap between classification and segmentation to lessen. Still, the increased annotation cost of segmentation given available tools should be weighed against the expected benefits (detailed above) of using segmentation-for-classification on a given dataset.
Historically, segmentation networks have been employed only when required due to their high annotation cost. Given the increasing ease of training segmentation networks, we aim to better flesh out the decision space between training a segmentation vs. classification network. In Section 2, we develop intuition and formally analyze why segmentation and classification networks can perform differently. In Section 3 and 4, we explore methods to implement segmentation-for-classification and empirically evaluate segmentation-for-classification over many datasets and training conditions. Finally, in Section 5, we discuss the benefits and drawbacks of segmentation-for-classification. We show that segmentation-for-classification can lead to more performant, robust, and interpretable models. In clinical settings, one may prefer more these upsides despite the added annotation cost.
Broadly, these results underline the potential of deep learning-based segmentation as a general tool for automating image analysis. Segmentation data provides an information-rich environment for deep neural networks to learn from, which we've shown results in separable, robust embedding spaces. In turn, segmentation network outputs capture detailed information about images, which can support many downstream analysis tasks.
In future work, we are interested in how to leverage segmentation for targets that don't fit the task specification we study here, perhaps utilizing anatomical or coarser forms of segmentation to still reap some of the benefits of training segmentation models. Additionally, we are interested in reader studies focused on quantifying the value of the location information contained in segmentation masks for radiologists working with AI algorithms. Finally, we are interested in how to leverage segmentation to benefit other machine learning applications in radiology, extending our study of how task framing impacts a model's properties.
We thank Dr. Leslie Smith and Dr. Darwin Yi for their discussion and ideas related to this topic.
Sarah Hooper is supported by the Fannie and John Hertz Foundation, the National Science Foundation Graduate Research Fellowship under Grant No. DGE-1656518, and as a Texas Instruments Fellow under the Stanford Graduate Fellowship in Science and Engineering. Additionally, we gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under No. W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare. Figures in this paper were created with BioRender. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
[1] Phillip M Cheng, Emmanuel Montagnon, Rikiya Yamashita, Ian Pan, Alexandre Cadrin-Chênevert, Francisco Perdigón Romero, Gabriel Chartrand, Samuel Kadoury, and An Tang. Deep learning: an update for radiologists. Radiographics, 41(5):1427-1445, 2021.
[2] Simeng Zhu, Marissa Gilbert, Indrin Chetty, and Farzan Siddiqui. The 2021 landscape of fda-approved artificial intelligence/machine learning-enabled medical devices: an analysis of the characteristics and intended use. International journal of medical informatics, 165:104828, 2022.
[3] S Niyas, SJ Pawan, M Anand Kumar, and Jeny Rajan. Medical image segmentation with 3d convolutional neural networks: A survey. Neurocomputing, 493:397-413, 2022.
[4] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
[5] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9650-9660, 2021.
[6] Mark Hamilton, Zhoutong Zhang, Bharath Hariharan, Noah Snavely, and William T Freeman. Unsupervised semantic segmentation by distilling feature correspondences. arXiv preprint arXiv:2203.08414, 2022.
[7] Shanghua Gao, Zhong-Yu Li, Ming-Hsuan Yang, Ming-Ming Cheng, Junwei Han, and Philip Torr. Large-scale unsupervised semantic segmentation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[8] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselet, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
[9] Alexander J Ratner, Christopher M De Sa, Sen Wu, Daniel Selsam, and Christopher Ré. Data programming: Creating large training sets, quickly. Advances in neural information processing systems, 29, 2016.
[10] Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E Hinton. Big self-supervised models are strong semi-supervised learners. Advances in neural information processing systems, 33:22243-22255, 2020.
[57] Nimit Sohoni, Jared Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. Advances in Neural Information Processing Systems, 33:19339-19352, 2020.
[58] Evan Z Liu, Behzad Haghloo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pages 6781-6792. PMLR, 2021.
[59] Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein. nnu-net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods, 18(2):203–211, 2021.
[60] William Falcon et al. Pytorch lightning. https://github.com/PyTorchLightning/pytorch-lightning, 3, 2019.
[61] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[62] Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3431-3440, 2015.
[63] Gautam Machiraju, Sylvia Plevritis, and Parag Mallick. A dataset generation framework for evaluating megapixel image classifiers and their explanations. In European Conference on Computer Vision, pages 422-442. Springer, 2022.
A1 Related work
Segmentation to benefit classification. Our study is inspired by past works that employ segmentation networks in classification problems and report various benefits of doing so. Improved overall accuracy. For example, [16] uses a segmentation network to generate segmentation masks of different tissue types in lung ultrasound images, then trains a sequence of global average pooling and fully connected layers to convert the tissue masks into classification labels. The authors report higher aggregate performance compared to training a classification model from scratch for their lung ultrasound classification task. In [17], the authors expand this approach by including diagnostic tissue segmentation targets (i.e., segmenting benign vs. malignant tissue as different segmentation classes), then converting the segmentation masks into class labels via a likelihood score over the segmented tissue types. Similarly, [18] jointly segments and classifies liver lesions in CT scans by segmenting the lesion as one of multiple possible lesion types and finds that the segmentation approach results in higher classification performance than traditional classification or multitask learning. Training with fewer labels. Other reports focus on segmentation networks' ability to learn from limited labeled data. For instance, [19] converts segmentation outputs into classification labels via a likelihood score and compares the aggregate performance of the segmentation network against classification and multitask learning networks when trained with varying dataset sizes. The authors find that the segmentation network achieves higher classification scores in the limited data regime. Similarly, [20] uses features from a UNet pretrained on a segmentation task to fine tune a classifier and again finds higher classification performance using the pretrained segmentation weights in the limited data regime. Improved robustness. Finally, [21] shows that training a segmentation network instead of a classification network results in predictors more robust to spurious correlations.
Multitask learning. Other works have combined segmentation and classification via multitask learning, in which the classification and segmentation networks share the same backbone but different task heads trained concurrently [39, 40, 41, 31, 42]. The multitask approach differs from the approaches discussed above, which each train a segmentation network, freeze it, then use the segmentation network's embeddings or outputs to produce a classification label. In multitask frameworks, the classification task can impact the weights learned for the segmentation task via the shared backbone, and training schedules and task weighting must be managed.
Classification to benefit segmentation. Other studies focus on the reverse of what we are interested in by first training a classification network then using that classifier to generate segmentation information. These approaches typically aim either to generate segmentation labels for "free" by extracting location information from the classification models or to use this location information to improve a classification network's interpretability. Examples of this class of related work include using saliency maps [33], class activation maps [34, 35], occlusion maps [43], and classifiers with built-in attention mechanisms [44].