[AI 논문 리뷰] Segment Anything

1. 논문이 나온 배경

자연어 처리 분야에서는 GPT와 같은 대규모 언어 모델들이 웹에서 수집한 방대한 텍스트 데이터로 사전학습된 후, 프롬프트 엔지니어링을 통해 다양한 downstream 태스크를 zero-shot 또는 few-shot 방식으로 해결하는 foundation model 패러다임이 확립되었습니다. CLIP과 ALIGN 같은 모델들은 이미지와 텍스트를 연결하는 vision-language 분야에서도 성공을 거두었습니다. 하지만 순수한 컴퓨터 비전, 특히 이미지 세그멘테이션 분야에는 이러한 범용 foundation model이 존재하지 않았습니다.

 

세그멘테이션 분야가 foundation model을 갖지 못한 가장 큰 이유는 데이터 부족 문제였습니다. 텍스트 데이터는 웹에 자연스럽게 풍부하게 존재하지만, 세그멘테이션에 필요한 마스크 주석은 웹에서 쉽게 구할 수 없습니다. 기존의 가장 큰 세그멘테이션 데이터셋인 Open Images조차 100만 이미지에 270만 마스크만을 포함하고 있었고, 수동 라벨링은 막대한 시간과 비용이 소요되었습니다. 또한 기존의 세그멘테이션 모델들은 instance segmentation, semantic segmentation, panoptic segmentation 등 특정 태스크에만 특화되어 있어, 학습 시 보지 못한 새로운 도메인이나 태스크에 일반화하는 능력이 부족했습니다.

2. 무엇을 개선할 것인지

연구팀은 이미지 세그멘테이션을 위한 foundation model을 구축하기 위해 세 가지 핵심 질문에 답해야 했습니다. 첫째, zero-shot generalization을 가능하게 하는 태스크는 무엇인가? 둘째, 그에 상응하는 모델 아키텍처는 무엇인가? 셋째, 이 태스크와 모델을 지원할 데이터를 어떻게 확보할 것인가?

 

이에 대한 해답으로 연구팀은 "promptable segmentation"이라는 새로운 태스크를 정의했습니다. 이는 어떤 프롬프트가 주어지더라도 유효한(valid) 세그멘테이션 마스크를 반환하는 것을 목표로 합니다. 프롬프트는 foreground/background 포인트, 박스, 마스크, 자유 형식 텍스트 등 세그멘테이션할 대상을 지정하는 모든 정보를 포함할 수 있습니다. 여기서 "유효한" 마스크란 프롬프트가 모호하여 여러 객체를 가리킬 수 있는 경우에도, 출력이 그 중 적어도 하나의 합리적인 마스크여야 함을 의미합니다.

 

모델은 유연한 프롬프팅을 지원하고, 대화형 사용을 위해 실시간으로 마스크를 계산할 수 있어야 하며, 모호성을 인식할 수 있어야 했습니다. 또한 기존 데이터셋의 한계를 극복하기 위해 11억 개의 마스크를 포함하는 대규모 데이터셋을 구축해야 했고, 이를 위해 효율적인 모델을 활용한 "data engine"을 개발해야 했습니다.

3. 어떤 원리로 개선되었는지

모델 아키텍처

모델의 개략적인 구조

모델의 세부적인 구조

 

SAM(Segment Anything Model)은 세 가지 주요 컴포넌트로 구성됩니다. 첫 번째는 이미지 인코더로, MAE로 사전학습된 Vision Transformer(ViT-H)를 사용합니다. 이 인코더는 1024×1024 해상도의 입력 이미지를 64×64×256 차원의 임베딩으로 변환하며, 이미지당 한 번만 계산되므로 비용이 분산됩니다.

 

