來(lái)源：公眾號(hào)“HiLab實(shí)驗(yàn)室”

原文鏈接：https://mp.weixin.qq.com/s/0qp3IN0hdHGTzjn3BAUOBg?scene=1&click_id=45

1.為什么少樣本醫(yī)學(xué)圖像分割重要？

醫(yī)學(xué)圖像分割通常需要醫(yī)生或?qū)I(yè)標(biāo)注人員進(jìn)行逐像素標(biāo)注，這一過(guò)程不僅耗時(shí)，而且成本很高。對(duì)于罕見(jiàn)病灶、新器官結(jié)構(gòu)或新成像協(xié)議下的數(shù)據(jù)，收集大規(guī)模標(biāo)注集往往并不現(xiàn)實(shí)。因此，如何讓模型在僅有少量標(biāo)注樣本的情況下快速適應(yīng)新類別，是醫(yī)學(xué)圖像分析中的關(guān)鍵問(wèn)題。

FSMIS正是為了解決這一問(wèn)題而提出：給定少量圖像-標(biāo)注對(duì)，模型需要對(duì)查詢圖像中相同類別的結(jié)構(gòu)進(jìn)行分割。這種設(shè)定更貼近真實(shí)臨床場(chǎng)景，因?yàn)榕R床中經(jīng)常出現(xiàn)標(biāo)注稀缺、類別變化和數(shù)據(jù)分布漂移的問(wèn)題。

2.現(xiàn)有FSMIS方法的局限

已有方法主要沿著兩條路線發(fā)展：

一類是原型匹配方法，即從支持集中提取類別原型，再與查詢圖像特征進(jìn)行匹配；另一類是支持-查詢交互方法，通過(guò)注意力、雙分支網(wǎng)絡(luò)或多尺度對(duì)齊機(jī)制增強(qiáng)支持圖像和查詢圖像之間的聯(lián)系。

這些方法雖然在標(biāo)準(zhǔn)設(shè)置下取得了不錯(cuò)效果，但仍然存在一個(gè)核心問(wèn)題：它們通常是在有限醫(yī)學(xué)數(shù)據(jù)上訓(xùn)練得到的任務(wù)特定模型，因此學(xué)到的視覺(jué)先驗(yàn)較窄。當(dāng)測(cè)試數(shù)據(jù)來(lái)自新的模態(tài)、新設(shè)備或新的解剖分布時(shí)，模型容易出現(xiàn)明顯性能退化。

3.為什么考慮Stable Diffusion？

Stable Diffusion這類大規(guī)模擴(kuò)散模型在海量數(shù)據(jù)上訓(xùn)練，具備較強(qiáng)的紋理、形狀和上下文建模能力。雖然它最初主要用于圖像生成任務(wù)，但近年來(lái)已有研究發(fā)現(xiàn)，擴(kuò)散模型內(nèi)部特征也可以遷移到語(yǔ)義對(duì)應(yīng)、像素級(jí)預(yù)測(cè)、開(kāi)放詞匯分割等視覺(jué)理解任務(wù)中。

本文的出發(fā)點(diǎn)是：如果Stable Diffusion已經(jīng)具備豐富的通用視覺(jué)先驗(yàn)，那么是否可以通過(guò)較小的適配成本，將這些先驗(yàn)遷移到少樣本醫(yī)學(xué)圖像分割中。

SD-FSMIS 模型總體

SD-FSMIS的核心思想是：不再?gòu)牧阍O(shè)計(jì)一個(gè)少樣本醫(yī)學(xué)分割網(wǎng)絡(luò)，而是把預(yù)訓(xùn)練Stable Diffusion當(dāng)作一個(gè)具有強(qiáng)視覺(jué)先驗(yàn)的基礎(chǔ)模型，再把它改造成少樣本分割模型。傳統(tǒng)FSMIS方法通常依賴支持-查詢特征匹配、原型匹配或注意力融合，而SD-FSMIS認(rèn)為Stable Diffusion已經(jīng)在大規(guī)模圖文數(shù)據(jù)上學(xué)到了豐富的形狀、紋理和上下文先驗(yàn)，這些先驗(yàn)可以幫助模型在少樣本醫(yī)學(xué)圖像中更好地泛化。論文的方法主要由兩個(gè)關(guān)鍵模塊組成：Support-Query Interaction（支持-查詢交互模塊，SQI）和Visual-to-Textual Condition Translator（視覺(jué)到文本條件轉(zhuǎn)換模塊，VTCT）。

