[논문리뷰] Segment Anything

이번 글에선, ICCV 2023에 소개된 meta의 "Segment Anything" 논문에 대해 리뷰할 예정이다.

SAM(Segment Anything Model)이라고 불리며, 아래와 같이 엄청난 성능으로 모두를 놀라게 한 논문으로 유명하다.

 

리뷰 순서는 아래와 같다.

 

1. Introduction

2. Segment Anything Task

3. Segment Anything Model

    - 3.1 Image encoder

    - 3.2 Prompt encoder

    - 3.3 Mask decoder 

    - 3.4 Losses and training

4. Segment Anything Dataset

5. Experiments

6. Conclusion

 

그림 1. SAM을 통한 output.

 

1. Introduction

방대한 양의 데이터셋으로 학습한 LLM의 등장으로, zero-shot과 few-shot generalization 등 기존에 어려웠던 task를 해결하고 있다. 이처럼, large dataset으로 학습한 거대 모델로 보지 못한 task 및 data 분포를 일반화하는 모델을 이른바 "foundation model"이라고 칭한다. Foundation model은 방대한 양의 데이터셋의 영향도 있지만, prompt engineering이 차지하는 역할이 굉장히 크다. 그중, hand-crafted text (특정 task를 염두하고 설계.)를 이용하여 주어진 작업에 대한 올바른 응답을 유도하는 과정은 강력한 성능을 내는데 큰 역할을 한다. (즉, 학습 후 downstream task를 처리하는 BERT 같은 모델이 우수한 성능을 낼 수밖에 없다는 의미로 해석됨.) 

결국, 본 논문의 목적은 언어 모델과 같이 prompt를 입력받아, zero-shot 상황에서도 강건한 foundation model for image segmentation을 구축하는 것을 목표로 함.

 

2. Segment Anything Task

본 논문에서 다루는 task는 NLP에서 영감을 받아, prompt를 통한 segmentation task를 다룬다.

 

그림 2. Promptable segmentation task의 예시.

언어 모델이 "애매한 프롬프트"에 대해 일관된 응답을 출력해야하는 것과 비슷하게, 적어도 하나의 object에 대해 의미있는 mask를 출력해야 한다. 이러한 task를 채택한 이유는 다음과 같다.

1) natural pre-training algorithm : 특정한 방법을 강요하지 않고도, image segmentation model을 효과적으로 학습시키는 데 적합한 알고리즘을 제안하기 위해서.

2) general method for zero-shot transfer : 모델이 사전 학습에서 보지 못한 새로운 image segmentation task에도 추가적인 학습 없이도 적절히 대응할 수 있게 된다.

 

그림 3. Prompt에 따른 valid output.

즉, SAM task는 그림 3 과 같이, 각 prompt (해당 그림에선 point)에서 valid한 output을 출력할 수 있어야 한다. 여기서 valid output이란, 의미론적으로 유효한 output을 의미한다. 예를 들어, 첫 번째 컬럼의 prompt를 살펴보면, 타조 전체, 타조 몸통, 타조 머리 모두 의미론적으로 유효하다고 해석할 수 있다. 

 

3. Segment Anything Model

그림 4. Segment Anything Model (SAM) overview.

 

Model은 그림 4 와 같이 구성되며, image encoder, prompt encoder, mask decoder 3가지 components가 존재한다. 

 

3.1 Image encoder

Image encoder는 masked auto encoder (MAE)를 사용하여 학습된 ViT의 image encoder를 이용하였다. 여기서 MAE란, input 일부에 mask를 적용하여 mask를 복원하면서 image encoder의 표현력을 높이기 위해 많이 사용된다. 

 

그림 5. MAE 학습 방법의 예시.

 