두 번째는 프롬프트 인코더입니다. 포인트와 박스 같은 sparse 프롬프트는 positional encoding과 학습된 임베딩의 합으로 표현되고, 마스크 같은 dense 프롬프트는 convolution을 통해 임베딩되어 이미지 임베딩과 element-wise로 합쳐집니다. 텍스트 프롬프트의 경우 CLIP의 텍스트 인코더를 활용합니다.

 

세 번째는 경량 마스크 디코더로, Transformer 디코더 블록을 수정한 2-layer 구조를 사용합니다. 각 디코더 레이어는 네 단계를 거칩니다: 토큰에 대한 self-attention, 토큰을 쿼리로 한 이미지 임베딩과의 cross-attention, point-wise MLP, 이미지 임베딩을 쿼리로 한 토큰과의 cross-attention입니다. 이렇게 업데이트된 이미지 임베딩은 4배 업샘플링되고, 업데이트된 출력 토큰은 MLP를 거쳐 동적 선형 분류기로 각 위치의 마스크 foreground 확률을 계산합니다.

 

모호성을 해결하기 위해 SAM은 단일 프롬프트에 대해 세 개의 마스크를 동시에 예측합니다. 이는 중첩된 마스크가 주로 세 단계 깊이(whole, part, subpart)를 갖는다는 관찰에 기반합니다. 학습 중에는 세 개의 예측 마스크와 ground truth 간의 loss 중 최소값만 backpropagation하며, 각 마스크에 대한 IoU 신뢰도 점수도 예측하여 마스크 순위를 결정합니다.

학습 전략

SAM은 promptable segmentation 태스크를 위해 interactive segmentation을 시뮬레이션하는 방식으로 학습됩니다. 각 학습 샘플에 대해 첫 번째로 동일한 확률로 foreground 포인트 또는 박스가 무작위로 선택됩니다. 포인트는 ground truth 마스크에서 균등하게 샘플링되고, 박스는 ground truth 마스크의 bounding box에 각 좌표마다 박스 변 길이의 10%(최대 20픽셀)의 표준편차를 가진 노이즈를 추가하여 생성됩니다.

 

첫 번째 프롬프트로 예측을 한 후, 이후 포인트들은 이전 마스크 예측과 ground truth 마스크 간의 오류 영역에서 균등하게 선택됩니다. 각 새 포인트는 오류 영역이 false negative인지 false positive인지에 따라 foreground 또는 background로 설정됩니다. 이전 반복의 마스크 예측도 추가 프롬프트로 제공되는데, 최대 정보를 제공하기 위해 이진화된 마스크 대신 threshold 적용 전의 마스크 logit을 사용합니다.

 

총 11번의 반복이 수행됩니다: 1번의 초기 프롬프트 샘플링, 8번의 반복적 포인트 샘플링, 그리고 새로운 외부 정보 없이 자체 마스크 예측을 개선하도록 학습하는 2번의 반복입니다. 이렇게 상대적으로 많은 반복 횟수가 가능한 이유는 경량 마스크 디코더가 이미지 인코더 계산량의 1% 미만만을 요구하기 때문입니다.

 

Loss function으로는 focal loss와 dice loss를 20:1 비율로 선형 결합하여 사용하고, IoU 예측 헤드는 예측된 마스크와 ground truth 마스크 간의 IoU를 mean-square-error로 학습합니다. AdamW optimizer를 사용하며, 초기 learning rate는 8e-4이고, 90k iteration 동안 학습합니다(SA-1B의 약 2 epoch).

Data Engine

연구팀은 1.1B 마스크 데이터셋을 수집하기 위해 세 단계로 진화하는 "data engine"을 구축했습니다.

 

