作者 | 上海交通大學、上海人工智慧實驗室
編輯 | ScienceAI
近日，上海交通大學與上海人工智慧實驗室聯合團隊釋出3D醫學影像分割大模型SAT（Segment Anything in radiology scans, driven by Text prompts），在3D醫學影像（CT、MR、PET）上，基於文字提示實現對人體497種器官/病灶的通用分割。所有資料和程式碼、模型均已開源。
論文連結：https://arxiv.org/abs/2312.17183
程式碼連結：https://github.com/zhaoziheng/SAT
資料連結：https://github.com/zhaoziheng/SAT-DS/
研究背景
醫學影像分割在診斷、手術規劃和疾病監測等一系列臨床任務中都有重要作用。然而，傳統的研究針對每個特定的分割任務訓練「專用」模型，導致每個「專用」模型的應用範圍都相對有限，無法高效便捷地滿足廣泛多樣的醫療分割需求。
與此同時，大語言模型最近在醫療領域取得了巨大成功，而要進一步推動通用醫療人工智慧的發展，構建一個可以連線語言和定位能力的醫學分割工具變得十分必要。
為了克服這些挑戰，來自上海交通大學和上海人工智慧實驗室的研究者們提出了第一個基於知識增強，使用文字提示的3D醫療影像通用分割模型，名為SAT（Segment Anything in radiology scans, driven by Text prompts），並做出了以下三點主要貢獻：
1. 該研究首次探索將人體解剖學知識注入文字編碼器，以精準編碼解剖學術語，實現了透過文字提示的放射學影像通用醫療分割模型。
2. 該研究構建了第一個包含6K+人體解剖學概念的多模態醫療知識圖譜。同時，構建了當前最大規模的3D醫學影像分割資料集，名為SAT-DS，彙集了72個公開資料集，來自CT、MR和PET三種模態的22K+影像，302K+分割標註，涵蓋了人體8個主要部位中的497個分割目標。
3. 基於SAT-DS，該研究訓練了兩款不同大小的模型：SAT-Pro（447M引數）和SAT-Nano（110M引數），並設計實驗從多個角度驗證了SAT的價值：SAT的效能與72個nnU-Nets專家模型相當（在每個資料集上單獨調參和最佳化，共約2.2B引數），並在域外資料上表現出更強的泛化能力；SAT可以作為一個基於大規模資料預訓練的基礎分割模型，透過下游微調遷移到特定任務時，可以表現出比nnU-Nets更好的效能；此外，與基於box提示的MedSAM對比，SAT基於文字提示可以實現更精準、更高效的分割；最後，在域外的臨床資料上，研究團隊展示了SAT可以被用作大語言模型的代理工具，在報告生成等任務中直接賦予後者定位和分割的能力。
接下來將從資料、模型與實驗結果三個方面介紹原文細節。
資料構建
多模態知識圖譜：為了實現精準編碼解剖學術語，研究團隊首先收集了一個包含6K+人體解剖學概念的多模態知識圖譜，其內容來自於三個來源：
1. Unified Medical Language System（UMLS）是由美國國家醫學圖書館構建的生物醫學字典。研究團隊從中提取了近230K的生物醫學概念與定義，以及涵蓋1M+條相互關係的知識圖譜。
2. 網路上的權威解剖學知識。研究團隊篩選了6502個人體解剖學概念，並藉助檢索增強的大語言模型從網路上檢索相關資訊，獲取了6K+概念與定義，涵蓋了38K+解剖結構相互關係的知識圖譜。
3. 公開的分割資料集。研究團隊收集了大規模的公開3D醫學影像分割資料集，將分割區域透過解剖學概念（類別標籤）與上述文字知識庫中的知識對應連線，提供視覺知識對照。
SAT-DS：為了訓練通用分割模型，研究團隊構建了領域內最大規模的3D醫學影像分割資料集合SAT-DS。特別的，72個多樣的公開分割資料集被收集並整理，總計包括22186例3D影像，302033個分割標註，來自CT、MR和PET三種模態，以及涵蓋人體8個主要區域的497個分割類別（解剖學結構或病灶）。
為了儘可能降低異構資料集間的差異，研究團隊對不同資料集間的方向、體素間距、灰度值等影像屬性進行了標準化，用統一的解剖學術語系統命名了不同資料集中的分割類別。
圖 3：SAT-DS是一個大規模、多樣化的3D醫學影像分割資料集合，涵蓋人體8個主要區域共497個分割類別。
模型架構
知識注入：為了構建可以精準編碼解剖學術語的提示編碼器，研究團隊首先將多模態解剖學知識用對比學習的方式注入到文字編碼器。
如下圖a所示，用解剖學概念將多模態知識連線成對，隨後使用視覺編碼器（visual encoder）和文字編碼器（text encoder）分別編碼視覺和文字知識，透過對比學習在特徵空間中將解剖學結構的視覺特徵和文字知識對齊，並構建解剖學結構之間的關係，從而學習到對解剖學概念的更好編碼，作為提示引導視覺分割模型的訓練。
基於文字提示的通用分割：研究團隊進一步設計了基於文字提示的通用分割模型框架，如下圖b所示，包含文字編碼器、視覺編碼器、視覺解碼器與提示解碼器。
其中，考慮到同一解剖結構在不同影像中存在差異，提示解碼器（query decoder）使用視覺編碼器輸出的影像特徵增強解剖學概念特徵，即分割提示。最後，在分割提示與視覺解碼器輸出的畫素級特徵間計算點積，得到分割預測結果。
模型測評
該研究將SAT與兩個代表性的方法進行對比，即「專用」模型nnU-Nets和互動式通用分割模型MedSAM。測評包含了域內資料集測試（綜合分割效能）與zero-shot域外資料集測試（跨中心資料遷移能力）兩方面，評測結果從資料集、類別和人體區域三個層面進行了整合：
類別：不同資料集之間相同類別的分割結果進行彙總、平均；
區域：基於類別結果，將同一人體解剖區域內的類別結果進行彙總、平均；
資料集：傳統的分割模型評估方式，同一資料集內的分割結果進行平均；
與專用模型nnU-Nets的對比實驗
為了最大化nnU-Nets的效能，該研究在每個單獨的資料集上訓練nnU-Nets並與SAT對比，具體設定如下：
1. 在域內測試中，使用SAT-DS中的全部72個資料集進行測試和對比。對於SAT，使用72個訓練集的總和進行訓練，並在72個測試集上進行測試；對於nnU-Nets，彙總72個nnU-Nets在各自測試集上的結果作為一個整體。
2. 在域外測試中，進一步劃分72個資料集，使用其中的49個資料集（命名為SAT-DS-Nano）的訓練集訓練SAT-Nano，在10個域外的測試集上zero-shot測試；對於nnU-Nets，使用49個nnU-Nets在10個域外測試集上測試並彙總結果。
域內測試結果：從表1可以看到，SAT-Pro在域內測試中表現出與72個nnU-Nets十分接近的效能，並在多個區域上超越nnU-Nets。需要注意的是，SAT可以僅用一個模型完成72個分割任務，並在模型尺寸上遠小於nnU-Nets的集合（如下圖c所示）。
Fine-tuning遷移測試結果：該研究進一步將SAT-Pro在每個資料集上單獨fine-tune後進行測試，命名為SAT-Pro-Ft。從表1中可以看出，SAT-Pro-Ft相比SAT-Pro在所有區域上都有了明顯的效能提升，並在總體效能上超過了nnU-Nets。
域外測試結果：如表2所示，SAT-Nano在10個資料集20個指標中的19項上超過了nnU-Nets，表現出整體更強的遷移能力。
表 2：SAT-Nano與nnU-Nets、MedSAM的域外測試對比，結果以資料集為單位呈現。
與互動式分割模型MedSAM的對比實驗
該研究直接使用MedSAM的公開checkpoint進行測試和SAT對比，具體設定如下：
1. 在域內測試中，從72個資料集中進一步篩選了32個MedSAM訓練中使用過的資料集進行對比。
2. 在域外測試中，篩選了5個MedSAM訓練中沒有使用過的資料集進行對比。
對於MedSAM，考慮兩種不同的Box提示: 使用包含ground truth分割的最小矩形（Oracle Box） ，記為MedSAM（Tight）；在Oracle Box基礎上加入隨機偏移，記為MedSAM（Loose）。同時測試Oracle Box直接作為預測的效果。對於SAT，直接沿用nnU-Nets對比實驗中的模型在這些資料集上測試，不重新訓練。
域內測試結果：如表3所示，SAT-Pro幾乎在所有的區域上都比MedSAM表現更好，且綜合表現來看SAT-Pro和SAT-Nano的效能都優於MedSAM。儘管在病灶上SAT-Pro表現不如MedSAM，Oracle Box本身作為預測在病灶上的表現也足夠好，甚至在DSC上超越了MedSAM。這表明MedSAM在病灶的分割上的更優效能很可能來自於Box提示的強先驗資訊。
表 3：SAT-Pro、SAT-Nano與MedSAM的域內測試對比，結果以區域或病灶為單位整合。
定性比較：圖6從域內測試的結果中挑選了兩個典型例子進行視覺化展示，進一步對比SAT和MedSAM。如圖6上所示，在對心肌的分割中，Box提示很難區分心肌和被心肌包裹的心室，因此MedSAM也錯誤將兩者一起分割了出來，這表明Box提示在類似的複雜空間關係中很容易有歧義，導致分割不精準。
相比之下，基於文字提示（直接輸入解剖結構的名字）的SAT可以精準地區分心肌和心室。此外，在圖6下展示的腸道腫瘤分割中可以看到，Oracle Box對於病灶目標來說已經是很好的預測結果，而MedSAM的分割結果可能不會比所得到的Box提示更好。
圖 6：SAT-Pro與MedSAM（Tight）的定性比較。其中MedSAM使用Oracle Box作為提示，Box用藍色標識。第一行展示了一個心肌分割的例子；第二行展示了一個腸道腫瘤分割的例子。
域外測試結果：如表2所示，與MedSAM（Tight）相比，SAT-Nano在5個資料集共10個指標中的5個上超過了MedSAM。而MedSAM（Loose）在所有指標上都有明顯的效能下降，表明MedSAM對於使用者輸入的Box提示的偏移比較敏感。
消融實驗
在設計SAT時，視覺骨幹網路和文字編碼器是兩個關鍵的部分，該研究嘗試在SAT框架中使用不同的視覺網路結構或文字編碼器，通用消融實驗來探討它們的影響。
為了節省實驗的開銷，消融實驗中的所有SAT模型訓練和測試都在包含49個資料集的SAT-DS-Nano上進行，它包含了13303個3D影像，151461個分割標註，以及429個分割類別。
視覺骨幹網路：在SAT-Nano的框架下，該研究選用了三種主流的分割網路結構進行對比，即U-Net（110M引數），SwinUNETR（107M引數）和U-Mamba（114M引數）。為了公平對比，該消融實驗中控制它們的引數量大致接近。同時為了計算開銷，省略了知識注入的步驟，直接使用MedCPT（MedCPT是基於PubMed文獻，使用225M私有使用者點選資料訓練的文字編碼器，在一系列醫療語言任務中取得了最好的效能）作為文字編碼器生成提示。三種變體分別被記為U-Net-CPT、SwinUNETR-CPT和U-Mamba-CPT。
從圖7中可以看到，使用U-Net與U-Mamba作為視覺骨幹網路，最終的分割效能比較接近，其中U-Net略好於U-Mamba；而使用SwinUNETR時的分割效能有明顯下降。最後，研究團隊選擇U-Net作為SAT的視覺骨幹網路。
文字編碼器：在SAT-Nano的框架下，該研究選擇了三種代表性的文字編碼器進行對比：使用上文提出的知識注入方式訓練的文字編碼器（記為Ours），使用先進的醫療文字編碼器MedCPT，使用沒有針對醫療資料微調的文字編碼器BERT-base。
為了公平，該消融實驗統一採用U-Net作為視覺網路。三種變體分別被記為U-Net-Ours，U-Net-CPT和U-Net-BB。如圖8所示，總體而言，使用MedCPT相比使用BERT-base對於分割效能有微小的提升，表明領域知識對提供好的分割提示有一定的幫助；而使用該研究提出的文字編碼器在所有的類別上都取得了最好的效能，表明構建多模態人體解剖學知識庫和知識注入對於分割模型有明顯的幫助。
長尾分佈是分割資料集的一個明顯特徵。如圖9中a圖和b圖所示，研究團隊調研了用於消融實驗的SAT-DS-Nano中429個類別的標註數量分佈。如果把標註數最多的10個類（前2.33%）定義為頭部，標註數最少的150個類別（後34.97%）定義為尾部類，可以發現尾部類的標註數量僅佔總體標註數的3.25%。
該研究進一步探究文字編碼器對長尾分佈中不同類別的分割結果影響。如圖9中c圖所示，研究團隊提出的編碼器在頭部、尾部和中間類別上都取得了最好的效能，其中尾部類上的提升比頭部類更加明顯。同時，MedCPT在頭部類上表現略低於BERT-base，但在尾部類上效果更好。這些結果都表明，領域知識，特別是多模態人體解剖學知識的注入，對於長尾類別的分割有明顯幫助。
與大語言模型的結合
由於SAT可以基於文字提示進行分割，它可以被直接用作大語言模型的代理工具，提供分割能力。為了展示應用場景，研究團隊選擇了4個多樣的真實臨床資料，使用GPT4從報告中提取分割目標並呼叫SAT進行zero-shot分割，結果展示在圖10中。
可以看到，GPT-4可以很好地檢測到報告中重要的解剖學結構，並且呼叫SAT，在不需要任何資料fine-tune的情況下，在真實臨床影像上很好地分割出這些目標。
研究價值
作為首個基於文字提示的3D醫療影像通用分割大模型，SAT的價值體現在許多方面：
SAT構建了高效靈活的通用分割：SAT-Pro僅用一個模型，在廣泛的分割任務上展現出與72個nnU-Nets相當的效能，並且有更少的模型引數量。這表明相比傳統的醫療分割方法需要配置、訓練和部署一系列專用模型，SAT-Pro作為通用分割模型是更靈活高效的解決方法。同時，研究團隊也證明SAT-Pro在域外資料上泛化效能更好，能更好地滿足跨中心遷移等臨床需求。
SAT是基於大規模分割資料預訓練的基礎模型：SAT-Pro在大規模的分割資料集上進行訓練後，當透過微調遷移到特定資料集上時，表現出了明顯的效能提升，並比nnU-Nets總體效能更好。這表明SAT可以被視作強大的基礎分割模型，可以透過微調遷移在特定任務上表現更好，從而平衡通用分割與專用分割的臨床需求。
SAT實現了基於文字提示的準確、魯棒分割：相比基於Box提示的互動式分割模型，SAT基於文字提示可以取得更準確和對提示魯棒的分割結果，並且可以節省使用者圈畫Box的大量時間，從而實現自動的、可批量化的通用分割。
SAT可以作為大語言模型的代理工具：研究團隊在真實臨床資料上展示了SAT可以和大語言模型無縫銜接，透過文字作為橋樑，直接為任何大語言模型提供分割和定位能力。這對進一步推動通用醫療智慧（Generalist Medical Artificial Intelligence）的發展有重要價值。
模型尺寸對於分割的影響：透過訓練兩種不同大小的模型：SAT-Nano和SAT-Pro，該研究觀測到在域內測試中SAT-Pro相比SAT-Nano有明顯的提升。這暗示著在大規模資料集上訓練通用分割模型時，scaling-law依然適用。
領域知識對於分割的影響：研究團隊提出了第一個多模態的人體解剖學知識庫，並探索用知識增強提升通用分割模型的效能，特別是對長尾類別的分割。考慮到分割標註，特別是長尾類別上的標註，相對稀缺，這一探索對於構建通用的分割模型有重要意義。