SuperSAM: Crafting a SAM Supernetwork via Structured Pruning and Unstructured Parameter Prioritization
^††thanks: Github: Here is a link to the repository.

Vision Transformers (ViTs) Dosovitskiy (2020) have transformed the landscape of computer vision by leveraging the self-attention mechanism originally developed for natural language processing. The Segment Anyting Model (SAM) Kirillov et al. (2023a), is a ViT-based foundation model for image segmentation. SAM was trained on the SA1B dataset Kirillov et al. (2023a) which consists of 11M images and 1.1B mask annotations. SAM’s model architecture is composed of three major components, a large ViT-based image encoder, a lightweight prompt encoder and a lightweight mask decoder. Given the size of image encoder and its computationally expensive attention mechanism, several works have been conducted in compressing the model for deployment in resource constrained environments Zhang et al. (2023); Xiong et al. (2023); Fu et al. (2024).

Neural Architecture Search (NAS) techniques are advanced methods used to automatically discover efficient architectures for deep learning models. By automating the architectural design process, NAS methods aim to identify architectures that offer state-of-the-art performance while minimizing human intervention and computational cost. Recently, one-shot NAS methods have proven to be more efficient compared to traditional NAS techniques White et al. (2023). This is because one-shot NAS trains a single weight-sharing supernetwork that acts as a container of other subnetworks rather than train new architectures from scratch. As a result, the subnetworks directly inherit their weights from the same overparametrized supernetwork. This approach provides a means of training an exponential number of architectures for linear computation cost.

A major challenge in NAS involves designing of the the architecture search space Muñoz et al. (2024). The search space is the set of all possible sub-architecture configurations (subnetworks) that could be derived by pruning the supernetwork (i.e. which is first initialized as the pre-trained SAM model). The supernetwork acts as an over-parameterized architecture that contains all possible sub architectures White et al. (2023). Unsurprisingly, the search space is combinatorial making it quite infeasible to enumerate and optimize for the NAS task. In general, a NAS search space is first designed to bound the pool of potential architectures that can be derived. But even after bounding the search space, it could potentially contain millions of possible subnetwork configurations. Hence, an ideal search space not only reduces the number of candidate subnetworks but also contains promising subnetworks that require less resources to train.

The NAS training process iteratively samples subnetworks from the search space and optimizes them on every batch of training data. In particular, our approach follows the ‘sandwich rule’ Yu and Huang (2019), where for a single batch of training data, the maximal subnetwork (supernetwork), the minimal subnetwork and a randomly selected subnetwork are sequentially sampled and optimized. After optimization, gradients of the sampled subnetworks are pushed back into the supernetwork (main model) for updating its parameters.

Several search space design strategies have been proposed in the past. For instance, Cai et al. (2019) proposes Once-for-All, where 4-dimensional elasticity is applied to train a weight-sharing supernetwork for CNN architectures. They also propose progressive shrinking, a training strategy where the size of sampled subnetworks shrinks as the NAS training progresses. Recently, EFTNAS Muñoz et al. (2024) proposed performance aware search-space design coupled with first-order weight-reordering for transformer based models. NASViTGong and Wang (2022) designs efficient small and medium-sized models. Its search space is inspired by LeViTGraham et al. (2021), which utilizes a hybrid architecture combining convolutions and transformers. For each CNN block, the search focuses on optimizing channel widths, block depths, expansion ratios, and kernel sizes. For each transformer block, it explores the optimal number of windows, hidden feature dimensions, depths, and MLP expansion ratios. The BigNASYu et al. (2020) search space includes multiple dimensions such as kernel sizes, channel numbers, input resolutions, and network depths. These dimensions are simultaneously searched to identify optimal child models.

