1 Introduction↩︎

In recent years, the deployment of complex deep neural networks have a lot of applications [1]–[10] particularly in resource-constrained environments, which has driven extensive research into model efficiency and compression techniques, such as network pruning [11], [12], quantization, parameter-efficient tuning [13] and knowledge distillation [14]–[20]. Building upon these broad advances in efficient AI, optimizing large-scale models for specific clinical domains remains a critical challenge.

Although deep learning models, such as U-Net [21] and Transformer-based architectures [22]–[29], have achieved remarkable success in automated medical image segmentation [30], their massive parameter counts and heavy computational burdens severely hinder the deployment on platforms such as mobile health devices and edge computing nodes. To alleviate this, structured network pruning [31], [32] is widely adopted to remove entire structural units, maintaining compatibility with general-purpose hardware.

However, significant gaps persist in applying structured pruning to medical multimedia analysis. First, general approaches like HRank [33] and CHEX [34] evaluate network units in isolation, overlooking intrinsic inter-channel correlations. In medical imaging, anatomical structures induce high feature similarity, and neglecting these relationships risks discarding essential fine-grained semantics. Second, existing medical-specific pruning strategies [35], [36] typically prioritize the direct elimination of “redundant” filters. Because redundancy often contributes to model robustness against imaging noise, there is a critical need for frameworks that integrate representation similarity [37] to fuse correlated features, rather than simply deleting them, to preserve collective representational power during compression.

To address these challenges, we propose a novel Hierarchical Clustering Compression method (MIS-HCC) for efficient medical image segmentation. Instead of a binary keep-or-drop decision, MIS-HCC leverages the Wasserstein distance to rigorously quantify the similarity between channel distributions. We then employ hierarchical clustering to group redundant channels and fuse their parameters. This approach—applicable to both CNN and Transformer backbones for 2D and 3D data—maintains the original network topology for seamless hardware acceleration while significantly reducing model size.

The main contributions of this paper are summarized as follows:

• We propose MIS-HCC, a structured compression framework that introduces hierarchical clustering to medical multimedia model pruning, replacing simple filter removal with intelligent parameter fusion.

• We utilize the Wasserstein distance to construct a robust channel similarity matrix, explicitly modeling inter-channel correlations to preserve critical visual features during compression.

• Extensive experiments on multiple medical segmentation benchmarks demonstrate that MIS-HCC achieves superior compression-accuracy trade-offs compared to state-of-the-art pruning methods.

2 Related Works↩︎

2.1 Medical Image Segmentation↩︎

Deep learning has been applied across various fields [38]–[54]. While deep learning architectures, including CNNs (e.g., U-Net [21]) and Transformers (e.g., TransUNet [22]), have become cornerstones of medical image segmentation, their clinical deployment is often bottlenecked by computational inefficiency. Advanced variants incur substantial parameter overhead; for instance, U-Net++ [55] increases the parameter count by 30% over the original U-Net, demanding extensive GPU memory for high-resolution volumetric data. These heavy computational burdens preclude their direct application in resource-limited clinical environments, motivating the need for efficient model compression.

2.2 Network Pruning↩︎

Network pruning mitigates computational overhead by removing redundant structures [56]–[58]. While unstructured pruning sparsifies weights, structured pruning eliminates entire filters or channels, ensuring compatibility with general-purpose hardware. Prominent methods such as HRank [33] and CHEX [34] rely on proxy metrics to identify and discard less important channels. In the medical domain, tailored approaches like Sauron U-Net [35] have been proposed to reduce parameters while preserving critical anatomical features. Furthermore, coupling pruning with knowledge distillation has proven highly effective in maintaining robust representational capacity under aggressive compression Your_Paper_Citation_Here?. Despite these advances, most existing techniques treat network channels independently. By prioritizing the sheer deletion of “redundant” channels, they fail to exploit the intrinsic structural similarities highly prevalent in medical imagery.

2.3 Representation Similarity↩︎

Analyzing representation similarity is crucial for identifying redundant structures and guiding intelligent network compression [37]. Traditional metric approaches often struggle to evaluate meaningful similarities within high-dimensional feature spaces. However, recent advances in neural clustering [59] suggest that feature extraction can be viewed as selecting and fusing representatives from the underlying data distribution. Motivated by this perspective, rather than evaluating channels in isolation, our study leverages representation similarity to guide hierarchical clustering, selectively fusing functionally overlapping channels to achieve context-aware compression for medical image segmentation.

