CVPR 2026 I SD-FSMIS：把Stable Diffusion改造成小樣本醫(yī)學(xué)分割利器

來源：公眾號“HiLab實驗室”

原文鏈接：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)注，這一過程不僅耗時，而且成本很高。對于罕見病灶、新器官結(jié)構(gòu)或新成像協(xié)議下的數(shù)據(jù)，收集大規(guī)模標(biāo)注集往往并不現(xiàn)實。因此，如何讓模型在僅有少量標(biāo)注樣本的情況下快速適應(yīng)新類別，是醫(yī)學(xué)圖像分析中的關(guān)鍵問題。

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

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

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

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

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

3.為什么考慮Stable Diffusion？

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

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

SD-FSMIS 模型總體

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

可以把整個方法拆成四個階段：

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

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

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

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

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

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

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

Stable Diffusion作為分割骨干

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

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

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

圖2:SD-FSMIS 總體流程

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

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

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

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

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

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

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

SQI包含兩個部分：

●Support Information Injection，SII

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

●Query Enhancement，QE

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

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

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

圖3:修改后的 BasicTransformerBlock

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

而在SD-FSMIS中，作者在這個結(jié)構(gòu)里額外加入了支持信息注入模塊。它的作用是讓查詢圖像的潛特征先與支持圖像的潛特征進(jìn)行交互，再接受由視覺到文本條件轉(zhuǎn)換模塊生成的條件信息。這個改動很關(guān)鍵，因為少樣本分割的核心提示并不是自然語言中的類別名稱，而是支持圖像和支持標(biāo)注本身。也就是說，模型需要從支持樣本中理解“要分割什么”，再把這個信息傳遞給查詢圖像的分割過程。

核心模塊二：

Visual-to-Textual Condition Translator（VTCT）

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

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

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

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

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

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

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

推理流程：單步預(yù)測查詢圖像的分割結(jié)果

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

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

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

方法總結(jié)

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

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

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

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

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

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

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

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

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

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

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