In this work, we introduce SuperSAM, by transforming SAM into an ‘elastic’ supernetwork, enabling the derivation of subnetworks with varying architectures tailored to a wide range of resource constraints. The subnetworks derived from SuperSAM exhibit comparable performance to the pre-trained SAM with just a fraction of its parameters. To do so, we propose combining probabilistic layer-wise pruning along side row/column-wise Wanda Sun et al. (2023b) parameter prioritization to design the search space and guide the subnetwork selection method. This approach first ranks the transformer layers and maintains the most important ones while the rest are assigned pruning probabilities. Moreover, the Wanda row/column-wise prioritization performs row/column reordering and slicing of MLP blocks within remaining transformer layers to retain the most salient parameters. This proposed dual elasticity acts as a hierarchical mechanism where the structured pruning allows the discovery of a few high-performing subnetworks whereas the reordering and slicing operations expand the intermediate search space to discover a wide range of robust architectures.

In summary:

• •

  We propose a novel algorithm for search-space design for transformer-based NAS.

• •

  We utilize this scheme to convert the SAM model into a weight-sharing supernetwork.

• •

  We train supernetworks on multiple datasets demonstrating a higher training efficiency.

• •

  We apply opentuner for subnetwork discovery and extract efficient subnetworks with comparable performance to the supernetwork.

The remainder of the paper is organized as follows: section II discusses related work, section III presents the methodology, section IV presents the evaluation. And we finally conclude the paper in section V.

Neural Architecture Search:

The search space, consisting of a set of neural network architectures, is explored by NAS methods, which apply search and evaluation strategies to identify high-performing architectures that are often smaller and more efficient than human-crafted ones. Elsken et al. (2019) Traditional NAS methods require training each architecture from scratch, which is costly. In contrast, one-shot approachesWhite et al. (2023), introduced in 2022, train all architectures in the search space simultaneously by using a single ”supernetwork,” an over-parameterized model containing all possible subnetworks. In one-shot weight-sharing approaches, it has been demonstrated through careful experimental analysis that it is possible to efficiently identify promising architectures from a complex search space Bender et al. (2018). After training, search algorithms, such as reinforcement learning, evolutionary algorithms, or gradient-based methods, can be used to explore the vast space for possible architectures under certain constraints.

Various techniques exist for training the generated super-network. For instance, Once-for-All (OFA)Cai et al. (2019) introduces progressive shrinking, which enforces a training order that starts with large subnetworks and gradually moves to smaller ones. NASVITGong and Wang (2022) offers a set of methods, including a gradient projection algorithm, switchable layer scaling design, and a streamlined approach to data augmentation and regularization, all of which significantly enhance the convergence and performance of subnetworks. BigNAS Yu et al. (2020) challenges the conventional view that post-processing of weights is required for good prediction accuracy. Without additional retraining or post-processing, it trains a single set of shared weights. A sandwich sampling rule with inplace knowledge distillation (KD) is used to simultaneously optimize the supernet and sub-networks for each mini-batch, stabilizing training and improving convergence. EFTNAS Muñoz et al. (2024) proposes using first-order weight reordering to improve the search space design. In particular, a column-wise reordering of attention and MLP layers is used to improve results of elasticity operation.

Network Pruning:

Pruning is a popular technique for model compression. Generally, pruning can be categorized as structured or unstructured depending on the size of network components removed by the pruning operation. Structured pruning involves the removal of large network componentsCheng et al. (2024). For example, in ShortGPTMen et al. (2024), redundant layers are removed based on a metric that measures layer importance. LLM-PrunerMa et al. (2023) adopts a structured pruning approach by selectively removing non-critical coupled structures using gradient information. Similarly, BlockPrunerZhong et al. (2024) splits each transformer layer into Multi-Head Attention (MHA) and MLP blocks, evaluates their importance using perplexity measures, and employs a heuristic search for iterative pruning to optimize model efficiency. Additionally, Isomorphic PruningFang et al. (2024) demonstrates its effectiveness across various network architectures, such as Vision Transformers and CNNs, and consistently delivers competitive performance across models of different sizes, further showcasing the versatility of structured pruning techniques.

