[논문 리뷰] (Meta AI Research, FAIR) Segment Anything

https://arxiv.org/abs/2304.02643

Segment Anything

 

Segment Anything은 컴퓨터 비전 그 중에서도 Segmentation Task를 위한 Foundation Model 이다.

비슷한 예시로는 CLIP과 ALIGN으로 텍스트와 이미지 서로 다른 두 모달리티를 정렬하도록 학습된 후

이를 Engineered text prompt로 zero-shot이 가능하며 이미지 생성과 같은 Downstream Task에 사용할 수 있다.

 

Segment Anything Model(SAM)은 조립부품같은 역할로써 여러 Downstream Task의 일부로 사용할 수 있는 Foundation Model이라고 생각하면 된다.

 

CLIP(Contrastive Language–Image Pre-training): OpenAI에서 개발한 모델로, 인터넷에서 수집한 이미지와 해당 이미지에 대한 캡션(설명 문장)을 사용해 학습. 이미지 인코더와 텍스트 인코더를 각각 학습해서, 주어진 이미지와 문장이 서로 얼마나 잘 일치하는지를 판단하는데 사용.


ALIGN (A Larger-Scale Image and Noisy-text pre-training): CLIP과 비슷하지만 훨씬 더 큰 데이터셋으로 학습. 
노이즈가 있는(불완전한) 데이터도 적극적으로 활용.

 

특히 3가지 핵심 요소를 목표로 두었다.

1. Zero-shot generalization을 가능하게 할 작업은 무엇인가?

2. 이에 적합한 모델 아키텍처는 무엇인가?

3. 이 작업과 모델을 학습하기 위해 필요한 데이터는 무엇인가?

 

이를 위해 사전학습으로 다양한 Downstream Task 지원할 수 있는 Promptable Segmentation Task를 정의했다.

Promptable Segmentation Task는 Prompt가 주어졌을 때 이에 맞는 Segmentation Mask를 출력하여 상호작용적 사을 가능하게 해주는 Task이다.

 

그러므로 사전학습을 위해 대규모 데이터 소스가 필요하고 이를 위해 데이터 엔진을 구축하였다.

 

Task

 

Promptable Segmentation Task는 어떠한 segmentation prompt가 주어지더라도 유효한 segmentation mask를 반환하는 것을 목표로 한다.

 

이때, Prompt는 단순히 이미지에서 어떤 대상을 분할할 것인지를 지정하는 역할을 한다.

Prompt는 공간 정보(spatial information)나 텍스트 정보(text information)를 포함하여 특정 대상을 지정할 수 있는데

이때, 반드시 유효한 segmentation mask를 출력해야한다.

 

유효한 segmentation mask는 prompt가 모호하거나 여러 객체를 참조할 수 있는 경우에도 적어도 하나의 객체에 대해 타당한 mask를 출력해야 한다는 의미다.

예를들어 셔츠 위에 포인트가 지정(공간 정보)되면, 셔츠일 수도 있고 입고 있는 사람일 수도 있지만, 이 중 하나에 대해 합리적인 mask를 출력해야 한다.

 

그리고 이렇게 출력한 mask를 정답과 비교하는 방식으로 사전학습을 한다. (Supervised)

 

이렇게 사전학습을 한 모델은 추론시 어떠한 prompt에도 적절하게 반응을 할 수 있으며,

적절한 prompt를 설계함으로써 다양한 Downstream Task에 적용할 수 있다.

 

가능한 Task로는

 

• Interactive Segmentation → 사용자와의 상호작용(포인트, 박스, 클릭)을 통해 Segmentation
• Edge Detection → 객체의 경계선 검출
• Super Pixelization → 유사한 픽셀을 작은 영역으로 묶음
• Object Proposal Generation → 객체가 있을 법한 영역을 제안
• Foreground Segmentation → 관심 있는 객체만을 남기고, 배경은 제거
• Semantic Segmentation → 픽셀을 클래스별로 분류 (객체 개별 구분 X)
• Instance Segmentation → 픽셀을 클래스별로 분류하고, 객체 개별 구분
• Panoptic Segmentation → 객체 개별 구분 + 배경 영역 통합

 

등이 있으며 다른 Task에도 일부 구성요소로 사용이 가능하다.

핵심은 더 큰 시스템의 구성 요소로 작동하면서 새로운 작업을 처리하는 것이 가능하다.

예를 들어, 기존 객체 탐지기(object detector)와 promptable segmentation 모델을 결합하면 instance segmentation 작업을 수행할 수 있다.

Model

 

모델은 반드시 3가지 조건을 만족해야한다.

 

• 유연한 prompt를 지원해야 하고,
• 상호작용적 사용을 위해 실시간으로 마스크를 계산해야 하며,
• 모호성(ambiguity)을 인식할 수 있어야 한다.

논문에서는 간단한 모델 설계로 이 조건을 만족했다.

 

• Image Encoder: 강력한 이미지 인코더가 이미지 임베딩(image embedding)을 생성한다.
• Prompt Encoder: Prompt 인코더는 prompt 정보를 임베딩한다.
• Mask Decoder: 이 두 정보를 경량화된 mask decoder에서 결합하여 segmentation mask를 예측한다.