첫 번째 assisted-manual 단계에서는 전문 라벨러 팀이 SAM의 도움을 받아 브라우저 기반 대화형 세그멘테이션 도구를 사용하여 마스크를 라벨링했습니다. 라벨러들은 foreground/background 포인트를 클릭하고 pixel-precise "brush"와 "eraser" 도구로 마스크를 정제할 수 있었습니다. 의미론적 제약은 부과하지 않았고, 라벨러들이 이름 붙이거나 설명할 수 있는 객체를 자유롭게 라벨링하도록 했습니다. 이 단계에서 SAM은 공개 세그멘테이션 데이터셋으로 초기 학습되었고, 충분한 데이터가 모인 후에는 새로 수집된 마스크만으로 재학습되었습니다. 총 6번의 재학습을 거쳤고, 마스크당 평균 라벨링 시간은 34초에서 14초로 감소했으며, 이미지당 평균 마스크 수는 20개에서 44개로 증가했습니다. 이 단계에서 120k 이미지에서 4.3M 마스크를 수집했습니다.

 

두 번째 semi-automatic 단계에서는 마스크 다양성을 증가시키는 데 초점을 맞췄습니다. 먼저 SAM으로 confident 마스크를 자동으로 감지하고, 이 마스크들이 미리 채워진 이미지를 라벨러에게 제시하여 추가로 주석이 달리지 않은 객체만 라벨링하도록 했습니다. Confident 마스크를 감지하기 위해 첫 번째 단계의 모든 마스크를 사용하여 generic "object" 카테고리로 bounding box detector를 학습시켰습니다. 이 단계에서 180k 이미지에서 추가로 5.9M 마스크를 수집했으며(총 10.2M 마스크), 새로운 데이터로 모델을 5번 재학습했습니다. 자동 마스크를 제외한 마스크당 평균 라벨링 시간은 34초로 돌아갔는데, 이는 더 어려운 객체들을 라벨링했기 때문입니다. 자동 마스크를 포함한 이미지당 평균 마스크 수는 44개에서 72개로 증가했습니다.

 

세 번째 fully automatic 단계는 두 가지 주요 개선 덕분에 가능했습니다. 첫째, 이 단계 시작 시점에는 이전 단계의 다양한 마스크를 포함하여 충분한 마스크를 수집했습니다. 둘째, ambiguity-aware 모델을 개발하여 모호한 경우에도 유효한 마스크를 예측할 수 있게 되었습니다. 구체적으로 32×32 정규 그리드의 포인트로 모델에 프롬프트하고, 각 포인트에 대해 유효한 객체에 해당할 수 있는 마스크 세트를 예측했습니다. Ambiguity-aware 모델 덕분에 포인트가 part나 subpart에 있으면 모델은 subpart, part, whole object를 모두 반환합니다. 모델의 IoU 예측 모듈을 사용하여 confident 마스크를 선택하고, stable 마스크만 선택했습니다(0.5-δ와 0.5+δ에서 threshold한 마스크가 유사한 경우 stable로 간주). 마지막으로 confident하고 stable한 마스크를 선택한 후 NMS를 적용하여 중복을 제거했습니다. 작은 마스크의 품질을 더욱 향상시키기 위해 여러 겹치는 확대 이미지 crop도 처리했습니다. 이 fully automatic 마스크 생성을 11M 이미지 전체에 적용하여 총 1.1B 개의 고품질 마스크를 생성했습니다.

4. 실험 결과 및 성능

데이터셋 구조

SA-1B 데이터셋은 11M 개의 다양한 고해상도 라이센스 이미지와 1.1B 개의 고품질 세그멘테이션 마스크로 구성됩니다. 원본 이미지는 평균 3300×4950 픽셀이지만, 접근성과 저장 문제를 고려하여 짧은 쪽이 1500픽셀이 되도록 다운샘플링하여 공개됩니다. 얼굴과 차량 번호판은 개인정보 보호를 위해 블러 처리되었습니다. 마스크의 99.1%는 fully automatic으로 생성되었으며, 데이터셋은 자동 생성된 마스크만 포함합니다.

 

