Segment Anything

Hướng dẫn đọc

Vấn đề và động lực

Các foundation model được tiền huấn luyện trên dữ liệu quy mô web đã tạo nên cuộc cách mạng trong NLP và căn chỉnh thị giác - ngôn ngữ (ví dụ: CLIP), nhưng lĩnh vực thị giác máy tính tổng quát vẫn thiếu một hệ thống nền tảng có khả năng tổng quát hóa zero-shot rộng rãi trên nhiều tác vụ phân vùng khác nhau. Các tập dữ liệu hiện có cho phân vùng ảnh rất nhỏ so với các kho ngữ liệu ảnh-văn bản quy mô web, và các kiến trúc phân vùng truyền thống thường bị ràng buộc vào các lớp ngữ nghĩa cố định hoặc các thiết lập tương tác chuyên biệt.

Để xây dựng một vision foundation model cho phân vùng, các tác giả đồng thời giải quyết ba yếu tố cốt lõi:

1. Tác vụ (Task): Mục tiêu nào cho phép chuyển giao zero-shot tổng quát?
2. Mô hình (Model): Kiến trúc nào hỗ trợ nhận prompt linh hoạt theo thời gian thực đồng thời xử lý được tính nhập nhằng (ambiguity) cố hữu của thị giác?
3. Dữ liệu (Data): Làm thế nào để thu thập các mask phân vùng ở quy mô web khi dữ liệu internet thiếu các nhãn mask dày đặc (dense masks)?

Các ý tưởng cốt lõi

• Tác vụ phân vùng bằng prompt (Promptable Segmentation Task): Mô hình phải trả về một mask phân vùng hợp lệ khi nhận bất kỳ prompt nào—chẳng hạn như các điểm tiền cảnh/hậu cảnh, bounding box, dense mask hoặc văn bản tự do. Khi prompt có tính nhập nhằng (ví dụ: một điểm duy nhất trên áo sơ mi), mô hình phải xuất ra nhiều mask hợp lý theo phân cấp (bộ phận phụ, bộ phận, toàn thể).
• Kiến trúc bất đối xứng giảm tải tính toán (Asymmetric Amortized Architecture - SAM): Một image encoder nặng về tính toán chỉ chạy một lần cho mỗi hình ảnh để tạo ra embedding, trong khi một prompt encoder gọn nhẹ và Transformer decoder hai lớp sẽ xử lý các prompt trong khoảng $\sim 50\text{ ms}$ trên CPU bên trong trình duyệt web.
• Quy trình dữ liệu bánh đà (Flywheel Data Engine): Quy trình gán nhãn ba giai đoạn:
  1. Thủ công có hỗ trợ (Assisted-manual): SAM hỗ trợ người gán nhãn trong trình duyệt, giảm thời gian gán nhãn từ 34 giây xuống 14 giây cho mỗi mask (thu thập 4,3 triệu mask).
  2. Bán tự động (Semi-automatic): SAM phát hiện các mask có độ tin cậy cao để hướng người gán nhãn chú ý vào các vật thể chưa được gán nhãn và ít nổi bật hơn (thu thập thêm 5,9 triệu mask).
  3. Tự động hoàn toàn (Fully automatic): SAM được cung cấp prompt trên một lưới điểm đều $32\times 32$ kết hợp với các vùng ảnh cắt phóng to (zoomed-in image crops), tạo ra 1,1 tỷ mask trong SA-1B.

Phương pháp

Kiến trúc mô hình

SAM bao gồm ba thành phần:

1. Image Encoder: Một Vision Transformer được tiền huấn luyện bằng MAE (mặc định là ViT-H/16) được điều chỉnh với windowed attention và bốn khối global attention, xử lý đầu vào kích thước $1024\times 1024$ để tạo ra embedding kích thước $64\times 64$ với 256 kênh.
2. Prompt Encoder:
   • Sparse prompts (prompt thưa): Điểm và bounding box được ánh xạ thành các vector 256 chiều bằng cách sử dụng positional encoding cộng với learned type embedding. Văn bản tự do được mã hóa bằng text encoder của CLIP.
   • Dense prompts (prompt dày đặc): Mask đầu vào (ở độ phân giải thấp hơn $4\times$) được giảm kích thước qua các lớp tích chập $2\times 2$ với stride 2 và được cộng theo từng phần tử (element-wise) vào image embedding.
