CVPR 2024 | 分割一切模型SAM泛化能力差?域適應策略給解決了

机器之心發表於2024-04-09
第一個針對「Segment Anything」大模型的域適應策略來了!相關論文已被CVPR 2024 接收。

引言

語言模型(LLMs)的成功激發了計算機視覺領域探索分割基礎模型的興趣。這些基礎分割模型通常透過 Prompt Engineer 來進行 zero/few 影像分割。其中,Segment Anything Model(SAM)是最先進的影像分割基礎模型。

CVPR 2024 | 分割一切模型SAM泛化能力差?域適應策略給解決了

圖 SAM 在多個下游任務上表現不佳

但是最近的研究表明,SAM 在多種下游任務中並非具有很強的魯棒性與泛化性,例如在醫學影像、偽裝物體、新增干擾的自然影像等領域表現較差。這可能是由於訓練資料集與下游的測試資料集之間存在較大的域差異(Domain Shift)所致。因此,一個非常重要的問題是,如何設計域自適應方案,使 SAM 在面對現實世界和多樣化的下游任務中更加魯棒?

將預訓練好的 SAM 適應到下游任務主要面臨三個挑戰:

  • 首先,傳統的無監督域自適應正規化需要源資料集和目標資料集,由於隱私和計算成本較為不可行。
  • 其次,對於域適應,更新所有權重通常效能更好,同時也受到了昂貴的記憶體成本的限制
  • 最後,SAM 可以針對不同種類、不同顆粒度的提示 Prompt,展現出多樣化的分割能力,因此當缺乏下游任務的提示資訊時,無監督適應將非常具有挑戰性。

圖片

圖 1 SAM 在大規模資料集上進行預訓練,但存在泛化性問題。我們使用弱監督的方式在各種下游任務上對 SAM 進行自適應

為了應對上述挑戰,我們提出了一種具有錨點正則化低秩微調弱監督自訓練架構,以提高自適應的魯棒性和計算效率

具體而言,我們首先採用無源域的自訓練策略,從而避免對源資料的依賴。自訓練產生偽標籤,用於監督模型的更新,但是容易受到錯誤偽標籤的影響,我們引入凍結的 source model 作為錨定網路,以規範模型的更新。

為了進一步減少更新完整模型權重的高計算成本,我們對編碼器應用低秩權重分解,並透過低秩快捷路徑進行反向傳播。

最後,為了進一步提高無源域自適應的效果,我們在目標域引入了弱監督(weak supervise),例如稀疏的點註釋,以提供更強的域適應資訊,同時這種弱監督與 SAM 中的提示編碼器自然相容。

藉助弱監督作為 Prompt,我們獲得了更區域性、更明確的自訓練偽標籤。經過調整的模型在多個下游任務上表現出了更強的泛化能力。

我們總結本工作的貢獻如下:

1. 我們受到 SAM 在下游任務中泛化問題的啟發,提出了一種與任務無關且無需源資料的解決方案,透過自訓練來適應 SAM。
2. 我們利用弱監督,包括 box、point 等標籤,以提高自適應效果。這些弱監督標籤與 SAM 的提示編碼器完全相容。
3. 我們對 5 種型別的下游例項分割任務進行了大量實驗,證明了所提出的弱監督自適應方法的有效性。

圖片

  • 論文地址:https://arxiv.org/pdf/2312.03502.pdf
  • 專案地址:https://github.com/Zhang-Haojie/WeSAM
  • 論文標題:Improving the Generalization of Segmentation Foundation Model under Distribution Shift via Weakly Supervised Adaptation

方法

方法介紹分為四個部分:

  • Segment Anything 模型
  • 基於自訓練的自適應框架
  • 弱監督如何幫助實現有效的自訓練
  • 低秩權重更新

1.Segment Anything Model

SAM 主要由三個元件構成:影像編碼器(ImageEncoder)、提示編碼器(PromptEncoder)、和解碼器(MaskDecoder)

