SAM: Segment Anything

Abstract

저자들은 Segment Anything 프로젝트를 통해 promptable segmentation task, 이를 수행하는 모델인 SAM, 11M 이미지와 1B개 이상의 마스크로 구성된 SA-1B를 함께 제안한다. SAM은 다양한 prompt에 따라 적절한 마스크를 예측할 수 있도록 설계되었으며, 새로운 데이터와 작업에도 zero-shot transfer가 가능하다. 실험 결과 SAM은 여러 과제에서 강한 일반화 성능을 보였고, 저자들은 이를 바탕으로 모델과 데이터셋을 공개하여 segmentation foundation model 연구를 확장하고자 하였다. 

Introduction

최근 NLP에서는 대규모 데이터로 학습된 모델이 prompt를 통해 다양한 작업을 수행하는 방식이 널리 사용되고 있고, 컴퓨터 비전에서도 CLIP이나 ALIGN과 같은 모델이 등장하면서 비슷한 흐름이 나타나고 있다(대규모 사전학습 + prompt 기반 일반화). 하지만 이미지 segmentation 분야에서는 아직 이러한 범용 모델이 충분히 등장하지 않았다는 한계가 있다.

 

이러한 문제를 해결하기 위해 논문은 이미지 segmentation을 위한 foundation model을 만드는 것을 목표로 한다. 이를 위해 세 가지 핵심 질문을 제시한다. 먼저 어떤 형태의 segmentation task가 다양한 상황에 일반화될 수 있는지, 그 task를 수행하기 위한 모델 구조는 무엇인지, 그리고 이를 학습하기 위한 대규모 segmentation 데이터를 어떻게 확보할 수 있는지를 고민한다.

 

그리해 저자들은 Segment Anything 프로젝트를 제안한다. 먼저 point, bounding box, mask, text 등 다양한 형태의 입력을 prompt로 받아 객체의 segmentation mask를 생성하는 promptable segmentation task를 정의한다. 이를 통해 하나의 모델이 다양한 segmentation 문제를 prompt 방식으로 해결할 수 있도록 한다.

 

SAM은 image encoder, prompt encoder, 그리고 lightweight mask decoder로 구성된다. image encoder는 이미지 임베딩을 생성하고, prompt encoder는 point, box, mask 등의 prompt 정보를 임베딩한다. 이후 lightweight mask decoder가 이 두 정보를 결합하여 segmentation mask를 예측하며, 이러한 구조 덕분에 이미지 임베딩을 한 번만 계산하고 여러 prompt에 재사용할 수 있기 때문에 interactive하게 segmentation을 수행할 수 있다.

 

마지막으로 segmentation 데이터 부족 문제를 해결하기 위해 data engine을 구축하고, 이를 통해 SA-1B라는 대규모 데이터셋을 생성한다. 이 데이터셋은 약 1,100만장의 이미지와 10억 개 이상의 segmentation mask로 이루어져 있으며, 기존 segmentation 데이터셋보다 훨씬 큰 규모를 가진다.

 

Segment Anything Task

해당 섹션은 다양한 segmentation 문제에 일반화할 수 있는 새로운 task를 정의한다. 이 task는 point, bounding box, mask, text 등 다양한 형태의 prompt를 입력으로 받아 해당 객체의 segmentation mask를 예측하는 방식이다. 예를 들면, 사용자가 어떤 객체를 분할할지 prompt로 지정하면 모델이 이에 대응하는 segmentation 결과를 생성한다.

하나의 point prompt가 여러 객체를 가리킬 수 있는 모호한 상황에서 SAM이 어떻게 동작하는지를 보여줌; 하나의 prompt에 대해 여러 개의 가능한 segmentation mask를 동시에 예측하도록 설계됨

 

또한 prompt가 모호할 수 있기 때문에 모델은 가능한 여러 객체 중 하나에 대해 유효한 mask를 반환하도록 설계되어 있다. 이러한 prompt 기반 방식은 하나의 모델이 다양한 segmentation 작업에 활용될 수 있도록 하며, 예를 들어 객체 탐지 모델이 예측한 bounding box를 prompt로 사용하면 instance segmentation과 같은 다른 작업도 수행할 수 있다.

 

Segment Anything Model

 

