source: Segment Anything
이번에 살펴볼 논문은 SAM: Segment Anything이다.

사실 SAM을 VL-SAM, GenerateU 등 OED 모델들을 리뷰하기 전에 먼저 살펴봤어야하는데, 순서가 약간 뒤바뀌었다..😵‍💫

본 포스트에서는 SAM의 모든 부분을 리뷰하지는 않고, 기존 Object detection의 proposal 개념을 SAM 방식으로 재해석하는 것에 초점을 맞추고자 한다.

SAM은 이미지 segmentation을 위한 하나의 Foundation Model이다.

SAM은 RPN처럼 region proposal module을 학습하는 모델인가, 아니면 promptable segmentation model을 proposal generator처럼 사용하는 것인가?

위 질문을 가지고 SAM을 알아보고자 한다.

저자들은 SAM이 Promptable하게 만들어져서 zeero-shot transfer가 가능하다고 한다.
Promptable이 대체 무엇이길래 이것이 가능할까?

promptable하다는 것은 "무엇을 segment할지"를 모델 파라미터 안에 고정하지 않고, inference 때 prompt로 지정할 수 있다는 뜻이다.

그래서 새로운 task를 새로 학습하지 않고도 prompt 형태로 바꿔 넣으면 쓸 수 있고, 이게 zero-shot transfer와 연결된다.

기존 segmentaiton/detection 모델은 보통 task가 고정되어있다.

입력 이미지 -> COCO 80개 class에 대한 mask 예측
입력 이미지 -> LVIS class별 instance segmentation
입력 이미지 -> semantic segmentation map

이런 모델은 출력 공간이 훈련 때 정해져있어서 새로운 데이터셋, 새로운 클래스, 새로운 task가 오면 보통 fine-tuning이 필요하다.

반면 SAM은 다음과 같이 학습된다.
이미지 + prompt -> prompt가 가리키는 valid mask

여기서 prompt는 point, box, mask, text 등이 될 수 있다. 위에 첨부한 Figure 1에서 (a)를 보면 알 수 있다.

즉, 모델이 배우는 것은 '고양이 클래스를 찾아라'가 아니라,

이 점이 가리키는 영역을 segment해라
이 box 안의 object를 segment해라
이 mask를 refine해라

와 같은 조건부 mask 생성 능력이다.

그래서 zero-shot이 가능해진 이유는

"새로운 task를 모델 head를 바꾸거나 재학습하지 않고, prompt를 설계해서 기존 Promptable segmentation task로 변환할 수 있기 때문이다."

근데, promptable하다고 자동으로 zero-shot이 보장되는 것은 아니다.
새 task를 prompt 형태로 잘 바꿀 수 있을 때 zero-shot transfer가 가능하다.

저자들은 SAM을 디자인할 때 광범위한 downstream application을 가능하게 할 만큼 충분히 일반적인 Promptable segmentation 작업을 정의한다 라고 말한다.

"광범위한 downstream application을 가능하게 한다"?

이것은 SAM을 특정 하나의 segmentation 모델로만 쓰는 것이 아닌, 여러 다른 vision task의 부품처럼 가져다 쓸 수 있다는 뜻이다.

여기서 downstream application은 "SAM을 학습한 다음에 실제로 적용하는 하위/응용 task들"을 말한다.

그럼 promptable segmentation은 무엇이지?
이것은 말 그래도 이미지와 prompt를 같이 입력하면, 그 prompt가 가리키는 대상의 segmentation mask를 출력하는 task. 인 것이다.

---

우리가 주목할 것은

2. Segment Ahything Task

이다. 그 전에, Introduction을 아주 짧게 요약하자면 다음과 같다.

Introduction 요약

SAM 논문은 NLP의 Foundation model처럼, segmentation에서도 다양한 task와 데이터 분포에 일반화되는 모델을 만들고자 한다. 이를 위해 저자들은 특정 class나 task에 고정된 segmentation이 아니라, point/box/text 같은 prompt를 받아 어떤 대상이든 mask로 반환하는 promptable segmentation task를 제안하고, 이를 학습하기 위한 모델 SAM과 대규모 데이터셋 SA-1B를 함께 구축했다.

Section 2: Segment Anything Task는 SAM 논문의 "문제 정의" 파트이다.

여기서 말하는 핵심은
"Segmentation을 특정 class를 맞히는 문제가 아니라, prompt가 주어졌을 때 그 prompt가 가리키는 valid mask를 출력하는 문제로 다시 정의하자" 이다.

Introduction에서 저자들은 NLP의 Foundation Model에서 영감을 얻었다고 말했다.
Segmentation에서도 이러한 모델을 만들고자, 유사한 기능을 가진 작업을 정의하는데, Task는 다음과 같다.

