💡Foundation model : Foundation models are models built on unlabeled data using self-supervision Label이 지정되지 않은 광범위한 dataset에 대해 훈련된 대규모 인공 지능 모델로, 광범위한 다운스트림 작업에 적용할 수 있는 AI 모델 유형
그래서 저자들은 promptable model을 개발하고 강력한 generalization이 가능하도록 광범위한 dataset으로 pre-train을 시도하려고 했다.
Foundation model을 통해 연구원들은 수많은 downstream segmentation problem을 새로운 data distribution에 대해 가능할 수 있도록 하려고 했다.
이 연구계획의 성공 여부에는 3가지 요소에 달려있다고 한다 : task, model, data
What task will enable zero-shot generalization?
What is the corresponding model architecture?
What data can power this task and model?
Defining a promptable segmentation task
General enough to provide a powerful pre-training objective and to enable a wide range of downstream applications
이 task는 interactive한 사용을 위해 flexible prompting이 가능하고 real-time으로 segmentation mask들을 output으로 내보낼 수 있는 모델을 필요로 한다.
그러나 segmentation에는 web-scale의 dataset이 없기 때문에 연구원들은 “data engine”이라는 것을 만들었다.
Segment Anything Task
💡Foundation models are a promising development that can perform zero-shot and few-shot learning for new datasets and tasks often by using “prompting” techniques
Promptable segmentation task : To return a valid segmentation mask given any segmentation prompt
Pre-training
연구원들의 목표는 어떠한 prompt가 주어지더라도 심지어 prompt가 ambiguous 할 때도 valid mask를 예측하도록 만드는 것이다.
이로 인해 pre-trained model이 ambiguity가 있을 때 효과적이다 라는 것이 보장되게 된다.
Zero-shot transfer
SAM의 pre-training task가 model이 inference하는 동안 어떠한 prompt에 대해서도 적절하게 반응하는 능력을 부여해 주었고, 결과적으로 downstream task들이 engineering prompt들에 의해 해결될 수 있게 되었다.
Related tasks
Interactive Segmentation
Edge Detection
Super Pixelization
Object Proposal Generation
Foreground Segmentation
Semantic Segmentation
Instance Segmentation
Panoptic Segmentation
SAM의 promptable segmentation task의 궁극적인 목표는prompt engineering을 통해 기존 또는 새로운 segmentation task에 적용할 수 있는 foundation model을 만드는 것이다. [Task Generalization]
Segment Anything Model
The model must support flexible prompts, need to compute masks in amortized real-time to allow interactive use, and must be ambiguity-aware
Three constraints of simple design
A powerful image encoder computes an image embedding
A prompt encoder embeds prompts
Two information sources are combined in a lightweight mask decoder that predicts segmentation masks
Image encoder
MAE pre-trained ViT를 사용해서 최소한으로 process에 high-resolution input을 적용하려고 함. Image encoder의 경우 image당 1번 적용 되며, model을 prompting하기 위해 prior로 적용될 수 있다.
output = C X H X W image embedding
Prompt encoder
Prompt의 경우 sparse[points, boxes, text]와 dense [masks]로 나뉜다. 저자들은 point와 box들을 positional encoding들로 표현했고 CLIP의 text encoder로 학습한 embedding들을 합쳤다.
Dense prompt들은 convolution으로 embedding되었고 image embedding과 함께 element-wise sum으로 합쳐졌다.
Point는 point의 위치와 point가 foreground인지 background인지를 나타내는 2개의 학습된 embeddinge들의 positional encoding의 합으로 표현된다.
Box는 embedding pair로 표현이 된다.
Top-left corner positional encoding
Bottom right corner positional encoding
free-form text를 표현하기 위해 CLIP으로부터 text encoder를 사용한다.
Mask decoder
Mask decoder는 효율적으로 image embedding, prompt embedding, output token들을 mask에 mapping 시킨다. decoder의 design은 DETR에 영감을 받았고, Transformer의 decoder block을 dynamic mask prediction head에 이어 수정해서 만들었다고 한다.
수정된 decoder block은 prompt self-attention과 cross-attention을 두 방향으로 사용하면서 모든 embedding들을 update한다.
2 block들을 진행한 뒤 image embedding을 upsampling하고 MLP가 output token을 dynamic linear classifier에 mapping 시키는데 그 이후에 각각의 image 위치에서 mask foreground probability를 계산한다.
Resolving ambiguity
만약 ambiguous prompt가 주어지게 되면 모델은 하나의 output에 대해 여러개의 valid mask들을 생성하게 될 것이다.
이를 해결하기 위해 single prompt에 대한 여러 output mask들을 예측하도록 model을 수정했다. 그 과정으로 가장 흔한 case를 해결하기 위해 3개의 mask output이 충분하다는 것을 발견했다.
Training 동안 mask에 대해 최소한의 backpropagation 연산을 진행했고, mask들의 순위를 매기기 위해 각각의 mask에 대해 confidence score [estimated IOU]를 예측하게 만들었다고 한다.
Efficiency
전체적인 model의 design은 efficiency가 중심이 된다. 주어진 precomputed image embedding에 대해 prompt encoder와 mask encoder는 CPU의 web browser에서 동작을 하고 ~50ms가 걸린다고 한다. 이 runtime 성능은 seamless, real-time interactive prompting을 가능하게 해준다.
Losses and training
Focal loss와 Dice loss의 linear combination으로 mask prediction을 진행하였고, geometric prompt들의 mixture를 사용해서 promptable segmentation task를 위해 training을 시킴
Data engine
새로운 data distribution에 강력한 generalization을 보이기 위해 SAM을 크고 다양한 mask set에 대해 training을 시켜야 겠다고 생각함. 이에 대한 해결책이 “data engine”을 만드는 것.
Three stages of data engine
Assisted-manual : SAM assists annotators in annotating masks, similar to a classic interactive segmentation setup
Semi-automatic : SAM can automatically generate masks for a subset of objects by prompting it with likely object locations and annotators focus on annotating the remaining objects, helping increase mask diversity
Aim to increase the diversity of masks in order to improve SAM’s ability to segment anything
Fully automatic : Prompt SAM with a regular grid of foreground points, yielding on average ~100 high-quality masks per image
Annotation was fully automatic
Collect enough masks to greatly improve the model $\rightarrow$ developed the ambiguity-aware model
Apply NMS [Non-Maximal Suppression] to filter duplicates
Process multiple overlapping zoomed-in image crops
SA-1B : data engine의 final stage를 사용해서 fully automatic하게 data를 수집해서 기존의 segmentation dataset들보다 400배 이상 많은 mask들을 얻을 수 있었음
이로 인해 SAM을 training하는 것이 robust하고 general 해짐
Experiments
Extensive use
SAM produces high-quality masks from a single foreground point
Consistently strong quantitative and qualitative results on a variety of downstream tasks under a zero-shot transfer protocol using prompt engineering
Discussion
Foundation models
Pre-trained models have been adapted to downstream tasks since the early days of machine learning. This paradigm has become increasingly important in recent years with a growing emphasis on scale, and such models have recently been [re-]branded as “foundation models”
즉, Foundation Model = Broad data로 training되어 다양한 downstream task에 적용할 수 있는 모델
SAM의 경우 self-supervised technique [MAE]로 시작이 되었지만, 대부분의 SAM의 capability는 large-scale supervised training으로부터 온다.
Compositionality
Pre-trained model들은 training의 순간에도 능력을 발휘할 수 있다.
이 논문의 목표는 이런 종류의 SAM과의 직접적인 composition을 만드는 것이다. 이를 위해 SAM이 다양한 segmentation prompt들에 대해 valid한 mask를 예측할 수 있도록 하는 것을 타겟으로 삼았다고 한다.
결과적으로 SAM이 ego-centric image와 같은 새로운 domain에 잘 일반화하는 능력 덕분에 추가적인 training 없이도 잘 동작하게 되었다고 한다.
Limitations
SAM이 일반적으로 잘 동작은 하지만, 단점도 존재한다.
SAM의 경우 fine structure를 놓칠 수 있고 때때로 small disconnected component들을 환각[hallucinate]할 수도 있다 또한 zoom-in 같이 computationally intensive 한 방식으로 boundary를 만들어내지 못할 수 있다.
일반적으로 많은 point들이 주어져있을 때 SAM을 뛰어넘는 interactive segmentation 방식을 기대하지만, 이런 방식과 달리 SAM은 generality와 광범위한 사용을 목표로 설계가 되어 있다.
SAM이 많은 task들에 동작할 수 있지만, 어떻게 semantic, panoptic segmentation에 동작하는지는 불분명하다. 결론적으로 domain-specific한 방식들이 각각의 domain에서 SAM을 능가하기를 기대하게 되는 것이다.
Conclusion
SAM은 image segmentation을 foundation model의 수준으로 높여준 연구이다.
이 논문을 통해 task [promptable segmentation], model [SAM], dataset [SA-1B]를 등장시켜서 이를 가능하게 했다.
