基于一種用于醫(yī)學(xué)圖像分割的方法

這次我們要解讀的工作發(fā)表在 IPMI 2023（IPMI全名 Information Processing in Medical Imaging，兩年一屆，是醫(yī)學(xué)影像分析處理領(lǐng)域公認(rèn)的最具特色的會(huì)議），同時(shí)也是 Test Time Adaptation 系列的文章，之前的 TTA 論文解決在：

CoTTA

EcoTTA

DIGA

對(duì) TTA 不了解的同學(xué)可以先看上面這幾篇新工作?；?a href="http://www.ttokpm.com/v/tag/448/" target="_blank">深度學(xué)習(xí)的醫(yī)學(xué)成像解決方案的一個(gè)主要問(wèn)題是，當(dāng)一個(gè)模型在不同于其訓(xùn)練的數(shù)據(jù)分布上進(jìn)行測(cè)試時(shí)，性能下降。將源模型適應(yīng)于測(cè)試時(shí)的目標(biāo)數(shù)據(jù)分布是解決數(shù)據(jù)移位問(wèn)題的一種有效的解決方案。以前的方法通過(guò)使用熵最小化或正則化等技術(shù)將模型適應(yīng)于目標(biāo)分布來(lái)解決這個(gè)問(wèn)題。

在這些方法中，模型仍然通過(guò)使用完整測(cè)試數(shù)據(jù)分布的無(wú)監(jiān)督損失反向傳播更新。但是在現(xiàn)實(shí)世界的臨床環(huán)境中，實(shí)時(shí)將模型適應(yīng)于新的測(cè)試圖像更有意義，并且需要避免在推理過(guò)程中由于隱私問(wèn)題和部署時(shí)缺乏計(jì)算資源的情況。TTA 在遇到來(lái)自未知領(lǐng)域的圖像時(shí)，有望提高深度學(xué)習(xí)模型的魯棒性。本文要介紹的工作屬于 Fully Test Time Adaptation，既在推理時(shí)需要網(wǎng)絡(luò)要做完整的反向傳播。下表簡(jiǎn)單列出幾種常見(jiàn) settings，其中 Fully TTA 只需要 target data 和 test loss。

請(qǐng)?zhí)砑訄D片描述

現(xiàn)有的 TTA 方法性能較差，原因在于未標(biāo)記的目標(biāo)域圖像提供的監(jiān)督信號(hào)不足，或者受到源域中預(yù)訓(xùn)練策略和網(wǎng)絡(luò)結(jié)構(gòu)的特定要求的限制。這篇工作目標(biāo)是將源域中的預(yù)訓(xùn)練與目標(biāo)域中的適應(yīng)分開(kāi)，以實(shí)現(xiàn)高性能且更具一般性的 TTA，而不對(duì)預(yù)訓(xùn)練策略做出假設(shè)。

概述

這篇工作（被叫做 UPL-TTA）提出了一種用于醫(yī)學(xué)圖像分割的 “Uncertainty-aware Pseudo Label guided Fully Test Time Adaptation" 方法，該方法不要求預(yù)訓(xùn)練模型在適應(yīng)到目標(biāo)領(lǐng)域之前在源域中進(jìn)行額外的輔助分支訓(xùn)練或采用特定策略。

如下圖所示，以嬰兒腦部 MRI 腫瘤分割為例，從 HASTE 跨域到 TrueFISP，UPL-TTA 的效果要比不做任何適應(yīng)的結(jié)果好很多。

請(qǐng)?zhí)砑訄D片描述

UPL-TTA 首先引入了 Test Time Growing（TTG），也就是說(shuō)，在目標(biāo)領(lǐng)域中多次復(fù)制源模型的預(yù)測(cè)頭部。并為它們的輸入圖像和特征映射添加一系列隨機(jī)擾動(dòng)（例如，Dropout，空間變換），以獲得多個(gè)不同的分割預(yù)測(cè)。然后，通過(guò)這些預(yù)測(cè)的集成來(lái)獲取目標(biāo)領(lǐng)域圖像的偽標(biāo)簽。為了抑制潛在不正確的偽標(biāo)簽的影響，引入了集成方法和 MC dropout 不確定性估計(jì)來(lái)獲得可靠性 Map。可靠像素的偽標(biāo)簽用于監(jiān)督每個(gè)預(yù)測(cè)頭部的輸出，而不可靠像素的預(yù)測(cè)則通過(guò)平均預(yù)測(cè)圖上的熵最小化進(jìn)行規(guī)范化。

請(qǐng)?zhí)砑訄D片描述

具體實(shí)現(xiàn)

Pre-trained Model from the Source Domain

