[논문리뷰] SA(Segment Anything)

본 게시물에서는 2023년도 Meta AI 연구진이 발표한 SA(Segment Anything) 이라는 모델에 대해 리뷰를 하고자 합니다.

 

💡 Background

Segment Anything Demo

 

ChatGPT와 같은 대용량 언어 모델(LLM)은 strong zero-shot(이전 학습 내용을 기반으로 새로운 Task 및 데이터 예측) 및 few-shot(몇 개의 샘플을 통해 새로운 Task에 적용) 을 통해 NLP분야에서 새로운 패러다임을 이끌어 가고있습니다. 이러한 "Foundation Model"은 새로운 데이터 분포에서 다양한 Task를 수행할 수 있으며, 이를 기반으로 hand-crafted 입력을 즉각적으로 출력할 수 있는 Prompt Engineering이 가능하다는 것이 가장 큰 장점입니다. 

 

컴퓨터 비전의 모델은 주로 이미지만을 학습하여 성능을 향상 시켜왔습니다. 하지만 이러한 접근 방식은 모델의 강건함과 일반화의 한계라는 단점을 가지고 있습니다. 이에 CLIP과 ALIGN 알고리즘을 통해 웹 상의 텍스트와 이미지를 학습하고 새로운 개념 및 데이터 분포에 대한 Zero-Shot 일반화가 가능하다는 장점을 가지고 있어, 이를 컴퓨터 비전 초기의 Foundation Model이라고 칭했습니다. 이처럼 컴퓨터 비전에서 많은 진전이 있었지만, 방대한 학습 데이터가 존재하지 않다는 문제점이 있었습니다. 

Segment Anything (Task/Model/Data)

 

이에 본 논문의 저자는 대용량 언어 모델과 비슷한 학습 방법으로 Image Segmetation 전용의 Foundation Model를 구축하는 것이 목표입니다. 하지만 이러한 목표를 이루기 위해서는 3가지의 문제점을 해결해야 했습니다. 

1. What Task will enable zero-shot generalization? 
2. What is the corresponding Model architecture? 
3. What Data can power this taks and model?

본 논문의 저자들은 복잡하고 다양한 문제에 대응하기 위해 강력한 사전 학습 목적 함수를 제공하고 하위 어플리케이션을 활성화할 수 있는 Promptable Segmentation Task를 정의하는 것으로 시작합니다. 이 task에서는 flexible prompt를 지원하고 상호 작용하도록 메시지가 표시될 때 segmentation mask를 실시간으로 출력할 수 있는 모델이 필요합니다. 하지만 모델을 학습시킬려면 다양하고 방대한 양의 데이터가 필요합니다만, 안타깝게도 segmentation를 위한 웹 규모 데이터가 존재하지 않습니다. 이에 11M개의 초고화질 이미지와 1.1B개의 segmentation mask를 포함한 "데이터 엔진"을 구축하여 문제를 해결하였습니다. 이러한 접근을 통해 다양하고 방대한 데이터 소스를 학습시키는 것을 반복하여, 효율적인 모델을 구축하게 됩니다. 

---

💡 Segment Anything Task

컴퓨터 비전과 NLP 분야에서 prompting 기술을 기반으로 한 few-shot과 zero-shot을 학습은 최근에 큰 각광을 받고 있습니다. 이러한 방법은 작은 양의 레이블 또는 아예 레이블이 없는 데이터로도 모델을 효과적으로 학습이 가능하며, 다양한 하위 task를 해결할 수 있게 되었습니다. 이에 본 논문의 저자들은 다음 토큰 예측 task가 foundation model 사전 학습에 사용이 되는점을 고려하여, Segmentation을 위한 foundation model 구축을 목표로 삼았습니다. 

 

이에 본 논문의 저자는 Promptable Segmentation Task는 어떠한 프롬프트(Box, Point, mask)가 주어지면 타당한 segmentation mask를 반환할 수 있다고 합니다. 여기에서 타당하다는 것은 프롬프트가 애매모호하여 여러 객체에 대해 검출하였을 때, 그 객체 중 하나에 대한 합리적인 마스크여야 함을 의미합니다. 이 요구 사항은 모호한 프롬프트에 대해 일관된 응답을 출력하는 언어 모델을 기대하는 것과 유사합니다.

 

또한, pre-training을 통해서 어떠한 prompt가 입력되어도 유효한 마스크를 생성할 수 있는 사전 학습 모델이 있기 때문에, zero-shot transfer이 가능합니다. 예를 들어, 고양이에 대한 boundary box가 존재할 경우, 개별적으로 학습하지 않고 출력 값을 모델에 프롬프트로 제공하여 고양이 instance segmentation 가능하다고 합니다. 

---

💡 Segment Anything Model