마스크 품질을 평가하기 위해 500개 이미지(약 50k 마스크)를 무작위로 샘플링하여 전문 라벨러에게 모든 마스크의 품질을 개선하도록 요청했습니다. 자동 예측 마스크와 전문가가 수정한 마스크 쌍 간의 IoU를 계산한 결과, 94%의 쌍이 90% 이상의 IoU를 보였고, 97%의 쌍이 75% 이상의 IoU를 보였습니다. 참고로 선행 연구에서는 라벨러 간 일관성을 85-91% IoU로 추정했습니다.

 

평가를 위해 연구팀은 23개의 다양한 세그멘테이션 데이터셋으로 구성된 새로운 벤치마크를 구축했습니다. 이 데이터셋들은 자기중심적(egocentric), 현미경, X-ray, 수중, 항공, 시뮬레이션, 자율주행, 회화 이미지 등 다양한 도메인을 포괄합니다. 효율적인 평가를 위해 15k 이상의 마스크를 가진 데이터셋은 서브샘플링하여 각각 약 10k 마스크를 포함하도록 했습니다.

성능 평가지표와 결과

Zero-Shot Single Point Valid Mask 평가에서는 단일 foreground 포인트에서 유효한 마스크를 생성하는 SAM의 핵심 목표를 테스트했습니다. 이 태스크는 하나의 포인트가 여러 객체를 가리킬 수 있어 본질적으로 ill-posed합니다. 자동 메트릭(mIoU)과 인간 평가를 모두 수행했습니다.

23개 데이터셋에 대한 mIoU 평가에서 SAM은 가장 강력한 단일 포인트 세그멘터인 RITM과 비교하여 16개 데이터셋에서 더 높은 결과를 얻었으며, 최대 약 47 IoU 차이를 보였습니다. "Oracle" 결과(SAM의 3개 마스크 중 ground truth와 가장 관련있는 것을 선택)에서는 모든 데이터셋에서 RITM을 능가했습니다. 이는 단일 포인트의 모호성이 자동 평가에 미치는 영향을 보여줍니다.

인간 평가 연구에서는 LVIS, VISOR, DRAM, IBD, NDD20, OVIS, iShape 데이터셋에 대해 1-10점 척도로 마스크 품질을 평가했습니다. SAM은 모든 데이터셋에서 RITM보다 일관되게 높은 평가를 받았으며, 평균 점수가 7-9 사이였습니다. 이는 "높은 점수(7-9): 객체를 식별할 수 있고 오류가 작고 드물다"는 정성적 평가 가이드라인에 해당합니다. 흥미롭게도 DRAM과 IBD처럼 자동 메트릭에서 SAM이 더 낮은 점수를 받은 데이터셋에서도 인간 평가에서는 일관되게 더 높은 점수를 받았습니다.

포인트 수를 1개에서 9개로 늘리면서 평가한 결과, 포인트가 늘어날수록 방법 간 격차가 줄어드는 것을 관찰했습니다. 이는 태스크가 더 쉬워지고, SAM이 매우 높은 IoU 영역에 최적화되지 않았기 때문입니다. 기본 center 포인트 샘플링을 random 포인트 샘플링으로 바꾸면 SAM과 baseline 간의 격차가 더 벌어졌고, SAM은 두 샘플링 방법에서 비슷한 결과를 달성했습니다.

 

Zero-Shot Edge Detection 평가는 BSDS500 데이터셋에서 수행되었습니다. SAM에 16×16 정규 그리드의 foreground 포인트로 프롬프트하여 768개의 마스크를 예측했고(포인트당 3개), NMS로 중복을 제거한 후 Sobel 필터링과 표준 edge NMS를 적용하여 edge map을 생성했습니다. SAM은 edge detection을 위해 학습되지 않았음에도 합리적인 edge map을 생성했습니다. Ground truth와 비교하여 SAM은 더 많은 edge를 예측하는데, 여기에는 BSDS500에 주석이 달리지 않은 합리적인 edge도 포함됩니다. 정량적으로 ODS 0.768, OIS 0.786, AP 0.794, R50 0.928을 달성하여, BSDS500에서 학습된 선구적인 딥러닝 방법인 HED(ODS 0.788)에 근접했고, 이전의 zero-shot transfer 방법들보다 훨씬 우수했습니다.

 