第一步我們需要優(yōu)化一個(gè)源域的預(yù)訓(xùn)練模型，即 Fig. 2 的 A 部分， 和 分別表示編碼器和解碼器的初始權(quán)重：

Test-Time Growing for Adaptation

對(duì)于 TTG 的過(guò)程，首先對(duì)于一張圖像 首先需要一個(gè)空間變換 ，包括隨即翻轉(zhuǎn)，旋轉(zhuǎn) π/2，π 和 3π/2。然后經(jīng)過(guò) Dropout（特征級(jí)別的擾動(dòng)）后在輸入到 decoder中，再進(jìn)行空間變換逆過(guò)程，得到概率圖。

上面概述里我們提到過(guò)，通過(guò)在目標(biāo)領(lǐng)域中多次復(fù)制源模型的預(yù)測(cè)頭部，并為它們的輸入圖像和特征映射添加一系列隨機(jī)擾動(dòng)（例如，Dropout，空間變換），以獲得多個(gè)不同的分割預(yù)測(cè)。最后，對(duì) K 個(gè)頭進(jìn)行集成：

Supervision with Reliable Pseudo Labels

這一步我們關(guān)注如何獲得一個(gè)可靠的偽標(biāo)簽。Fig. 2 中的 Reliable map 簡(jiǎn)單理解為一個(gè) Mask，用于優(yōu)化偽標(biāo)簽。設(shè)定一個(gè) 的閾值，我們通過(guò)對(duì)概率圖的值大小確定 的每個(gè)像素的值，只保留偽標(biāo)簽中較高可信度的像素：

到這里我們會(huì)得到三個(gè)目標(biāo)，一個(gè)是 K 個(gè)頭輸出的預(yù)測(cè)圖，第二個(gè)是偽標(biāo)簽，還有用于優(yōu)化偽標(biāo)簽的 Mask：

Mean Prediction-Based Entropy Minimization

熵最小化是 TTA 中很常用的手段，但是在 UPL-TTA 中，我們有 K 個(gè)集成。假設(shè)一種情況，第 K-1 個(gè)頭的的預(yù)測(cè)概率是 0，第 K 個(gè)頭的預(yù)測(cè)概率是 1，這時(shí)兩個(gè)頭的熵值都是最小的，但是一旦平均下來(lái)之后，0.5 對(duì)應(yīng)的熵就是大的。所以我們需要同時(shí)熵最小化 K 個(gè)頭：

Adaptation by Self-training

最后，我們的優(yōu)化目標(biāo)是下面兩個(gè)損失：

實(shí)驗(yàn)

數(shù)據(jù)集：

請(qǐng)?zhí)砑訄D片描述

和其他 SOTA 方法的對(duì)比：

請(qǐng)?zhí)砑訄D片描述

可視化結(jié)果：

請(qǐng)?zhí)砑訄D片描述

下圖是自訓(xùn)練中不同訓(xùn)練步驟的偽標(biāo)簽。Epoch 0 表示“僅源域”（自適應(yīng)之前），n 表示目標(biāo)域驗(yàn)證集上的最佳輪數(shù)。在（c）-（g）中，只有可靠的偽標(biāo)簽用顏色編碼。

請(qǐng)?zhí)砑訄D片描述

消融實(shí)驗(yàn)：

請(qǐng)?zhí)砑訄D片描述

總結(jié)

這篇 IPMI 2023 工作提出了一種完全測(cè)試時(shí)間自適應(yīng)的方法，該方法能夠在不知道源模型的訓(xùn)練策略的情況下，將源模型適應(yīng)到未標(biāo)記的目標(biāo)域。在沒(méi)有訪問(wèn)源域圖像的情況下，提出的基于不確定性的偽標(biāo)簽引導(dǎo)的 TTA 方法通過(guò)測(cè)試時(shí)間增長(zhǎng)（TTG）為目標(biāo)域中的同一樣本生成多個(gè)預(yù)測(cè)輸出。它生成高質(zhì)量的偽標(biāo)簽和相應(yīng)的可靠性Map，為未標(biāo)記的目標(biāo)域提供有效的監(jiān)督。具有不可靠偽標(biāo)簽的像素通過(guò)對(duì)復(fù)制的頭部的平均預(yù)測(cè)進(jìn)行熵最小化進(jìn)一步規(guī)范化，這也引入了隱式的一致性規(guī)范化。在胎兒腦分割的雙向跨模態(tài) TTA 實(shí)驗(yàn)中，優(yōu)于幾種最先進(jìn)的 TTA 方法。未來(lái)，實(shí)現(xiàn)該方法的 3d 版本并將其應(yīng)用于其他分割任務(wù)是很有興趣的方向。

審核編輯：彭菁