可以把整個(gè)方法拆成四個(gè)階段：

●潛空間編碼階段：使用Stable Diffusion自帶的VAE編碼器，把支持圖像、支持標(biāo)注和查詢圖像都編碼到潛空間。

●支持-查詢交互階段：通過(guò)SQI模塊，把支持集中的目標(biāo)類別信息注入查詢的處理過(guò)程。

●視覺(jué)條件翻譯階段：通過(guò)VTCT模塊，把支持圖像中的目標(biāo)區(qū)域視覺(jué)特征轉(zhuǎn)換成類似文本嵌入的條件向量，用來(lái)指導(dǎo)Stable Diffusion的U-Net。

●掩膜預(yù)測(cè)階段：U-Net在支持信息和隱式文本條件的共同引導(dǎo)下，直接預(yù)測(cè)查詢掩膜潛特征，再由VAE解碼器解碼得到最終分割結(jié)果。

任務(wù)設(shè)定：少樣本醫(yī)學(xué)圖像分割

SD-FSMIS采用標(biāo)準(zhǔn)的基于任務(wù)單元的少樣本分割設(shè)定。每個(gè)任務(wù)單元包含一個(gè)支持集和一個(gè)查詢集。支持集中有少量帶掩膜的樣本，查詢集中有需要預(yù)測(cè)掩膜的圖像。論文主要采用單類別單樣本設(shè)定，即每次任務(wù)只分割一個(gè)目標(biāo)類別，并且只有一個(gè)標(biāo)注支持樣本可用。模型需要從支持圖像和支持標(biāo)注中理解“這次要分割什么結(jié)構(gòu)”，然后在查詢圖像上找出同類結(jié)構(gòu)。

這個(gè)設(shè)定和普通醫(yī)學(xué)圖像分割不同。普通分割通常在固定類別上訓(xùn)練和測(cè)試，例如一直分割肝臟或腎臟；而少樣本分割更強(qiáng)調(diào)泛化能力，即模型在測(cè)試階段可能遇到訓(xùn)練時(shí)沒(méi)有顯式學(xué)習(xí)過(guò)的新器官或新域數(shù)據(jù)。因此，SD-FSMIS的關(guān)鍵不是單純擬合某個(gè)器官類別，而是學(xué)習(xí)如何利用支持樣本動(dòng)態(tài)地指導(dǎo)查詢分割。

Stable Diffusion作為分割骨干

SD-FSMIS使用Stable Diffusion v1.5作為基礎(chǔ)模型。原本的Stable Diffusion是一個(gè)文本生成圖像模型，輸入文本條件后，通過(guò)U-Net在潛空間中逐步去噪生成圖像。SD-FSMIS并不是直接用它生成醫(yī)學(xué)圖像，而是重新利用它的潛擴(kuò)散結(jié)構(gòu)，把“圖像生成”任務(wù)改造成“掩膜預(yù)測(cè)”任務(wù)。論文使用的是Stable Diffusion的VAE、U-Net和交叉注意力機(jī)制，其中VAE保持凍結(jié)，用于把圖像和掩膜映射到潛空間，再?gòu)臐摽臻g解碼回像素空間。

由于Stable Diffusion的VAE原本接收3通道RGB圖像，而醫(yī)學(xué)圖像和二值分割掩膜通常是單通道，作者采用了一個(gè)簡(jiǎn)單處理：把單通道醫(yī)學(xué)圖像和掩膜復(fù)制成3通道偽RGB輸入，并把像素值歸一化到Stable Diffusion需要的范圍。推理時(shí)，VAE解碼器輸出3通道掩膜，作者再對(duì)三個(gè)通道取平均，得到最終單通道分割掩膜。

