導讀

利用變分推斷進行分割置信度的預測。

在過去的十年里，深度學習在一系列的應用中取得了巨大的成功。然而，為了驗證和可解釋性，我們不僅需要模型做出的預測，還需要知道它在做出預測時的置信度。這對于讓醫學影像學的臨床醫生接受它是非常重要的。在這篇博客中，我們展示了我們在韋洛爾理工學院進行的研究。我們使用了一個基于變分推理技術的編碼解碼架構來分割腦腫瘤圖像。我們比較了U-Net、V-Net和FCN等不同的主干架構作為編碼器的條件分布采樣數據。我們使用Dice相似系數(DSC)和IOU作為評價指標來評價我們在公開數據集BRATS上的工作。

醫學圖像分割

在目前的文獻中主要利用兩種技術成功地解決了醫學圖像的分割問題，一種是利用全卷積網絡(FCN)，另一種是基于U-Net的技術。FCN體系結構的主要特點是在最后沒有使用已成功用于圖像分類問題的全連接層。另一方面，U-Net使用一種編碼器-解碼器架構，在編碼器中有池化層，在解碼器中有上采樣層。

貝葉斯神經網絡

這是一種可擴展的避免神經網絡過擬合的方法，同時也給了我們一個不確定性的度量。神經網絡學習給定的數據集的后驗分布的權重，而不是基于點的估計，如下面的公式所示。

預測分布可以通過逼近積分來計算，如下式所示。

變分推斷

變分推斷通過最大化證據下界來尋找分布的參數。ELBO由前后分布的Kullback-Leibler (KL)散度和負對數似然(NLL)兩項之和構成。需要最小化的KL散度項如下式所示。

KL散度定義如下式。

由于上述方程中的積分在本質上是難以處理的，它可以寫成另一種形式。該方程可轉化為優化問題，如下式所示。

隨機不確定性和認知不確定性

有兩種類型的不確定性 —— 隨機不確定性和認知不確定性，其中方差是兩者的總和。對于最終的預測，單個的均值和方差可以估計，如下兩個方程所示。

方差中的第一項表示隨機不確定性，而第二項表示認知不確定性。

網絡結構

先驗分布有助于整合網絡上的權值學習。我們的模型使用了與VAEs中使用的類似的編碼器解碼器體系結構，編碼器的輸入來自預先訓練好的圖像分割結構。輸入到編碼器只需要表示置信度的條件分布的標準差向量的均值，以此來正確預測像素點。參數經過編碼器后，被轉換為一個潛在表示，再采樣的平均值和標準偏差向量。解碼器隨后將其恢復到原始分布。采用傳統的反向傳播算法進行梯度下降模型的訓練。本工作中使用的模型架構如圖1所示：

圖1：模型結構

算法

下面是基于隨機梯度下降的訓練網絡的算法。

數據集

為了評估我們的網絡性能，我們使用BRATS18腦瘤分割數據集。它包含175名惡性膠質瘤和低級別惡性膠質瘤患者的MRI掃描。圖像分辨率為240×240×155像素。ground truth標簽是由神經放射學專家創建的。數據集的一個示例如圖2所示。

圖2：MRI切片的例子以及分割的ground truth

評估指標

評價指標為Dice相似系數(DSC)，也稱F1-score和IoU。對應的方程如下所示。

損失函數

采用二元交叉熵和dice損失相結合的方法對網絡進行訓練。第一部分二元交叉熵是分類問題中常用的損失函數，如下式所示:

二元交叉熵損失的問題在于它沒有考慮到類的不平衡，因為背景是占主導地位的類。dice損失解決了這個問題，可以寫成如下公式。

這兩個損失項被合并在一個項中，并給予dice損失項更多的權重，因為它能更好地處理類別不平衡問題。這是用下面的公式定義的。

結果

分割所涉及的不確定性如圖3所示。深的顏色表示更自信，而淺的顏色表示模型在這些區域不太自信。

圖3：與ground truth分割相比，測試樣本上的模型預測示例。第一列:輸入圖像，第二列:真值分割，第三列:預測分割，第四列:隨機不確定性，第五列:認知不確定性

總結

在這個博客中，我們提出了一種在醫學圖像分割中量化不確定性的方法。我們的模型基于一個類似于VAEs所使用的編碼器解碼器框架。網絡的權值代表分布而不是點估計，從而在進行預測的同時給出了一種原則性的測量不確定性的方法。編碼器的輸入來自于預訓練的骨干架構，如U-Net, V-Net, FCN，這些架構都是從條件分布中采樣的，代表了像素被正確標記的置信度。我們在公開數據集BRATS上評估我們的結果，使用DSC和IOU指標，我們的網絡優于以前的SOTA的結果。

責任編輯：lq