CVPR 2024 | 分割一切模型SAM泛化能力差?域適應(yīng)策略給解決了
引言
大語言模型(LLMs)的成功激發(fā)了計(jì)算機(jī)視覺領(lǐng)域探索分割基礎(chǔ)模型的興趣。這些基礎(chǔ)分割模型通常通過 Prompt Engineer 來進(jìn)行 zero/few 圖像分割。其中,Segment Anything Model(SAM)是最先進(jìn)的圖像分割基礎(chǔ)模型。
圖 SAM 在多個(gè)下游任務(wù)上表現(xiàn)不佳
但是最近的研究表明,SAM 在多種下游任務(wù)中并非具有很強(qiáng)的魯棒性與泛化性,例如在醫(yī)學(xué)圖像、偽裝物體、添加干擾的自然圖像等領(lǐng)域表現(xiàn)較差。這可能是由于訓(xùn)練數(shù)據(jù)集與下游的測試數(shù)據(jù)集之間存在較大的域差異(Domain Shift)所致。因此,一個(gè)非常重要的問題是,如何設(shè)計(jì)域自適應(yīng)方案,使 SAM 在面對(duì)現(xiàn)實(shí)世界和多樣化的下游任務(wù)中更加魯棒?
將預(yù)訓(xùn)練好的 SAM 適應(yīng)到下游任務(wù)主要面臨三個(gè)挑戰(zhàn):
- 首先,傳統(tǒng)的無監(jiān)督域自適應(yīng)范式需要源數(shù)據(jù)集和目標(biāo)數(shù)據(jù)集,由于隱私和計(jì)算成本較為不可行。
- 其次,對(duì)于域適應(yīng),更新所有權(quán)重通常性能更好,同時(shí)也受到了昂貴的內(nèi)存成本的限制。
- 最后,SAM 可以針對(duì)不同種類、不同顆粒度的提示 Prompt,展現(xiàn)出多樣化的分割能力,因此當(dāng)缺乏下游任務(wù)的提示信息時(shí),無監(jiān)督適應(yīng)將非常具有挑戰(zhàn)性。
圖 1 SAM 在大規(guī)模數(shù)據(jù)集上進(jìn)行預(yù)訓(xùn)練,但存在泛化性問題。我們使用弱監(jiān)督的方式在各種下游任務(wù)上對(duì) SAM 進(jìn)行自適應(yīng)
為了應(yīng)對(duì)上述挑戰(zhàn),我們提出了一種具有錨點(diǎn)正則化和低秩微調(diào)的弱監(jiān)督自訓(xùn)練架構(gòu),以提高自適應(yīng)的魯棒性和計(jì)算效率。
具體而言,我們首先采用無源域的自訓(xùn)練策略,從而避免對(duì)源數(shù)據(jù)的依賴。自訓(xùn)練產(chǎn)生偽標(biāo)簽,用于監(jiān)督模型的更新,但是容易受到錯(cuò)誤偽標(biāo)簽的影響,我們引入凍結(jié)的 source model 作為錨定網(wǎng)絡(luò),以規(guī)范模型的更新。
為了進(jìn)一步減少更新完整模型權(quán)重的高計(jì)算成本,我們對(duì)編碼器應(yīng)用低秩權(quán)重分解,并通過低秩快捷路徑進(jìn)行反向傳播。
最后,為了進(jìn)一步提高無源域自適應(yīng)的效果,我們?cè)谀繕?biāo)域引入了弱監(jiān)督(weak supervise),例如稀疏的點(diǎn)注釋,以提供更強(qiáng)的域適應(yīng)信息,同時(shí)這種弱監(jiān)督與 SAM 中的提示編碼器自然兼容。
借助弱監(jiān)督作為 Prompt,我們獲得了更局部、更明確的自訓(xùn)練偽標(biāo)簽。經(jīng)過調(diào)整的模型在多個(gè)下游任務(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. 我們對(duì) 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
方法
方法介紹分為四個(gè)部分:
- Segment Anything 模型
- 基于自訓(xùn)練的自適應(yīng)框架
- 弱監(jiān)督如何幫助實(shí)現(xiàn)有效的自訓(xùn)練
- 低秩權(quán)重更新
1.Segment Anything Model
SAM 主要由三個(gè)組件構(gòu)成:圖像編碼器(ImageEncoder)、提示編碼器(PromptEncoder)、和解碼器(MaskDecoder)。
圖像編碼器使用 MAE 進(jìn)行預(yù)訓(xùn)練,整個(gè) SAM 在擁有 11 億標(biāo)注的訓(xùn)練集 SA-1B 上進(jìn)一步進(jìn)行微調(diào),訓(xùn)練時(shí)使用了 Focal loss 和 Dice loss 的組合。推理時(shí),測試圖片 x 首先由圖像編碼器進(jìn)行編碼,然后給定提示 Prompt,輕量級(jí)的解碼器將進(jìn)行三個(gè)級(jí)別的預(yù)測。
2.Source-Free 域適應(yīng)自訓(xùn)練
圖 2 所提出的具有錨定網(wǎng)絡(luò)正則化和對(duì)比損失正則化的自訓(xùn)練架構(gòu)
針對(duì)未提供標(biāo)記的目標(biāo)數(shù)據(jù)集 DT={xi} 和預(yù)訓(xùn)練的分割模型。我們采用了 student-teacher 架構(gòu)進(jìn)行自訓(xùn)練。如圖 2 所示,我們維護(hù)三個(gè)編碼器網(wǎng)絡(luò),即 anchor model、student model、teacher model,其中 student 和 teacher model 共享權(quán)重。
具體來說,對(duì)于每個(gè)樣本 xi,應(yīng)用一個(gè)隨機(jī)的弱數(shù)據(jù)增強(qiáng)作為 anchor 和 teacher model 的輸入,應(yīng)用一個(gè)隨機(jī)的強(qiáng)數(shù)據(jù)增強(qiáng)作為 student model 的輸入,三個(gè)編碼器網(wǎng)絡(luò)編碼產(chǎn)生三個(gè)特征圖。
在解碼器網(wǎng)絡(luò)中,給定一定數(shù)量 Np 的提示 prompt,例如 box、point 或 coarse mask,將推理出一組實(shí)例分割的 masks。
基于以上知識(shí),我們下面詳細(xì)闡述用于自訓(xùn)練的三組優(yōu)化目標(biāo)。
1) Student-Teacher 自訓(xùn)練
我們首先使用與訓(xùn)練 SAM 時(shí)相同的損失函數(shù)作為自訓(xùn)練優(yōu)化目標(biāo)來更新 student/teacher model。自訓(xùn)練廣泛應(yīng)用于半監(jiān)督學(xué)習(xí),最近還被證明了對(duì)無源域自適應(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ù)測的錯(cuò)誤偽標(biāo)簽積累的影響,即所謂的確認(rèn)偏差。觀察也表明,僅使用自訓(xùn)練長時(shí)間迭代后性能會(huì)下降。現(xiàn)有的無源域自適應(yīng)方法通常采用額外的約束來防止自訓(xùn)練的負(fù)面影響,例如對(duì)預(yù)測進(jìn)行均勻分布。
我們通過 anchor 損失來進(jìn)行正則化,如公式 3 所示,分別最小化了 anchor model 與 student/teacher model 之間的 Dice loss。凍結(jié)的 anchor model 作為從源域(source domain)繼承的知識(shí),不鼓勵(lì)源模型和自訓(xùn)練更新模型之間出現(xiàn)過大的偏差,可以防止模型崩潰。
3) 對(duì)比損失正則化編碼器特征空間
圖 3 兩個(gè)分支下的對(duì)比損失
以上兩個(gè)訓(xùn)練目標(biāo)是在解碼器的輸出空間中執(zhí)行的。實(shí)驗(yàn)部分揭示出,更新編碼器網(wǎng)絡(luò)是適應(yīng) SAM 最有效的方法,因此有必要直接對(duì)從編碼器網(wǎng)絡(luò)輸出的特征應(yīng)用正則化。具體如圖 3 所示,我們根據(jù) anchor 和 teacher 分支中預(yù)測 mask 從特征圖中裁剪出每個(gè)實(shí)例的特征
。
我們進(jìn)一步定義對(duì)比損失中的正負(fù)樣本對(duì),正樣本對(duì)是由兩個(gè)分支中使用相同的 prompt 對(duì)應(yīng)的實(shí)例特征構(gòu)建,而負(fù)樣本對(duì)是由不同 prompt 對(duì)應(yīng)的實(shí)例特征來構(gòu)建的。最終的對(duì)比損失如下所示,其中
是溫度系數(shù)。
4) 總損失
我們將上述三個(gè)損失函數(shù)組合成最終的 Source-Free 自適應(yīng)損失。
3. 自訓(xùn)練的 Prompt 生成
SAM 分割需要 Prompt 輸入來指示出待分割的目標(biāo)對(duì)象,但可能會(huì)存在顆粒度模糊的問題。Prompt 工程可以以完全自動(dòng)化的方式實(shí)現(xiàn),也可以通過人工交互實(shí)現(xiàn)。
1) 完全自動(dòng)生成 Prompt
我們首先使用網(wǎng)格密集采樣點(diǎn)作為 prompt 輸入,通過 Anchor model 生成初始階段分割的 masks,剔除 IoU 和穩(wěn)定性得分低的 mask,然后進(jìn)行非極大值抑制來獲得分割結(jié)果。接下來從最終的 masks 中產(chǎn)生一組固定的 prompts,作為所有三個(gè)分支的 prompt 輸入。因此,三個(gè)網(wǎng)絡(luò)分割輸出的 mask 長度相同,并且具有精確的一對(duì)一對(duì)應(yīng)關(guān)系。
2) 弱監(jiān)督作為 Prompt
盡管可以通過在圖像上使用網(wǎng)格采樣獲得 prompts,并過濾掉質(zhì)量低和重復(fù)的 mask 來進(jìn)行自動(dòng)分割。但這些分割質(zhì)量相對(duì)較差,可能包含許多誤報(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ì)算友好的低秩更新方法。對(duì)于編碼器網(wǎng)絡(luò)中的每個(gè)權(quán)重 θ,我們使用低秩近似 ω = AB,并設(shè)定一個(gè)壓縮率 r。只有 A 和 B 通過反向傳播進(jìn)行更新以減少內(nèi)存占用。在推理階段,通過將低秩近似和原始權(quán)重組合來重構(gòu)權(quán)重,即 θ = θ + AB。
實(shí)驗(yàn)
在實(shí)驗(yàn)中,我們提供了與最先進(jìn)方法的詳細(xì)比較和定性結(jié)果。最后,我們分析了各個(gè)部分的有效性以及網(wǎng)絡(luò)的具體設(shè)計(jì)。
1. 數(shù)據(jù)集
在這項(xiàng)工作中,我們對(duì)五種不同類型的下游分割任務(wù)進(jìn)行評(píng)估,其中一些與 SA-1B 存在明顯的分布偏移。數(shù)據(jù)集涵蓋了清晰的自然圖像、添加干擾的自然圖像、醫(yī)學(xué)圖像、偽裝物體和機(jī)器人圖像,總計(jì) 10 種。
數(shù)據(jù)劃分:每個(gè)下游數(shù)據(jù)集被劃分為互不重疊的訓(xùn)練集和測試集。
表 1 中列出了每種類型下游任務(wù)所評(píng)估的數(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)練和評(píng)估階段的 Prompt 輸入均是由從實(shí)例分割 GT mask 計(jì)算而來,模擬人類交互作為弱監(jiān)督。
具體來說,我們從整個(gè) GT mask 的最小邊界框中提取 box。Point 是通過在 GT mask 內(nèi)隨機(jī)選擇 5 個(gè)正樣本點(diǎn)和 5 個(gè) 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)和額外分析
我們?cè)?COCO 數(shù)據(jù)集上分析了三個(gè)自訓(xùn)練優(yōu)化目標(biāo)各自的有效性,如表 7 所示。表 7 中,我們還分析了所提出方法在不使用任何弱監(jiān)督信息時(shí)進(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ù)中仍會(huì)存在性能不佳的情況。我們研究了 Segment-Anything 模型在多個(gè)下游圖像分割任務(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ī)器之心