這樣做的好處是最大程度復(fù)用Stable Diffusion的原始組件，不需要重新訓(xùn)練一個(gè)醫(yī)學(xué)圖像專用VAE，也避免了少樣本場(chǎng)景下模型參數(shù)過(guò)多導(dǎo)致的過(guò)擬合問(wèn)題。

圖2:SD-FSMIS 總體流程

圖2展示了SD-FSMIS的整體流程。首先，支持圖像、支持標(biāo)注和查詢圖像都會(huì)經(jīng)過(guò)凍結(jié)的VAE編碼器，被映射到潛空間中，得到對(duì)應(yīng)的潛空間表示。

接著，支持圖像潛特征和支持標(biāo)注潛特征會(huì)在通道維度上進(jìn)行拼接，形成帶有類別提示的信息輸入；查詢圖像潛特征則經(jīng)過(guò)查詢?cè)鰪?qiáng)模塊處理，得到增強(qiáng)后的查詢潛特征。

隨后，這些潛空間特征會(huì)被送入修改后的Stable Diffusion U-Net。這里的U-Net不再用于生成自然圖像，而是在支持信息和VTCT生成的隱式文本條件共同引導(dǎo)下，預(yù)測(cè)查詢圖像對(duì)應(yīng)的掩膜潛特征。

最后，VAE解碼器將預(yù)測(cè)得到的掩膜潛特征解碼回圖像空間，得到最終的分割結(jié)果。

核心模塊一：Support-Query Interaction（SQI）

SQI是SD-FSMIS中負(fù)責(zé)少樣本信息傳遞的核心模塊。它的作用是讓查詢圖像在分割過(guò)程中能夠顯式利用支持圖像和支持標(biāo)注提供的類別信息。

傳統(tǒng)少樣本分割方法通常通過(guò)原型匹配或特征相關(guān)性計(jì)算來(lái)建立支持與查詢之間的聯(lián)系。不同的是，SD-FSMIS將這種支持-查詢交互直接嵌入到Stable Diffusion的U-Net注意力模塊中。也就是說(shuō)，模型不是在外部單獨(dú)設(shè)計(jì)一個(gè)匹配模塊，而是在擴(kuò)散模型原有的U-Net結(jié)構(gòu)內(nèi)部完成支持樣本到查詢樣本的信息傳遞，從而更充分地利用預(yù)訓(xùn)練Stable Diffusion中的視覺(jué)先驗(yàn)。

SQI包含兩個(gè)部分：

●Support Information Injection，SII

Stable Diffusion的U-Net中包含基礎(chǔ)Transformer模塊，原始結(jié)構(gòu)通常由自注意力、交叉注意力和前饋網(wǎng)絡(luò)組成。SD-FSMIS在自注意力之后額外加入了一個(gè)支持到查詢的交叉注意力模塊。具體來(lái)說(shuō)，模型將查詢圖像的潛空間特征作為查詢項(xiàng)，將支持圖像的潛空間特征作為鍵和值，使查詢特征能夠主動(dòng)關(guān)注支持中與目標(biāo)類別相關(guān)的區(qū)域。這樣，支持標(biāo)注提供的類別提示就可以在U-Net內(nèi)部傳遞到查詢分支，從而指導(dǎo)最終的分割預(yù)測(cè)。

●Query Enhancement，QE

僅僅讓查詢圖像關(guān)注支持圖像還不夠，因?yàn)獒t(yī)學(xué)圖像中不同器官可能形態(tài)差異大、邊界模糊，而且灰度分布相近。因此，QE模塊進(jìn)一步借鑒原型學(xué)習(xí)的思想，對(duì)查詢圖像的潛空間特征進(jìn)行增強(qiáng)。

具體來(lái)說(shuō)，模型首先根據(jù)支持標(biāo)注，在支持圖像的潛空間特征上進(jìn)行掩膜平均池化，得到目標(biāo)區(qū)域的前景原型。然后，模型計(jì)算查詢圖像特征與該前景原型之間的余弦相似度，從而找到查詢圖像中可能屬于目標(biāo)類別的區(qū)域。作者設(shè)定相似度閾值為0.7，篩選出高相似度區(qū)域，并對(duì)這些區(qū)域的特征求平均，得到查詢?cè)?。最后，將查詢?cè)蛿U(kuò)展到與原始特征相同的空間尺寸，并與原始查詢特征拼接，形成增強(qiáng)后的查詢表示。

