Distribution Aligned Diffusion and Prototype-guided network for Unsupervised Domain Adaptive Segmentation

The exploration of latent feature representations from generative models has been widely investigated in many dense prediction tasks [21, 24]. However, GANs are prone to collapse and out-of-distribution issues. On the other hand, the diffusion model has attracted much research attention and has demonstrated superior performance in many vision tasks due to its noise robustness.

Instead of taking labels as input in the original DPM [1], we adapt the vanilla Diffusion [7], which is originally used to generate an image. This allows us to obtain latent feature representations from the intermediate activations of DPM, which contain rich semantic information [2]. In order to align the feature distributions between the source and target domains, we propose Distribution Aligned Diffusion (DADiff), as shown in Fig. 1. The generated images are discarded, and only the activations features of DPM are utilized. DADiff introduces a sequence of Gaussian noises to the input image in the forward pass and reconstructs the image in the reverse pass to predict the noise. Specifically, given a sample $x_{0}$, the forward process $q$ is:

where $\beta_{t}$ denotes the variance condition at the step $t$. $I$ is the identity matrix. $\mathcal{N}$ is a distribution. The noisy sample $x_{t}$ at the time step $t$ can be obtained as:

where $\alpha_{t}=1-\beta_{t}$. $\bar{\alpha_{t}}=\prod_{n=1}^{t}\alpha_{n}$. Meanwhile, the reverse diffusion process $p_{\theta}$ is parametered by $\theta$ and can be described as:

Instead of directly obtaining the distribution $\mathcal{N}$ of Eq. 3, our DADiff follows [12] to predict the noise via a UNet $\epsilon_{\theta}$:

where $\sigma_{t}$ denotes the variance scheme learnt by $\epsilon_{\theta}$. $\epsilon_{\theta}$ is often assign with a UNet. We follows [7] to compute the noise estimation loss to train the UNet $\epsilon_{\theta}$.

We utilize the trained UNet to extract latent diffusion features $F^{S}$ from the source image $x_{0}^{S}$ and $F^{T}$ from the target image $x_{0}^{T}$. Specifically, in the reverse diffusion process, our DADiff follows the Markov step to collect the latent activations $F^{S}$ and $F^{T}$, respectively. To reduce the domain discrepancy between these two domains, we implement a domain discriminator $D$ with a gradient reversal layer (GRL) [11] to ensure inter-domain indistinguishability, as shown in Fig. 1. Specifically, we devise the following loss $\mathcal{L}_{D}$ to learn the domain discrimintor:

where ${l}^{S}=1$ and ${l}^{T}=0$ denote the source and target domain labels, respectively. The coordinate in the diffusion feature map is denoted as $(x,y)$.

Given the latent diffusion features generated by DADiff in the first stage, we introduce a Prototype-guided Consistency Learning (PCL) module in the second stage to further narrow the gap between the source and target domains, where the detail is shown in Fig. 2. In the UDA setting, we have access to the source domain features $\{F_{i}^{S}\}^{N_{S}}_{i=1}$ and their corresponding segmentation annotations $\{Y_{i}^{S}\}^{N_{S}}_{i=1}$, as well as the target domain features $\{F_{i}^{T}\}^{N_{T}}_{i=1}$ without any label. To perform domain adaptation, we first pass the source domain features through a shared segmentation decoder $G$ to obtain segmentation results, on which we perform supervised learning to improve the segmentation performance on the source domain. Similarly, we generate pseudo labels $\{\hat{Y}_{i}^{T}\}^{N_{T}}_{i=1}$ for the target domain based on a probability threshold selection $\mathbbm{1}[p^{T}\geq\gamma]$. With inputting $F^{S}$ and $F^{T}$, one can obtain the output feature $f^{S}$ and $f^{T}$ from the last convolution layer of the segmentation decoder. And we compute the object prototype $z_{obj}^{S}$ at the source domain and the object prototype $z_{obj}^{T}$ at and target domain as:

where $i$ denotes the index of pixel and $N_{obj}$ is the number of pixels of object classes. However, the prototype of the target domain may be unreliable due to false predictions produced by $G$. To address it, we follow [4] and estimate the uncertainty to denoise the coarse pseudo labels. Using Monte Carlo Dropout [10], we subject the decoder to dropout during $K$ stochastic forward passes on each sample, generating a series of coarse pseudo labels $\{\hat{Y}_{k}^{T},k=1,...,K\}$. We compute the uncertainty map $m$ by taking the standard deviation $std(\hat{Y}_{1}^{T},...,\hat{Y}_{k}^{T})$ of the $K$ forward predictions. We then obtain a refined pseudo label $\acute{z}_{obj}^{T}$ by applying an uncertainty threshold $\eta$ to generate a stable prediction.

After refining the pseudo labels in the target domain, we propose a prototype-guided consistency learning loss $\mathcal{L}_{pcl}$ to further narrow the gap between the source and target domains for domain adaptation. This loss aims to reduce the distance between the prototypes (i.e., $z_{obj}^{S}$ and $\acute{z}_{obj}^{T}$) from the source and target domains. The definition of $\mathcal{L}_{pcl}$ is given by:

By adding $\mathcal{L}_{pcl}$ with the supervised learning loss $\mathcal{L}_{seg}$ at the source domain, the total loss $\mathcal{L}_{M}$ of the stage 2 is computed by:

where the weight $\lambda$ is empirically set as $\lambda=0.5$ in our experiments.

We conduct ablation study experiments to investigate major components of our framework. We first construct a baseline (denoted as “M1”) by only extracting diffusion features to predict the segmentation results. It is equal to remove the domain discriminator $\mathcal{L}_{D}$ at Eq. 5 and the prototype-guided consistency loss $\mathcal{L}_{pcl}$ at Eq. 7 of the PCL module. Then, we add the domain discriminator $\mathcal{L}_{D}$ into “M1” to construct another baseline network (“M2”). After that, we add a DPL [4] module into “M2” to build a baseline network (“M3”), and the prototype-guided consistency loss $\mathcal{L}_{pcl}$ into “M2” to build “M4”. Apparently, “M4” is equal to the full setting of our network.

Table 2 reports the Dice score of the optic disc and cup segmentation of our method and baseline networks in terms of Drishti-GS and RIM-ONE-r3 datasets. Apparently, “M2” has larger Dice scores than “M1” on the optic disc and cup segmentation for both two datasets. It indicates that the domain discriminator $\mathcal{L}_{D}$ of Eq. 5 on diffusion features can reduce the domain gap of the source and target domains, thereby improving the UDA performance. Then, “M3” and our method (i.e., “M4”) outperform “M2” for the optic disc and cup segmentation on the two datasets. It demonstrates that exploring the prototype information to compute consistency loss can further enhance the domain adaptation performance. More importantly, our method has a superior Dice score over “M3”. It indicates that our PCL can better reduce the domain gap than DPL [4]. Specifically, compared to “M3”, our method improves the Dice score of the optic disc segmentation from 0.963 to 0.966, and the Dice score of the optic cup segmentation from 0.876 to 0.844 on the Drishti-GS dataset. For the RIM-ONE-r3 dataset, our method improves the Dice score of the optic cup segmentation from 0.829 to 0.852. And the Dice scores of the optic disc segmentation for our method and “M3” are very close, and they are 0.916 and 0.913, respectively.