ICCV 2023：CLIP 驅(qū)動的器官分割和腫瘤檢測通用模型
目錄
-
前言 -
概述 -
Text branch -
Vision branch -
Masked back-propagation -
實驗 -
總結(jié) -
參考
前言
這次要介紹的文章屬于 CLIP 在醫(yī)學(xué)圖像上的一個應(yīng)用，思路上不算是創(chuàng)新。CLIP（Contrastive Language-Image Pre-training）是一種多模態(tài)模型，這意味著它可以同時處理文本和圖像數(shù)據(jù)。它的目標(biāo)是將文本描述和圖像內(nèi)容關(guān)聯(lián)起來，使得模型能夠理解文本描述與圖像之間的語義關(guān)系。它通過學(xué)習(xí)大量的文本和圖像來獲得對于語義理解的通用知識，這種通用知識可以在各種具體任務(wù)中進(jìn)行微調(diào)，使得模型可以適應(yīng)不同領(lǐng)域的任務(wù)。CLIP 使用對比學(xué)習(xí)的方法來訓(xùn)練模型。它要求模型將相關(guān)的文本描述和圖像匹配在一起，而將不相關(guān)的文本描述和圖像分開。這樣，模型可以學(xué)習(xí)如何捕捉文本和圖像之間的語義相似性。
在這篇解讀中，后文出現(xiàn)的 CLIP embedding 指的是由 CLIP 模型生成的表示文本和圖像之間語義關(guān)系的特征向量。這些嵌入向量是 CLIP 模型的核心輸出之一。它們編碼了文本描述和圖像內(nèi)容之間的關(guān)聯(lián)信息，使得模型能夠理解文本與圖像之間的語義相似性。CLIP 嵌入向量是在大規(guī)模數(shù)據(jù)上進(jìn)行預(yù)訓(xùn)練的，因此它們可以用于各種不同的任務(wù)，而不需要從頭開始訓(xùn)練一個新的模型。這種通用性使得 CLIP 嵌入在多個領(lǐng)域都有廣泛的應(yīng)用潛力。
對于下游任務(wù)，CLIP 會重新利用預(yù)訓(xùn)練獲得的能力，即預(yù)測某個圖像和某段文本是否為一對的能力。具體來說，對一個分類數(shù)據(jù)集，我們將所有 class 的名字分別作為 text，讓 CLIP 預(yù)測輸入圖像與哪個最有可能是一對。當(dāng)然，text 也可以進(jìn)一步加上 prompt，例如下圖（2）的 “A photo of [class]”。顯然這個過程是一種 zero-shot 遷移。
關(guān)于 CLIP 更多的細(xì)節(jié)，本篇解讀就不做更多介紹了。
我們回到這篇文章中，目前，在醫(yī)學(xué)圖像領(lǐng)域中公共數(shù)據(jù)集越來越多。然而，由于每個數(shù)據(jù)集的規(guī)模較小且都是部分標(biāo)記問題，以及對不同類型腫瘤的有限探究，由此產(chǎn)生的模型通常僅限于分割特定器官/腫瘤，并忽略了解剖結(jié)構(gòu)的語義，也不能擴(kuò)展到新領(lǐng)域。為了解決這些問題，這篇論文提出了基于 CLIP 學(xué)習(xí)的文本嵌入的通用模型，將其融入分割模型中。這種基于 CLIP 的標(biāo)簽編碼捕捉了解剖關(guān)系，使模型能夠?qū)W習(xí)結(jié)構(gòu)化特征嵌入并分割 25 個器官和 6 種腫瘤。所提出的模型是從 14 個數(shù)據(jù)集的組合中開發(fā)而來，使用了共計 3,410 個 CT 掃描進(jìn)行訓(xùn)練，然后在來自 3 個額外數(shù)據(jù)集的 6,162 個外部 CT 掃描上進(jìn)行評估。
概述
用于腹部影像的公開可用數(shù)據(jù)集主要關(guān)注不同的器官和腫瘤，例如，用于 4 個器官分割的 AbdomenCT-1K 數(shù)據(jù)集，用于 16 個器官分割的 WORD 數(shù)據(jù)集，以及用于 104 個解剖結(jié)構(gòu)分割的 TotalSegmentor 數(shù)據(jù)集（我們不需要關(guān)注這些數(shù)據(jù)集，這里只是舉個例子）。由于這些數(shù)據(jù)集的標(biāo)簽分類不一致，當(dāng)在這些數(shù)據(jù)集的組合上訓(xùn)練人工智能模型時，會出現(xiàn)部分標(biāo)簽問題。對于部分標(biāo)簽問題，我們舉一個例子，比如對于一張圖像切片，同時有肝臟和脾兩個類別，但是在一個肝臟分割小數(shù)據(jù)集上，只給了肝臟的標(biāo)注。同時，我們需要組合這些個小數(shù)據(jù)集，所以導(dǎo)致了部分標(biāo)簽問題出現(xiàn)。
為了充分利用部分標(biāo)簽，已經(jīng)研究了幾種方法，旨在構(gòu)建一個能夠執(zhí)行器官分割和腫瘤檢測的單一模型。這些研究存在以下限制：（1）由于數(shù)據(jù)集組合規(guī)模較小，組合數(shù)據(jù)集的潛力不太明顯。它們的性能與特定數(shù)據(jù)集模型相似，并未在官方 benchmark 上進(jìn)行評估。（2）由于采用了一位有效標(biāo)簽，器官和腫瘤之間的語義關(guān)系被丟棄。論文中做了相關(guān)實驗，使用 prompt 的性能要比使用 one-hot 編碼（存在正交性）更好。正交性指的是在"one-hot"編碼中，每個類別之間的表示是互相獨(dú)立的，沒有重疊或交互。具體來說：對于一個具有N個不同類別的分類問題，使用 one-hot 編碼后，每個類別都由一個長度為 N 的向量表示，其中只有一個元素的值為 1（代表當(dāng)前類別），而其余元素的值都為 0。這確保了每個類別的表示都是正交的，不受其他類別的影響。正交性使得模型可以明確地區(qū)分不同的類別，因為每個類別的表示不會與其他類別的表示發(fā)生干擾。這在許多分類任務(wù)中是有用的，例如圖像分類或自然語言處理中的詞匯分類。然而，正交性也具有一些限制，特別是在涉及到一些復(fù)雜的關(guān)系和語義信息的任務(wù)中。在某些情況下，one-hot 編碼可能無法捕捉類別之間的相關(guān)性或語義關(guān)系，因為它將每個類別都視為彼此獨(dú)立的。這在某些機(jī)器學(xué)習(xí)任務(wù)中可能會限制模型的性能。
如下圖是這篇文章提出的基于 CLIP 的通用模型，用于腹部器官分割和腫瘤檢測。為了解決標(biāo)簽不一致和正交性問題，將 CLIP 嵌入與分割模型整合在一起，從而實現(xiàn)了一個靈活而強(qiáng)大的分割器。該模型可以有效地從部分標(biāo)簽的數(shù)據(jù)集中學(xué)習(xí)，并取得了高性能。
Text branch
如概述圖中的上半部分，某一類的 CLIP embedding，它是由 CLIP 中預(yù)訓(xùn)練的文本編碼器和一個醫(yī)學(xué) prompt（例如，“[CLS] 的計算機(jī)斷層掃描”，其中 [CLS] 是具體的類名）生成的。首先將 CLIP embedding 和全局圖像特征（fimage）連接在一起，然后輸入到一個 MLP，即基于文本的控制器（已經(jīng)有其他工作這么做了），以生成參數(shù)（θk），即 θk = MLP(wk ⊕ f)，其中 ⊕ 表示連接。盡管 CLIP embedding 明顯優(yōu)于一位有效標(biāo)簽，但醫(yī)學(xué)提示模板的選擇至關(guān)重要。
最終在 Text branch 中，我們會得到 k 個生成參數(shù) θ，分別和所有要分割的類別對應(yīng)，形成 k 個分支，解決標(biāo)簽正交性的問題。
進(jìn)一步展示了嵌入空間的 t-SNE 可視化，分別對比了 one-hot 編碼和 CLIP 編碼，如上圖所示。可以看到，CLIP 編碼的解碼器嵌入展現(xiàn)出更好的特征聚類和解剖結(jié)構(gòu)。例如，通用模型中的右腎和左腎特征在嵌入空間中更接近（肝和肝腫瘤也是接近的，這樣就避免了 one-hot 的正交性），這與下圖中顯示的 CLIP embedding 之間的余弦相似度高度匹配。這驗證了 CLIP 基礎(chǔ)的編碼可以幫助模型捕捉解剖關(guān)系并學(xué)習(xí)結(jié)構(gòu)化的特征嵌入。
Vision branch
首先對圖像進(jìn)行預(yù)處理，使用等間距和標(biāo)準(zhǔn)化強(qiáng)度比例，以減小不同數(shù)據(jù)集之間的域差異（減小 domain gap），后由視覺編碼器進(jìn)行處理。令 F 表示由視覺編碼器提取的圖像特征。為了處理 F，使用了三個連續(xù)的卷積層，卷積核為 1×1×1，即文本驅(qū)動的解碼器。前兩層具有 8 個通道，最后一層具有 1 個通道。k 類的預(yù)測計算為 Pk = Sigmoid (((F ? θk1) ? θk2) ? θk3)，其中 θk = {θk1，θk2，θk3} 在文本分支中計算，* 表示卷積。對于每個類別 k，我們生成表示每個類別的前景的預(yù)測 Pk ∈ ，以一對多的方式進(jìn)行計算（即使用 Sigmoid 而不是 Softmax，因為每個像素可以同時屬于多個類別）。
Masked back-propagation
為了解決標(biāo)簽不一致性問題，這篇文章提出了掩碼反向傳播技術(shù)。使用 BCE 損失函數(shù)進(jìn)行監(jiān)督。屏蔽了不包含在對應(yīng)類別的損失項，并且只對準(zhǔn)確的監(jiān)督進(jìn)行反向傳播以更新整個框架。掩碼反向傳播解決了部分標(biāo)簽問題中的標(biāo)簽不一致性。具體來說，部分標(biāo)簽的數(shù)據(jù)集將一些其他器官標(biāo)記為背景，導(dǎo)致現(xiàn)有的訓(xùn)練方案失效。
實驗
在 MSD 數(shù)據(jù)集上的比較如下表，可以發(fā)現(xiàn)，通用模型的性能，已經(jīng)超過了 nnUNet：
一些腫瘤檢測 cases 的可視化如下圖：
文章中還驗證了通用模型的 FLOPS，因為通用模型只是利用了 CLIP 的 text encoder，backbone 是不受限制的，所以參數(shù)量可以很低。文中使用 CLIP 的預(yù)訓(xùn)練文本編碼器 ViTB/32 作為文本分支?？梢蕴崛〔⒋鎯ξ谋咎卣?，以減少在訓(xùn)練和推理階段由文本編碼器帶來的開銷，因為 CLIP embedding 僅依賴于固定的字典。
總結(jié)
總的來說，這篇文章介紹了一種使用 CLIP embedding 的通用模型，用于醫(yī)學(xué)圖像分割和腫瘤檢測。該模型在多個方面展示了出色的性能，解決了醫(yī)學(xué)圖像分割任務(wù)中的一些挑戰(zhàn)，特別是在捕捉解剖結(jié)構(gòu)和處理標(biāo)簽不一致性和正交性方面取得了顯著進(jìn)展。文章強(qiáng)調(diào)了 CLIP embedding 相對于傳統(tǒng)的獨(dú)熱編碼和其他預(yù)訓(xùn)練嵌入的優(yōu)勢，尤其是在醫(yī)學(xué)圖像分割任務(wù)中。CLIP embedding 能夠更好地捕捉圖像和文本之間的關(guān)系。
參考
-
https://arxiv.org/pdf/2301.00785.pdf -
https://github.com/ljwztc/CLIP-Driven-Universal-Model