可以理解為，SII負(fù)責(zé)“讓查詢看支持”，QE負(fù)責(zé)“讓查詢自己內(nèi)部先找到疑似目標(biāo)區(qū)域”。前者強(qiáng)調(diào)支持-查詢交互，后者強(qiáng)調(diào)查詢的目標(biāo)區(qū)域增強(qiáng)。

圖3:修改后的 BasicTransformerBlock

圖3展示了作者對(duì)Stable Diffusion中U-Net基礎(chǔ)Transformer模塊的修改。原始Stable Diffusion主要用于根據(jù)文本條件生成圖像，因此其中的交叉注意力模塊通常用來(lái)讓圖像潛特征接收文本特征的引導(dǎo)。

而在SD-FSMIS中，作者在這個(gè)結(jié)構(gòu)里額外加入了支持信息注入模塊。它的作用是讓查詢圖像的潛特征先與支持圖像的潛特征進(jìn)行交互，再接受由視覺(jué)到文本條件轉(zhuǎn)換模塊生成的條件信息。這個(gè)改動(dòng)很關(guān)鍵，因?yàn)樯贅颖痉指畹暮诵奶崾静⒉皇亲匀徽Z(yǔ)言中的類別名稱，而是支持圖像和支持標(biāo)注本身。也就是說(shuō)，模型需要從支持樣本中理解“要分割什么”，再把這個(gè)信息傳遞給查詢圖像的分割過(guò)程。

核心模塊二：

Visual-to-Textual Condition Translator（VTCT）

VTCT是SD-FSMIS的另一個(gè)關(guān)鍵創(chuàng)新。Stable Diffusion原本擅長(zhǎng)根據(jù)文本嵌入進(jìn)行條件生成，但在醫(yī)學(xué)少樣本分割中，支持集提供的條件并不是文字，而是一張醫(yī)學(xué)圖像及其對(duì)應(yīng)的分割標(biāo)注。如果直接使用空文本嵌入，模型就無(wú)法明確當(dāng)前要分割的目標(biāo)結(jié)構(gòu)。因此，作者設(shè)計(jì)了VTCT，將支持圖像中目標(biāo)區(qū)域的視覺(jué)信息轉(zhuǎn)換成Stable Diffusion能夠利用的“隱式文本條件”。

VTCT的流程可以概括為三步。首先，模型使用一個(gè)凍結(jié)的預(yù)訓(xùn)練圖像編碼器提取支持圖像的視覺(jué)特征，論文中采用的是DINOv2-small。然后，利用支持標(biāo)注對(duì)支持圖像特征進(jìn)行掩膜平均池化，只聚合目標(biāo)區(qū)域的前景特征，從而得到與當(dāng)前類別相關(guān)的視覺(jué)原型。最后，這個(gè)視覺(jué)原型會(huì)被送入一個(gè)可學(xué)習(xí)的MLP，并映射到Stable Diffusion中U-Net交叉注意力所需的文本嵌入空間，生成用于引導(dǎo)分割的隱式文本條件。

這個(gè)設(shè)計(jì)的意義在于，模型不需要人工編寫(xiě)提示詞，例如“肝臟”或“腎臟”，也不依賴類別名稱是否準(zhǔn)確。它直接從支持圖像中提取“這次要分割的目標(biāo)長(zhǎng)什么樣”，再把這個(gè)視覺(jué)信息轉(zhuǎn)換成Stable Diffusion原本熟悉的條件格式。也就是說(shuō)，VTCT相當(dāng)于一個(gè)從視覺(jué)到語(yǔ)義的橋梁，把醫(yī)學(xué)圖像中的類別提示轉(zhuǎn)換成擴(kuò)散模型可以利用的條件信號(hào)。

訓(xùn)練目標(biāo)：從查詢圖像潛特征預(yù)測(cè)分割掩膜潛特征