影像編碼器使用 MAE 進行預訓練,整個 SAM 在擁有 11 億標註的訓練集 SA-1B 上進一步進行微調,訓練時使用了 Focal loss 和 Dice loss 的組合。推理時,測試圖片 x 首先由影像編碼器進行編碼,然後給定提示 Prompt,輕量級的解碼器將進行三個級別的預測。

2.Source-Free 域適應自訓練

圖片

圖 2 所提出的具有錨定網路正則化和對比損失正則化的自訓練架構

針對未提供標記的目標資料集 DT={xi} 和預訓練的分割模型。我們採用了 student-teacher 架構進行自訓練。如圖 2 所示,我們維護三個編碼器網路,即 anchor model、student model、teacher model,其中 student 和 teacher model 共享權重

具體來說,對於每個樣本 xi,應用一個隨機的弱資料增強作為 anchor 和 teacher model 的輸入,應用一個隨機的強資料增強作為 student model 的輸入,三個編碼器網路編碼產生三個特徵圖。

在解碼器網路中,給定一定數量 Np 的提示 prompt,例如 box、point 或 coarse mask,將推理出一組例項分割的 masks。

基於以上知識,我們下面詳細闡述用於自訓練的三組最佳化目標。

1) Student-Teacher 自訓練

我們首先使用與訓練 SAM 時相同的損失函式作為自訓練最佳化目標來更新 student/teacher model。自訓練廣泛應用於半監督學習,最近還被證明了對無源域自適應非常有效。具體而言,我們使用 teacher model 產生的預測結果,作為偽標籤(Pseudo label),並使用 Focal loss 和 Dice loss 來監督 student 的輸出。

圖片

2) Anchor 損失用於魯棒正則化

僅使用自訓練損失進行網路訓練容易受到 teacher 網路預測的錯誤偽標籤積累的影響,即所謂的確認偏差。觀察也表明,僅使用自訓練長時間迭代後效能會下降。現有的無源域自適應方法通常採用額外的約束來防止自訓練的負面影響,例如對預測進行均勻分佈。

我們透過 anchor 損失來進行正則化,如公式 3 所示,分別最小化了 anchor model 與 student/teacher model 之間的 Dice loss。凍結的 anchor model 作為從源域(source domain)繼承的知識,不鼓勵源模型和自訓練更新模型之間出現過大的偏差,可以防止模型崩潰。

圖片

3) 對比損失正則化編碼器特徵空間

圖片

圖 3 兩個分支下的對比損失

以上兩個訓練目標是在解碼器的輸出空間中執行的。實驗部分揭示出,更新編碼器網路是適應 SAM 最有效的方法,因此有必要直接對從編碼器網路輸出的特徵應用正則化。具體如圖 3 所示,我們根據 anchor 和 teacher 分支中預測 mask 從特徵圖中裁剪出每個例項的特徵圖片

圖片

我們進一步定義對比損失中的正負樣本對,正樣本對是由兩個分支中使用相同的 prompt 對應的例項特徵構建,而負樣本對是由不同 prompt 對應的例項特徵來構建的。最終的對比損失如下所示,其中 圖片是溫度係數。

圖片

4) 總損失

我們將上述三個損失函式組合成最終的 Source-Free 自適應損失。

圖片

3. 自訓練的 Prompt 生成

SAM 分割需要 Prompt 輸入來指示出待分割的目標物件,但可能會存在顆粒度模糊的問題。Prompt 工程可以以完全自動化的方式實現,也可以透過人工互動實現。

1) 完全自動生成 Prompt

我們首先使用網格密集取樣點作為 prompt 輸入,透過 Anchor model 生成初始階段分割的 masks,剔除 IoU 和穩定性得分低的 mask,然後進行非極大值抑制來獲得分割結果。接下來從最終的 masks 中產生一組固定的 prompts,作為所有三個分支的 prompt 輸入。因此,三個網路分割輸出的 mask 長度相同,並且具有精確的一對一對應關係。

2) 弱監督作為 Prompt

儘管可以透過在影像上使用網格取樣獲得 prompts,並過濾掉質量低和重複的 mask 來進行自動分割。但這些分割質量相對較差,可能包含許多誤報預測,並且顆粒度不明確。由此產生的 prompt 質量參差不齊,使得自訓練效果較差。

