SGAV-MP: Smaller Segmentation Models That Keep Their Uncertainty

Anyone who has tried to ship a segmentation model on a hospital workstation or a small inference box knows the next request after “make it accurate” is “make it smaller”. The usual answer is post-training quantization: convert the FP32 weights to INT8 or lower, rerun the validation set, check that Dice barely moved, and ship. Sagar Kumar, Shailja Sharma, Tirthajit Baruah and Punit Rathore of the Indian Institute of Science show why that check is incomplete in arXiv:2610.25460, posted on 18 October 2026. Their method, SGAV-MP (score-guided aggregate-validated mixed precision), picks a precision for each layer and accepts a change only after the whole compressed model still matches the original on Dice, calibration, uncertainty maps and failure detection.

Why a viewer or platform team should care

Many segmentation products now show the model’s confidence next to the mask, as an uncertainty overlay in the viewer, a per-case flag that sends a scan to a radiologist, or a threshold that holds a volume out of a report. All of that depends on the probability map, not on the hard mask. Quantization perturbs the logits. The argmax label can survive that perturbation while the probabilities flatten, so the mask looks the same and the confidence overlay changes. The authors prove this for a simple two-class case and then show it in real models.

Table 1 of the paper shows the gap. On the PROMISE12 prostate MRI set, a ResNet50 segmenter quantized uniformly to INT4 kept Dice at 0.881 (from 0.886), while the rank correlation between its uncertainty map and the FP32 model’s map fell from 1.0 to 0.324. If your QA sheet only tracks Dice, that model passes. If your viewer draws its uncertainty overlay, the overlay is now mostly different from the one you validated.

What SGAV-MP does

The input is a trained FP32 segmentation model and a labeled calibration set that is kept apart from the test data. No weights are retrained. The method works through a list of candidate actions, where each action sets one layer or block to FP16, INT8, INT4 or INT2.

• Score every action. For each candidate, it estimates how much storage the change saves and how much damage that layer alone does to the reliability metrics when quantized. It adds a structural risk term that penalises layers with many connections or layers that sit in skip or multi-scale fusion paths, since errors there spread further.
• Drop dominated actions. Any action that saves less and hurts more than another one is removed before the search.
• Greedy search with a full-model check. The highest-scoring action is applied, and the complete mixed-precision model is run on the calibration set. The change is kept only if four conditions all hold relative to FP32: Dice within a tolerance, expected calibration error within a tolerance, uncertainty-map rank correlation above a floor, and the area under the risk-coverage curve (a failure-detection measure) within a tolerance. Otherwise the action is rejected and the next one is tried.

The output is a per-layer precision policy. The score only decides the order; acceptance always comes from evaluating the whole model, so the method does not assume that layer-by-layer damage adds up neatly.

What it was tested on

In 2D, the authors used PROMISE12 (prostate MRI) and MosMedPlus (COVID-19 lung infection on CT) with segmentation networks built on five encoders: ResNet18, ResNet50, MobileNetV2, EfficientNet-B0 and DenseNet121. In 3D, they used SegResNet on AMOS22 (abdominal organs), BraTS (brain tumour) and BTCV (abdominal CT). Baselines included uniform FP16, INT8 and INT4, AdaRound, a Hessian-guided greedy method, a knapsack-style mixed-precision method (MCKP) and EfficientQ, a medical post-training quantization method.

On the same PROMISE12 ResNet50 model from Table 1, SGAV-MP reached 7.29 times storage compression with Dice 0.8858 and an uncertainty-map correlation of 0.97, close to what uniform INT4 gives in size but without the collapse in the uncertainty map. Uniform INT8 stayed reliable too but stopped at about 4 times. The 3D results follow the same pattern: uniform INT4 on BraTS dropped both Dice and uncertainty agreement, while SGAV-MP kept both near FP32.

Storage compression is not the same as speed

Mixed INT4 and INT2 policies are reported as storage compression only, because standard runtimes do not execute those mixes natively. The authors ran separate TensorRT 10.3 measurements on an NVIDIA RTX 6000 Ada with FP16 and INT8 only. In that setting the gains are smaller. On the 2D models, the TensorRT-compatible SGAV-MP policies compressed less than uniform INT8 but ran slightly faster and kept the uncertainty map closer to FP32. On a 3D UNet3D for AMOS22, the existing MedPTQ pipeline was both smaller and faster than SGAV-MP, while SGAV-MP kept the uncertainty map a little closer to FP32.

If your product reads uncertainty, SGAV-MP gives you a way to choose how far to compress while keeping that signal intact. If you only need a faster runtime, uniform INT8 on a supported backend is still a strong default.

How it would fit into a pipeline

1. Freeze the FP32 model you already validated, and keep its predictions and entropy maps on a held-out calibration set as the reference.
2. Pick the tolerances that match your product. If the viewer shows uncertainty, the correlation floor matters most. If you route cases by confidence, the risk-coverage term matters most.
3. Run the search offline on a GPU. It is not quick. The 3D SegResNet runs took about 8,900 to 28,300 seconds (roughly 2.5 to 8 hours) with 3.4 to 6.0 GB of peak GPU memory, and the 2D ResNet50 run on PROMISE12 took about 2,450 seconds with close to 20 GB of peak memory.
4. Export the policy that your runtime can execute (for TensorRT today that means FP16 and INT8 layers), then rerun the same four checks on a separate test set before release.

Where it falls short

• No public code. The paper does not link a repository or weights, and none was public when we checked on 25 October 2026. You would need to reimplement the search yourself.
• It needs labeled calibration data. Dice and the risk-coverage term require ground truth, so a site without labels cannot run the checks as described.
• It inherits the FP32 model’s flaws. The goal is to match the original model’s uncertainty. If that uncertainty was poorly calibrated to begin with, the compressed model will faithfully copy the problem.
• Calibration-set guarantees only. The authors state that passing the checks on the calibration set does not guarantee reliability on unseen data. A cross-dataset test (calibrate on PROMISE12, evaluate on MosMedPlus) held up reasonably, but stronger domain shift was not studied.
• Limited model coverage. The models are CNN encoders and SegResNet. Foundation-style segmenters, other modalities and real INT4 or INT2 runtime backends are listed as future work.

Related reading on PYCAD

• Hippocampal Segmentation Ensembles That Flag Their Own Failures
• SAT3D Puts Uncertainty-Guided Tumour Segmentation in 3D Slicer
• MID: Score Segmentation Batches Without Per-Case Labels

Sources

• Sagar Kumar, Shailja Sharma, Tirthajit Baruah, Punit Rathore. Reliability-Preserving Compression for Deployable Medical Image Segmentation. arXiv:2610.25460, 2026. Abstract and PDF. Licensed CC BY 4.0.
• Featured image: Fig. 1, Table 1 and Algorithm 1 of the paper, cropped and resized only.