이 모델을 Segment Anything Model(SAM)이라고 부른다.

그리고 이미지를 한번 임베딩하면 다양한 prompt에 반복 재사용이 가능해서 계산비용을 줄일 수 있다.

 

SAM은 주로 포인트(point), 박스(box), 마스크(mask) prompt에 중점을 두고 있고, 자유형 텍스트 프롬프트(free-form text prompt)에 대한 초기 결과도 제시한다.

 

또한, 유효한 segmentation mask를 출력하기 위해 모호성을 인식하도록 1개의 prompt에 대해 여러개의 mask를 예측하도록 만들었다.

 

1. Image Encoder

이미지 인코더는 CxHxW 형태의 이미지 임베딩을 출력하는 네트워크라면 어떤 것이든 사용이 가능하다.

본 논문에서는 MAE(Masked Autoencoder)로 사전 학습된 Vision Transformer(ViT)를 사용했다.

이미지당 한 번만 작동하며, prompt가 입력되기 전에 사전 처리로 작동된다.

 

구체적으로 Exploring plain vision transformer backbones for object detection. ECCV, 2022을 따라 14x14 windowed attention과 4개의 equally-spaced global attention을 가진 ViT-H/16을 사용했다.

 

이미지 인코더의 출력은 입력 이미지의 16x downscaled 임베딩이다.

 

전처리는 Simple copy-paste is a strong data augmentation method for instance segmentation. CVPR, 2021을 따라 짧은 쪽을 패딩하여 1024x1024로 resize를 했다.

따라서, 이미지 임베딩의 크기는 16배 작은 64x64 크기가 된다.

 

이후 채널 수를 줄이기 위해 Exploring plain vision transformer backbones for object detection. ECCV, 2022을 따라 1x1 컨볼루션으로 채널 수를 256개로 줄이고 채널 수를 유지한채 3x3  컨볼루션을 한번 더 적용했다.

각 컨볼루션 뒤에는 Layer Normalization을 적용했다.

 

2. Prompt Encoder

사용자의 프롬프트 정보를 임베딩하는 모델로,

프롬프트로는 Sparse Prompts(포인트, 박스, 자유형 텍스트), Dense Prompts(마스크)가 있다.

 

Sparse Prompts는 256차원 벡터 임베딩으로 변환된다.

포인트는 Fourier features let networks learn high frequency functions in low dimensional domains. NeurIPS, 2020을 따라 포인트 위치에 대한 positional encoding과 foreground 또는 background를 나타내는 학습된 임베딩을 더한다.

 

박스는 2개의 임베딩 쌍으로 표현된다. top-left corner의 positional encoding과 top-left corner를 나타내는 학습된 임베딩을 더한다. 마찬가지로 bottom-right corner도 positional encoding과 학습된 임베딩을 더한다.

 

자유형 텍스트는 CLIP의 텍스트 인코더로 임베딩했다.

 

Dense Prompts는 이미지와 spatial correspondence를 가진다.

입력 이미지보다 4배 낮은 해상도로 mask를 입력했고, 이후 2x2 stride-2 컨볼루션을 2번 적용하여 4채널, 16채널 순으로 한 번 더 4배 축소했다. 마지막으로 1x1 컨볼루션을 통해 채널 수를 256으로 맞추었다.

각 층은 GELU 활성화 함수와 Layer Normalization을 적용하여 구분했다.

mask와 이미지 임베딩은 요소별로 더했다.

만일, mask prompt가 없는 경우 "no mask"를 나타내는 학습된 임베딩을 각 이미지 임베딩 위치에 추가했다.

 

3. Lightweight Mask Decoder

 

마스크 디코더는 이미지 임베딩, 프롬프트 임베딩 세트를 출력 마스크로 매핑한다.

DETR과 Perpixel classification is not all you need for semantic segmentation. NeurIPS, 2021로 부터 영감을 받아 표준 Transformer 디코더 블록을 수정했다.

 

디코더 적용 전에 프롬프트 임베딩 세트에 학습된 출력 토큰 임베딩을 삽입했다. (ViT의 [class] 토큰과 유사한 역할)

이 출력 토큰 하나가 하나의 mask를 예측한다고 보면된다.

여기서부터는 이미지 임베딩을 제외한 모든 임베딩을 토큰이라고 부른다.

 

각 디코더 레이어는 4가지 단계를 수행한다.

1. Self-attention on the tokens
토큰에 대한 self-attention
2. Cross-attention from tokens (as queries) to the image embedding
토큰을 쿼리로 사용하여 이미지 임베딩에 대한 cross-attention
3. A point-wise MLP updates each token
point-wise MLP를 사용해 각 토큰을 업데이트
4. Cross-attention from the image embedding (as queries) to tokens
이미지 임베딩을 쿼리로 사용하여 토큰에 대한 교차 어텐션 수행

4단계에서는 이미지 임베딩이 프롬프트 정보를 통해 업데이트된다.

 

cross-attention에서 이미지 임베딩은 64x64 크기의 256차원 벡터 집합으로 처리되고,각 self/cross-attention과 MLP는 residual connection, layer normalization, dropout을 포함한다.

 

