圖 SAM 在多個下游任務上表現不佳 但是最近的研究表明,SAM 在多種下游任務中並非具有很強的魯棒性與泛化性,例如在醫學影像、偽裝物體、新增干擾的自然影像等領域表現較差。這可能是由於訓練資料集與下游的測試資料集之間存在較大的域差異(Domain Shift)所致。因此,一個非常重要的問題是,如何設計域自適應方案,使 SAM 在面對現實世界和多樣化的下游任務中更加魯棒? 將預訓練好的 SAM 適應到下游任務主要面臨三個挑戰:
圖 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)繼承的知識,不鼓勵源模型和自訓練更新模型之間出現過大的偏差,可以防止模型崩潰。