Breaking the Memory Barrier: Efficient Multi-Class 3D Segmentation for Hundreds of Classes

Olivier Jaubert, William Traynor, Shadia Mikhael, John H. Hipwell, Sonia Dahdouh
Proceedings of The 9th International Conference on Medical Imaging with Deep Learning, PMLR 315:2369-2387, 2026.

Abstract

Medical image segmentation has transformed clinical routine by providing fast and accurate methods for the automated measurement of biomarkers and lesions. While foundation models promise broad generalization across hundreds of anatomical structures, they often under-perform compared to task-specific deep learning methods like nnUNet. However, these specialized models face scalability challenges when segmenting large numbers of classes in 3D images. We introduce a class scalable 3D segmentation method combining a low rank basis and projection operator with a chunked cross entropy and Dice loss. This design decouples the number of classes and the peak memory requirements enabling the segmentation of hundreds of classes in 3D. Integrated into the nnUNet framework, the proposed method supports state-of-the-art training and architectures. Scalability of our framework was demonstrated by creating and obtaining high Dice scores ($>0.95$) on a novel synthetic 3D “Toy Dataset” with up to 1000 different classes. Performance on the TotalSegmentator dataset (117 classes) was assessed showing comparable mean Dice scores between the proposed method and the multi-model TotalSegmentator baseline ($0.913$ vs $0.928$) and outperforming VISTA3D ($0.803$). These results highlight a practical path toward a unified, scalable foundation model for comprehensive 3D medical image segmentation of thousands of classes.

Cite this Paper

BibTeX
@InProceedings{pmlr-v315-jaubert26a, title = {Breaking the Memory Barrier: Efficient Multi-Class 3D Segmentation for Hundreds of Classes}, author = {Jaubert, Olivier and Traynor, William and Mikhael, Shadia and Hipwell, John H. and Dahdouh, Sonia}, booktitle = {Proceedings of The 9th International Conference on Medical Imaging with Deep Learning}, pages = {2369--2387}, year = {2026}, editor = {Huo, Yuankai and Gao, Mingchen and Kuo, Chang-Fu and Jin, Yueming and Deng, Ruining}, volume = {315}, series = {Proceedings of Machine Learning Research}, month = {08--10 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v315/main/assets/jaubert26a/jaubert26a.pdf}, url = {https://proceedings.mlr.press/v315/jaubert26a.html}, abstract = {Medical image segmentation has transformed clinical routine by providing fast and accurate methods for the automated measurement of biomarkers and lesions. While foundation models promise broad generalization across hundreds of anatomical structures, they often under-perform compared to task-specific deep learning methods like nnUNet. However, these specialized models face scalability challenges when segmenting large numbers of classes in 3D images. We introduce a class scalable 3D segmentation method combining a low rank basis and projection operator with a chunked cross entropy and Dice loss. This design decouples the number of classes and the peak memory requirements enabling the segmentation of hundreds of classes in 3D. Integrated into the nnUNet framework, the proposed method supports state-of-the-art training and architectures. Scalability of our framework was demonstrated by creating and obtaining high Dice scores ($>0.95$) on a novel synthetic 3D “Toy Dataset” with up to 1000 different classes. Performance on the TotalSegmentator dataset (117 classes) was assessed showing comparable mean Dice scores between the proposed method and the multi-model TotalSegmentator baseline ($0.913$ vs $0.928$) and outperforming VISTA3D ($0.803$). These results highlight a practical path toward a unified, scalable foundation model for comprehensive 3D medical image segmentation of thousands of classes.} }
Endnote
%0 Conference Paper %T Breaking the Memory Barrier: Efficient Multi-Class 3D Segmentation for Hundreds of Classes %A Olivier Jaubert %A William Traynor %A Shadia Mikhael %A John H. Hipwell %A Sonia Dahdouh %B Proceedings of The 9th International Conference on Medical Imaging with Deep Learning %C Proceedings of Machine Learning Research %D 2026 %E Yuankai Huo %E Mingchen Gao %E Chang-Fu Kuo %E Yueming Jin %E Ruining Deng %F pmlr-v315-jaubert26a %I PMLR %P 2369--2387 %U https://proceedings.mlr.press/v315/jaubert26a.html %V 315 %X Medical image segmentation has transformed clinical routine by providing fast and accurate methods for the automated measurement of biomarkers and lesions. While foundation models promise broad generalization across hundreds of anatomical structures, they often under-perform compared to task-specific deep learning methods like nnUNet. However, these specialized models face scalability challenges when segmenting large numbers of classes in 3D images. We introduce a class scalable 3D segmentation method combining a low rank basis and projection operator with a chunked cross entropy and Dice loss. This design decouples the number of classes and the peak memory requirements enabling the segmentation of hundreds of classes in 3D. Integrated into the nnUNet framework, the proposed method supports state-of-the-art training and architectures. Scalability of our framework was demonstrated by creating and obtaining high Dice scores ($>0.95$) on a novel synthetic 3D “Toy Dataset” with up to 1000 different classes. Performance on the TotalSegmentator dataset (117 classes) was assessed showing comparable mean Dice scores between the proposed method and the multi-model TotalSegmentator baseline ($0.913$ vs $0.928$) and outperforming VISTA3D ($0.803$). These results highlight a practical path toward a unified, scalable foundation model for comprehensive 3D medical image segmentation of thousands of classes.
APA
Jaubert, O., Traynor, W., Mikhael, S., Hipwell, J.H. & Dahdouh, S.. (2026). Breaking the Memory Barrier: Efficient Multi-Class 3D Segmentation for Hundreds of Classes. Proceedings of The 9th International Conference on Medical Imaging with Deep Learning, in Proceedings of Machine Learning Research 315:2369-2387 Available from https://proceedings.mlr.press/v315/jaubert26a.html.

Related Material