다음 디코더 레이어는 이전 레이어에서 업데이트된 토큰과 이미지 임베딩을 입력으로 사용한다.논문에서는 2개의 레이어를 사용했다.

 

attention layer마다 업데이트된 토큰에 원래 프롬프트 토큰(positional encoding 포함)을 다시 더했다.

 

디코더를 통과한 후 2개의 Transposed 컨볼루션 레이어를 사용하여 이미지 임베딩을 4배 업샘플링하였다.

그리고 토큰을 한번 더 이미지와 cross-attention을 수행했고,

이렇게 업데이트된 출력 토큰 임베딩을 3-layer MLP를 거쳐 업샘플링된 이미지 임베딩과 동일한 채널 차원을 갖는 벡터를 출력한다.

마지막으로, 업샘플된 이미지 임베딩과 MLP 출력 벡터의 point-wise product를 통해 최종 mask를 예측했다.

 

모호성 문제를 해결하기 위해 하나의 프롬프트가 3개의 mask(전체, 부분, 하위 부분)를 출력하도록 수정하였다.

학습 중에는 가장 작은 손실의 mask만으로 역전파를 수행했다.

실제로는 예측된 mask의 순위를 매기기 위해, 출력 토큰에 작은 head를 추가해서 각 예측된 마스크와 해당 객체 간의 IoU를 예측하도록 설계했다. (이때 , IOU loss는 MSE 사용)

다만, 다중 프롬프트가 주어질 경우 모호성이 거의 발생하지 않기 때문에 불필요한 손실 계산을 줄이기 위해 단일 mask만 예측하도록했다. 이를 위해, 추가적인 mask 예측을 위한 4번째 출력 토큰을 추가했다. (다중 프롬프트에서만 출력됨)

이부분을 요약하면 이전에 디코더 적용 전에 프롬프트 임베딩 세트에 학습된 출력 토큰 임베딩을 삽입했다고 했는데

여기서 1, 2, 3번 토큰은 모호성을 해결하기 위해 3개의 mask를 예측하고 4번 토큰은 1개의 명확한 mask만 예측한다.

이때, 다중 프롬프트면 1, 2, 3번은 사용하지 않고 4번 토큰만 사용한다.

1, 2, 3번은 단일 프롬프트에서 모호성 해결용이고 4번은 다중 프롬프트에서 효율성을 높이기 위한 단일 mask 전용이다.

 

Loss는 focal loss와 dice loss의 선형 결합을 사용하였다.

Data Engine & Dataset

 

그리고 대규모이면서도 다양한 Mask 데이터로 학습을 하기 위해 데이터 엔진을 구축했다.

모델과 데이터셋 주석을 동시에 개발하는 방식으로 3단계로 구성된다.

1. Assisted-Manual (보조 수작업 단계)
2. Semi-Automatic (반자동 단계)
3. Fully Automatic (완전 자동 단계)

보조 수작업 단계에서 SAM은 annotator(주석자)를 보조하여 segmentation mask를 주석하도록 지원한다.

주석자는 SAM이 제안한 Mask를 검토하고 수정한다.

(기존의 Interactive Segmentation 방식과 유사함)

 

반자동 단계에서 SAM은 객체의 위치가 유력한 영역에 대해 자동으로 mask를 생성한다.

주석자는 남은 객체만을 수동으로 주석하여 mask의 다양성을 높인다.

 

완전 자동 단계에서는 SAM에 대해 32x32 정규 격자 방식으로 foreground 포인트를 prompt로 제공하여,

각 포인트에 대해 자동으로 mask를 생성한다. (총 32x32=1,024개의 포인트)

1개의 포인트가 부분, 세부에 위치할 경우 전체, 부분, 세부의 mask 3개를 모두 출력한다.

이후 신뢰도가 가장 높은 mask를 선택하고,

NMS(Non-Maximal Suppression)을 통해 중복된 마스크를 제거했다.

 

이렇게 구축한 최종 데이터셋을 SA-1B라고 부른다.

SA-1B는 데이터 엔진의 마지막 완전 자동 단계를 통해 수집했고, 지리적, 경제적으로 다양한 국가에서 수집했다.

이미지 수는 11M(1,100 만 장)에 mask 수는 11.B (11억 개)이다.

 

특히, 고해상도 이미지로 평균 3300x4950 크기이다. (실제 공개에는 저장 공간때문에 다운샘플링함)

그리고 99.1%는 완전 자동으로 생성된 마스크이다.

Conclusion

Segment Anything 프로젝트는 이미지 세분화를 Foundation Model의 시대로 끌어올리려는 시도이다.

주요 기여는 새로운 작업(promptable segmentation), 모델(SAM), 데이터셋(SA-1B)이다.

SAM이 진정한 Foundation Model로 자리 잡을 수 있을지는 커뮤니티에서의 활용에 달려 있다.

그럼에도 불구하고, 본 연구의 관점, 10억 개 이상의 마스크 공개, 프롬프트 세분화 모델이 향후 발전의 토대가 되기를 기대한다.