Zero-Shot Object Proposals 평가는 LVIS v1에서 AR@1000 메트릭으로 수행되었습니다. SAM의 자동 마스크 생성 파이프라인을 수정하여 64×64 포인트 그리드와 0.9 NMS threshold로 이미지당 평균 약 900개의 마스크를 생성했습니다. Baseline인 cascade ViTDet-H와 비교하여, SAM은 전체 AR에서는 59.3(vs 63.0)으로 뒤처졌지만, medium 객체(81.6 vs 80.8)와 large 객체(86.9 vs 87.0)에서는 ViTDet을 능가했습니다. 또한 common(63.9 vs 63.3)과 rare(65.8 vs 58.3) 카테고리에서도 우수한 성능을 보였습니다. SAM은 small 객체와 frequent 객체에서만 ViTDet에 뒤처졌는데, 이는 ViTDet이 LVIS에서 학습되어 데이터셋 특유의 주석 편향을 쉽게 학습할 수 있었기 때문입니다.

 

Zero-Shot Instance Segmentation 평가에서는 SAM을 fully-supervised ViTDet-H의 출력 박스로 프롬프트하여 세그멘테이션 모듈로 사용했습니다. COCO와 LVIS v1에서 mask AP 메트릭을 비교한 결과, SAM은 COCO에서 46.5 AP(vs 51.0), LVIS에서 44.7 AP(vs 46.6)를 기록하여 격차를 보였습니다. 그러나 출력을 시각화한 결과 SAM 마스크가 종종 ViTDet보다 질적으로 더 우수하고 경계가 더 선명했습니다.

이를 조사하기 위해 추가 인간 평가 연구를 수행했습니다. LVIS ground truth 박스에 적용된 ViTDet 마스크와 SAM 마스크를 1-10점 척도로 평가한 결과, SAM이 일관되게 ViTDet을 능가했습니다(SAM 8.1±0.07 vs ViTDet 7.9±0.08). LVIS ground truth는 8.6±0.06, COCO ground truth는 7.6±0.12를 받았습니다. 연구팀은 COCO에서 마스크 AP 격차가 더 크고 ground truth 품질이 상대적으로 낮은 점(인간 평가에서 확인)을 고려할 때, ViTDet이 COCO 마스크의 특정 편향을 학습한다고 가설을 세웠습니다. Zero-shot 방법인 SAM은 이러한 편향을 활용할 수 없습니다. LVIS 데이터셋은 더 높은 품질의 ground truth를 가지고 있지만, 여전히 특정한 특이성(예: 마스크에 구멍이 없고 단순한 다각형으로 구성)과 modal vs amodal 마스크에 대한 편향이 있습니다.

 

Zero-Shot Text-to-Mask는 SAM이 텍스트 프롬프트를 처리할 수 있는 능력에 대한 개념 증명입니다. 이 실험만을 위해 SAM의 학습 절차를 수정하여 텍스트를 인식하도록 했지만, 새로운 텍스트 주석은 필요하지 않았습니다. 면적이 100² 픽셀보다 큰 각 수동 수집 마스크에 대해 CLIP 이미지 임베딩을 추출했고, 학습 중에는 이 CLIP 이미지 임베딩을 첫 번째 상호작용으로 SAM에 프롬프트했습니다. CLIP의 이미지 임베딩이 텍스트 임베딩과 정렬되도록 학습되었기 때문에, 이미지 임베딩으로 학습하고 텍스트 임베딩을 추론에 사용할 수 있었습니다. 정성적 결과에서 SAM은 "a wheel"과 같은 간단한 텍스트 프롬프트뿐만 아니라 "beaver tooth grille"과 같은 구문도 세그먼트할 수 있었습니다. SAM이 텍스트 프롬프트만으로 올바른 객체를 선택하지 못할 때, 추가 포인트를 제공하면 예측이 수정되는 경우가 많았습니다.

 

