SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Semantic segmentation. A great deal of work has focused on improving semantic segmentation methods [45, 46], which underlie the segmentation-for-classification methods we explore here. Segmentation-for-classification frameworks contain an additional step converting a semantic segmentation network or output into an image-level classification label, particularly for tasks which are not typically framed as segmentation tasks. Additionally, many recent works have advanced our ability to train semantic segmentation networks with limited labeled data, including using self-supervised networks, semi-supervision, and foundation models [4, 5, 6, 7, 47, 8], many of which have been adapted specifically to medical imaging [48, 49, 50]. These methods make training segmentation-for-classification models less labor-intensive.
Transfer learning. The segmentation-for-classification framing can be viewed as a type of transfer learning. While many transfer learning studies have focused on settings where the task type stays consistent (e.g., transferring one classification network to a different classification task), some studies have looked at transferring classification weights to segmentation tasks or vice versa [51, 52, 53]. Some segmentation-for-classification methods can be considered transfer learning, where
Algorithm 1 Compute binary label
Require: ProbSegMask  $\in [0,1]^{H\times W\times D},t\in [0,1],\tau \in \mathbb{Z}^*$  1:BinSegMask  $\leftarrow$  ProbSegMask  $\geq t$  2:if sum(BinSegMask)  $\geq \tau$  then 3:return 1 4:else 5:return 0 6:end if
the segmentation network is trained first and much of the segmentation network is frozen before transferring the learned weights to a classification task.
Label granularity. A related line of work explores the right label granularity for training classification models (e.g., training "herbivore vs. carnivore" compared to training "horse vs. deer vs. cat vs. dog vs...") [54, 55]. Similar to our findings, these works find that more granular labels can lead to improved network optimization and generalization [54].
Modified classification networks. Finally, many works have modified straightforward supervised classification training schemes. For example, self- and semi-supervised learning aim to improve performance with limited labeled data [56], and past work has proposed various modifications to classification training to improve robustness [57, 58].
In this manuscript, we use straightforward classification as we are interested in exploring fundamental differences between classification and segmentation. We observed that many performance benefits can be achieved simply by using a segmentation model, improving performance along many axes without specialized tooling. Conversely, classification modifications are typically different for different objectives (e.g., methods for reducing reliance on spurious correlations are different from methods that handle class imbalance are different from methods that improve interpretability). We believe the fact that segmentation results in many benefits without specialized tooling for each is a major strength of the proposed approach.
Our contributions. We first analyze why we see performance differences, supporting and helping to explain the empirical results that we and others have observed as well as the conditions under which we can expect to see benefits. Next, we implement both previously proposed and new methods for using segmentation-for-classification and compare these approaches across multiple datasets using varying amounts of training data. We perform this head-to-head comparison to build up a recommendation of best practices for using segmentation networks in classification problems, including: recommending a summarizing function; showing an off-the-shelf segmentor can achieve improved classification performance; showing an existing semi-supervised training method can improve performance; and adapting that semi-supervised method specifically to the segmentation-for-classification setting. Finally, we aggregate and expand on the benefits of switching from classification to segmentation. We expand on benefits by showing improved performance with low class prevalence, rare subtypes, and location spurious correlations, and we discuss these quantitative benefits in conjunction with qualitative benefits like improved interpretability and model evaluation. We use this consolidated evidence to make the broader case that many classification problems would benefit from a segmentation framing.
A2 Summarizing functions
Below we provide additional details on the summarizing functions we evaluate in the main text.
A2.1 Rule-based summarizing functions
In Section 3.1, we introduce a rule-based summarizing function using thresholds to transform the segmentation output into a classification label. Here we give further details of that summarizing function. Specifically, in Algorithm 1, we give the algorithm we use to compute a binary, image-level label from a probabilistic segmentation mask. Given a probabilistic segmentation mask—which
Algorithm 2 Compute probabilistic label
Require: ProbSegMask  $\in [0,1]^{H\times W\times D},t\in [0,1],$  BinaryLabel  $\in \{0,1\}$  1: if BinaryLabel  $= 1$  then   
2: PosProbs  $\leftarrow$  ProbSegMask[ProbSegMask  $\geq t]$    
3: return mean(PosProbs)   
4: else   
5: NegProbs  $\leftarrow$  ProbSegMask[ProbSegMask  $<  t]$    
6: return mean(NegProbs)   
7: end if
contains the segmentation network output, the pixelwise probabilities indicating where the class of interest was found—we first binarize the segmentation mask using threshold  $t$ . At default,  $t = 0.5$ . We then sum all pixel values in the binary segmentation mask; if there are more than threshold  $\tau$  positive pixels in the segmentation mask, we return a positive binary image-level label. In this work, we set  $\tau = 100$ .
In Algorithm 2 we provide the method we use to compute a probabilistic image-level label from a probabilistic segmentation mask. We use these probabilistic labels to compute the AUROC. Given a probabilistic segmentation mask, we first compute the binary label using Algorithm 1. If the binary label is positive, we compute the probabilistic image-level label as the mean of pixel values that are greater than threshold  $t$ . Similarly, if the binary label is negative, the probabilistic image-level label is the mean of pixel values less than  $t$ . At default,  $t$  takes value 0.5.
A2.2 Trained summarizing functions
In Section 3.2, we introduce trained summarizing functions to transform the segmentation output into a classification label. We give further details of those summarizing functions here.
Summarizing functions that operate over segmentation outputs. After first training the segmentation network, we generate predicted segmentation masks  $\hat{S}$  for all training and validation images. Note that we do not binarize these masks, but retain the pixel-wise probabilities. Then, using the same dataset splits used for training the segmentation network, we train the summarizing function  $g(\cdot)$  to take the predicted segmentation mask  $\hat{S}$  as input and produce a class label vector  $g(\hat{S}) = \hat{y}$  by minimizing the cross entropy loss  $CE(\hat{y},\hat{y})$ . At inference, the segmentation network first generates a predicted segmentation mask for the test image, then the predicted mask is fed into the summarizing function which produces class probabilities.
We consider three architectures for trained summarizing functions that operate over segmentation network outputs. The first and simplest architecture is a fully connected layer on the flattened predicted segmentation mask  $\hat{S}$ , which is used to produce class scores that are normalized with a softmax function.
The next summarizing function we consider is a global average pooling layer followed by a fully connected layer. We take the output segmentation mask  $\hat{S}$  and produce a vector with values representing the average of each class's segmentation mask. The fully connected layer then transforms these class averages into scores for each class, which are normalized with a softmax function. This method is similar in spirit to thresholding as described above but with a threshold learned from the values observed in the train set. This global average pooling to fully connected approach was previously proposed in [16].
Finally, we consider known image classification architectures as summarizing functions. At a high level, the task of transforming a segmentation mask into a classification label is exactly an image classification task: does the "image" represented by the predicted segmentation mask correspond to a positive or negative class label? So, we investigate if existing image classification architectures are suited for learning mappings from predicted segmentation masks to class labels. In this work, we consider SqueezeNet [27] and ResNet50 [28] as summarizing functions on top of the segmentation output.
Model	Deep embedding layer name	Shallow embedding layer name
models segmentation.fcn_resnet50	backbone(layer2	backbone(layer4
Generic_UNet [59]	conv_blocks_localization[2]	conv_blocks_localization[4]
Table A1: Layer names we pull embeddings from for the trained summarizing functions. These names are from public networks available in Pytorch.
Summarizing functions on pretrained segmentation embeddings. We next consider summarizing functions that operate over segmentation network embeddings. We first train the segmentation network  $f(\cdot)$ , then store embeddings for all training and validation images. We train the summarizing function  $g(\cdot)$  to take the segmentation embeddings as input and produce a class label vector  $\hat{y}$  as output by minimizing the cross entropy loss  $CE(\hat{y}, y)$ . At inference, the image is passed through the pretrained segmentation network and embeddings from the segmentation network are stored. These embeddings are then processed with the summarizing function to produce class probabilities.
We consider two types of embeddings, shallow and deep. The intuition for exploring different embedding depths is that different depths contain different information about the input image, which may vary in their usefulness for the classification task. We give the name of the layer in the Pytorch model that we pull embeddings from in Table A1. We also consider two different architectures to classify the segmentation network's embeddings. First, we define a simple network architecture consisting of a global average pooling layer followed by three fully connected layers, which produces class scores that we normalize with a softmax function. Second, we consider a more complex classification head proposed previously in [20]. The classification head takes the segmentation embeddings as input, then processes the embeddings with three blocks of convolution-batch normalization-ReLU-maximum pooling operations. The convolutional kernels are size  $3 \times 3$  and each produce 16 channels; the maximum pooling operation operates over a  $2 \times 2$  window. The resulting embedding is flattened. Then, we process the flattened embedding with two blocks of linear transform-batch normalization-ReLU-dropout. The linear layer produces 100 features and the dropout layer has a probability of dropout set to 0.25. Finally, the resulting feature map is processed with a single linear layer to produce the final output, a vector equal in length to the number of classes, which can be normalized with the softmax function.
A3 Supporting analysis
A3.1 Proof of Proposition 1
Proof. We can write  $\operatorname*{Pr}(X_k|y = 1)$  as
$$
\begin{array}{l} \operatorname * {P r} (X _ {k} | \exists j s. t. S _ {j} = 1) = \frac {\operatorname * {P r} (\exists j s . t . S _ {j} = 1 , X _ {k})}{\operatorname * {P r} (\exists j s . t . S _ {j} = 1)} = \frac {\sum_ {S _ {1} , \ldots , S _ {L} : \sum S _ {k} \geq 1} \operatorname * {P r} (S _ {1} , \ldots , S _ {L}) \operatorname * {P r} (X _ {k} | S _ {1} , \ldots , S _ {L})}{\operatorname * {P r} (\exists j s . t . S _ {j} = 1)} (Al) \\ = \frac {\sum_ {S _ {1} , \dots , S _ {L} : \sum S _ {k} \geq 1} \Pr \left(S _ {1} , \dots , S _ {L}\right) \Pr \left(X _ {k} \mid S _ {k}\right)}{\Pr (\exists j s . t . S _ {j} = 1)} (A2) \\ = \frac {\sum_ {S _ {- k}} \Pr \left(S _ {k} = 1 , S _ {- k}\right) \Pr \left(X _ {k} \mid S _ {k} = 1\right) + \sum_ {S _ {- k}: \sum S _ {- k} \geq 1} \Pr \left(S _ {k} = 0, S _ {- k}\right) \Pr \left(X _ {k} \mid S _ {k} = 0\right)}{\Pr (\exists j s . t . S _ {j} = 1)} (A3) \\ = \frac {\Pr \left(S _ {k} = 1\right) \Pr \left(X _ {k} \mid S _ {k} = 1\right) + \Pr \left(X _ {k} \mid S _ {k} = 0\right) \sum_ {S _ {- k}: \sum S _ {- k} \geq 1} \Pr \left(S _ {k} = 0, S _ {- k}\right)}{\Pr (\exists j s . t . S _ {j} = 1)} (A4) \\ = \lambda \Pr \left(X _ {k} \mid S _ {k} = 1\right) + (1 - \lambda) \Pr \left(X _ {k} \mid S _ {k} = 0\right), (A5) \\ \end{array}
$$
where  $\lambda = \frac{\operatorname*{Pr}(S_k = 1)}{\operatorname*{Pr}(\exists j s.t. S_j = 1)}$ . Therefore, we have written the probability of a pixel value given the image-level label,  $\operatorname*{Pr}(X_k|y = 1)$ , as a mixture of the probability of the pixel value given the pixel-level label,  $\operatorname*{Pr}(X_k|S_k = 1)$  and  $\operatorname*{Pr}(X_k|S_k = 0)$ . Next, we can write  $\operatorname*{Pr}(X_k|y = 0)$  as  $\operatorname*{Pr}(X_k|S_1,\ldots ,S_L = 0) = \operatorname*{Pr}(X_k|S_k = 0)$  by our data generating process.
Putting together  $\operatorname*{Pr}(X_k|y = 1)$  and  $\operatorname*{Pr}(X_k|y = 0)$ , the right hand side of the inequality in Proposition 1 is
$$
\begin{array}{l} D _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid y = 0\right) \mid \mid \Pr \left(X _ {k} \mid y = 1\right)\right) = D _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid S _ {k} = 0\right) \mid \mid \lambda \Pr \left(X _ {k} \mid S _ {k} = 1\right) + (1 - \lambda) \Pr \left(X _ {k} \mid S _ {k} = 0\right)\right) (A6) \\ \leq \lambda D _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid S _ {k} = 0\right) \mid \mid \Pr \left(X _ {k} \mid S _ {k} = 1\right)\right) + (1 - \lambda) D _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid S _ {k} = 0\right) \mid \mid \Pr \left(X _ {k} \mid S _ {k} = 0\right)\right) (A7) \\ = \lambda D _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid S _ {k} = 0\right) \mid \mid \Pr \left(X _ {k} \mid S _ {k} = 1\right)\right) \leq D _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid S _ {k} = 0\right) \mid \mid \Pr \left(X _ {k} \mid S _ {k} = 1\right)\right) (A8) \\ \end{array}
$$
where we have obtained the desired inequality due to convexity of the KL divergence. Note that by keeping  $\lambda$  in the inequality, we can make the stronger claim
$$
D _ {\mathrm {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 _ {\mathrm {K L}} \left(\Pr \left(X _ {k} \mid y = 0\right) \mid \mid \Pr \left(X _ {k} \mid y = 1\right)\right). \tag {A9}
$$
Variable name	Description	Default value
nobj	Number of total objects in the image	50
ntrain	Number of training images	100,000
nval	Number of validation images	20,000
ntest	Number of test images	100,000
ctarget	Color of the target object	navy blue
rtarget	Radius of the target object	8
starget	Shape of the target object	circle
ltargetx	Possible x coordinates for the target object	~ unif{0,224}
ltargety	Possible y coordinates for the target object	~ unif{0,224}
cspurious	Color of the spurious object	pink
rspurious	Radius of the spurious object	8
spurious	Shape of the spurious object	square
lspuriousx	Possible x coordinates for the spurious object	~ unif{0,224}
lspuriousy	Possible y coordinates for the spurious object	~ unif{0,224}
cbackground	Colors of the background objects	∈{yellow, green, pink, red, orange, purple, eight shades of blue}
rbackground	Radii of the background objects	8
sbackground	Shapes of the background objects	∈{circle, rectangle, diamond, octagon, star, square}
lbackgroundx	Possible x coordinates for the background object	~ unif{0,224}
lbackgroundy	Possible y coordinates for the background object	~ unif{0,224}
ppos	Probability of a positive class (i.e., class balance). At ppos=0 the dataset contains no positive samples; at ppos=1 the dataset contains only positive samples.	0.5
pspur	Strength of spurious correlation. At pspur=1 the spurious object always and only co-occurs with the target object; at pspur=0 the spurious object never co-occurs with the target object, but appears every time the target object is absent; at pspur=0.5, the spurious object has no dependence on the target object.	0.5
Table A2: Variables that define a synthetic dataset. All experiments are run with datasets constructed from the default values unless specified otherwise.
Experiment	Adjusted variable
Number of training images	ntrain is varied from 1,000 to 100,000.
Class balance	Class balance is changed in the train data by varying ppos ∈ [0.01, 1]. Note that class balance in the validation and test data is not changed from default.
Size bias	Instead of all target objects having the same radius, we draw the radius from a distribution increasingly biased towards larger radii. Note that the validation and test datasets only show small objects.
Spurious correlation	Spurious correlation strength is changed in the train data by varying spur ∈ [0, 1] and the spurious feature size is changed to 28. Note that spurious correlation strength in the validation and test data is not changed from default.
Location bias	The location of the target object is changed in the train data by varying ltargetx and ltargety from ~ unif{112, 112} to ~ unif{0, 224}. Note that location distributions are not changed from default in the validation or test data.
Target size	The size of the target and background objects are swept from 8 to 40. We change ntrain to 10,000.
Number of objects	The difficulty of the classification task is changed by sweeping the number of unique background objects from 1 to 10; the total number of objects nobj remains set to 50. We change ntrain to 25,000.
Table A3: Experiments we run on synthetic datasets. For each experiment, we adjust the variable values indicated above (leaving the other variables at default) and assess how classification and segmentation-for-classification performances are impacted.
A4 Synthetic dataset and experiments
We conduct experiments on a synthetic dataset to evaluate how dataset and task characteristics impact the performance of classification and segmentation-for-classification models. Here we provide additional details on our synthetic dataset and experiments. The synthetic dataset is visualized in Figure 4.
A4.1 Synthetic dataset
The synthetic dataset is made up of images containing a gray background and many objects of varying shapes and colors. The task is to classify if a given image has a certain colored shape present. Each synthetic dataset is defined by the set of parameters given in Table A2.
To construct a synthetic image, we initialize a gray background of size (224, 224, 3). We then draw the class label  $y \sim Bernoulli(p_{pos})$ . We sample the spurious label  $s$  according to the strength of the spurious correlation  $p_{spur}$ . Specifically, we define and sample the random variable match  $\sim Bernoulli(p_{spur})$ . If match  $= 1$ , we set the spurious label  $s = y$ ; else, we set  $s = 1 - y$ . If  $s = 1$ , we create the spurious object specified by  $(c_{spurious}, r_{spurious}, s_{spurious})$ —the object's color, radius, and shape—and place the spurious object at the coordinates drawn from the distributions  $(l_{spurious_x}, l_{spurious_y})$ . We follow similar processes to generate and place the  $n_{obj} - y - s$  background objects and target object. We place the target object last to ensure it is placed in the foreground.
To run experiments on the synthetic dataset, we vary one dataset variable, train the classification and segmentation-for-classification models, and evaluate performance. The variable we adjust for each experiment is described in Table A3. Additionally, we visualize the training dataset for each experiment in Figure A1.
A4.2 Training details for synthetic datasets
We train a ResNet50 to classify the synthetic data and a ResNet50 with a fully convolutional segmentation head to segment the synthetic data. We train the classifier and segmentation networks by minimizing the cross entropy loss between the predicted and ground truth labels. We do not use data augmentation on the synthetic data. We train each network with an Adam optimizer and a learning rate of 1e-4, tuned from [1e-6, 1e-5, 1e-4, 1e-3]. We checkpoint the networks on the maximum balanced accuracy on the validation set. We use the threshold-based summarizing function.
Figure A1: Visualization of the synthetic dataset task and experiments. The task is to identify if the image contains a navy blue circle. Each experiment changes a characteristic of the training dataset to observe how classification vs. segmentation-for-classification performance is impacted.
Figure A2: Visualization of the annotation quantity experiment. The input image is a grid of colored squares; classification supervision receives a class label per image; segmentation supervision with subsampling receives one class label per colored square; segmentation supervision without subsampling receives a class label per pixel.
Figure A3: Results from annotation quantity experiment. This graph shows that providing segmentation supervision for every pixel is the most sample efficient, though providing supervision for every colored square still improves sample efficiency compared to providing one label for the entire image. This graph supports that it is both annotation quantity and the differences in class distribution leading to segmentation-for-classification's improved performance in the limited data regime.
A4.3 Additional synthetic experiments
In the main text, we focus on the divergence between positive and negative classes as to why classification and segmentation perform differently. We do not explicitly model the benefit from denser annotation during training, instead taking it as a given that the denser annotation should help segmentation learn with fewer images. To show that annotation quantity is not the only component contributing to segmentation's improved performance, we run the following experiment.
We study a simplified setting visualized in Figure A2. Specifically, we divide a  $112 \times 112$  grid into sixteen isotropic squares. The task is to identify if any of the squares are navy blue. Class balance is set to 0.5, meaning an image is as likely to have the target class as not. The remaining square colors are drawn randomly. In this simplified setting, we still expect divergence between segmentation classes to be greater than divergence between classification classes, following our analysis in Section 2.
Classification supervision is applied at the image-level, per usual. To understand the influence of annotation quantity, we evaluate segmentation-for-classification performance in two settings. First, we provide segmentation supervision with only one pixel per square. This level of supervision results in 16 annotated pixels per image, but an equal number of positive class annotations as the classification supervision. Second, we provide segmentation supervision over all pixels (this is the usual segmentation-for-classification setting), in effect providing 12,544 annotated pixels per image and  $784 \times$  as many positive class annotations. The supervision for each of these models is also visualized in Figure A2.
In effect, this experiment separates the impact of just divergence vs. divergence+higher quantity of (positive) annotations. We sweep the number of training images and report results in Figure A3.
We see that segmentation without subsampling outperforms segmentation with subsampling outperforms classification. This result confirms that both the denser annotation and greater divergence during training benefit segmentation-for-classification performance.