SD-FSMIS的訓(xùn)練目標(biāo)比較直接。模型輸入支持圖像、支持標(biāo)注和查詢圖像，輸出查詢掩膜的潛空間預(yù)測(cè)結(jié)果。訓(xùn)練時(shí)的監(jiān)督信號(hào)來(lái)自真實(shí)查詢掩膜經(jīng)過(guò)VAE編碼器后得到的掩膜潛特征。作者使用MSE損失約束預(yù)測(cè)掩膜潛特征和真實(shí)掩膜潛特征之間的差異。

這和普通醫(yī)學(xué)圖像分割常用的交叉熵?fù)p失或Dice損失不同。SD-FSMIS并不是直接在像素空間預(yù)測(cè)每個(gè)像素屬于前景或背景的概率，而是在Stable Diffusion的潛空間中學(xué)習(xí)如何生成分割掩膜。這樣做可以更自然地復(fù)用Stable Diffusion原有的U-Net結(jié)構(gòu)和潛擴(kuò)散表示能力。

在具體實(shí)現(xiàn)中，論文還采用單步DDIM調(diào)度器，并將時(shí)間步設(shè)置為999。這樣，模型可以用一步預(yù)測(cè)完成掩膜潛特征的生成，而不需要像標(biāo)準(zhǔn)擴(kuò)散采樣那樣進(jìn)行多輪迭代去噪。因此，SD-FSMIS本質(zhì)上是把擴(kuò)散模型的潛空間生成能力簡(jiǎn)化為一次性的掩膜潛特征預(yù)測(cè)，從而更適合分割任務(wù)。

推理流程：?jiǎn)尾筋A(yù)測(cè)查詢圖像的分割結(jié)果

推理階段不需要重新訓(xùn)練模型。給定一個(gè)支持圖像-標(biāo)注對(duì)和一個(gè)查詢圖像，模型首先通過(guò)VAE編碼器將它們映射到潛空間，得到對(duì)應(yīng)的潛特征。

隨后，支持圖像潛特征和支持標(biāo)注潛特征會(huì)在通道維度上拼接，并作為帶有類別提示的信息輸入到U-Net中。與此同時(shí)，查詢潛特征會(huì)經(jīng)過(guò)QE模塊進(jìn)一步強(qiáng)化目標(biāo)相關(guān)區(qū)域；VTCT則從支持圖像前景區(qū)域中提取目標(biāo)視覺(jué)信息，并生成隱式文本條件。

最后，U-Net在支持樣本信息和隱式文本條件的共同引導(dǎo)下，直接預(yù)測(cè)查詢掩膜潛特征。預(yù)測(cè)結(jié)果再經(jīng)過(guò)VAE解碼器解碼回像素空間，并通過(guò)通道平均得到最終的二值分割掩膜。整體來(lái)看，SD-FSMIS的推理過(guò)程可以理解為：支持樣本告訴模型“要分割什么”，查詢圖像提供“在哪里找目標(biāo)”，擴(kuò)散模型則利用已有視覺(jué)先驗(yàn)完成最終分割。

方法總結(jié)

SD-FSMIS的方法貢獻(xiàn)可以概括為三點(diǎn)。

第一，它將Stable Diffusion從文本生成圖像模型改造成少樣本醫(yī)學(xué)圖像分割模型，證明擴(kuò)散模型中學(xué)到的視覺(jué)先驗(yàn)不僅可以用于圖像生成，也可以用于醫(yī)學(xué)圖像中的密集預(yù)測(cè)任務(wù)。

第二，它提出了SQI模塊，在Stable Diffusion的U-Net潛空間注意力結(jié)構(gòu)中建立支持圖像和查詢圖像之間的信息交互，使support mask中的類別提示能夠有效傳遞到query圖像的分割過(guò)程中。

第三，它提出了VTCT模塊，將支持圖像前景區(qū)域的視覺(jué)信息轉(zhuǎn)換成隱式文本條件，使模型能夠繼續(xù)利用Stable Diffusion原有的條件生成機(jī)制，而不是簡(jiǎn)單使用空文本提示或完全丟棄文本條件分支。

