Pruning of MedSAM’s ViT-B Encoder Across Medical Imaging Modalities

Abstract

Unlike modality or structure-specific segmentation models such as U-Net, foundation models like Medical Segment Anything Model (MedSAM) segment across imaging modalities and anatomical structures without retraining, at a fixed cost. MedSAM pairs a Vision Transformer Base Model (ViT-B) image encoder with a prompt encoder and a mask decoder. The encoder dominates, accounting for 95.66% of the parameters and, at 966 giga Floating Point Operations (FLOPs), for almost all of the computation of a single forward pass, incurred for every image regardless of the modality segmented. This work characterises that cost analytically, deriving a block’s FLOPs from its retained attention heads and Multi-Layer Perceptron (MLP) neurons rather than by profiling, and then asks how much of the encoder is redundant for medical image segmentation. Redundancy is studied at two levels. At the level of the individual weight, redundant weights are zeroed while the weight-matrix shapes are left unchanged, so no computation is saved. At the level of the whole computational structure, removing an attention head or an MLP neuron physically reduces the weight count. It is the only level at which inference is accelerated. Sparsity is induced during training by a penalty on the loss: the Lasso ℓ1 penalty drives individual weights towards zero, paired with unstructured pruning, whilst Sparse Group Lasso (SGL) groups the weights of each head and neuron and drives whole groups towards zero, paired with structured pruning. As both penalties are non-smooth at zero, they are imposed via Proximal Gradient Descent (PGD), in which the gradient step optimises the task loss and a closed-form shrinkage sets small weights to exactly zero. Both pruned regimes are compared against a dense task-loss-fine-tuned reference and a variant with the encoder removed entirely, the latter serving as a floor
that tests whether the encoder is needed at all for a given target. Evaluated across seven imaging modalities and eleven evaluation subsets, the encoder is found to hold substantial but unequal redundancy. At the individual-weight level, 72.49% of the prunable encoder weights can be zeroed while maintaining accuracy comparable to the baseline, but the irregular sparsity yields no computational savings. As whole structures, 91.7% of the prunable encoder weight volume can be removed, reducing the model from 93.74 to 15.82 million parameters and the encoder cost from 966.32 to 81.08 giga FLOPs, which does come at a cost of the Dice Similarity Coeffi-
cient (DSC). This cost is not uniform across structures. Some targets retain usable accuracy with the encoder fully removed, such as the ultrasound foetal head and the fundus optic disc, both above 92% DSC, whilst others depend heavily on it, most notably the Chest X-ray
(CXR) lung fields, whose median DSC falls to 54.84%. The per-structure cost ranges from
the optic cup, which improves by 3.08 points under compression, to GlaS (testB), which
falls by 8.60. The MedSAM encoder therefore carries redundancy that can be removed for
real savings in size and speed. However, the extent of removal depends on the structure
and modality being segmented.