Abstract
Many-class 3D medical image segmentation models suffer from a memory bottleneck caused by the output logit tensor, which scales linearly with class count and limits both training and inference. We propose a fused-kernel approach that eliminates this bottleneck. Using the Triton language and a custom ONNX Runtime operator, we implemented fused kernels combining the final linear projection with binary cross-entropy and Dice loss (training) or argmax (inference), avoiding any actual instantiation of the full logit tensor. We evaluated on the TotalSegmentator v2 dataset (1228 CT scans, 117 classes) with a 3D U-Net, comparing against the official TotalSegmentator 3 mm baseline. Segmentation accuracy, memory usage, and inference time on consumer hardware (32\,GiB RAM, no GPU) were assessed. The fused kernel increased the feasible training batch size 4-fold and reduced peak inference memory by 14.6\,% in patch-based mode. This enabled whole-scan (patch-free) inference, reducing median inference time by 74.3\,% (IQR [62.4\,%, 82.1\,%]) compared to the baseline. Segmentation performance was comparable (mean Dice 0.880.18 vs. 0.890.14). We show that fused-kernel computation eliminates the output logit bottleneck in many-class 3D segmentation, enabling substantially faster whole-scan inference on consumer hardware with comparable accuracy.
Links to Paper and Supplementary Materials
Main Paper (Open Access Version): https://papers.miccai.org/miccai-2026-sat/paper/EMA_027.pdf
SharedIt Link: Not yet available
SpringerLink (DOI): Not yet available
Supplementary Material: Not Submitted
Link to Open Review
Open Review Page: Not Available
BibTex
@InProceedings{ThiFel_Segmentation_MICCAISAT2026,
author = { Thielke, Felix AND Meine, Hans},
title = { { Segmentation Without Full Logits: Fused Kernels for Memory-Efficient Many-Class Models } },
booktitle = {Medical Image Computing and Computer Assisted Intervention -- MICCAI 2026 Workshops and Challenges},
year = {2026},
publisher = {Springer Nature Switzerland},
volume = {LNCS 17265},
month = {pending},
page = {pending}
}