Segment Anything Model(SAM)은 image encoder, prompt encoder, lightweight mask decoder로 구성된 구조로, 다양한 형태의 prompt를 이용해 segmentation mask를 예측하도록 설계되었다.

 

Image Encoder
Image encoder는 입력 이미지를 특징 임베딩으로 변환하는 역할을 한다. SAM에서는 MAE로 사전학습된 Vision Transformer(ViT-H/16)를 사용하며, 입력 이미지는 1024×1024 해상도로 처리된다. 이를 통해 생성된 이미지 임베딩은 입력 이미지보다 16배 다운샘플링된 64×64 크기의 feature map 형태로 표현된다. 이 임베딩은 이미지마다 한 번만 계산되고 이후 여러 prompt에 대해 재사용되기 때문에 전체 시스템의 효율성을 높인다. 

 

Prompt Encoder
Prompt encoder는 사용자가 입력한 prompt 정보를 모델이 이해할 수 있는 형태로 변환하는 모듈이다. SAM에서는 prompt를 크게 sparse prompt와 dense prompt 두 가지로 나누어 처리한다. 

 

Sparse prompt는 객체의 위치나 대상을 간단히 지정하는 입력이다. 대표적으로 point, bounding box, text가 있다.

• Point prompt는 점의 위치와 함께 해당 점이 foreground(관심 객체)인지 background인지를 나타내는 정보를 함께 사용한다.
• Box prompt는 bounding box의 왼쪽 위와 오른쪽 아래 좌표를 이용해 객체의 영역을 지정한다.
• Text prompt는 CLIP의 text encoder를 이용해 임베딩으로 변환된다.

Dense prompt는 mask처럼 이미지와 같은 공간 구조를 가지는 입력이다.
mask prompt는 convolution을 이용해 임베딩으로 변환된 뒤 image embedding과 element-wise로 결합되어 모델이 해당 영역 정보를 활용할 수 있도록 한다. 또한 mask prompt가 없는 경우에는 이를 나타내는 “no mask” 임베딩이 대신 사용된다. 

 

Lightweight Mask Decoder

 

Lightweight mask decoder는 image encoder와 prompt encoder에서 생성된 정보를 이용해 최종 segmentation mask를 생성하는 모듈이다. 이 모듈은 Transformer 기반 구조를 사용하여 이미지 특징과 prompt 정보를 결합한다.

먼저 prompt token들 사이에서 self-attention을 수행하는데, 이 과정에서 각 prompt token은 다른 prompt token들과의 유사도를 계산하고 그 결과를 가중치로 사용해 서로의 정보를 섞는다. 이를 통해 여러 prompt 사이의 관계가 정리되고 필요한 정보가 통합된다.

이후 prompt token이 image embedding을 참고하는 cross-attention을 수행하는데, 이 과정에서는 prompt가 이미지 특징과의 유사도를 계산하여 해당 prompt가 이미지의 어느 영역과 가장 관련이 있는지 찾게 된다(prompt → image). 반대로 image embedding도 prompt 정보를 참고하도록 cross-attention을 수행해 이미지 특징에 prompt 정보가 반영되도록 한다(image → prompt). 이러한 과정을 통해 이미지 정보와 prompt 정보가 점차 결합되면서 최종적으로 segmentation 대상 영역을 찾게 된다.

 

이 과정을 통해 이미지 특징과 prompt 정보가 결합된 뒤, decoder의 출력은 업샘플링을 통해 해상도를 높이고 작은 MLP를 이용해 객체의 특징을 나타내는 벡터를 생성한다. 이후 이 벡터와 image embedding을 위치별로 결합하여 각 위치가 객체에 속할 확률을 계산하고, 이를 통해 최종 segmentation mask가 생성된다. SAM에서는 이 decoder를 lightweight하게 설계하여 이미지 특징은 한 번만 계산하고, 이후 다양한 prompt에 대해 빠르게 segmentation 결과를 생성할 수 있도록 했다.

 

Training and Efficiency
SAM은 계산 효율성을 고려해 설계되었으며, 이미지 임베딩은 한 번만 계산되고 이후 prompt encoder와 mask decoder를 통해 빠르게 결과를 생성한다. 실제로 prompt 기반 mask 예측은 웹 브라우저 환경에서도 약 50ms 수준으로 수행될 수 있다. 학습에는 focal loss와 dice loss를 조합해 사용하며, 여러 개의 mask 중 정답과 가장 가까운 mask에 대해서만 loss를 계산하는 방식으로 학습이 진행된다.

 

