Most 3D segmentation models that score well on public benchmarks are heavy. SegMamba-V2, a Mamba-based network, has about 139 million parameters and needs 3,730.7 GFLOPs for one forward pass on a BraTS patch. A group at the University of Glasgow asked how much of that you can throw away. Rahim Ullah, Hasan Abbas and Syed Waqar Nabi distilled SegMamba-V2 into TinyUNet3D, a plain 3D U-Net with 0.957 million parameters, and tested it on brain tumor MRI, hippocampus MRI and spleen CT. The paper is arXiv:2611.10033 and the training code is on GitHub under MIT.
The small student lands within about one Dice point of the big teacher, and on BraTS 2023 it matches it. The second finding is less flattering to distillation and more useful to anyone planning a project. With all the labels, distillation barely changes the student. It helps when the student has few labels and the teacher has many.
What TinyUNet3D is
TinyUNet3D is a four-level encoder and decoder. Every convolution is depthwise-separable: a 3 × 3 × 3 depthwise convolution followed by a 1 × 1 × 1 pointwise convolution, then instance norm and LeakyReLU. The decoder upsamples with trilinear interpolation, which adds no parameters, and concatenates skip connections the usual way. With a base width of 64 (channels 64, 128, 256 and 512) the network has 956,912 parameters. In float32 that is roughly a 4 MB weights file.
The repository default is base 32 (0.262M parameters). The paper’s model is base 64, and the README warns you to pass STUDENT_BASE=64 to the trainers or you will train a different network.
How the distillation works
The teacher, SegMamba-V2, is trained first on the same fixed split, then frozen. The student sees each 3D patch, and its loss has two parts. One is the usual DiceCE loss against the ground-truth mask. The other is DIST, a distillation loss from Huang and colleagues (NeurIPS 2022) that compares the relational structure of the two networks’ class probabilities instead of forcing the student to copy the teacher’s probabilities voxel by voxel.
DIST uses a Pearson-correlation distance in two directions. Inside each voxel it compares the student’s and teacher’s probability vectors across classes. Inside each class it compares the two probability maps across all voxels. Both terms and the overall distillation weight are set to 1. In the BraTS data-fraction runs the per-class term is computed over foreground voxels only.
The authors retrained the teacher themselves instead of using the released BraTS checkpoint, because that checkpoint was trained on cases that overlap their test split. Their reproduced teacher scores 88.88 Dice on BraTS 2023, below the 91.32 reported for SegMamba in its own paper under its own protocol, so keep in mind that the teacher here is not the strongest published version.
What they found
On fixed held-out test sets (students averaged over three seeds), the distilled student reaches 96.10 Dice on MSD Spleen against the teacher’s 96.84, 87.59 on MSD Hippocampus against 88.68, and 88.96 on BraTS 2023 against 88.88. That is 99.2%, 98.8% and 100.1% of the teacher’s Dice. Compute drops too, though less dramatically than the parameter count: the student needs 211.5 GFLOPs per BraTS patch against the teacher’s 3,730.7, and 235.7 against 1,566 on Spleen.
A student trained with labels alone, no teacher at all, scores 88.95 on BraTS, 95.86 on Spleen and 87.88 on Hippocampus. So with full labels distillation adds 0.01, adds 0.24 and loses 0.29 points. Most of the student’s accuracy comes from the architecture and supervised training.
The authors then retrained the student on 5%, 10%, 25%, 50% and 100% of the 876 BraTS training cases. The teacher stays the same in every run, and the training script loads the one teacher trained on all 876 cases. With 5% of the labels, distillation lifts test Dice from 0.849 to 0.868, about 1.8 points. With 10% the gain is about 1.7 points. At 25% and 50% it is around half a point, and at 100% it is gone. Their recommendation is to switch distillation on when labels are scarce and leave it off otherwise.
Why a viewer or clinic team should care
If you ship segmentation inside a DICOM viewer, a hospital workstation or a browser, model size decides a lot: download time, memory on a shared GPU, whether CPU inference is tolerable, whether you can run several models side by side. A 4 MB 3D network is easy to bundle and easy to version. The paper’s Fig. 3 shows the student’s masks next to the ground truth on one test slice per dataset. They look nearly identical, though one slice per dataset is an illustration and not a test.
Read the data-fraction sweep carefully before you plan around it. The student trained on 5% of BraTS learned from a teacher that had seen all 876 labeled cases, so the extra 1.8 points come from labels the student never saw directly. That setup fits a real situation: you have a large model trained on a big labeled set and want a small model for deployment, perhaps fine-tuned on a few local cases. The paper does not test a teacher trained on the same small set. Once the student has hundreds of labeled cases, the extra teacher pipeline adds almost nothing, and you can train the small network directly.
What it does not tell you yet
- No timings. The paper reports parameters and FLOPs, not latency or memory on real hardware. Depthwise 3D convolutions cut FLOPs, but how fast they run depends on kernel support in your runtime. The authors list inference speed and memory on constrained clinical hardware as future work.
- No weights. Datasets and checkpoints are not distributed, so you have to train the teacher and the student yourself. The scripts also import the SegMamba-V2 code base, which is not bundled, and they include SLURM launchers for a GPU cluster.
- Narrow evidence. The data-fraction sweep covers BraTS 2023 only, with one teacher and one student design. MSD Spleen has 41 labeled cases and 8 test cases, so the spleen differences sit within noise, as the README says.
- Not like-for-like comparisons. The nnU-Net, UNETR, CoTr and Swin UNETR numbers in Table 1 come from their own papers and protocols. Read them as a rough guide, not a ranking.
- Small inconsistencies. The paper says teachers trained for 400 epochs and students for 250. The README says 1,000 epochs for the BraTS teacher and 300 for BraTS students. Check the scripts before you plan GPU time.
How to try it
Clone 2802990u/tinyunet3d-dist-kd, then clone the SegMamba-V2 code base next to it and copy the student model and the BraTS scripts into its dataset folder as the README describes. python models/student_unet3d.py prints the parameter count for each base width, and python losses/dist_loss.py smoke-tests the loss. If you only want a small 3D network for your own structure, the model file is under 100 lines of PyTorch and trains without the teacher.
For a different take on lean 3D segmentation, with released weights and a 3D Slicer extension, see our post on SlimFormer++.
Source and license
Paper: Rahim Ullah, Hasan Abbas and Syed Waqar Nabi, “Ultra-Lightweight 3D Medical Image Segmentation via Knowledge Distillation in Label-Scarce Domains”, School of Computing and James Watt School of Engineering, University of Glasgow, arXiv:2611.10033 (cs.CV, 5 November 2026). The README describes it as an ISBI paper under review. Code: github.com/2802990u/tinyunet3d-dist-kd, MIT license. No pretrained weights are released. BraTS 2023 and the Medical Segmentation Decathlon keep their own licenses. Featured image: Fig. 3 of the paper, panels cropped and scaled.