總體而言，SD-FSMIS的核心不是“用擴(kuò)散模型生成更多醫(yī)學(xué)數(shù)據(jù)”，而是直接把Stable Diffusion適配成一個(gè)少樣本分割模型。它的優(yōu)勢(shì)在于充分利用大規(guī)模預(yù)訓(xùn)練擴(kuò)散模型中的通用視覺(jué)先驗(yàn)，并通過(guò)少量結(jié)構(gòu)修改，將這些先驗(yàn)遷移到醫(yī)學(xué)少樣本分割任務(wù)中。

圖5：BUSGen 用于乳腺癌預(yù)后相關(guān)指標(biāo)預(yù)測(cè)。圖5a 展示 TNBC 與 non-TNBC 分子亞型分類結(jié)果；圖5b 展示腋窩淋巴結(jié)轉(zhuǎn)移狀態(tài)分類結(jié)果；圖5c 和圖5d 展示不同模型的特征分布；圖5e 和圖5f 展示模型在預(yù)測(cè)不同預(yù)后指標(biāo)時(shí)關(guān)注的關(guān)鍵影像區(qū)域。

總體來(lái)看，SD-FSMIS的價(jià)值不只是將Stable Diffusion引入少樣本醫(yī)學(xué)圖像分割，更重要的是展示了擴(kuò)散基礎(chǔ)模型在數(shù)據(jù)高效醫(yī)學(xué)分割中的更大潛力。

首先，Stable Diffusion有望進(jìn)一步發(fā)展為統(tǒng)一的生成—判別框架。當(dāng)前方法主要利用其預(yù)訓(xùn)練視覺(jué)先驗(yàn)完成查詢圖像掩膜預(yù)測(cè)，而擴(kuò)散模型本身還具備強(qiáng)大的數(shù)據(jù)生成能力。未來(lái)可以探索讓模型一方面生成結(jié)構(gòu)合理的合成醫(yī)學(xué)圖像及對(duì)應(yīng)掩膜，擴(kuò)展少樣本訓(xùn)練數(shù)據(jù)；另一方面直接承擔(dān)分割任務(wù)。相比傳統(tǒng)數(shù)據(jù)增強(qiáng)，這種方式能夠在潛空間中生成更具解剖合理性和任務(wù)相關(guān)性的樣本，從而為少樣本分割提供更豐富、更可控的監(jiān)督信號(hào)。

其次，SD-FSMIS也啟發(fā)我們重新思考支持集—查詢集交互中“位置先驗(yàn)”和“語(yǔ)義信息”的作用。醫(yī)學(xué)圖像中的器官通常具有相對(duì)穩(wěn)定的解剖位置，但真正支持跨模態(tài)、跨中心泛化的，可能是更穩(wěn)定的目標(biāo)語(yǔ)義與形態(tài)表征。未來(lái)如果能夠?qū)⑽恢眉s束與語(yǔ)義表示進(jìn)一步解耦，一部分模塊學(xué)習(xí)跨模態(tài)一致的器官語(yǔ)義，另一部分模塊建?？勺兊慕馄士臻g關(guān)系，將有助于提升模型在更復(fù)雜模態(tài)遷移場(chǎng)景中的適應(yīng)能力。

此外，該工作也為弱監(jiān)督和偽標(biāo)簽適配提供了新的思路。醫(yī)學(xué)圖像的像素級(jí)密集標(biāo)注成本高、獲取難，而具有強(qiáng)視覺(jué)先驗(yàn)的基礎(chǔ)模型有望降低對(duì)高質(zhì)量人工標(biāo)注的依賴。結(jié)合偽標(biāo)簽生成、可靠性篩選、少量人工修正和自訓(xùn)練策略，未來(lái)少樣本醫(yī)學(xué)圖像分割可能從“少量精確標(biāo)注樣本”進(jìn)一步走向“強(qiáng)基礎(chǔ)模型先驗(yàn)+少量標(biāo)注+可控偽標(biāo)簽”的實(shí)用流程。

因此，SD-FSMIS的意義不僅在于提升少樣本醫(yī)學(xué)圖像分割性能，更在于提示了一條新的研究方向：通過(guò)生成能力、語(yǔ)義—位置解耦和弱監(jiān)督適配，將擴(kuò)散基礎(chǔ)模型構(gòu)建為更加通用、魯棒且數(shù)據(jù)高效的醫(yī)學(xué)圖像分割框架。