*Focal loss는 가중치를 조절해 어려운 픽셀에 더 집중하도록 만들어 클래스 불균형 문제를 완화하며, 특히 segmentation에서 대부분의 픽셀이 배경이기 때문에 모델이 배경에만 집중하는 문제를 줄여준다. 반면 Dice loss는 예측한 mask와 실제 mask의 겹치는 정도를 직접 측정하여 전체 segmentation 형태가 잘 맞도록 학습을 돕는다. SAM은 이 두 loss를 함께 사용해 픽셀 단위 정확도와 전체 mask 품질을 동시에 향상시킨다.

 

Segment Anything Data Engine

SAM을 학습하기 위해서는 대규모 segmentation 데이터가 필요하지만, 기존에는 이러한 규모의 mask 데이터셋이 충분하지 않았다. 이를 해결하기 위해 논문에서는 Data Engine(데이터 수집 방식)을 제안한다. Data Engine은 모델을 활용해 데이터를 생성하고, 생성된 데이터를 이용해 모델을 다시 학습하는 과정을 반복하는 방식으로 구성된다. 이 과정은 크게 세 단계로 이루어진다:

1. Assisted-manual stage
   - 사람이 직접 segmentation을 수행하지만 SAM이 이를 보조한다. 사용자가 foreground나 background를 클릭하면 모델이 mask를 생성하고, annotator가 이를 수정하며 데이터를 구축한다.
2. Semi-automatic stage
   - 모델이 일부 객체의 mask를 자동으로 생성하고, 사람은 남은 객체만 추가로 annotation한다.
3. Fully automatic stage
   - 모델이 이미지 위에 여러 point를 생성하고 각 point에 대해 segmentation mask를 예측하여 데이터를 완전히 자동으로 생성한다.

이러한 과정을 통해 최종적으로 SA-1B 데이터셋이 구축되었으며, 약 1,100만 장의 이미지와 10억 개 이상의 segmentation mask로 이루어진 대규모 데이터가 만들어졌다.

 

Segment Anything Dataset

SA-1B는 약 1,100만장의 이미지와 11억 개 이상의 segmentation mask로 구성된 대규모 데이터셋으로, SAM 학습을 위해 구축되었다. 라이선스가 확보된 고해상도 이미지로 구성되며 기존 segmentation 데이터셋보다 훨씬 큰 규모를 가진다. 이미지는 전문 사진 제공 업체로부터 수집된 고해상도 이미지로 평균 해상도는 약 3300×4950이다. 공개 시에는 접근성과 저장 효율을 위해 short side 1500px로 다운샘플링된 버전이 제공된다. 데이터 엔진을 통해 총 1.1B개의 mask가 생성되었으며, 99.1%가 완전 자동 방식으로 생성되었다. 자동 생성 마스크의 품질이 충분히 높다고 판단되어 최종 데이터셋에는 자동 마스크만 포함된다. 무작위로 샘플링한 약 5만 개의 마스크를 인간 주석과 비교한 결과, 94%가 IoU 90% 이상을 기록하여 높은 품질을 보였다. SA-1B는 기존 데이터셋보다 더 많은 마스크와 다양한 객체 크기를 포함하며, 한 이미지당 평균 약 100개의 마스크가 존재한다. 또한 마스크와 형태 분포 역시 기존 segmentation 데이터셋과 유사한 특성을 보인다.

Segment Anything RAI Analysis

SA-1B 이미지의 지리적 분포를 분석한 결과, 유럽과 아시아, 오세아니아 지역의 비중이 높은 반면 아프리카와 저소득 국가의 데이터는 상대적으로 적게 포함되어 있다. 다만 모든 지역에서 충분한 수의 마스크가 존재하도록 데이터가 구성되었다. SAM이 사람을 분할할 때 성별 표현, 연령대, 피부색에 따른 성능 차이를 평가하였다. 실험 결과 대부분의 그룹에서 성능 차이가 크지 않았으며, 전반적으로 유사한 segmentation 성능을 보였다. 연구에서는 현재 분석이 일부 지표에 한정되어 있으며, SAM이 다른 시스템의 구성 요소로 사용될 경우 새로운 편향이 발생할 수 있음을 언급한다. 따라서 실제 응용에서는 이러한 잠재적 편향을 지속적으로 검토할 필요가 있다고 설명한다.

 