因此,借鑑先前的弱監督域自適應工作,我們提出使用三種弱監督方式,包括邊界框 box、稀疏點標註 point 和粗分割多邊形 coarse mask。在 SAM 中,這些弱監督方式與 prompt 輸入完美匹配,可以無縫整合弱監督以適應 SAM。

4. 低秩權重更新

基礎模型龐大的編碼器網路使得更新所有模型的權重變得異常困難。然而,許多現有研究表明,更新編碼器網路權重是調整預訓練模型的有效方法。

為了能夠更加有效且低成本地更新編碼器網路,我們選擇了一種計算友好的低秩更新方法。對於編碼器網路中的每個權重 θ,我們使用低秩近似 ω = AB,並設定一個壓縮率 r。只有 A 和 B 透過反向傳播進行更新以減少記憶體佔用。在推理階段,透過將低秩近似和原始權重組合來重構權重,即 θ = θ + AB。

實驗

在實驗中,我們提供了與最先進方法的詳細比較和定性結果。最後,我們分析了各個部分的有效性以及網路的具體設計。

1. 資料集

在這項工作中,我們對五種不同型別的下游分割任務進行評估,其中一些與 SA-1B 存在明顯的分佈偏移。資料集涵蓋了清晰的自然影像、新增干擾的自然影像、醫學影像、偽裝物體和機器人影像,總計 10 種。

資料劃分:每個下游資料集被劃分為互不重疊的訓練集和測試集。

表 1 中列出了每種型別下游任務所評估的資料集,以及訓練和測試資料集的劃分。

圖片

2. 實驗細節

Segment-Anything 模型:由於記憶體限制,我們採用 ViT-B 作為編碼器網路。採用標準提示編碼器和 mask 解碼器。

Prompt 生成:訓練和評估階段的 Prompt 輸入均是由從例項分割 GT mask 計算而來,模擬人類互動作為弱監督。

具體來說,我們從整個 GT mask 的最小邊界框中提取 box。Point 是透過在 GT mask 內隨機選擇 5 個正樣本點和 5 個 mask 外的負樣本點建立的。Coarse mask 是透過將多邊形擬合到 GT mask 來模擬的。

3. 實驗結果

表 2、3、4、5 分別是在新增干擾的自然影像、清晰的自然影像、醫學影像、偽裝物體資料集上的測試結果,完整的實驗結果可以在論文中找到。實驗證明了我們的方案在幾乎所有的下游分割資料集上都優於預訓練的 SAM 和最先進的域適應方案。

圖片

圖片

圖片

圖片

4. 視覺化結果

部分視覺化結果如圖 4 所示,更多的視覺化結果可以在論文中找到。

圖片

圖片

圖片

圖片

圖 4 部分例項的視覺化結果

5. 消融實驗和額外分析

我們在 COCO 資料集上分析了三個自訓練最佳化目標各自的有效性,如表 7 所示。表 7 中,我們還分析了所提出方法在不使用任何弱監督資訊時進行自適應的效果。

圖片

我們分析了訓練和測試使用不同類別的 prompt 的效能差異,如表 8 所示。實驗表明我們的方案在 cross-prompt 條件下依然表現良好。

圖片

此外,我們還分析了最佳化不同模組,包括解碼器、LayerNorm 和不同的 finetune 方案以及他們的組合的實驗結果,實驗證明了 finetune 編碼器的 LoRA 方案效果最佳。

圖片

總結

儘管視覺基礎模型可以在分割任務上表現出色,但其在下游任務中仍會存在效能不佳的情況。我們研究了 Segment-Anything 模型在多個下游影像分割任務中的泛化能力,並提出了一種基於錨點正則化和低秩微調的自訓練方法。該方法無需訪問源資料集、記憶體成本低、與弱監督自然相容,可以顯著提升自適應效果。經過廣泛的實驗驗證,結果表明我們提出的域適應方法可以顯著改善 SAM 在各種分佈遷移下的泛化能力。

相關文章