Segment Anything Model은 입력 이미지와 Prompt에 대한 특징을 추출할 수 있는 두 개의 인코더와 이를 통해 도출된 표현벡터를 결합해 원하는 객체만 추출할 수 있는 디코더로 구성되어 있습니다. 또한, 본 논문의 저자들은 실시간 성능에 대한 특정 trade-off가 있는 Transformer Vision 기반의 모델을 구축하였습니다. 

 

📌 모델 구성 

 

Image Encoder : MAE라는 사전 학습 방식을 사용해 학습된 Encoder 사용하며, 이미지당 한번 실행되며 모델을 프롬프트하기 전에 적용할 수 있습니다. 

Prompt Encoder : 종류는 크게 Sparse(Box, Point, Text)와 Dense(Mask)으로 나뉩니다. 
① box와 point - Positional Encoding (해당 점의 위치 + 피사체와 배경을 구별하게 학습된 임베딩) 추출
② text - CLIP의 text encoder 활용 
③ mask - convolution 연산 후, 이미지 임베딩과 element-wise로 더함

Mask Decoder : Image/Prompt embedding과 Ouput token를 받아서 mask를 생성하게 됩니다. 전체 Embedding을 업데이트 하기 위해 양방향으로 self-attention과 cross attention을 적용하였습니다. 이후 MLP를 통해 Masking probability를 반환하게 됩니다.

 

📌 모델 설정

 

Resolving Ambiguity : 3개의 prediction mask(whole/part/subpart)을 통해 애매모호한 프롬프트를 처리할 수 있습니다. 학습 중에는 마스크 중 Minimum Loss만 backprop하며, 신뢰도 점수(ex. estimated IoU)를 기반으로 마스크 순위를 매깁니다. 

Losses and Training : Focal Loss와 Dice Loss를 linear하게 합쳐서 사용하였습니다. 불균형한 데이터셋에서 소수 클래스를 더 잘 예측하면서도 전체 세그멘테이션의 정확도를 유지하거나 개선할 수 있습니다. 
      * Focal Loss: Cross Entropy의 클래스 불균형 문제를 개선한 개념이며, 쉽게 오분류되는 케이스에 더 많은 가중치 
      * Dice Loss : 두 개의 영역에 대해 얼마나 겹치는 나타나는 지표 (IoU) 

---

💡 Segment Anything Data Engine

Image Segmentation을 위한 데이터가 존재 하지않아, SA-1B 라는 데이터 엔진을 구축하였습니다. 이는 세 단계로 구성됩니다.

 

Assitsted-manual stage : Interative segmentation 도구(erase/brush)을 통해 전경/배경 개체 지점을 클릭하여 마스크를 생성하는 단계입니다. 개체에 레이블을 지정하는데 제약을 부과하지 않았으며, stuff와 things를 모두 자유롭게 레이블을 지정하게 하였습니다. 눈에 띄는 순서대로 개체에 레이블을 지정하였으며, 마스크를 주석하는데 30초 이상이 걸리면 다음 이미지로 넘어가도록 하였습니다. 이에 Common Public 데이터 셋을 사용을 기반으로 학습하였습니다. 또한 충분하게 주석이 되었다면, 새로운 주석이 달린 마스크를 SAM을 기반으로 재학습 되었습니다. 점차적으로 방대한 양의 데이터가 수집되면서, 이미지 인코다가 ViT-B에서 ViT-H로 확장되었습니다. 이렇게 총 6번을 재학습 하였으며, 120K 이미지에서 4.3M 마스크가 생성되었다고 합니다. 

 

Semi-automatic stage: 이 단계에서는 SAM의 전반적인 성능을 향상시키기 위해 눈에 잘 띄지 않는 개체에 집중을 합니다. 주석가에게 SAM으로 어느정도 레이블링이 되어있는 이미지를 제공하고, annotataion이 없다면 추가하는 작업을 진행하게 됩니다. 신뢰할 수 있는 마스크를 감지하기 위해 첫번째 단계 마스크에서 Faster-RCNN을 학습하였습니다. 이에 180K 이미지에서 5.9M 마스크를 추가적으로 생성하였습니다.

Fully automatic stage : 위의 두과정을 통해 모델을 크게 개선할 수 있는 충분한 마스크를 수집하였으며, 이번 단계에서는 모든 것을 모델이 예측한 mask로 학습하는 단계입니다. 32x32 grid point에 point prompt를 입력으로 넣어 각 포인트에 대해 valid object를 segmentation를 진행하였으며, 이를 계기로 ambiguity-aware model를 만드는데 도움이 되었다고 합니다. 이후 subpart, part, whole 객체의 마스크를 출력하도록 하였습니다. IoU prediction module를 사용하여 stable mask를 선택하였으며, NMS를 통해 최종적인 마스크를 출력하였습니다.