Zero-Shot Transfer Experiments

SAM의 핵심 목표는 promptable segmentation 능력을 기반으로 새로운 작업과 데이터 분포에서도 별도의 학습 없이 적용되는 것이다. 이를 검증하기 위해 edge detection, object proposal, instance segmentation 등 다양한 작업을 prompt engineering 기반 zero-shot 방식으로 평가하였다. 모델은 SA-1B 데이터로 학습된 MAE 기반 ViT-H encoder를 사용하며, 훈련 데이터에 포함되지 않은 여러 데이터셋에서 실험이 수행된다.

 

Zero-Shot Single Point Valid Mask Evaluation

이 실험은 단일 foreground point prompt만으로 객체 segmentation이 가능한지 평가한다. 한 개의 점은 여러 객체를 가리킬 수 있기 때문에 문제 자체가 모호하며, 이를 고려하여 mIoU와 함께 인간 평가도 수행하였다. 실험은 23개의 다양한 segmentation 데이터셋에서 진행되었으며, 강력한 interactive segmentation 모델인 RITM과 비교하였다. 

 

Figure 9(a)에서 SAM은 23개 데이터셋 중 16개에서 RITM보다 높은 IoU를 기록하였으며, 예측된 3개 마스크와 ground truth 간 IoU의 최대값을 사용하는 oracle evaluation에서는 모든 데이터셋에서 RITM보다 높은 성능을 보였다. 또한 Figure 9(b)의 인간 평가에서도 SAM의 마스크 품질이 기존 방법보다 더 높게 평가되었으며, 이를 통해 SAM이 단일 foreground point prompt만으로도 유효하고 품질 높은 segmentation mask를 생성할 수 있음을 보여준다. 이어서 Figure 9(c)에서는 center point를 사용할 때 point 수가 증가할수록 다른 interactive segmentation 기법들과의 성능 차이가 줄어들지만, 1-point 설정에서는 SAM이 강한 성능을 보인다. 마지막으로 Figure 9(d)에서는 center point 대신 random point를 사용한 경우에도 SAM이 기존 방법들보다 더 안정적이고 우수한 성능을 보이며, 점의 위치가 달라져도 강한 일반화 성능을 유지함을 보여준다.

 

Zero-Shot Edge Detection

 

이 실험에서는 SAM을 edge detection에 직접 학습시키지 않고, 자동 mask generation 절차를 변형해 BSDS500에서 zero-shot edge detection 성능을 평가하였다. 구체적으로는 regular grid point prompt로 다수의 mask를 생성한 뒤, mask probability map으로부터 edge map을 추출하는 방식이다. 

Figure 10은 SAM이 생성한 edge map 예시를 보여주며, 학습 없이도 비교적 괜찮은 경계를 복원함을 시각적으로 확인할 수 있다. 또한   Table 3에 따르면 SAM은 전용 edge detector보다는 낮지만, Sobel이나 Canny 같은 전통적 zero-shot 방법보다 우수한 성능을 보인다.

 

- 평가 지표

ODS(Optimal Dataset Scale): 데이터셋 전체에 하나의 고정 threshold를 사용했을 때의 최적 F-score

OIS(Optimal Image Scale): 이미지마다 가장 좋은 threshold를 따로 선택했을 때의 최적 F-score

AP(Average Precision): threshold를 바꿔가며 precision-recall를 계산한 평균 정밀도

R50(Recall at 50% Precision): precision이 50%일 때의 recall

 

Zero-Shot Object Proposals

 

이 실험에서는 SAM을 object proposal generator로 사용하였다. 저자들은 automatic mask generation pipeline을 약간 수정하여 SAM이 생성한 mask들을 object proposal로 활용하고, LVIS v1에서 평균 재현율(AR; Average Recall; 전체 정답 객체 수 중 찾아낸 정답 객체 수)로 평가하였다. Table 4에 따르면 fully supervised detector 기반 방법인 ViTDet-H가 전체적으로 가장 높지만, SAM은 zero-shot 방식임에도 불구하고 medium/large object와 rare/common category에서 경쟁력 있는 성능을 보였다. 특히 ambiguity-aware한 원래 SAM이 단일 출력만 내는 변형보다 일관되게 더 좋은 결과를 보여, 다중 mask 예측 구조가 object proposal 생성에도 유효함을 확인할 수 있다. 

 