조금 더 자세하게 설명하면, ViT의 경우 image를 patch 단위로 나누어 input으로 사용하게 된다. 여기서, patch에 mask를 씌우고, 해당 image token을 encoder의 input으로 사용한다. 이후 decoder에 image embedding을 넣어 원본을 복원하는 방식으로 iamge encoder를 학습하게 된다. 이러한 과정은, mask로 인해 제한된 정보를 통해 image encoder가 더 높은 표현력을 가지도록 유도했다고 해석할 수 있다.

 

3.2 Prompt encoder

여기서 prompt는 두 가지 set으로 나뉜다. Sparse (points, boxes, text)와 dense (masks)로 나뉘게 된다. 간략히 말해서, 간소화된 정보를 가진 points, boxes, text와 보다 확장된 범위의 정보를 가진 masks 두 가지 prompt에 대해 고려하게 된다. 

 

그림 6. 각 prompt에 대한 encoding 예시.

 

그림 6 과 같이, dense prompt의 경우, 일반적인 conv filter를 통과하여 image embedding과 더해지는 방식을 사용한다. 반면, sparse prompts 중 points와 box는 positional encoding을 더해주어 각 prompt가 어떤 위치에 있는지 모델에게 알기 쉽게 전달해준다고 한다. 또한, text의 경우 CLIP의 text encoder를 그대로 사용하여 encoding을 진행했다고 한다. (논문의 표현에 따르면, with an off-the-shelf라고 함.)

 

3.3 Mask decoder 

그림 7. Mask decoder overview.

Mask decoder의 경우, DETR의 decoder 부분에 영감을 받아 dynamic mask prediction이 가능하도록 설계했다고 한다. 이를 약간 수정해서, 그림 7 과 같은 구조로 변경하였다. Prompt token 간의 self attention과 prompt-image 간의 cross attention을 거쳐 embedding을 업데이트하게 된다. 여기서 빨간색 box를 보면, point-wise MLP를 사용했다고 한다. 코드를 살펴보니, 2 linear로 구성되어 차원을 늘린 후 다시 원래 embedding dim으로 줄이는 방식을 이야기한다. 이러한 과정을 통해, 기존 크게 압축되어있던 임베딩에서 조금 더 크게 자유도를 준 후, 찾지 못한 특성을 업데이트 해줬다 라고 해석할 수 있을 것 같다. 논문에 직접적으로 나와있지 않지만, DETR에서 영감을 받았다는 의미가 object query와 feed forward network (FFN)을 사용하여 이미지 내의 object를 찾도록 하는 DETR 구조와, prompt token과 point-wise MLP를 사용한 것이 비슷한 점이 아닐까 한다.  

 

3.4 Resolving ambiguity

그림 8. Model의 ambiguity를 해결하기 위한 방법.

모호한 prompt와, 이로 인한 모호한 output을 해결하기 위해, model이 여러 output을 내도록 수정하였다. 논문에 따르면, 3개의 output을 내었을 때, 일반적인 모든 상황을 해결할 수 있다고한다. 따라서, 3개의 output 중 가장 작은 loss를 가지는 mask만 backprob하는 방식을 채택하여 모델을 학습하였다.

 

3.4 Losses and training

Loss의 경우, focal loss와 DICE loss를 20:1 비율로 선형 결합하여 사용하였다. 모델을 학습할 때는 다음과 같은 프로세스를 따르게 된다. Foreground와, background라는 용어를 사용하는데, 각각 object 혹은 배경 영역을 의미한다고 보면 됨.

 

1. 이미지에서 random하게 foreground point, box가 선택.

2. 해당 prompt를 통해 mask를 예측.

3. Error region에서 그다음 prompt를 sampling. 여기서 error region이란, predict mask와 GT간의 차이를 나타냄.

4. Sampling된 prompt가 background에 존재하면 FP, foreground에 존재하면 FN.

 

4. Segment Anything Dataset

그림 9. SA-1B dataset 특성.

