Beyond one-hot encoding? Journey into compact encoding for large multi-class segmentation

TL;DR

This work addresses the memory and compute challenges of large-scale multi-class semantic segmentation in medical imaging by exploring binary-encoding strategies that aim to reduce output complexity from $N_C$ to roughly $O(\log N_C)$. It systematically evaluates vanilla binary, ECOCs (e.g., Hamming) with hard/soft decoding, a binary-tree head with sequential conditioning, and class-to-codeword assignment (random and graph-based) on a BraTS 2021 subset with $N_C=108$, using the nnU-Net backbone. The key finding is that all binary-encoding approaches underperform the conventional one-hot baseline (Dice ~ $82.4$) across structures, with the best binary results around $72.7$–$73.8$ and the binary-tree method dropping to ~ $39.3$, indicating substantial accuracy losses despite memory reductions. These informative negative results suggest that naive compact encodings introduce boundary and small-structure errors and that future work should focus on ensuring cross-head consistency or leveraging more nuanced class relationships to recover segmentation quality in large-scale medical imaging tasks.

Abstract

This work presents novel methods to reduce computational and memory requirements for medical image segmentation with a large number of classes. We curiously observe challenges in maintaining state-of-the-art segmentation performance with all of the explored options. Standard learning-based methods typically employ one-hot encoding of class labels. The computational complexity and memory requirements thus increase linearly with the number of classes. We propose a family of binary encoding approaches instead of one-hot encoding to reduce the computational complexity and memory requirements to logarithmic in the number of classes. In addition to vanilla binary encoding, we investigate the effects of error-correcting output codes (ECOCs), class weighting, hard/soft decoding, class-to-codeword assignment, and label embedding trees. We apply the methods to the use case of whole brain parcellation with 108 classes based on 3D MRI images. While binary encodings have proven efficient in so-called extreme classification problems in computer vision, we faced challenges in reaching state-of-the-art segmentation quality with binary encodings. Compared to one-hot encoding (Dice Similarity Coefficient (DSC) = 82.4 (2.8)), we report reduced segmentation performance with the binary segmentation approaches, achieving DSCs in the range from 39.3 to 73.8. Informative negative results all too often go unpublished. We hope that this work inspires future research of compact encoding strategies for large multi-class segmentation tasks.

Figures (3)

• Figure 1: Binary tree representing the hierarchy of codeword bits. Weight vectors $\mathbf{w}_i$ for $i>0$ are determined by preceding bits which are determined by the ground truth during training and by the model prediction during inference. The dashed nodes at the bottom of the tree represent the class associated with the path taken from the root node.
• Figure 2: Comparison of predicted segmentations with one-hot encoding and vanilla binary encoding. Both encoding strategies yield sensible predictions. However, the white arrows point out examples of misclassified pixels at structure boundaries as a result of the binary encoding strategy.
• Figure 3: Comparison between one-hot encoding and vanilla binary encoding. The difference in average DSC between both encoding strategies shows that small structures experience the largest drop in DSC.