掃碼關注下方公眾號："Python編程與深度學習"，領取配套學習資源，並有不定時深度學習相關文章及代碼分享。
今天分享一篇發表在CVPR 2020上的論文：LT-Net: Label Transfer by Learning Reversible Voxel-wise Correspondence for One-shot Medical Image Segmentation (原文鏈接：[1])。
1 研究背景
近年來隨着深度學習的快速發展，深度卷積神經網絡 (DCNNs)在許多分割任務上取得很好的性能。但是對於3D醫學圖像分割任務，獲得3D空間中的體素標注是困難的，因此模型需要學習如何從一個或者少量幾個標注樣本中進行有效地學習。
對於這個問題，傳統的解決方案是atlas-based的分割方法。而在這篇文章中，作者直接采用深度學習模型來模擬atlas-based的經典分割方法，提出了LT-Net。LT-Net以atlas有標注圖像和無標注圖像作為輸入，預測前者到后者的對應映射關系，這樣子就可以利用映射關系將atlas圖像上的分割圖轉移到無標注的圖像上。
這篇文章主要有三點貢獻：
- 為了解決缺少標注的問題，借助經典的atlas-based方法的分割思想來解決one-shot分割問題
- 以端到端方式將對應映射關系的學習擴展到one-shot分割框架中，其中前向和反向構成的循環一致性 (forward-backward cycle-consistency)在圖像、轉換和標簽空間中起到了額外監督的重要作用 （注：這里的前向和反向不同於深度學習的前向傳播和反向傳播，具體見后面章節的描述）
- 通過充分的實驗證明所提出方法的有效性
2 方法
2.1 整體流程
假設$(l,l_s)$表示一對具有標注的圖像對，其中$l\in \mathbb{R}^{h\times w\times c}$是atlas圖像，$l_s\in \mathbb{R}^{h\times w\times c}$是對應的分割圖。對於所有輸入圖像，他們都被定義在3D空間$\Omega\in\mathbb{R}^3$中，其中無標注圖像集合可以表示為$\{u^{(i)}|u^{(i)}\in\mathbb{R}^{h\times w\times c}\}$，后面將采用$u$來表示一副無標注的圖像。
在這篇文章中，作者希望模型不單單能夠學習從$l$到$u$的對應映射關系，也應該學習從$u$到$l$的對應映射關系，對於前者稱為前向過程 (forward)，對於后者稱為反向過程 (backward)。
首先采用$\Delta_{p_F}$表示有標注的atlas圖像$l$和無標注圖像$u$之間的對應映射關系（前向過程），則將$\Delta_{p_F}$作用在$l$上的過程可以表示為：
$$\bar{u}=l\circ\Delta_{p_F},$$
其中$\circ$表示變換操作 (warp operation)，$\bar{u}$表示從$l$變換出來的atlas圖像，那么其對應的分割圖可以由$l_s$經過同樣的映射關系變換得到：
$$\bar{u}_s=l_s\circ\Delta_{p_F},$$
然后采用$\Delta_{p_B}$表示無標注圖像$u$到有標注的atlas圖像$l$之間的對應映射關系（反向過程），則將$\Delta_{p_B}$作用在$\bar{u}$上的過程可以表示為：
$$\bar{l}=\bar{u}\circ\Delta_{p_B},$$
對應的分割圖也可以進行變換：
$$\bar{l}_s=\bar{u}_s\circ\Delta_{p_B},$$
其中我們期望$\bar{l}$以及$\bar{l}_s$應該分別和$l$以及$l_s$是一致的，因此上述的這兩個過程就構成了一個循環 (cycle)。接下來的兩小節將詳解怎么構造模型來對$\Delta_{p_F}$和$\Delta_{p_F}$進行預測。
2.2 前向對應映射關系的學習
文中采用DCNN（例如VoxelMorph）構成的生成網絡$G_F$來學習對應的映射關系$\Delta_{p_F}$，其輸入是$l$和$u$，輸出是他們之間的對應關系$\Delta_{p_F}$，網絡的優化過程是通過最小化兩個無監督損失函數：
- 圖像相似度損失函數$\mathcal{L}_{sim}(u,\bar{u})$：采用局部歸一化互相關損失 (locally normalized cross-correlation (CC) loss)
- 轉換平滑損失函數$\mathcal{L}_{smooth}(\Delta_{p_F})=\sum_{t\in\Omega}\left\|\bigtriangledown(\Delta_{p_F}(t))\right\|_2$
除了上述兩個損失函數之外，文中還加入了輔助GAN損失方式來提供額外的監督。GAN子網絡由$G_F$和判別器$D$構成，目的是使得通過$G_F$生成的$\bar{u}$能夠在判別器$D$下“蒙混過關”。這部分的損失函數定義為：
$$\mathcal{L}_{GAN}(l,u,\bar{u})=\mathbb{E}_{u\sim p_d(u)}[\left\|D(u)\right\|]_2+\mathbb{E}_{l\sim p_d(l),u\sim p_d(u)}[\left\|D(\bar{u})-\mathbf{1}\right\|_2],$$
其中$G_F$和$D$在目標函數下交替訓練，在兩者的最小-最大博弈中進行學習：$$min_{G_F}max_D\mathcal{L}_{GAN}(G_F,D).$$
2.3 反向對應映射關系的學習
文中把上述2.1小節中介紹的反向過程中的對應映射關系$\Delta_{p_B}$稱為反向對應映射關系，同樣采用DCNN（例如VoxelMorph）構成的生成網絡$G_B$來進行學習，有$\Delta_{p_B}=G_B(\bar{u},l)$。因此對於上面2.2小節中由$G_F$生成的$\bar{u}$，可以再通過$G_B$生成$l$的重構atlas圖像$\bar{l}$，表示如下：
$$\bar{l}=\bar{u}\circ\Delta_{p_B},$$
結合上面2.2小節，總的轉換平滑損失函數可以表示為：
$$\mathcal{L}_{smooth}=\mathcal{L}_{smooth}(\Delta_{p_F})+\mathcal{L}_{smooth}(\Delta_{p_B}),$$
此外，如Figure 2所示，$G_F$和$G_B$構成的循環中可以加入三個監督信號：
- 圖像空間上的循環一致性監督：$l$根據$G_F$的預測結果轉換成$\bar{u}$，$\bar{u}$再根據$G_B$的預測結果轉換為$\bar{l}$，其中期望$\bar{l}$與$l$是一致的，因此加入L1損失函數進行監督學習，有：$$\mathcal{L}_{cyc}(l,\bar{l})=\mathbb{E}_{l\sim p_d(l)}[\left\|\bar{l}-l\right\|_1],$$
- 轉換空間上的循環一致性監督：對於某個位置，通過前向過程轉換之后，再通過反向過程的轉換應該回到原來的位置，這部分的損失可以表示為：$$\mathcal{L}_{trans}(\Delta_{p_F},\Delta_{p_B})=\sum_{t\in \Omega}\rho(\Delta_{p_F}(t)+\Delta_{p_B}(t+\Delta_{p_F}(t))),\rho(x)=(x^2+\epsilon^2)^{\gamma},$$
- 標簽空間上的循環一致性監督（第一部分）：$l_s$根據$G_F$的預測結果轉換成$\bar{u}_s$，$\bar{u}_s$再根據$G_B$的預測結果轉換為$\bar{l}_s$，其中期望$\bar{l}_s$與$l$是一致的，加入Dice損失函數進行監督，有：$$\mathcal{L}_{anatomy\_cyc}(l_s,\bar{l}_s)=1-\frac{2\sum_{t\in\Omega}l_s(t)\bar{l}_s(t)}{\sum_{t\in\Omega}l_s^2(t)+\sum_{t\in\Omega}\bar{l}_s^2(t)},$$
- 標簽空間上的循環一致性監督（第二部分）：由於模型的目標是學習可用於將altas圖像上的分割圖轉移到每個未標注圖像的對應關系，因此文中還額外提出了差異一致性損失，以間接地規范化合成的分割圖的質量，有：$$\mathcal{L}_{diff\_cyc}(l_s,\bar{u}_s,\bar{l}_s)=\sum_{t\in\Omega}\rho(\left |l_s(t)-\bar{u}_s(t)\right |-\left |\bar{u}_s(t)-\bar{l}_s(t)\right |).$$
2.4 損失函數
綜合上述2.2小節和2.3小節，整個模型的優化目標可以表示為：
$$\mathcal{L}=\mathcal{L}_{GAN}+\mathcal{L}_{sim}+\lambda_1\mathcal{L}_{cyc}+\lambda_2(\mathcal{L}_{anatomy\_cyc}+\mathcal{L}_{smooth}+\mathcal{L}_{trans}+\mathcal{L}_{diff\_cyc}),$$
其中$\lambda_1=10,\lambda_2=3.$
3 實驗結果
這里我只給出論文中的部分實驗結果，具體的實驗結果分析以及實驗和參數的設置請看原文。
4 參考資料
[1] http://openaccess.thecvf.com/content_CVPR_2020/papers/Wang_LT-Net_Label_Transfer_by_Learning_Reversible_Voxel-Wise_Correspondence_for_One-Shot_CVPR_2020_paper.pdf