Task

먼저 SAM에서의 prompt들을 살펴보면 다음과 같다.
foreground / background point, rough box, rough mask, text, any information indicating what to segment. 그래서 task는 단순하다.

Image + prompt -> Valid segmentation mask
여기서 중요한 것이 바로 Valid mask이다.
논문의 저자들은 prompt가 애매할 수 있다고 본다. 예를 들어 한 점을 찍었는데 그 점이 사람의 셔츠 위라면, prompt가 의미하는 대상은 셔츠, 상반신, 사람 전체 모두 의미할 수 있다.

이때 SAM은 반드시 "정답 하나"만을 맞히는 게 아니라, 그중 하나라도 합리적인 mask를 내면 된다고 정의한다. 이게 SAM이 Ambiguity, 모호함을 다루는 출발점이다.

왜 이런 Task를 정의했을까?

저자들은 NLP의 prompt 방식을 segmentation으로 가져오고 싶어 한다.
언어모델이 다양한 질문에 prompt로 반응하듯, segmentation 모델도 다양한 segmentation 요구를 prompt로 저리하게 만들자는 의미이다.

즉, 기존 방식이 Image -> 정해진 class들의 segmentation이었다면, SAM은
Image + prompt -> prompt가 지시하는 대상의 segmentation 으로 바꾼 것이다.

Pre-training은 어떻게 연결되는가?

SAM은 학습할 때 interactive segmentation처럼 prompt sequence를 시뮬레이션 한다.

예를 들자면,

처음 point 또는 box prompt 제공
-> mask 예측
-> 이전 예측 mask와 GT mask가 다른 영역(Error Region)에 추가 point 제공
-> 다시 mask 예측

참고로 Error Region은 두 종류로, FP: Foregound Point로 추가, FN: Background point로 추가한다.

SAM 학습에서는 이전 예측과 GT가 불일치하는 Error Region에서 다음 point를 뽑는다. 그 point가 FN에 있으면 Foreground point, FP에 있으면 Background point로 사용한다.

다시 말해서, 이전 mask 예측이 GT와 다른 지점에 point를 찍어서, 그 위치가 target object에 포함되어야 하는지 아닌지를 prompt로 알려준다.

하지만, 기존 interactive segmentation과도 차이가 있다. 기존 방식은 "사람이 충분히 클릭하면 결국 맞추자"에 가깝고, SAM은 Prompt가 하나만 있어도 가능한 valid mask를 내야 한다는 쪽이다.

Zero-shot transfer와의 연결

이 Task 정의 덕분에 downstream task를 prompt로 바꿔서 풀 수 있다.

• Object proposal: Grid point prompt를 많이 넣기
• Instance segmentation: Detector가 만든 box를 prompt로 넣기
• Interactive segmentation: 사람의 click을 point promptfh sjgrl
• Mask refinement: Rough mask를 prompt로 넣기

즉, SAM은 각각의 task에 대해 새로 학습되는 것이 아닌, task를 prompt 형태로 바꿔서 수행한다. 이게 Section 2에서 말하는 zero-shot transfer의 핵심이다.

섹션의 Core idea

SAM은 segmentation을 class-specific prediction문제가 아니라 prompt-conditioned mask generation 문제로 재정의한다. 이때 prompt가 모호하더라도 하나 이상의 합리적인 valid mask를 출력하도록 task를 설계했기 때문에, 다양한 downstream segmentation task를 prompt enginnering으로 zero-shot 수행할 수 있다.

3. Segment Anything Model

이번 파트는 Sectino 2의 task 정의를 실제 모델 구조로 구현한 부분이다.

핵심을 먼저 말하자면

"SAM은 이미지를 한 번 무겁게 인코딩하여 image embedding을 만들고, 이후 point/box/mask/text prompt를 가볍게 인코딩해서 빠른 mask decoder로 mask를 뽑는는 prompt-conditioned segmentation model이다."

구조는 크게 3개이다.

• Image Encoder
• Prompt Encoder
• Mask Decoder

전체 흐름은 다음과 같이 표현할 수 있다.

image
 -> image encoder
 -> image embedding

prompt: point / box / mask / text
 -> prompt encoder
 -> prompt embedding

image embedding + prompt embedding
 -> mask decoder
 -> masks + predicted IoU score

Image Encoder

이미지를 feature로 바꾸는 부분이다.

SAM에서는 ViT 기반 Image Encoder를 사용한다. 논문에서는 MAE로 pre-trained된 Vision Transformer를 사용한다고 설명한다.

역할은 단순하다. 입력 이미지 -> dense image embedding

중요한 점은 image encoder가 무겁다라는 점이다. 대신 이미지를 한 번만 처리한다.

