Task-Aware Low-Rank Adaptation of Segment Anything Model

TL;DR

The Task-Aware Low-Rank Adaptation (TA-LoRA) method is proposed, which enables SAM to work as a foundation model for multi-task learning and leverages a low-rank tensor decomposition method to incorporate both task-shared and task-specific information.

Abstract

The Segment Anything Model (SAM), with its remarkable zero-shot capability, has been proven to be a powerful foundation model for image segmentation tasks, which is an important task in computer vision. However, the transfer of its rich semantic information to multiple different downstream tasks remains unexplored. In this paper, we propose the Task-Aware Low-Rank Adaptation (TA-LoRA) method, which enables SAM to work as a foundation model for multi-task learning. Specifically, TA-LoRA injects an update parameter tensor into each layer of the encoder in SAM and leverages a low-rank tensor decomposition method to incorporate both task-shared and task-specific information. Furthermore, we introduce modified SAM (mSAM) for multi-task learning where we remove the prompt encoder of SAM and use task-specific no mask embeddings and mask decoder for each task. Extensive experiments conducted on benchmark datasets substantiate the efficacy of TA-LoRA in enhancing the performance of mSAM across multiple downstream tasks.

Figures (4)

• Figure 1: Comparison between the outputs of (a) the original SAM and (b) the mSAM proposed in this paper. It can be observed that the original SAM outputs segmentation results at three different levels with an identical number of channels. In contrast, the mSAM with TA-LoRA can produce results for different tasks with varying numbers of channels.
• Figure 2: An overview of the original SAM and the proposed mSAM with TA-LoRA. Compared to the original SAM, we froze the heavyweight image encoder while using TA-LoRA to fine-tune the update parameter tensors, and generate task-specific image embeddings for each task. Additionally, mSAM does not utilize the prompt encoder from the original SAM. Instead, we introduce trainable "no mask" embeddings with the corresponding number of output channels. Through these modifications, we can adapt to tasks with varying numbers of output channels. Moreover, mSAM has a task-specific mask decoder for each task.
• Figure 3: Comparison between (a) LoRA and (b) TA-LoRA. Both LoRA and TA-LoRA employ low-rank approximation operations on the update parameter matrices. LoRA uses separate low-rank matrices to approximate each task's update parameter matrix, while TA-LoRA concatenates the update parameter matrices of each task into an update parameter tensor and applies low-rank tensor decomposition for approximation.
• Figure 4: Generation quality of mSAM fine-tuned by LoRA-HPS, LoRA-STL, and TA-LoRA on NYUv2 dataset. It can be observed that our method can generate outputs of higher quality, particularly in the regions in the white bounding box.