Ablation Study에서는 23개 데이터셋 suite에서 단일 center 포인트 프롬프트 프로토콜로 여러 ablation을 수행했습니다. Data engine 단계별 누적 데이터로 학습한 SAM의 성능을 플롯한 결과, 각 단계가 mIoU를 증가시켰습니다. 세 단계 모두로 학습할 때 자동 마스크가 수동 및 반자동 마스크를 압도적으로 많이 차지하므로, 학습 중 수동 및 반자동 마스크를 10배 오버샘플링하는 것이 최상의 결과를 제공했습니다. 또한 자동 생성된 마스크만 사용하는 네 번째 설정을 테스트한 결과, 모든 데이터를 사용하는 것보다 약 0.5 mIoU만 낮았습니다. 따라서 기본적으로 학습 설정을 단순화하기 위해 자동 생성된 마스크만 사용합니다.

데이터 볼륨의 영향을 조사하기 위해 11M 이미지를 포함하는 전체 SA-1B를 1M과 0.1M으로 균등하게 서브샘플링했습니다. 0.1M 이미지에서는 모든 설정에서 큰 mIoU 감소가 관찰되었습니다. 그러나 전체 데이터셋의 약 10%인 1M 이미지를 사용하면 전체 데이터셋 사용과 비슷한 결과를 얻었습니다. 이 데이터 regime은 여전히 약 100M 마스크를 포함하며, 많은 사용 사례에 실용적인 설정일 수 있습니다.

마지막으로 ViT-B, ViT-L, ViT-H 이미지 인코더로 결과를 비교했습니다. ViT-H는 ViT-B에 비해 상당히 개선되었지만, ViT-L에 비해서는 미미한 개선만 보였습니다. 추가적인 이미지 인코더 스케일링은 현재로서는 유익하지 않은 것으로 보입니다.

5. 향후 연구 방향성 및 미해결 과제

SAM은 일반적으로 잘 작동하지만 완벽하지는 않습니다. 미세한 구조를 놓칠 수 있고, 때때로 작은 분리된 컴포넌트를 환각하며, "zoom-in" 방식을 사용하는 더 계산 집약적인 방법들만큼 경계를 선명하게 생성하지 못합니다. 일반적으로 많은 포인트가 제공될 때는 전용 interactive segmentation 방법들이 SAM을 능가할 것으로 예상됩니다. SAM과 달리 이러한 방법들은 일반성과 사용 폭보다는 높은 IoU interactive segmentation을 위해 설계되었습니다. 게다가 SAM은 프롬프트를 실시간으로 처리할 수 있지만, heavy 이미지 인코더를 사용할 때 전체 성능은 실시간이 아닙니다.

텍스트-투-마스크 태스크로의 진출은 탐색적이며 완전히 견고하지 않지만, 더 많은 노력으로 개선될 수 있다고 믿습니다. SAM이 많은 태스크를 수행할 수 있지만, semantic과 panoptic segmentation을 구현하는 간단한 프롬프트를 설계하는 방법은 불분명합니다. 마지막으로 의료 영상 등 특정 도메인에 특화된 도구들이 해당 도메인에서 SAM을 능가할 것으로 예상됩니다.

연구팀은 Segment Anything 프로젝트가 이미지 세그멘테이션을 foundation model 시대로 끌어올리려는 시도라고 밝혔습니다. SAM이 foundation model의 지위를 달성할지 여부는 커뮤니티에서 어떻게 사용되는지에 따라 결정될 것이지만, 이 작업의 관점, 10억 개 이상의 마스크 공개, promptable segmentation 모델이 앞으로 나아갈 길을 열어줄 것으로 기대한다고 결론지었습니다.