이미지 1장에 대해 Image Embedding을 미리 계산하여 여러 prompt가 들어와도 Image Encoder를 다시 돌리지 않는다.

이것이 바로 SAM이 Interactive하게 빠르게 동작할 수 있는 이유이다.
사용자가 점을 찍거나 box를 바꿀 때마다 전체 ViT를 다시 돌리는 것이 아닌, 이미 만들어둔 Image embedding에 prompt만 새로 넣어서 Decoder를 돌린다.

SAM의 general visual knowledge는 대부분 Image Encoder 안에 들어있다. prompt가 바뀌어도 Image Embedding은 고정된다.

Prompt Encoder

prompt를 Embedding으로 바꾸는 부분이다.

논문은 prompt를 크게 두 종류로 나눈다. Sparse prompt: point, box, text와 Dense prompt: mask.

point prompt는 위치 정보와 foreground/background 정보를 함께 가진다.
(Foreground point: 이 위치는 target object 안쪽이다. background point: 이 위치는 target object 바깥쪽이다.)

box prompt는 top-left corner와 bottom-right corner를 Embedding으로 표현한다.

mask prompt는 이미 spatial한 형태라서 convolution을 통해 image embedding과 같은 공간적 feature처럼 바꿔준다.

즉, prompt encoder의 역할은 사람이 준 클릭, 박스, rough mask 같은 입력을 mask decoder가 이해할 수 있는 token/embedding으로 바꿔주는 것이다.

여기서 중요한 건 Prompt가 class label이 아니라는 점이다.
dog, car, person 같은 class를 고르는 것이 아니라, 어디를 segment할지 알려주는 조건 정보인 것이다.

Mask Decoder

SAM에서 제일 중요한 부분이 바로 mask decoder이다.

역할은 image embedding + prompt embedding -> mask 로 말할 수 있겠다.

논문에서는 lightweight mask decoder라고 부르는데, image encoder에 비해 훨씬 가볍게 설계되어 있어서 그렇게 부른다. 이러한 장점 덕분에 prompt가 바뀔 때마다 빠르게 mask를 다시 만들 수 있다.

mask decoder의 내부는 Transformer decoder 계열 구조이다. 핵심 아이디어는 Image feature와 prompt feature가 attention으로 서로 정보를 주고받는 것이다.

쉽게 말하자면, prompt token은 "이 위치/박스/마스크가 관심 대상이야." 라고 알려주고, Image Embedding은 "이미지 전체의 시각 정보는 이렇게 생겼어" 라고 말해준다.

이 둘을 통해서 prompt가 가리키는 영역은 이미지에서 어디까지인가를 확인하고, mask foreground probability를 각 위치마다 예측한다.

Multi Masks: Ambiguity 해결

여기가 바로 SAM다운 부분이다.

한 prompt가 항상 명확하지는 않다. 예를 들어서 사람의 셔츠 위에 점 하나를 찍으면 가능한 답이 여러 개가 될 수 있기 때문이다.

만약 모델이 mask 하나만 출력하면, 여러 가능한 정답을 평균낸 이상한 mask를 만들 수 있다. 그래서 SAM은 Ambiguous prompt에 대해 여러 개의 valid mask를 출력하도록 설계된다.

논문에서는 기본적으로 3개 mask output이 충분하다고 본다.

• mask 1: part
• mask 2: object
• mask 3: whole object/context

정확히 이 세 이름으로 고정된 것은 아니지만, 직관적으로는 whole/part/subpart 구조를 커버하려는 것이다.

그리고 각 mask마다 score도 함께 출력한다.
mask candidates + predicted IoU scores
predicted IoU는 모델이 스스로 이 mask가 얼마나 좋을것 같은지 추정하는 confidence socre라고 보면 된다.

Predicted IoU?

이 predicted IoU의 평가 기준은 다음과 같다.
"이 mask가 GT mask와 얼마나 겹칠 것 같은지를 모델이 예측한 값"

즉, 기준은 일반적인 IoU이다. SAM은 학습할 때 GT mask를 알고 있으니까, 실제 IoU를 계산할 수 있다. 그래서 mask decoder가 mask 후보들을 만나면, 추가 Head가 각 mask에 대해 "이 mask의 실제 IoU가 대략 얼마일까?를 예측하도록 학습된다.

중요한 점은 Inference 때에는 GT가 없다는 것이다. 그래서 실제 IoU를 계산할 수 없는데, 대신 SAM이 학습해둔 IoU Prediction Head가 "품질 예상 점수"를 뽑는다.
mask 1: predicted IoU 0.91 \ mask 2: predicted IoU: 0.76.. 이런 식으로 말이다.

그래서 predicted IoU는 Class Confidence가 아니다.