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

CoTTA

EcoTTA

DIGA

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

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

請?zhí)砑訄D片描述

現(xiàn)有的 TTA 方法性能較差，原因在于未標(biāo)記的目標(biāo)域圖像提供的監(jiān)督信號不足，或者受到源域中預(yù)訓(xùn)練策略和網(wǎng)絡(luò)結(jié)構(gòu)的特定要求的限制。這篇工作目標(biāo)是將源域中的預(yù)訓(xùn)練與目標(biāo)域中的適應(yīng)分開，以實現(xiàn)高性能且更具一般性的 TTA，而不對預(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é)果好很多。

請?zhí)砑訄D片描述

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

請?zhí)砑訄D片描述

具體實現(xiàn)

Pre-trained Model from the Source Domain

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

Test-Time Growing for Adaptation

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

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

Supervision with Reliable Pseudo Labels

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

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

Mean Prediction-Based Entropy Minimization

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

Adaptation by Self-training

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

實驗

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

請?zhí)砑訄D片描述

和其他 SOTA 方法的對比：

請?zhí)砑訄D片描述

可視化結(jié)果：

請?zhí)砑訄D片描述

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

請?zhí)砑訄D片描述

消融實驗：

請?zhí)砑訄D片描述

總結(jié)

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

審核編輯：彭菁