Segmentation-for-classification is a methodology introduced and systematically evaluated by Hooper et al. (2023) in which a segmentation network is trained to produce pixel-level predictions, and these predictions are then converted into image-level classification labels via a summarizing function g(.).
The approach is motivated by an information-theoretic analysis showing that segmentation supervision creates more separable class distributions than classification supervision (via KL divergence bounds), particularly for tasks with small targets. This theoretical insight builds on the spatial specificity framework of Saab et al. (2022).
Summarizing Functions
Two families of summarizing functions convert segmentation outputs to classification labels:
Rule-based:
- Binarize the probabilistic mask using threshold t
- If the binarized mask contains over tau positive pixels, return a positive classification label
- Class probability is computed as the average pixel-wise probability
Trained:
- Operate on segmentation outputs: FC layer, pooling + FC, full classification network (e.g., SqueezeNet, ResNet50)
- Operate on segmentation embeddings: shallow or deep embeddings with simple or complex classification heads
Key Details
- Rule-based summarizing functions are recommended: they require no additional training, avoid overfitting, are fast at inference, and produce interpretable results
- Benefits are most pronounced with small datasets, low class prevalence, rare subgroups, and small targets
- Segmentation-for-classification improves robustness to spurious correlations because background features are less correlated with pixel-level labels
- Semi-supervised segmentation methods directly benefit segmentation-for-classification; a “boosted” variant uses classification labels to filter pseudo-masks
- The primary trade-off is higher per-image labelling cost for segmentation annotations
- Towards Certified Shortcut Unlearning in Medical Imaging uses average pixel pooling (a simple summarizing function) for its downstream classification evaluation