引言

大語言模型（LLMs）的成功激發(fā)了計(jì)算機(jī)視覺領(lǐng)域探索分割基礎(chǔ)模型的興趣。這些基礎(chǔ)分割模型通常通過 Prompt Engineer 來進(jìn)行 zero/few 圖像分割。其中，Segment Anything Model（SAM）是最先進(jìn)的圖像分割基礎(chǔ)模型。

圖  SAM 在多個下游任務(wù)上表現(xiàn)不佳

但是最近的研究表明，SAM 在多種下游任務(wù)中并非具有很強(qiáng)的魯棒性與泛化性，例如在醫(yī)學(xué)圖像、偽裝物體、添加干擾的自然圖像等領(lǐng)域表現(xiàn)較差。這可能是由于訓(xùn)練數(shù)據(jù)集與下游的測試數(shù)據(jù)集之間存在較大的域差異（Domain Shift）所致。因此，一個非常重要的問題是，如何設(shè)計(jì)域自適應(yīng)方案，使 SAM 在面對現(xiàn)實(shí)世界和多樣化的下游任務(wù)中更加魯棒？

將預(yù)訓(xùn)練好的 SAM 適應(yīng)到下游任務(wù)主要面臨三個挑戰(zhàn)：

• 首先，傳統(tǒng)的無監(jiān)督域自適應(yīng)范式需要源數(shù)據(jù)集和目標(biāo)數(shù)據(jù)集，由于隱私和計(jì)算成本較為不可行。
• 其次，對于域適應(yīng)，更新所有權(quán)重通常性能更好，同時也受到了昂貴的內(nèi)存成本的限制。
• 最后，SAM 可以針對不同種類、不同顆粒度的提示 Prompt，展現(xiàn)出多樣化的分割能力，因此當(dāng)缺乏下游任務(wù)的提示信息時，無監(jiān)督適應(yīng)將非常具有挑戰(zhàn)性。

圖 1 SAM 在大規(guī)模數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練，但存在泛化性問題。我們使用弱監(jiān)督的方式在各種下游任務(wù)上對 SAM 進(jìn)行自適應(yīng)

為了應(yīng)對上述挑戰(zhàn)，我們提出了一種具有錨點(diǎn)正則化和低秩微調(diào)的弱監(jiān)督自訓(xùn)練架構(gòu)，以提高自適應(yīng)的魯棒性和計(jì)算效率。

具體而言，我們首先采用無源域的自訓(xùn)練策略，從而避免對源數(shù)據(jù)的依賴。自訓(xùn)練產(chǎn)生偽標(biāo)簽，用于監(jiān)督模型的更新，但是容易受到錯誤偽標(biāo)簽的影響，我們引入凍結(jié)的 source model 作為錨定網(wǎng)絡(luò)，以規(guī)范模型的更新。

為了進(jìn)一步減少更新完整模型權(quán)重的高計(jì)算成本，我們對編碼器應(yīng)用低秩權(quán)重分解，并通過低秩快捷路徑進(jìn)行反向傳播。

最后，為了進(jìn)一步提高無源域自適應(yīng)的效果，我們在目標(biāo)域引入了弱監(jiān)督（weak supervise），例如稀疏的點(diǎn)注釋，以提供更強(qiáng)的域適應(yīng)信息，同時這種弱監(jiān)督與 SAM 中的提示編碼器自然兼容。

借助弱監(jiān)督作為 Prompt，我們獲得了更局部、更明確的自訓(xùn)練偽標(biāo)簽。經(jīng)過調(diào)整的模型在多個下游任務(wù)上表現(xiàn)出了更強(qiáng)的泛化能力。

我們總結(jié)本工作的貢獻(xiàn)如下：

1. 我們受到 SAM 在下游任務(wù)中泛化問題的啟發(fā)，提出了一種與任務(wù)無關(guān)且無需源數(shù)據(jù)的解決方案，通過自訓(xùn)練來適應(yīng) SAM。

2. 我們利用弱監(jiān)督，包括 box、point 等標(biāo)簽，以提高自適應(yīng)效果。這些弱監(jiān)督標(biāo)簽與 SAM 的提示編碼器完全兼容。

