SlowGuess/NeurIPS_2023_OCR huggingface dataset:
A greater number of annotated samples should improve most ML tasks, as has been extensively studied before. We focus on characterising the divergence as it is specific to segmentation-for-classification and less studied in prior work. Future work may study the relative benefits of these two aspects of supervision.
A5 Real datasets and experiments
We run experiments on three public medical imaging datasets, described below and fully detailed in their original studies. The medical datasets are visualized in Figure 4.
A5.1 Medical imaging datasets
CANDID: classifying pneumothorax in chest x-rays [29]
- Classification task. Our aim is to classify whether a chest x-ray contains a pneumothorax; ground truth classification labels are provided with the original dataset.
- Segmentation target. Our aim is to segment pneumothorax when it appears in a chest x-ray; ground truth segmentation masks are provided with the original dataset.
- Dataset. This dataset contains 19,237 chest x-rays. The images are two-dimensional, one-channel images of size  $1024 \times 1024$ , and we normalized each image to the range [0, 1]. We split this dataset randomly into  $60\%$  training images,  $20\%$  validation images, and  $20\%$  test images. In this dataset,  $16.6\%$  of images are positive for pneumothorax.
ISIC: classifying melanoma from photographs of skin lesions [30]
- Classification task. Our aim is to classify whether a skin lesion is melanoma or not; ground truth classification labels are provided with the original dataset.
- Segmentation target. Our aim is to segment the skin lesion when it is melanoma; ground truth segmentation masks are provided with the original dataset.
- Dataset. This dataset contains 2750 images of skin lesions. The images are two-dimensional, three-channel images that we resized to  $224 \times 224$ . We use the splits provided by the ISIC challenge organizers, resulting in 2000 training images, 150 validation images, and 600 test images.  $18.9\%$  of images are positive for melanoma.
SPINE: classifying cervical spine fractures in CT scans (RSNA 2022 Cervical Spine Fracture Detection Challenge)
- Classification task. Our aim is to classify whether there is a fracture in the cervical spine; ground truth classification labels are provided with the original dataset.
- Segmentation target. Our aim is to segment the vertabrae in which the fracture occurs. Unlike the previous two datasets, the ground truth segmentation masks were not available in this dataset for all training images. Instead, we used an off-the-shelf CT segmentation network [11] to generate segmentation masks of the vertebra. For each positive classification label, we create a corresponding segmentation mask which delineates the vertebrea in which the fracture occurs. This use case shows how off-the-shelf models can be used to convert some classification problems into segmentation problems without any additional labeling.
- Dataset. This dataset contains 2018 CT scans of the spine. The images are three-dimensional volumes with one channel. We applied a CT bone window then normalized each scan to the range [0, 1]. We resized each axial slice to  $128 \times 128$  voxels using bilinear interpolation and resampled the axial dimension to 128 slices around the spine, resulting in  $128 \times 128 \times 128$  image volumes.  $47.6\%$  of these scans contain a cervical spine fracture. We randomly split this dataset into  $60\%$  training,  $20\%$  validation, and  $20\%$  test.
Model	Dataset	Architecture	Learning rate	Weight decay
Segmentation	CANDID	ResNet50, SEResNet50, DeepLabv3, UNet	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Segmentation	SPINE	ResNet50 3d, UNet 3d, NNuNet 3d	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Segmentation	ISIC	ResNet50, SEResNet50, DeepLabv3, UNet	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Classification	CANDID	ResNet50, SEResNet50, DenseNet121, ConvNext	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Classification	SPINE	ResNet50 3d, UNet 3d encoder, NNuNet 3d encoder	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Classification	ISIC	ResNet50, SEResNet50, DenseNet121, ConvNext	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Multitask	CANDID	ResNet50	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Multitask	SPINE	ResNet50 3d	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Multitask	ISIC	ResNet50	1e-6, 1e-5, 1e-4, 1e-3	1e-6, 1e-4, 1e-2, 0
Table A4: Hyperparameters tested for the backbone networks; chosen hyperparameters are in bold.
Summarizing function	CANDID (low data)	SPINE (low data)	ISIC (low data)	CANDID (high data)	SPINE (high data)	ISIC (high data)
Seg. output → fully connected layer	1e-5	1e-6	1e-3	1e-5	1e-6	1e-4
Seg. output → avg. pool + fully connected layer	1e-2	1e-2	1e-3	1e-2	1e-2	1e-4
Seg. output → SqueezeNet	1e-5	1e-5	1e-5	1e-5	1e-6	1e-6
Deep seg. embeddings → simple head	1e-3	1e-3	1e-3	1e-3	1e-2	1e-3
Shallow seg. embeddings → simple head	1e-3	1e-4	1e-4	1e-4	1e-2	1e-3
Deep seg. embeddings → complex head	1e-4	1e-2	1e-3	1e-3	1e-4	1e-2
Shallow seg. embeddings → complex head	1e-3	1e-3	1e-4	1e-3	1e-3	1e-2
Table A5: Hyperparameters selected for training the summarizing functions.
A5.2 Training details for medical imaging datasets
In Table A4 we list each hyperparameter we evaluated for the backbone networks trained on the medical datasets. We chose hyperparameters that maximized the validation AUROC. Chosen hyperparameters are in bold.
For the trained summarizing functions, we only tuned the learning rate. We evaluated learning rates [1e-6, 1e-5, 1e-4, 1e-3, 1e-2] for each summarizing function and chose the learning rate that maximized the validation AUROC. In Table A5 we list the chosen learning rate for each summarizing function and dataset.
To evaluate performance, we report mean AUROC over the test split of each dataset. We binarize the probabilistic labels to report additional performance metrics by setting the operating point on the ROC curve that maximizes the Youden Index. We compute  $95\%$  confidence intervals using bootstrapping (1000 iterations). We conduct our experiments using Pytorch Lightning (Pytorch version 1.9.0, Lightning version 1.5.10) [60, 61].
A5.3 Ablations on model capacity
To report results in the main text, we tested multiple architectures for each task and dataset and selected the best-performing architecture for both segmentation and classification (Table A4). We did this to give each task and dataset the best performance, as it's not clear the best architecture for segmentation on a given dataset is the same that is best for classification.
However, segmentation networks are often higher capacity than classification networks due to their dense, pixel-wise outputs and corresponding decoders. To show that it is not simply a difference in model capacity that is leading to the observed performance differences between segmentation and classification, we perform the following ablations.
First, we use the same backbone for both segmentation and classification (Table A6). We use Resnet50 as the 2D backbone as it is a common backbone known to be useful for classification and segmentation. The Resnet50 classification head is a linear layer, while the Resnet50 segmentation head is the head described in the Fully Convolutional Network paper [62]. We use the UNet from the nnU-Net paper [59] for the 3D backbone. We see that segmentation shows the same or greater performance enhancement than what appeared in the original paper when using these standardized backbones. We further note that all of our synthetic experiments were run with a standardized backbone between classification and segmentation.
However, even when using the same backbone, segmentation models still typically have more parameters due to the convolutional decoder head. We next provide results comparing performance using
Backbone	Method	AUROC
CANDID	ResNet50	Classification	0.80
Seg-for-class	0.91
ISIC	ResNet50	Classification	0.56
Seg-for-class	0.67
SPINE	nnU-Net	Classification	0.60
Seg-for-class	0.62
Table A6: Performance comparison when classification and segmentation models are trained with the same network backbone. For the SPINE task, the nnU-Net encoder is used for the classification task.
Backbone	Method	AUROC
CANDID	ResNet101	Classification	0.82
ResNet50	Seg-for-class	0.91
ISIC	ResNet101	Classification	0.59
ResNet50	Seg-for-class	0.67
Table A7: Performance comparison when classification has a higher-capacity model (i.e., contains more parameters) compared to segmentation.
a higher capacity classification model—Resnet101, which has more parameters than the Resnet50 segmentation model—in Table A7. Again, we see improved performance with segmentation.
Finally, we emphasize that we matched training procedure for classification and segmentation networks, including matching the input data, augmentations, and codebase (which standardizes model checkpointing, loss function, hyperparameter tuning, etc.). Together, these results show it is not just model capacity or training procedure that lead to segmentation-for-classification's improved performance.
A5.4 Additional performance metrics
In Table A8 we report additional performance metrics using standard classification and segmentation-for-classification (including semi-supervised methods) on the CANDID dataset.
A5.5 Results on a natural image dataset
To show the proposed method extends to a natural image dataset, we include results classifying dog and cat breeds in the Oxford Pets dataset in the limited data regime (50 training images per class). We trained these models using ResNet50 architectures and the rule-based summarizing function. We observe segmentation-for-classification achieves an AUROC of 0.96 while standard classification achieves an AUROC of 0.81. These results show segmentation-for-classification improves average AUROC by  $18.5\%$  on a natural image dataset, mirroring our results on medical imaging datasets.
A5.6 Results on multiclass datasets
While we focus on single-class tasks in the main text, we show here that segmentation-for-classification can directly be used for multiclass tasks. Instead of having a binary mask output from the segmentation network, a user simply needs to specify multiple output channels. First, we note the Oxford Pets dataset contains 37 classes of dog and cat breeds; as shown in the previous subsection, we see segmentation outperforms classification by  $18.5\%$  on this multiclass dataset. Further, we extend one of our medical imaging datasets to the multiclass setting and perform 3-class classification on ISIC, classifying lesions as benign nevi, melanoma, or seborrheic keratosis. Again, we see improved performance with segmentation-for-classification, which achieves a balanced AUROC of 0.61 compared to classification's 0.50 in the limited data regime.
# images with seg and class labels	# images with only class labels	# images with no labels	AUROC	Bal. Accuracy	Recall	Precision	Specificity
Classification (limited)	0	198	0	0.74(0.72 - 0.76)	0.68(0.66 - 0.70)	0.60(0.56 - 0.64)	0.31(0.29 - 0.34)	0.76(0.74 - 0.77)
Seg-for-class (limited)	198	0	0	0.86(0.84 - 0.88)	0.78(0.76 - 0.80)	0.75(0.71 - 0.78)	0.42(0.39 - 0.45)	0.81(0.80 - 0.82)
Seg-for-class, semi-sup	198	0	3780	0.89(0.87 - 0.90)	0.81(0.79 - 0.83)	0.75(0.71 - 0.78)	0.51(0.48 - 0.55)	0.87(0.86 - 0.88)
Seg-for-class, boosted semi-sup	198	3780	0	0.90(0.89 - 0.92)	0.83(0.81 - 0.84)	0.73(0.70 - 0.77)	0.62(0.59 - 0.66)	0.92(0.91 - 0.93)
Classification (abundant)	0	3978	0	0.90(0.89 - 0.92)	0.82(0.81 - 0.84)	0.80(0.77 - 0.83)	0.49(0.46 - 0.52)	0.85(0.84 - 0.86)
Seg-for-class (abundant)	3978	0	0	0.96(0.96 - 0.97)	0.90(0.89 - 0.92)	0.89(0.86 - 0.92)	0.67(0.63 - 0.70)	0.92(0.91 - 0.93)
Table A8: Comparison of classification vs. semi-supervised segmentation-for-classification performance on the CANDID dataset in the most limited and abundant training data conditions. These models were trained with a balanced dataset. We report the average of each metric over the test set with 95% confidence intervals in parentheses.
A6 Limitations
A limitation of this work includes the maximum sizes of the datasets. With much larger datasets, we expect overall classification performance to improve. However, the dataset sizes we use in this manuscript (2,000–20,000 for medical datasets; 100,000 for synthetic dataset) are not uncommon medical imaging dataset sizes, particularly for new applications or datasets specific to an institution. Further, we expect many benefits to hold in the very large data regime (e.g., robustness to spurious correlations, performance on very rare subsets, human assessment, and model evaluation). As an additional limitation, the segmentation-for-classification methods explored here are most applicable to abnormalities that can be segmented. Images for which even experts do not know which regions directly correspond to the class labels (e.g., some histopathology images [63]) have less obvious segmentation targets. In these situations, classification in combination with methods like GradCAM can help with feature discovery, for which segmentation-for-classification may be less well-suited. Similarly, classification labels which are more global in nature (e.g., age prediction from radiographs, atrophy of the brain in a head CT) do not have clear segmentation targets and may not benefit as much from the segmentation signal.
A7 Broader impacts
This work focuses on automated medical image analysis. Automated image analysis algorithms have the potential to support radiologists by providing decision support, improving analysis efficiency, and improving analysis reproducibility. The work presented here helps to make progress towards these benefits by improving the performance of automated medical image analysis algorithms; improving the quality of model evaluation; and providing location information to the radiologist to facilitate human assessment. However, medical image analysis is a low failure tolerance application. The work presented here is meant as a contribution to shift the way we consider segmentation vs. classification algorithms in building medical image analysis tools; the specific neural networks trained here are not meant for clinical use.