본 논문에서는, SA-1B라는 dataset을 소개한다. 해당 dataset은 11M의 고해상도 이미지로 구성되어 있다. 특히, 해당 이미지를 사용하여 1.1B의 mask를 얻었고, 그중 99.1%가 fully automatic한 방식으로 얻을 수 있었다고 함. 사람의 cost가 적게 들어간 점은 굉장히 좋지만, label의 quality 또한 중요한 point 중 하나이다. 따라서, mask의 quality를 측정하기 위해 500개의 image를 sampling하여 전문 labelor에게 labeling한 것을 GT로 사용하여 IoU를 비교하였다. 그 결과, 94%가 0.9 이상의 IoU를 가지고 있는 것으로 확인되었다.

 

5. Experiments

 

그림 10. SAM의 zero-shot visualization

해당 논문에서 중요하게 생각하는 것은, segmentation을 잘하는 model보다, zero-shot transfer가 얼마나 가능한지가 중요한 point이다. 그림 10. 을 보면, small object, large object 상관없이 잘 분할하고 있는 것을 확인할 수 있다.

 

그림 11. SAM과 RITM model비교

그림 11. 에선, RITM이라는 model과의 성능 차이를 비교했다. RITM은 SAM과 마찬가지로, interactive segmentation 방식으로 학습하여 비교 대상으로 선정했다고 한다. 결과를 보면, 대부분의 dataset에서 우수한 성능 (주황색 bar)을 얻은 것을 볼 수 있다. 특정 dataset에서는 RITM보다 떨어지는 성능을 얻었는데, 논문에서는 아래와 같이 이야기하고 있다.

"Note that for datasets like DRAM and IBD, where SAM is worse on automatic metrics, it receives consistently higher ratings in the human study."

 

즉, 지표 상으로는 좋지 않을 수 있지만, human study에서는 지속적으로 높은 점수를 받았다고 한다. 따라서, 일부 dataset에서 성능이 낮게 나온 이유는 automatic metric과 human study 간의 차이에 기인할 수 있으며, 정확한 원인은 이 차이를 통해 다르게 해석될 수 있다고 할 수 있다.

 

그림 12. Text prompt를 통한 결과.

그림 12. 에선 text prompt를 통한 SAM의 output을 확인할 수 있다. Text prompt만을 주었을 때에도 해당 prompt가 지시하는 object를 잘 분할하는 것을 확인할 수 있다. Wiper의 경우, 만족스러운 결과물이 나오지 않는 것을 확인할 수 있다. 이러한 경우, text prompt와 sparse prompt (point)를 같이 사용하여 해결한 것을 볼 수 있다.

 

6. Conclusion

이처럼, 본 논문에서는 foundation model for image segmentation을 목표로 연구를 진행하였다. 기존의 LLM 방식에서 prompt를 사용한 것처럼, point, box, text를 prompt로 생각하여 성공적인 학습을 이루어낼 수 있었다. 또한, 이로 인한 SA-1B dataset은 fully automatic하게 수집하면서도 높은 quality의 label을 가지고 있는 것으로 확인되었다. 이처럼, SAM은 높은 performance를 가지고 있지만, 저자들은 다음과 같은 한계점이 있다고 지적한다. 

 

1. Fine structures를 놓칠 수 있으며, 경계를 정확하게 처리하지 못하는 경우가 발생할 수 있다고 한다.

해당 문제는 사실 어떠한 model에서도 그렇듯, 모든 model의 숙제로 남는 부분이 아닐까 생각함.

 

2. Text-to-mask의 한계점이 존재한다. 해당 task의 경우, 아직 연구의 초기 단계이기 때문에 안정적이지 않다고 언급한다.

 

3. 실시간 처리 성능 문제. SAM은 실시간으로도 가능하지만, 이는 image encoder의 lightweight함이 따라와야 한다. 아무래도 무거운 image encoder를 사용하게 되면 성능은 좋아지지만, 실시간성은 떨어지는 trade-off 관계를 가지기 때문에 실시간성을 위한 다른 방안이 존재해야 한다고 언급한다.