Unstructured pruning, also known as weight-wise pruning, targets individual weights by eliminating redundant connections in neural networks, setting the corresponding weights to zero. For example, magnitude pruningHan et al. (2015) reduces storage and computation requirements by learning and retaining only the most important connections, achieving this reduction without sacrificing accuracy. Wanda Sun et al. (2023a) removes weights with the smallest magnitudes, adjusted by the corresponding input activations, on a per-output basis. Similarly, SparseGPTFrantar and Alistarh (2023) approaches pruning as a layer-wise sparse regression problem, solving it approximately through a sequence of efficient Hessian updates and weight reconstructions, further optimizing neural network efficiency.

SAM:

Despite the impressive performance of the Segment Anything Model (SAM), its large ViT-based image encoder imposes a substantial inference cost, making it challenging to deploy in resource-constrained environments. Several methods have been proposed to address this issue. e.g., the Fast Segment Anything (FastSAM)Zhao et al. (2023) introduces FastSAM, which employs YOLOv8-seg, an object detector adapted for instance segmentation, to significantly reduce computational overhead. MobileSAMZhang et al. (2023) was developed by applying decoupled knowledge distillation from ViT-H image encoder to a new tiny-ViT image encoder. EfficientSAMsXiong et al. (2023) utilizes SAMI(masked image pretraining) to enhance visual representation learning by reconstructing features from SAM’s image encoder. By combining SAMI-pretrained lightweight encoders with a mask decoder, EfficientSAMs achieve both efficiency and effectiveness. The SAM-LighteningSonga et al. (2024) introduces Dilated Flash Attention, a re-engineered attention mechanism that increases inference speed by approximately 30 times compared to the original SAM.

Search Space design:

Following the intuition of previous compression attempts, we apply our proposed NAS technique on the image encoder. Unlike previous works that use direct distillation on a single architecture (eg. Zhang et al. (2023)), we first apply elasticity to the image encoder and conduct NAS training. In NAS literature, elasticity describes the potential variability a certain architectural component can have across subnetworks Muñoz et al. (2021). For instance, two subnetworks ($\alpha^{i}$ and $\alpha^{j}$) derived from the same supernetwork could have varying number of channels for the same layer $k$, i.e. $\alpha_{k}^{i}\neq\alpha_{k}^{j}$. This makes layer $k$ elastic.

We apply 2-dimensional elasticity to the SAM ViT image encoder. The first elastic dimension is applied to the transformer layers as a form of structured pruning (similar to Men et al. (2024)). In this case, we utilize few-shot evaluation to determine the importance of the transformer layers with respect to a performance metric. This allows us to identify a set of prunable layers $P\in\{0,1,\ldots,11\}$ that will be assigned a pruning probability $\theta_{i}$ during subnetwork selection. In particular, during NAS training, the pruning probability determines the selection of the corresponding transformer layer in the construction of the currently sampled subnetwork.

The second elasticity dimension applies weight reordering and slicing (windowing) of MLP-blocks in remaining (unpruned) layers. As shown in Fig. 2, MLP-blocks in remnant layers contain two linear layers. In the ViT-B version, these linear layers have a dimension of (3072 x 768) and (768 x 3072). Reordering and slicing is applied to the intermediate dimension (of size 3072). In particular, the rows of the first linear layer and the columns of the second linear layer are first ranked using the Wanda Sun et al. (2023a) metric. The row and column ranks are computed by aggregating Wanda scores of individual parameters in the rows and columns as shown in Eq. (1) below:

where $S_{i}$ is score for $i$-th column, $W_{ij}$ is the magnitude of the $ij$-th parameter and $||X_{j}||_{2}$ is the $l_{2}$ norm of the $j$-th input features aggregated across $N$ input batches with sequence length $L$.