3. 我們對 5 種類型的下游實(shí)例分割任務(wù)進(jìn)行了大量實(shí)驗(yàn)，證明了所提出的弱監(jiān)督自適應(yīng)方法的有效性。

• 論文地址：https://arxiv.org/pdf/2312.03502.pdf
• 項(xiàng)目地址：https://github.com/Zhang-Haojie/WeSAM
• 論文標(biāo)題：Improving the Generalization of Segmentation Foundation Model under Distribution Shift via Weakly Supervised Adaptation

方法

方法介紹分為四個部分：

• Segment Anything 模型
• 基于自訓(xùn)練的自適應(yīng)框架
• 弱監(jiān)督如何幫助實(shí)現(xiàn)有效的自訓(xùn)練
• 低秩權(quán)重更新

1.Segment Anything Model

SAM 主要由三個組件構(gòu)成：圖像編碼器（ImageEncoder）、提示編碼器（PromptEncoder）、和解碼器（MaskDecoder）。

圖像編碼器使用 MAE 進(jìn)行預(yù)訓(xùn)練，整個 SAM 在擁有 11 億標(biāo)注的訓(xùn)練集 SA-1B 上進(jìn)一步進(jìn)行微調(diào)，訓(xùn)練時使用了 Focal loss 和 Dice loss 的組合。推理時，測試圖片 x 首先由圖像編碼器進(jìn)行編碼，然后給定提示 Prompt，輕量級的解碼器將進(jìn)行三個級別的預(yù)測。

2.Source-Free 域適應(yīng)自訓(xùn)練

圖 2 所提出的具有錨定網(wǎng)絡(luò)正則化和對比損失正則化的自訓(xùn)練架構(gòu)

針對未提供標(biāo)記的目標(biāo)數(shù)據(jù)集 DT={xi} 和預(yù)訓(xùn)練的分割模型。我們采用了 student-teacher 架構(gòu)進(jìn)行自訓(xùn)練。如圖 2 所示，我們維護(hù)三個編碼器網(wǎng)絡(luò)，即 anchor model、student model、teacher model，其中 student 和 teacher model 共享權(quán)重。

具體來說，對于每個樣本 xi，應(yīng)用一個隨機(jī)的弱數(shù)據(jù)增強(qiáng)作為 anchor 和 teacher model 的輸入，應(yīng)用一個隨機(jī)的強(qiáng)數(shù)據(jù)增強(qiáng)作為 student model 的輸入，三個編碼器網(wǎng)絡(luò)編碼產(chǎn)生三個特征圖。

在解碼器網(wǎng)絡(luò)中，給定一定數(shù)量 Np 的提示 prompt，例如 box、point 或 coarse mask，將推理出一組實(shí)例分割的 masks。

基于以上知識，我們下面詳細(xì)闡述用于自訓(xùn)練的三組優(yōu)化目標(biāo)。

1) Student-Teacher 自訓(xùn)練

我們首先使用與訓(xùn)練 SAM 時相同的損失函數(shù)作為自訓(xùn)練優(yōu)化目標(biāo)來更新 student/teacher model。自訓(xùn)練廣泛應(yīng)用于半監(jiān)督學(xué)習(xí)，最近還被證明了對無源域自適應(yīng)非常有效。具體而言，我們使用 teacher model 產(chǎn)生的預(yù)測結(jié)果，作為偽標(biāo)簽（Pseudo label），并使用 Focal loss 和 Dice loss 來監(jiān)督 student 的輸出。

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

僅使用自訓(xùn)練損失進(jìn)行網(wǎng)絡(luò)訓(xùn)練容易受到 teacher 網(wǎng)絡(luò)預(yù)測的錯誤偽標(biāo)簽積累的影響，即所謂的確認(rèn)偏差。觀察也表明，僅使用自訓(xùn)練長時間迭代后性能會下降。現(xiàn)有的無源域自適應(yīng)方法通常采用額外的約束來防止自訓(xùn)練的負(fù)面影響，例如對預(yù)測進(jìn)行均勻分布。

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

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

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

