Segment Anything

논문 해설

문제점 및 동기

웹 스케일 데이터로 사전 학습된 파운데이션 모델은 NLP 및 비전-언어 정렬(예: CLIP)을 혁신적으로 변화시켰으나, 범용 컴퓨터 비전에는 다양한 분할 태스크에 걸쳐 광범위한 제로샷 일반화가 가능한 파운데이션 시스템이 부족했습니다. 이미지 분할을 위한 기존 데이터셋은 웹 스케일 이미지-텍스트 말뭉치에 비해 규모가 작으며, 전통적인 분할 아키텍처는 주로 고정된 시맨틱 클래스나 특화된 인터랙티브 설정에 종속되어 있습니다.

분할을 위한 비전 파운데이션 모델을 구축하기 위해 저자들은 다음 세 가지 핵심 요소를 동시에 해결합니다:

1. 태스크(Task): 일반적인 제로샷 전이를 가능하게 하는 목적 함수는 무엇인가?
2. 모델(Model): 고유한 시각적 모호성을 해결하면서 실시간으로 유연한 프롬프트를 지원하는 아키텍처는 무엇인가?
3. 데이터(Data): 인터넷 데이터에 조밀한(dense) 마스크 주석이 부족한 상황에서 어떻게 웹 스케일 분할 마스크를 확보할 수 있는가?

핵심 아이디어

• 프롬프트 지원 분할 태스크(Promptable Segmentation Task): 모델은 전경/배경 포인트, 바운딩 박스, 조밀한 마스크, 자유 형식 텍스트 등 임의의 프롬프트가 주어졌을 때 유효한 분할 마스크를 반환해야 합니다. 프롬프트가 모호한 경우(예: 셔츠 위의 단일 포인트), 모델은 그럴듯한 계층적 다중 마스크(하위 부분, 부분, 전체)를 출력해야 합니다.
• 비대칭 상각 아키텍처(Asymmetric Amortized Architecture, SAM): 계산량이 많은 이미지 인코더는 이미지당 한 번만 실행되어 임베딩을 생성하고, 가벼운 프롬프트 인코더와 2계층 Transformer 디코더는 웹 브라우저 내의 CPU에서도 약 $50\text{ ms}$만에 프롬프트를 평가합니다.
• 플라이휠 데이터 엔진(Flywheel Data Engine): 3단계 주석 파이프라인:
  1. 보조 수동(Assisted-manual): SAM이 브라우저에서 작업자를 보조하여 마스크당 주석 시간을 34초에서 14초로 단축(430만 개 마스크 수집).
  2. 반자동(Semi-automatic): SAM이 신뢰도 높은 마스크를 감지하여 작업자가 라벨링되지 않은 덜 두드러진 객체에 집중하도록 유도(590만 개 추가 마스크 수집).
  3. 완전 자동(Fully automatic): 균일한 $32\times 32$ 포인트 그리드와 확대된 이미지 크롭을 기반으로 SAM에 프롬프트를 입력하여 SA-1B의 11억 개 마스크를 생성.

방법론

모델 아키텍처

SAM은 세 가지 구성 요소로 이루어집니다:

1. 이미지 인코더(Image Encoder): 윈도우 어텐션과 4개의 글로벌 어텐션 블록으로 개조된 MAE 사전 학습 Vision Transformer(기본값 ViT-H/16)를 사용하여, $1024\times 1024$ 입력을 처리하고 256개 채널을 갖는 $64\times 64$ 임베딩을 생성합니다.
2. 프롬프트 인코더(Prompt Encoder):
   • 희소 프롬프트(Sparse prompts): 포인트와 바운딩 박스는 학습된 타입 임베딩과 더해진 위치 인코딩을 통해 256차원 벡터로 매핑됩니다. 자유 형식 텍스트는 CLIP의 텍스트 인코더로 인코딩됩니다.
   • 조밀한 프롬프트(Dense prompts): 입력 마스크($4\times$ 낮은 해상도)는 $2\times 2$ 스트라이드 2 합성곱을 통해 다운스케일된 후 이미지 임베딩에 요소별(element-wise)로 더해집니다.
3. 마스크 디코더(Mask Decoder): 2계층 수정된 Transformer 디코더입니다. 교차하는 프롬프트 셀프 어텐션과 양방향 크로스 어텐션(토큰-이미지 및 이미지-토큰)을 통해 프롬프트 토큰과 이미지 임베딩을 업데이트합니다. 이미지 피처는 전치 합성곱을 통해 $4\times$ 업스케일링되고, MLP는 출력 토큰을 투영하여 업스케일된 피처에 포인트 단위로 적용되는 동적 선형 분류기를 생성합니다.