Corresponding row and column ranks are averaged to compute a mean rank over the intermediate dimension (i.e. for both linear layers). For instance, the rank of the 3rd row in linear layer 1 is averaged with the rank of the 3rd column in linear layer 2 to compute the average rank of the third row-column pair. The rank averaging is computed across all corresponding row-column pairs to compute a mean rank over the intermediate dimension. The mean rank is then used to reorder corresponding rows and columns of linear layers 1 and 2 in descending order such that the reordered rows/columns align without distorting the forward/backward pass of the MLP-block. After reordering, a window $w_{i}\;\text{for}\;i\in\{1,2,\ldots,m\}$ is applied to the linear layers to slice the layers by retaining the first $w_{i}$ rows and columns. Note that slicing neither affects the dimension of input to linear layer 1 (i.e. 768), nor the output dimension of linear layer 2 (i.e. 768). Rather, it only shrinks the intermediate dimension (i.e. 3072). This approach effectively prioritizes and retains the most salient parameters in the linear layers significantly improving the quality of the search space.

The proposed 2D elasticity is not arbitrary. Instead, we observe that the layer-wise pruning serves as starting points (anchors) in the search space whereas windowing operations expand on those anchors. Together, the layer-wise pruning and windowing act as a hierarchical mechanism for designing the search space. This is because, as previous work Men et al. (2024) has shown, and as we observed from preliminary experiments, transformer layers are not equally important. While some are critical, others can be redundant so much so that pruning (removing) layers may result in either a dramatic drop or marginal decrease in performance as shown in Fig. 3. Layer pruning significantly reduces the model size facilitating the presence of a few high performing subnetworks. In contrast, when applying only windowing, the search space produces a substantial number of low-quality subnetworks. Nevertheless, we cannot solely depend on layer-wise pruning, as this method results in only a limited set of high-quality subnets (anchors), typically numbering just a few dozen. As such, it poses a limitation in exploring a broader search space that includes smaller subnets dispersed between the anchors. This phenomenon is demonstrated in section IV-B.

NAS Training:

We implement the sandwich rule Yu and Huang (2019) to train the supernetworks. In particular, for each training batch, we sample the maximal subnet, the smallest subnet and a randomly selected subnet for sequential optimization. Here, we compute the cost using dice loss as shown in Eq. 2 below.

where GT is ground truth mask and Pred is predicted mask.

In all three cases, we keep the prompt encoder and mask decoder architecturally intact across all subnetworks. i.e. Elasticity is only applied to the image encoder. This simplifies implementation and avoids slicing into smaller architectural components allowing all subnets to directly inherit the pre-trained architectures and weights. To train the SAM supernet, images are first resized to 1024 $\times$ 1024 and fed to the model along with prompts. SAM prompts can be one or more coordinate points and/or boxes on or around objects of interest Kirillov et al. (2023b). We freeze the mask decoder during training and jointly optimize the image encoder and mask decoder as shown in Fig. 1.

Deployment:

Once the SAM model is transformed into a supernetwork, its search space will have developed sufficiently to include subnetworks that perform on par with the supernetwork. As a result, a search algorithm can be used to identify subnets that meet specific resource constraints. For instance, we can query the search space to identify subnets with sizes under 60M parameters and mIoU performance $\geq$ 80% on a given task.

Once SuperSAM is trained, therefore, we are now able to deploy opentuner Ansel et al. (2014), a program auto-tuner to execute the constrained search. OpenTuner uses an ensemble of disparate search techniques simultaneously; dynamically allocating a larger proportion of tests to promising techniques. The ensemble techniques are themselves organized by a meta technique. For instance, an AUC Bandit meta technique can combine greedy mutation, differential evolution and hill climber instances to tackle the search problem. In the following section, we describe various experiments showcasing the efficacy of the proposed search space design strategy.