以上兩個訓(xùn)練目標(biāo)是在解碼器的輸出空間中執(zhí)行的。實(shí)驗(yàn)部分揭示出，更新編碼器網(wǎng)絡(luò)是適應(yīng) SAM 最有效的方法，因此有必要直接對從編碼器網(wǎng)絡(luò)輸出的特征應(yīng)用正則化。具體如圖 3 所示，我們根據(jù) anchor 和 teacher 分支中預(yù)測 mask 從特征圖中裁剪出每個實(shí)例的特征

。

我們進(jìn)一步定義對比損失中的正負(fù)樣本對，正樣本對是由兩個分支中使用相同的 prompt 對應(yīng)的實(shí)例特征構(gòu)建，而負(fù)樣本對是由不同 prompt 對應(yīng)的實(shí)例特征來構(gòu)建的。最終的對比損失如下所示，其中 

是溫度系數(shù)。

4) 總損失

我們將上述三個損失函數(shù)組合成最終的 Source-Free 自適應(yīng)損失。

3. 自訓(xùn)練的 Prompt 生成

SAM 分割需要 Prompt 輸入來指示出待分割的目標(biāo)對象，但可能會存在顆粒度模糊的問題。Prompt 工程可以以完全自動化的方式實(shí)現(xiàn)，也可以通過人工交互實(shí)現(xiàn)。

1) 完全自動生成 Prompt

我們首先使用網(wǎng)格密集采樣點(diǎn)作為 prompt 輸入，通過 Anchor model 生成初始階段分割的 masks，剔除 IoU 和穩(wěn)定性得分低的 mask，然后進(jìn)行非極大值抑制來獲得分割結(jié)果。接下來從最終的 masks 中產(chǎn)生一組固定的 prompts，作為所有三個分支的 prompt 輸入。因此，三個網(wǎng)絡(luò)分割輸出的 mask 長度相同，并且具有精確的一對一對應(yīng)關(guān)系。

2) 弱監(jiān)督作為 Prompt

盡管可以通過在圖像上使用網(wǎng)格采樣獲得 prompts，并過濾掉質(zhì)量低和重復(fù)的 mask 來進(jìn)行自動分割。但這些分割質(zhì)量相對較差，可能包含許多誤報(bào)預(yù)測，并且顆粒度不明確。由此產(chǎn)生的 prompt 質(zhì)量參差不齊，使得自訓(xùn)練效果較差。

因此，借鑒先前的弱監(jiān)督域自適應(yīng)工作，我們提出使用三種弱監(jiān)督方式，包括邊界框 box、稀疏點(diǎn)標(biāo)注 point 和粗分割多邊形 coarse mask。在 SAM 中，這些弱監(jiān)督方式與 prompt 輸入完美匹配，可以無縫集成弱監(jiān)督以適應(yīng) SAM。

4. 低秩權(quán)重更新

基礎(chǔ)模型龐大的編碼器網(wǎng)絡(luò)使得更新所有模型的權(quán)重變得異常困難。然而，許多現(xiàn)有研究表明，更新編碼器網(wǎng)絡(luò)權(quán)重是調(diào)整預(yù)訓(xùn)練模型的有效方法。

為了能夠更加有效且低成本地更新編碼器網(wǎng)絡(luò)，我們選擇了一種計(jì)算友好的低秩更新方法。對于編碼器網(wǎng)絡(luò)中的每個權(quán)重 θ，我們使用低秩近似 ω = AB，并設(shè)定一個壓縮率 r。只有 A 和 B 通過反向傳播進(jìn)行更新以減少內(nèi)存占用。在推理階段，通過將低秩近似和原始權(quán)重組合來重構(gòu)權(quán)重，即 θ = θ + AB。

實(shí)驗(yàn)

在實(shí)驗(yàn)中，我們提供了與最先進(jìn)方法的詳細(xì)比較和定性結(jié)果。最后，我們分析了各個部分的有效性以及網(wǎng)絡(luò)的具體設(shè)計(jì)。