- 평가 지표

AR@1000: proposal을 최대 1000개 사용할 때의 average recall

 

Zero-Shot Instance Segmentation

 

SAM은 ViTDet가 예측한 bounding box를 prompt로 받아 객체 mask를 생성하는 방식으로 zero-shot instance segmentation에 활용되었다. Table 5에서 SAM은 fully supervised ViTDet-H보다 AP는 낮았지만 전반적으로 근접한 성능을 보였고, 추가 인간 평가 결과를 제시한 Figure 11에서는 오히려 더 높은 mask 품질 점수를 기록하였다. 이는 SAM이 정량 지표에서는 다소 뒤처지더라도 실제 mask 경계 품질 측면에서는 더 우수할 수 있음을 시사한다.

 

- 평가 지표

AP: 전체 평균 성능, AP_S: 작은 물체에서 성능, AP_M: 중간 크기 물체에서 성능, AP_L: 큰 물체에서 성능

 

Zero-Shot Text-to-Mask

 

이 실험은 SAM이 자유 형식의 텍스트를 prompt로 처리할 수 있는지 확인하기 위해 수행되었다. 저자들은 별도의 텍스트 주석을 추가하지 않고, 수동 수집 마스크의 CLIP image embedding을 학습 prompt로 사용한 뒤, 추론 단계에서는 이를 CLIP text embedding으로 대체하는 방식으로 텍스트 입력을 다룰 수 있도록 하였다. Figure 12를 보면 SAM은 “a wheel”처럼 단순한 표현뿐 아니라 “beaver tooth grille”처럼 보다 구체적인 표현에도 적절한 mask를 생성한다. 반면 텍스트만으로는 대상 식별에 실패하는 경우도 있었고, 이 경우 point prompt를 함께 주면 예측이 보완되는 모습을 확인할 수 있다.

 

Ablations

 

저자들은 성능에 영향을 미치는 요인을 살펴보기 위해 데이터 엔진 단계, 학습 데이터 규모, image encoder 크기에 대한 ablation을 수행하였다. Figure 13(left)에서는 manual, semi-automatic, automatic 단계를 차례로 반영할수록 성능이 꾸준히 향상되었고, automatic 데이터만으로 학습한 경우에도 전체 단계를 모두 사용한 경우와 큰 차이가 없었다. Figure 13(middle)에서는 학습 이미지를 0.1M까지 줄이면 성능이 뚜렷하게 떨어졌지만, 1M만 사용한 경우에는 전체 11M를 사용했을 때와 비슷한 수준을 유지하였다. 또한 Figure 13(right)에서는 encoder의 규모가 커질수록 성능이 전반적으로 좋아졌으나, ViT-L에서 ViT-H로 확장할 때의 이득은 상대적으로 작아 성능 향상이 점차 포화되는 경향을 보였다.

 

Discussion

저자들은 SAM이 대규모 데이터와 promptable segmentation을 기반으로 다양한 downstream task에 적용될 수 있다는 점에서 segmentation 분야의 foundation model로서 가능성을 가진다고 본다. 또한 CLIP처럼 다른 시스템의 구성 요소로 쉽게 결합될 수 있어, 객체 분할을 필요로 하는 여러 응용으로 확장될 수 있다고 설명한다. 동시에 한계도 언급한다. SAM은 fine structures를 놓치거나 작은 분리된 영역을 이상하게 hallucination할 수 있고, 경계가 매우 정밀하지 않을 때도 있다. 또한 많은 point가 주어지는 고정밀 interactive segmentation에서는 전용 방법보다 불리할 수 있으며, 무거운 image encoder 때문에 전체 처리 속도가 완전히 실시간인 것은 아니다. 텍스트 기반 분할 역시 아직은 예비적 수준이고, semantic segmentation이나 panoptic segmentation을 위한 단순한 prompt 설계도 명확하지 않다고 지적한다.

 

 

 

•  

 

•