3 Method↩︎

In this study, as shown in Fig. 1, we propose a novel model compression method, named MIS-HCC, to enhance the efficiency of medical image segmentation. Our method follows three steps: First, we employ the Wasserstein distance to quantitatively measure representation similarity between input channels in a pre-trained network layer. Next, hierarchical clustering is applied on the derived similarity matrix to partition channels into clusters. Finally, channels within each cluster are averaged and fused to create a compact network, effectively reducing model size while preserving key features.

3.1 Representation Similarity↩︎

Modern CNN or Transformer blocks often contain redundant input channels that encode highly overlapping features. A principled pairwise similarity can expose such redundancy and guide structured pruning without task-specific heuristics. In order to compress the network, we aim to construct a pairwise channel similarity matrix that is applicable to layers with generic weight shapes, robust to scale, and suitable for downstream clustering.

Specifically, given a pre-trained network \(G\) with \(L\) layers \(\{l_1,\dots,l_L\}\), for any layer \(l_i\), let \(W_i\) denote the weight tensor with generic shape \(W_i \in \mathbb{R}^{\,N_{\mathrm{in}} \times N_{\mathrm{out}} \times *}\), where \(*\) denotes all remaining dimensions beyond input and output channels. This covers 1D/2D/3D convolutions and linear layers, and thus is suitable for both CNN and Transformers. Considering input-channel compression, we flatten all dimensions except the input-channel dimension to obtain the channel feature matrix \(M_1 \in \mathbb{R}^{\,N_{\mathrm{in}} \times D_i}\) where \(D_i\) is the product of all non-input-channel dimensions of \(W_i\). Let \(v_p\) and \(v_q\) denote the \(p\)-th and \(q\)-th channel in the channel feature matrix \(M_1\), we compute the Wasserstein distance for every channel pair \((p,q)\), where \(p, q \in [0, 1, 2, ..., N_{in}]\), forming the symmetric channel-pair distance matrix \([M_{\mathrm{dist}}]_{p,q} \;=\; W(v_p, v_q)\) as below: \[W(v_p, v_q) \;=\; \inf_{\gamma \in \Gamma(\mu_p, \mu_q)} \int_{X \times X} d(x,y)\, d\gamma(x,y), \label{eq:wd}\tag{1}\] where \(\mu_p,\mu_q\) are distributions induced by \(v_p,v_q\), \(X\) is the metric space, \(d(\cdot,\cdot)\) is a ground distance, and \(\Gamma(\mu_p,\mu_q)\) is the set of couplings with marginals \(\mu_p,\mu_q\). We convert distances to similarities by normalizing \(M_{\mathrm{dist}}\) to \([0,1]\) and defining the similarity matrix \(M_2 \in \mathbb{R}^{\,N_{\mathrm{in}} \times N_{\mathrm{in}}}\) as below: \[M_2 \;=\; \mathbf{1} - \widehat{M}_{\mathrm{dist}}, \quad M_2 \in \mathbb{R}^{\,N_{\mathrm{in}} \times N_{\mathrm{in}}},\] where larger \(M_2[p,q]\) indicates higher similarity.

Unlike cosine measures on raw vectors, the Wasserstein metric captures distributional discrepancies and is less sensitive to local mass shifts, making it suitable for channel features aggregated over spatial/temporal/kernel dimensions. Note that the same compression method is applied to the output channel dimension as well in practice.

3.2 Hierarchical Clustering↩︎

In a complex neural network, channels are not isolated from each other, and their internal relations and interdependence must be taken into account when simplifying the model, as pruning them independently risks removing complementary features. Thus, we propose to utilize hierarchical clustering to partition the \(N_{\mathrm{in}}\) input channels into \(K\leq N_{\mathrm{in}}\) clusters, where \(K\) is chosen by a desired compression ratio or a dendrogram cut threshold.

We perform agglomerative hierarchical clustering with Average Linkage on the Wasserstein-based affinities. We use the similarity matrix \(M_2\) to initialize hierarchical clustering, where each channel initially forms its own cluster. We then use Average Linkage algorithm to compute inter-cluster distances: \[d(U,V) \;=\; \frac{1}{|U|\,|V|}\sum_{i \in U}\sum_{j \in V} d(i,j),\] where \(U,V\) are two clusters and \(d(i,j)\) is taken from \(M_2\).

