## 3 Methodology
### Variational Autoencoder Stage  
Variational autoencoders are able to capture latent representations, which makes them ideal for use in generative settings [5; 30]. The evidence lower bound (ELBO), which is VAEs optimization objective, can be written as:  
\[\mathcal{L}_{VAE}(\mathbf{i},\mathbf{\hat{i}})=\mathcal{L}_{REC}(\mathbf{i},\mathbf{\hat{i}})+ \mathcal{K},\mathcal{L}[q_{\Lambda}(\mathbf{r}|\mathbf{i})||p(\mathbf{r})] \tag{9}\]  
where the term \(\mathcal{L}_{REC}(\mathbf{i},\mathbf{\hat{i}})\) denotes reconstruction loss, and can be further written as:  
\[\mathcal{L}_{REC}(\mathbf{i},\mathbf{\hat{i}})=-\mathbb{E}_{q_{\Lambda}|i}[log(p_{ \Delta}(\mathbf{i}|\mathbf{r}))] \tag{10}\]  
where \(\mathbf{\hat{i}}\) denotes the reconstructed input.  
The term \(\mathcal{K},\mathcal{L}[q_{\Lambda}(\mathbf{r}|\mathbf{i})||p(\mathbf{r})]\) from Equation (9) defines the \(KL\) divergence of the approximating variational density, which can be expressed as:  
\[q_{\Lambda}(\mathbf{r}|\mathbf{i})=\mathcal{N}(\mathbf{r};\mu_{\Lambda},\sigma_{\Lambda} ^{2}) \tag{11}\]  
The standard prior on the latent variable can be written as:  
\[p(\mathbf{r})=\mathcal{N}(\mathbf{r};0.1) \tag{12}\]  
where the aligned Gaussian \((\mu_{\Lambda},\sigma_{\Lambda}^{2})\) is expressed by the encoder network \(V_{\Lambda}(\cdot)\).  
Following this, the low dimensional representations of the input data \(\mathbf{i}\) can be obtained by introducing the latent random variable \(\mathbf{r}\). The input images are mapped to a low dimensional space using VAE encoder \(V_{\Lambda}(\cdot)\). After that, the output of the encoder of the segmentation network, \(E_{\Omega}(\cdot)\), takes samples from the latent space as shown in Figure 1. In this manner, segmentation encoder and VAE jointly share the decoder \(D_{\Delta}(\cdot)\), which can be also written as:  
\[\mathcal{L}(\mathbf{o},\mathbf{\hat{o}})=\mathcal{L}_{REC}(\mathbf{o},\mathbf{\hat{s}})+ \mathcal{KL}[q_{\Lambda}(\mathbf{r}|i)||p(\mathbf{r})] \tag{13}\]  
The final segmentation, \(\mathbf{\hat{o}}\), is obtained from the decoder using following expression:  
\[\mathbf{\hat{o}}=D_{\Delta}[E_{\Omega}(i)\circ V_{\Lambda}(i)] \tag{14}\]  
which can be written as:  
\[\mathbf{\hat{o}}=D_{\Delta}[\mathbf{h}\circ\mathbf{r}] \tag{15}\]  
where \(\circ\) denotes concatenation, \(H=E_{\Omega}(i)\) is the output of the segmentation encoder and \(r\sim q_{\Lambda}(r|i)\) is a sample from the latent space that is learnt by VAE.  
The loss function plays an important role in improving the models' performance. In this work, we employ total loss function that is the addition of soft dice loss, \(L2\) loss and standard VAE penalty term, and can be written as:  
\[L=L_{dice}+0.1\cdot L_{L2}+0.1\cdot L_{KL} \tag{16}\]  
The term that represents the soft dice loss, \(L_{dice}\)[31], can be written as:  
\[L_{dice}=\frac{2\cdot\sum p_{gt}\cdot p_{pred}}{\sum p_{gt}^{2}+\sum p_{pred}^{ 2}+\epsilon} \tag{17}\]  
where \(\epsilon\) denotes a small constant used for computational stability, i.e., to avoid zero division.  
The loss \(L2\) represents loss on the VAE encoder output and can be written as:  
\[L_{L2}=||I_{input}-I_{pred}||_{2}^{2} \tag{18}\]The standard VAE penalty term, \(L_{KL}\), represents \(KL\) divergence between a prior distribution \(N(0,1)\) and the estimated normal distribution \(N(\mu,\sigma^{2})\), which can be written as:  
\[L_{KL}=\frac{1}{N}\sum\mu^{2}+\sigma^{2}-log\sigma^{2}-1 \tag{19}\]  
where \(N\) represents the entire set of image voxels. Finally, the hyperparameter weight of 0.1 is empirically found to provide a good balance between VAE loss term and soft dice loss in Equation (16).