3. Mask Decoder: Một Transformer decoder hai lớp được sửa đổi. Nó cập nhật các prompt token và image embedding thông qua việc xen kẽ self-attention cho prompt và cross-attention hai chiều (tokens-to-image và image-to-tokens). Các đặc trưng hình ảnh được phóng to $4\times$ qua các phép tích chập chuyển vị (transposed convolutions), và một MLP sẽ chiếu output token để tạo ra một bộ phân loại tuyến tính động, được áp dụng theo từng điểm (point-wise) lên các đặc trưng đã phóng to.

Đầu ra nhận biết tính nhập nhằng và hàm mất mát

Một prompt đơn lẻ thường có thể mô tả nhiều cấu trúc lồng nhau. SAM giải quyết sự nhập nhằng bằng cách dự đoán đồng thời 3 mask cùng với các điểm IoU ước tính tương ứng. Trong quá trình huấn luyện, lan truyền ngược chỉ áp dụng cho tổn thất nhỏ nhất trong số các ứng viên: $${\mathcal{L}}_{\text{mask}}=\underset{i\in \{1,2,3\}}{\min }\left(20\cdot {\mathcal{L}}_{\text{focal}}(m_i,y)+{\mathcal{L}}_{\text{dice}}(m_i,y)\right)$$ Một head phụ trợ sẽ tối thiểu hóa sai số bình phương trung bình giữa IoU dự đoán và IoU ground truth: $${\mathcal{L}}_{\text{IoU}}=\Vert {\widehat{\text{IoU}}}_i-\text{IoU}(m_i,y){\Vert }^2$$ Khi nhiều prompt được cung cấp cùng lúc, SAM sẽ xuất ra một mask duy nhất không nhập nhằng bằng cách sử dụng một output token chuyên dụng thứ tư.

Kết quả

• Chuyển giao Zero-Shot Point-to-Mask: Trên 23 tập dữ liệu phân vùng đa dạng (bao gồm ảnh góc nhìn thứ nhất, ảnh dưới nước, ảnh chụp từ trên không và ảnh kính hiển vi), SAM vượt trội hơn RITM trên 16/23 tập dữ liệu chỉ với một prompt điểm ở trung tâm (với mức chênh lệch lên tới $+46.9$ IoU). Khi tính nhập nhằng được giải quyết thông qua bộ chọn oracle trên 3 dự đoán của SAM, SAM vượt trội hơn RITM trên toàn bộ 23 tập dữ liệu. Đánh giá chất lượng từ con người luôn xếp hạng SAM cao hơn đáng kể (từ 7 đến 9 trên 10).
• Zero-Shot Edge Detection (BSDS500): Được đánh giá zero-shot bằng lưới $16\times 16$ kết hợp lọc Sobel và NMS cạnh, SAM đạt ODS là 0,768 và R50 là 0,928, cạnh tranh tốt với các mô hình có giám sát như HED (ODS 0,788, R50 0,923).
• Zero-Shot Object Proposals (LVIS v1): SAM đạt 59,3 AR@1000 mà không cần giám sát theo lớp, vượt qua baseline ViTDet-H có giám sát hoàn toàn trên các vật thể kích thước trung bình (81,6 so với 80,8), vật thể lớn (86,9 so với 87,0) và các danh mục hiếm (65,8 so với 58,3).
• Zero-Shot Instance Segmentation: Khi được cung cấp prompt là các bounding box từ ViTDet, SAM đạt 46,5 AP trên COCO và 44,7 AP trên LVIS. Trong các đánh giá của con người, người gán nhãn liên tục đánh giá mask của SAM sắc nét hơn và có chất lượng cao hơn đầu ra của ViTDet có giám sát đầy đủ.

Hạn chế và các câu hỏi mở

• Chi tiết tinh vi và độ phóng đại cực lớn: SAM có thể bỏ sót các cấu trúc cực nhỏ hoặc tạo ra các phần nhỏ rời rạc so với các phương pháp tương tác tốn nhiều tài nguyên tính toán vốn phóng to rõ ràng vào các vùng ảnh cắt.
• Độ trễ toàn pipeline: Mặc dù giải mã prompt diễn ra theo thời gian thực ($\sim 50\text{ ms}$), việc chạy image encoder ViT-H rất nặng lại không đạt tốc độ thời gian thực.
• Hiểu biết ngữ nghĩa: SAM tạo ra các mask không phân biệt lớp (class-agnostic); nó không cung cấp nhãn danh mục ngữ nghĩa (semantic) hoặc toàn cảnh (panoptic) một cách nguyên bản nếu không kết hợp với các bộ phân loại bên ngoài.
• Nhận prompt text-to-mask: Pipeline nhận prompt bằng văn bản được căn chỉnh với CLIP ở dạng zero-shot vẫn còn ở mức sơ khởi và nhạy cảm hơn so với các prompt hình học.