Average Linkage balances both intra-cluster and inter-cluster distances, which leads to stable, consistent partitions of channels. Hierarchical clustering returns labels \(\{1,\dots,K\}\) and index sets \(\{I_c\}_{c=1}^K\) with \(I_c \subset \{1,\dots,N_{\mathrm{in}}\}\), ready for parameter fusion.

3.3 Parameter Fusion↩︎

If channels within a cluster are near-duplicates, replacing them by a single prototype preserves salient features while cutting parameters and FLOPs, without retraining or specialized inference kernels. Thus, we construct a fused weight tensor with \(K\) input channels by averaging channels within each cluster. For cluster \(c \in \{1,\dots,K\}\) with index set \(I_c\), we fuse weights by arithmetic averaging along the input-channel dimension: \[W^{(i)}_{\mathrm{fused}}[c,\,\cdot] \;=\; \frac{1}{|I_c|}\sum_{p \in I_c} W_i[p,\,\cdot], \label{eq:fuse}\tag{2}\] where \(W^{(i)}_{\mathrm{fused}}\) replaces the original input-channel dimension with size \(K\), and “\(\cdot\)” denotes all remaining indices.

The simple and naive fusion technique ensures closed-form, training-free compression, and is compatible with both convolutional layers and linear layers. For a layer with original parameters proportional to \(N_{\mathrm{in}}\), the fused layer uses \(K\) input channels, yielding a parameter and FLOP reduction factor of \(K/N_{\mathrm{in}}\). Stacking across layers produces the compressed model \(G'\) while retaining key representational characteristics.

4 Experiment↩︎

4.1 Experimental Settings↩︎

4.1.1 Backbones and Datasets↩︎

We evaluate our method on both 2D and 3D segmentation, with CNN backbones and Transformer backbones, respectively. For 2D segmentation, we conduct experiments on three publicly available datasets, using U-Net [21] and U-Net++ [55] as the backbone architectures. The first dataset is the Breast Ultrasound Images Dataset (BUSI) [60], which consists of 437 breast ultrasound images collected from women aged between 25 and 75 years. The second dataset is the Data Science Bowl (DSB) [61], containing 670 images of segmented nuclei from different cell types obtained under diverse experimental conditions. The third dataset is the International Skin Imaging Collaboration (ISIC) [62], which includes 2,594 dermoscopic images of skin lesions collected from multiple medical institutions worldwide. To ensure consistency across experiments, all images are uniformly cropped to a resolution of \(256 \times 256\) pixels. Each dataset is randomly divided into training and testing sets with an 80%–20% split.

For 3D segmentation, we conduct experiments on the MICCAI FLARE 2021 [63] dataset using the MedSAM2 [64] backbone. The FLARE 2021 challenge dataset consists of 511 fully annotated abdominal CT scans collected from 11 different medical centers. It provides segmentation masks for four abdominal organs: liver, kidney, spleen, and pancreas. For the backbone network, we adopt MedSAM2, a recently proposed promptable segmentation foundation model for 3D medical images and videos. MedSAM2 is developed by fine-tuning the Segment Anything Model 2 on a large-scale medical dataset comprising over 455,000 3D image–mask pairs and 76,000 frames across diverse modalities. In our implementation, the FLARE 2021 CT volumes are resampled to a consistent voxel spacing and cropped around the target organs’ bounding regions. The training and testing splits follow the official challenge protocol. For MedSAM2, we utilize bounding-box prompts on central slices with bidirectional propagation to generate volumetric predictions. For performance evaluation, we adopt Dice coefficient and Intersection over Union (IoU) metrics to assess both 2D and 3D segmentation performance.

4.1.2 Configurations↩︎

All experiments are implemented in the PyTorch (2.1.2) deep learning framework. For 2D segmentation tasks, training is conducted on a single NVIDIA RTX 4090 GPU, with the batch size set to 100. Each model is trained for 200 epochs using the Adam optimizer with an initial learning rate of 0.001 and a weight decay of 0.0001. The learning rate is scheduled following a cosine annealing strategy, with the minimum learning rate set to 1e-5. For 3D segmentation, due to GPU memory limitations, the batch size is reduced to 2 and the initial learning rate is adjusted to 1e-4, while all other training configurations remain consistent with the 2D setting.

To provide a comprehensive comparison, we evaluate our method against several representative pruning approaches. RP denotes random pruning, which serves as a model pruning baseline. PFEC [65] prunes filters with small weight magnitudes and requires fine-tuning to recover accuracy. HRank [33] identifies redundant filters based on the rank of feature maps, effectively reducing channel redundancy. CHEX [34] adopts a prune-and-regrow strategy during training, enabling channel exploration without the need for pretrained models. TPP [66] preserves the trainability of pruned networks by regularizing gradient flow during pruning, ensuring stable convergence after compression. These baselines cover both post-training and training-aware pruning paradigms, providing a diverse benchmark for evaluating our proposed method. Note that since CHEX [34] requires retraining during model pruning, we do not evaluate it on the 3D segmentation task with MedSAM2 [64] backbone.

4.2 Results and Analysis↩︎

4.2.1 Segmentation Results↩︎

4.2.1.1 2D Segmentation Results

Table 1 presents the performance comparison of our proposed MIS-HCC and other state-of-the-art methods on the BUSI, DSB, and ISIC datasets. We take U-Net and U-Net++ as the baseline networks and set a uniform pruning rate of 87.5%. To ensure a fair comparison, all methods were reproduced in the same experimental environment. Overall, MIS-HCC outperforms all other methods across the three datasets and maintains comparable performance to the baseline networks even after model compression. In contrast, other pruning methods lead to more evident accuracy reductions. For the U-Net model, before compression, the Dice/IoU scores are 0.7806 and 0.7035 on BUSI, 0.9196 and 0.8584 on DSB, and 0.9137 and 0.8551 on ISIC, respectively. After compression, random pruning (RP) exhibits the most significant relative decrease in accuracy, whereas TPP [66] shows a comparatively minor reduction. Our MIS-HCC achieves the best performance among all methods, with Dice/IoU scores of 0.7841 and 0.7072 on BUSI, 0.9226 and 0.8603 on DSB, and 0.9188 and 0.8588 on ISIC. On the BUSI dataset, MIS-HCC outperforms TPP (Dice = 0.7841 vs. 0.7781; IoU = 0.7072 vs. 0.6989). On DSB, it also achieves higher scores (Dice = 0.9226 vs. 0.9196; IoU = 0.8603 vs. 0.8559). On ISIC, MIS-HCC surpasses CHEX, with Dice = 0.9188 vs. 0.9124 and IoU = 0.8588 vs. 0.8513. For the U-Net++ model, the performance comparison across all methods shows similar trends. MIS-HCC again obtains the highest Dice and IoU scores, with performance changes comparable to those seen in U-Net. The only exception occurs on the ISIC dataset, where the improvements are slightly less pronounced compared to other baselines. We attribute this to the relatively large scale of the ISIC dataset. For example, the relative performance changes on Dice and IoU are 0.45% and 0.53% for our method, compared with 0.48% and 0.32%, and 0.33% and 0.22% observed in other methods (vs. 0.25% and 0.13%). These results demonstrate that combining similarity representation with hierarchical clustering can effectively improve network compression performance by better exploiting and fusing similarities within internal structures.

4.2.1.2 3D Segmentation Results

Table 2 demonstrates the performance of different pruning approaches on FLARE21, with MedSAM2 backbone. We adopt the same pruning setting as 2D segmentation, setting the pruning rate at 87.5%. Before compression, MedSAM2 achieves Dice and IoU scores of 0.8483 and 0.7807, respectively, providing a strong baseline. After pruning, most existing approaches suffer from considerable performance degradation. Specifically, RP leads to the most severe accuracy drop, with Dice/IoU reductions of 10.35% and 13.70%, respectively. HRank also shows notable performance loss, while PFEC alleviates the decline but still exhibits a Dice decrease of 3.87%. Among prior methods, TPP demonstrates relatively stable performance, reducing Dice by only 1.61% and IoU by 2.37%. However, our proposed MIS-HCC significantly outperforms all other pruning strategies, achieving Dice and IoU scores of 0.8471 and 0.7789, which are remarkably close to the uncompressed MedSAM2 (only 0.14% and 0.23% reduction, respectively). This observation confirms that MIS-HCC is particularly effective for large-scale 3D segmentation tasks, where preserving feature diversity is crucial. Taken together with the 2D segmentation results, these findings highlight the robustness and generalizability of MIS-HCC across different architectures and modalities. By explicitly modeling channel similarity and performing hierarchical clustering before fusion, MIS-HCC achieves superior compression efficiency while maintaining nearly lossless segmentation accuracy.

4.2.2 Results on Compression Performance↩︎

Table 3 demonstrates the Dice score, FLOPs, and parameter reductions for our proposed MIS-HCC and the state-of-the-art TPP method on the BUSI dataset. Both methods were applied to prune the original network U-Net, and with pruning rates set at 50%, 75%, and 87.5%. Overall, increasing the pruning rate leads to a certain degree of accuracy loss. Under all three pruning rates, MIS-HCC outperforms the TPP method in accuracy (0.7916 vs. 0.7911, 0.7855 vs. 0.7803, and 0.7841 vs. 0.7781). This is attributed to MIS-HCC’s ability to better preserve the original network features by fusing similar channels through similarity calculations. At pruning rates of 50% and 75%, compared to MIS-HCC, the TPP had higher reduction in FLOPs and parameters. However, at pruning rates of 87.5%, the MIS-HCC outperforms TTP in both FLOPS and parameter reductions (98.28% vs. 97.36% and 98.43% vs. 97.51%). These results indicate that our MIS-HCC method not only achieves better accuracy but also delivers superior overall performance in model compression.

4.3 Ablation Studies↩︎

In this section, we evaluate the performance of MIS-HCC method using different similarity calculation and clustering methods on the BUSI dataset, with U-Net as the original network. The U-Net_Wd means using Wasserstein distances (ticking), but with k-means clustering instead of Hierarchical Clustering. The U-Net_HC means using Hierarchical Clustering, but replaces Wasserstein distance with Cosine similarity.

As shown in Table 4, The performance of the original network improves when using Wasserstein distances alone (0.7826/0.7046 vs. 0.7806/0.7035). The improvement is even more significant when Hierarchical Clustering is applied on its own (0.7835/0.7051 vs. 0.7806/0.7035). The best performance is achieved when both Wasserstein distance and Hierarchical Clustering are combined in MIS-HCC. These results prove the effectiveness of our proposed entire pipeline of MIS-HCC in enhancing network performance.

4.4 Visual Results↩︎

Fig. 2 presents qualitative comparisons of different pruning strategies on both 2D and 3D medical image segmentation tasks. For the 2D datasets, compared with U-Net baseline, other pruning approaches such as RP, PFEC, HRank, CHEX, and TPP, exhibit issues of blurred lesion edges or redundant regions. In contrast, our proposed method provides segmentation maps that are more consistent with the ground truth, with sharper lesion contours and better preservation of fine details. Similarly, in the 3D medical segmentation task, compared with MedSAM2 baseline, compared pruning methods tend to under-segment small structures or produce significant false positives. Our method, however, achieves superior visual alignment with the ground truth, maintaining structural integrity and edge precision even under high pruning rates. These results demonstrate that our approach can effectively balance model compression and segmentation accuracy in both 2D and 3D scenarios, validating the robustness and generalizability of the proposed pruning framework.

5 Conclusion↩︎

In this study, we proposed a novel model compression algorithm for semantic segmentation of medical images, specifically designed for deployment on resource-constrained medical platforms. By utilizing Wasserstein distance to measure the representational similarity between channels in network layers and applying Hierarchical Clustering to divided the channels into different clusters, MIS-HCC effectively reduces model parameters and improved model efficiency. Experimental results on medical image benchmark datasets demonstrate that our method achieves excellent performance, making it a promising solution for the practical deployment and application of medical image segmentation models. In future work, we will extend these principles to other real-world scenarios, further improving model efficiency and applicability.

References↩︎

---

1. *Zequn Xie is the corresponding author.↩︎

2. \(^{1}\)Bo Zhao is with Yale University↩︎

3. \(^{2}\)Haoran Yu is with University of Florida↩︎

4. \(^{3}\)Lifei Liu is with Wichita State University↩︎

5. \(^{4}\)Zongcheng Chu is with Purdue University↩︎

6. \(^{5}\)Yining Liu is with University of California, Berkeley↩︎

7. \(^{6}\)Chang Liu is with the Institute of Computing Technology, Chinese Academy of Sciences↩︎

8. \(^{7}\)Szu-Yu Chen is with Stevens Institute of Technology↩︎

9. \(^{8}\)Zequn Xie is with Zhejiang University zqxie@zju.edu.cn↩︎