1. 數(shù)據(jù)集

在這項(xiàng)工作中，我們對五種不同類型的下游分割任務(wù)進(jìn)行評估，其中一些與 SA-1B 存在明顯的分布偏移。數(shù)據(jù)集涵蓋了清晰的自然圖像、添加干擾的自然圖像、醫(yī)學(xué)圖像、偽裝物體和機(jī)器人圖像，總計(jì) 10 種。

數(shù)據(jù)劃分：每個下游數(shù)據(jù)集被劃分為互不重疊的訓(xùn)練集和測試集。

表 1 中列出了每種類型下游任務(wù)所評估的數(shù)據(jù)集，以及訓(xùn)練和測試數(shù)據(jù)集的劃分。

2. 實(shí)驗(yàn)細(xì)節(jié)

Segment-Anything 模型：由于內(nèi)存限制，我們采用 ViT-B 作為編碼器網(wǎng)絡(luò)。采用標(biāo)準(zhǔn)提示編碼器和 mask 解碼器。

Prompt 生成：訓(xùn)練和評估階段的 Prompt 輸入均是由從實(shí)例分割 GT mask 計(jì)算而來，模擬人類交互作為弱監(jiān)督。

具體來說，我們從整個 GT mask 的最小邊界框中提取 box。Point 是通過在 GT mask 內(nèi)隨機(jī)選擇 5 個正樣本點(diǎn)和 5 個 mask 外的負(fù)樣本點(diǎn)創(chuàng)建的。Coarse mask 是通過將多邊形擬合到 GT mask 來模擬的。

3. 實(shí)驗(yàn)結(jié)果

表 2、3、4、5 分別是在添加干擾的自然圖像、清晰的自然圖像、醫(yī)學(xué)圖像、偽裝物體數(shù)據(jù)集上的測試結(jié)果，完整的實(shí)驗(yàn)結(jié)果可以在論文中找到。實(shí)驗(yàn)證明了我們的方案在幾乎所有的下游分割數(shù)據(jù)集上都優(yōu)于預(yù)訓(xùn)練的 SAM 和最先進(jìn)的域適應(yīng)方案。

4. 可視化結(jié)果

部分可視化結(jié)果如圖 4 所示，更多的可視化結(jié)果可以在論文中找到。

圖 4 部分實(shí)例的可視化結(jié)果

5. 消融實(shí)驗(yàn)和額外分析

我們在 COCO 數(shù)據(jù)集上分析了三個自訓(xùn)練優(yōu)化目標(biāo)各自的有效性，如表 7 所示。表 7 中，我們還分析了所提出方法在不使用任何弱監(jiān)督信息時進(jìn)行自適應(yīng)的效果。

我們分析了訓(xùn)練和測試使用不同類別的 prompt 的性能差異，如表 8 所示。實(shí)驗(yàn)表明我們的方案在 cross-prompt 條件下依然表現(xiàn)良好。

此外，我們還分析了優(yōu)化不同模塊，包括解碼器、LayerNorm 和不同的 finetune 方案以及他們的組合的實(shí)驗(yàn)結(jié)果，實(shí)驗(yàn)證明了 finetune 編碼器的 LoRA 方案效果最佳。

總結(jié)

盡管視覺基礎(chǔ)模型可以在分割任務(wù)上表現(xiàn)出色，但其在下游任務(wù)中仍會存在性能不佳的情況。我們研究了 Segment-Anything 模型在多個下游圖像分割任務(wù)中的泛化能力，并提出了一種基于錨點(diǎn)正則化和低秩微調(diào)的自訓(xùn)練方法。該方法無需訪問源數(shù)據(jù)集、內(nèi)存成本低、與弱監(jiān)督自然兼容，可以顯著提升自適應(yīng)效果。經(jīng)過廣泛的實(shí)驗(yàn)驗(yàn)證，結(jié)果表明我們提出的域適應(yīng)方法可以顯著改善 SAM 在各種分布遷移下的泛化能力。

本文轉(zhuǎn)自機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/KzQTWVGPuSScy_p2J6kqIQ??