모호성을 인식하는 출력 및 손실 함수

단일 프롬프트는 종종 여러 중첩된 구조를 나타낼 수 있습니다. SAM은 해당 추정 IoU 점수와 함께 3개의 마스크를 동시에 예측함으로써 모호성을 해결합니다. 학습 중 역전파는 후보군 중 가장 낮은 손실에만 적용됩니다: $${\mathcal{L}}_{\text{mask}}=\underset{i\in \{1,2,3\}}{\min }\left(20\cdot {\mathcal{L}}_{\text{focal}}(m_i,y)+{\mathcal{L}}_{\text{dice}}(m_i,y)\right)$$ 보조 헤드는 예측된 IoU와 실제 Ground Truth IoU 사이의 평균 제곱 오차를 최소화합니다: $${\mathcal{L}}_{\text{IoU}}=\Vert {\widehat{\text{IoU}}}_i-\text{IoU}(m_i,y){\Vert }^2$$ 여러 프롬프트가 주어지면, SAM은 전용 4번째 출력 토큰을 사용하여 모호하지 않은 단일 마스크를 출력합니다.

결과

• 제로샷 포인트 기반 마스크 전이(Zero-Shot Point-to-Mask Transfer): 23개의 다양한 분할 데이터셋(1인칭 시점, 수중, 항공 및 현미경 이미지 포함) 전반에서, SAM은 단일 중심점 프롬프트를 사용할 때 23개 중 16개 데이터셋에서 RITM을 능가했습니다(최대 $+46.9$ IoU 차이). SAM의 3개 예측 중 오라클 선택기(oracle selector)를 통해 모호성을 해결했을 때, SAM은 23개 모든 데이터셋에서 RITM보다 뛰어난 성능을 보였습니다. 사람이 평가한 품질 점수에서도 SAM은 일관되게 훨씬 높은 점수(10점 만점에 7~9점 사이)를 받았습니다.
• 제로샷 엣지 검출(Zero-Shot Edge Detection, BSDS500): Sobel 필터링과 엣지 NMS를 갖춘 $16\times 16$ 그리드를 사용해 제로샷으로 평가한 결과, SAM은 0.768의 ODS와 0.928의 R50을 달성하여 HED(ODS 0.788, R50 0.923)와 같은 지도 학습 모델과 대등한 성능을 기록했습니다.
• 제로샷 객체 제안(Zero-Shot Object Proposals, LVIS v1): SAM은 클래스 지도 없이 59.3 AR@1000을 달성하여 중간 크기 객체(81.6 대 80.8), 큰 객체(86.9 대 87.0), 희귀 카테고리(65.8 대 58.3)에서 완전 지도 학습된 ViTDet-H 베이스라인을 능가했습니다.
• 제로샷 인스턴스 분할(Zero-Shot Instance Segmentation): ViTDet 바운딩 박스를 프롬프트로 사용했을 때, SAM은 COCO에서 46.5 AP, LVIS에서 44.7 AP를 달성했습니다. 인간 평가에서 작업자들은 SAM의 마스크가 완전 지도 학습된 ViTDet 출력보다 일관되게 더 깔끔하고 품질이 높다고 평가했습니다.

한계점 및 미해결 과제

• 미세한 디테일 및 극단적 줌: SAM은 명시적으로 크롭을 확대하는 높은 연산량의 대화형 방식에 비해 극도로 미세한 구조를 놓치거나 작은 분리된 컴포넌트를 생성할 수 있습니다.
• 전체 파이프라인 지연 시간: 프롬프트 디코딩은 실시간(약 $50\text{ ms}$)이지만, 무거운 ViT-H 이미지 인코더를 실행하는 것은 실시간으로 작동하지 않습니다.
• 시맨틱 이해: SAM은 클래스에 구애받지 않는(class-agnostic) 마스크를 생성하므로, 결합 가능한 외부 분류기 없이는 기본적으로 시맨틱 또는 파놉틱 카테고리 라벨을 제공하지 못합니다.
• 텍스트-마스크 프롬프팅: 제로샷 CLIP 정렬 텍스트 프롬프팅 파이프라인은 기하학적 프롬프트에 비해 아직 기초적인 수준이며 민감합니다.