We evaluate the proposed approach by training supernetworks using the SAM ViT-B model on different downstream tasks. These tasks include the SA1B Kirillov et al. (2023b) dataset that was used to train the SAM model, as well as MS-COCO Lin et al. (2014), ADE20K Zhou et al. (2017) and Mitochondria Lucchi et al. (2013) datasets. SA1B was used to train the SAM model, it contains around 11 million images and over 1.1 billion masks. In our experiments, we utilize 0.01% of SA1B to train our supernetwork. i.e. 10,000 images for training and 100 for validation. We set a cutoff limit of 64 objects per image, by slicing the number of objects in the image if the objects exceed the cutoff limit and randomly repeating objects in case there are fewer objects. On the other hand, we utilize the entire MS-COCO dataset about 92K images for training and 5K for validation. Similar to the SA1B case, here, we set an object limit of 8 per image. Mitochondria is a small domain specific dataset containing gray-scale images of cells. We ‘patchify’ the images into 256x256 patches to generate around 800 images for training and 800 for validation. Finally we use the entire train set, 20K images for training and a subset of 100 images for validation when using the ADE20K dataset. We present example outputs of the smallest subnet in Figs. 6 and 7.

We train the supernets on the instance segmentation task where either a single point prompt per object or bounding box is applied to segment the objects of interest. In particular, while we utilize box prompts for Mitochondria, we use point prompts for the other datasets. To account for prompt noise, we generate the prompts by selecting a random positive point inside the ground truth. Similarly, in the case of box prompts, we generate the box prompt around the object with a randomly chosen padding size of [0-20] pixels. We train our subnets using the DiceLoss Sudre et al. (2017) cost function and apply a learning rate of $1e-5$. We also apply a lambda learning rate decay until the learning rate shrinks to 1% of its original value.

In designing the search space, we use a fraction of the SA1B, 100 images, to compare the importance of transformer layers of the model. Similar to observations made by ShortGPT Men et al. (2024), we notice that not all layers of the ViT-B image encoder are equally important. As shown in Fig 3, layer 0 exhibits highest importance such that removing it causes a $\approx$60% drop in mIoU. Moreover, after pruning layers [1,2,5,6,9] we notice that performance only drops by about $\approx$8% whereas the model size has shrunk to 62% of the original pre-trained model. With this insight, we design our search space by applying structured elasticity to layers [1,2,5,6,9]. Next, we apply column/row-wise elasticity windows of size [768, 1020, 1536, 2304], on remaining MLP blocks. Because of layer 0^th importance, we keep it intact by shielding it from both layer-wise and row/column-wise elasticity operations.

In this paper, we propose a new approach of search space design for training one-shot NAS in transformer-based vision models. We propose 2-dimensional elasticity where the first dimension applies probabilistic layer-wise structured pruning while the second dimension applies a Wanda-based row/column-wise reordering and slicing of MLP-Blocks of remaining layers. We show this hierarchical (2D) approach exploits the strengths of structured pruning to discover high performing subnetworks and utilizes windowing (slicing) to expand the search space to further discover robust subnetworks that lie in the intermediate space. We demonstrate the proposed technique on the Segment Anything Model (SAM), a newly developed foundation model for image segmentation. After training supernets on several tasks, we are able to discover and extract smaller networks with robust performance. In the SA1B case for instance, our subnetworks outperform the pretrained input model, SAM-ViT-B, despite the supernet being trained on just 0.1% of the dataset. Moreover, our proposed Wanda-based reordering and windowing technique shows superior convergence compared to ‘no reordering’, and ‘magnitude reordering’ cases. This is quite important as one of the challenges of NAS is the computation requirements required to train subnetworks in the search space.

This research is supported by the U.S. Department of Energy (DOE) through the Office of Advanced Scientific Computing Research’s “Orchestration for Distributed & Data-Intensive Scientific Exploration” and the “Cloud, HPC, and Edge for Science and Security” LDRD at Pacific Northwest National Laboratory. PNNL is operated by Battelle for the DOE under Contract DE-AC05-76RL01830. This research is also supported by the National Science Foundation (NSF) under Award Number 2211982. We thank Intel Labs and the Research IT team of Iowa State University for their support.