SuperSAM: Crafting a SAM Supernetwork via Structured Pruning and Unstructured Parameter Prioritization

TL;DR

This paper introduces SuperSAM, a novel method for designing a supernetwork from the Segment Anything Model (SAM) using structured pruning and parameter prioritization, enabling efficient subnetwork discovery that outperforms the original model.

Contribution

It presents a new search space design strategy for Vision Transformers by converting SAM into a supernetwork with automated pruning and prioritization, improving NAS efficiency.

Findings

Subnetworks are 30-70% smaller than original SAM.

Discovered subnetworks outperform pretrained models.

Automated search space design enhances NAS for ViT.

Abstract

Neural Architecture Search (NAS) is a powerful approach of automating the design of efficient neural architectures. In contrast to traditional NAS methods, recently proposed one-shot NAS methods prove to be more efficient in performing NAS. One-shot NAS works by generating a singular weight-sharing supernetwork that acts as a search space (container) of subnetworks. Despite its achievements, designing the one-shot search space remains a major challenge. In this work we propose a search space design strategy for Vision Transformer (ViT)-based architectures. In particular, we convert the Segment Anything Model (SAM) into a weight-sharing supernetwork called SuperSAM. Our approach involves automating the search space design via layer-wise structured pruning and parameter prioritization. While the structured pruning applies probabilistic removal of certain transformer layers, parameter prioritization performs weight reordering and slicing of MLP-blocks in the remaining layers. We train supernetworks on several datasets using the sandwich rule. For deployment, we enhance subnetwork discovery by utilizing a program autotuner to identify efficient subnetworks within the search space. The resulting subnetworks are 30-70% smaller in size compared to the original pre-trained SAM ViT-B, yet outperform the pretrained model. Our work introduces a new and effective method for ViT NAS search-space design.