SAM: Segment Anything

논문 원본: https://arxiv.org/abs/2304.02643

Segment Anything

 

서론

 

Segment Anything 논문은 이미지 안의 어떤 객체라도 다양한 프롬프트를 통해 실시간으로 분할할 수 있는 범용 세그멘테이션 모델(SAM)과 이를 학습시키기 위해 구축한 대규모 데이터셋 SA-1B를 제안한다. 최근 대규모 웹 데이터셋으로 사전 학습된 LLM이 zero-shot과 few-shot 일반화 성능을 통해 NLP 분야에서 혁신을 일으킨 것처럼, 이 연구는 컴퓨터 비전에서도 범용 모델을 구현하고자 한다. Zero-shot learning은 학습 시 보지 못한 작업을 추가 훈련 없이 수행하는 능력을 의미하며, few-shot learning은 극히 적은 예시만 보고도 새로운 작업을 잘 수행하는 능력을 의미한다. 이러한 범용 모델은 다양한 데이터와 연산 자원으로 미리 학습해 두며, 프롬프트 엔지니어링을 통해 새로운 작업에 적용 가능하다. 기존 컴퓨터 비전의 대표적인 예시로 CLIP과 ALIGN이 있으며, 이들은 텍스트-이미지 쌍을 contrastive learning 방식으로 학습해 zero-shot 전이가 가능하다. 그러나 이러한 접근은 모든 비전 문제를 포괄하지 못하며, 특히 세그멘테이션에서는 대규모 학습 데이터 부족이 문제로 남아 있다.

 

프롬프트 기반 세그멘테이션 과제 정의

SAM은 프롬프트 기반 세그멘테이션이라는 새로운 과제를 정의한다. 이 과제는 점, 박스, 마스크 등 어떤 형태의 프롬프트가 주어지든 유효한 세그멘테이션 마스크를 반환하는 것을 목표로 한다. 유효한 마스크란 프롬프트가 모호할 경우에도 가능한 객체 중 하나에 대해 합리적인 마스크를 생성하는 것을 의미한다. 프롬프트를 다양하게 설계하면 객체 인스턴스 세그멘테이션, 에지 검출, 오브젝트 프로포절 생성 등 다양한 다운스트림 작업으로 전환할 수 있다. 이는 다중 작업 학습과 달리 처음부터 모든 작업을 학습하지 않아도 새로운 작업에 바로 적용할 수 있다는 점에서 차별화된다. 프롬프트 기반 세그멘테이션은 사람뿐 아니라 다른 프로그램에서도 프롬프트를 제공할 수 있어 대규모 시스템 구성 요소로 활용될 수 있다.

 

모델 아키텍처

SAM의 구조는 이미지 인코더, 프롬프트 인코더, 마스크 디코더의 세 가지 핵심 모듈로 이루어진다. 이미지 인코더는 MAE로 사전 학습된 ViT를 사용하며, 이미지는 사전에 임베딩해두어 프롬프트가 들어올 때마다 재사용할 수 있다. 프롬프트 인코더는 점, 박스, 텍스트, 마스크를 각각의 방식으로 임베딩한다. 점과 박스는 위치 인코딩에 프롬프트 타입별 학습 임베딩을 더해 표현하고, 텍스트는 CLIP의 텍스트 인코더를 사용하며, 마스크는 컨볼루션을 통해 이미지 임베딩과 합산한다. 마스크 디코더는 이미지 임베딩과 프롬프트 임베딩을 입력받아 수정된 트랜스포머 디코더 블록과 동적 마스크 예측 헤드를 통해 전경 확률 마스크를 생성한다. 모호성을 해결하기 위해 하나의 프롬프트에 대해 세 개의 마스크를 예측하고, 학습 시 정답과 가장 유사한 마스크만으로 역전파를 수행하며, 각 마스크에는 IoU를 추정하는 신뢰도 점수를 부여한다.

 

효율성과 학습 방법

모델 효율성은 이미지 임베딩을 한 번만 계산하고 이후 프롬프트 인코더와 마스크 디코더만 실행하는 방식으로 확보한다. 이를 통해 웹 브라우저 CPU에서도 약 50ms 내에 결과를 출력할 수 있다. 학습 시에는 Focal Loss와 Dice Loss를 결합하고, 한 마스크에 대해 여러 랜덤 프롬프트를 생성해 사람이 클릭하며 마스크를 수정하는 인터랙티브 환경을 시뮬레이션한다. 이 방식은 SAM을 데이터 수집용 도구로도 활용 가능하게 한다.

 

데이터 엔진과 SA-1B 구축

데이터셋 SA-1B 구축을 위해 세 단계의 데이터 엔진을 설계한다. 첫 번째는 보조 수동 라벨링 단계로, 라벨러가 SAM 기반 툴을 사용해 점과 브러쉬로 마스크를 생성하고 수정한다. 초기에는 공개 데이터셋으로 학습한 SAM을 사용하고, 주기적으로 새로 라벨링한 데이터로 재학습하여 성능을 개선한다. 두 번째는 반자동 단계로, 1단계에서 생성된 데이터로 객체 탐지기를 학습해 자동 생성 마스크를 채워넣고, 라벨러는 나머지만 수동 라벨링한다. 세 번째는 완전 자동 단계로, 그리드 포인트를 프롬프트로 하여 다양한 범위의 마스크를 예측하고 IoU 기반 안정성 검증 및 NMS로 필터링해 최종 마스크를 생성한다. 이를 통해 총 1억 1천만 장 이미지에서 11억 개의 마스크를 수집한다.

 

데이터 분석과 공정성

RAI 분석에서는 SA-1B가 기존 데이터셋보다 지역과 소득 분포에서 더 다양한 커버리지를 갖지만, 여전히 저소득 국가와 아프리카 지역은 과소대표되는 경향이 있다고 보고한다. 사람 세그멘테이션에서 성별, 연령, 피부 톤에 따른 뚜렷한 성능 차이는 발견되지 않았으나, 다른 시스템과 결합할 경우 편향이 발생할 가능성은 남아 있다.

Zero-shot 전이 성능

Zero-shot 전이 실험에서 SAM은 학습 시 보지 못한 데이터셋과 도메인에서도 다양한 시각 작업을 수행할 수 있음을 보인다. 단일 포인트로 객체를 분할하는 실험에서는 다수의 데이터셋에서 기존 강력한 모델보다 우수한 성능을 보였으며, 에지 검출과 오브젝트 프로포절 생성에서도 합리적인 결과를 생성한다. 인스턴스 세그멘테이션에서는 경계 품질이 우수하며, 텍스트 프롬프트 기반 마스크 생성도 가능하나 성능은 아직 실험 단계이다.

논의와 한계

논문에서는 SAM이 대규모 데이터와 다양한 프롬프트 학습을 통해 여러 세그멘테이션 작업에 적응 가능하다고 결론짓는다. 그러나 세밀한 구조 감지나 작은 객체 분리에서는 한계가 있으며, 무거운 백본 사용 시 실시간 전처리에는 부담이 있다. 텍스트 기반 마스크 생성과 특정 도메인 특화 작업에서도 성능이 낮을 수 있다. 그럼에도 불구하고 프롬프트 기반 세그멘테이션이라는 새로운 과제 정의, SAM 모델, 그리고 11억 개 이상의 마스크를 포함한 SA-1B 데이터셋 공개는 향후 세그멘테이션 연구와 응용 확장에 크게 기여할 것으로 기대된다고 설명한다.