Fully Guided Neural Schrödinger bridge for Brain MR image synthesis

Recently, diffusion models have gained prominence in computer vision as an alternative to these GAN-based models [18, 19, 20]. There are two main operations of diffusion models: a forward process in which Gaussian noise of scheduled variance is applied to the image, and a reverse process in which the noise applied in the forward process is removed using a deep neural network. Diffusion models explicitly learn the likelihood due to its mathematical foundations, relatively straightforward training procedures, while effectively addressing the mode collapse issues commonly observed in GANs [17]. However, it still exists as a disadvantage that it needs to go through many sampling processes to generate a image [17].

DDGAN [17] addresses this limitation through an adversarial learning approach. In conventional diffusion models, the distribution of the reverse process is fundamentally assumed to be Gaussian. However, this theoretical assumption necessitates a large number of time steps in the reverse process to be mathematically precise [21, 22]. DDGAN implements significantly larger step sizes within each time step, thereby reducing the number of time steps in the reverse process. As a result, the reverse process distribution becomes non-Gaussian multimodal distribution, which is effectively approximated through an adversarial learning framework.

GAN-based models have great potential for many exciting real-world applications such as image generation, image processing, and super-resolution. However, The vast amount of required training data limit GAN-based models in real applications with only small image sets. In real-life scenarios, the data available to train a GAN can be minimal, such as medical data sets, a particular celebrity’s portrait set and a specific artist’s artworks.

The scarcity of training samples exacerbates the mode collapse of generator, presenting a significant challenge. Furthermore, discriminator overfitting to the limited data impedes the provision of meaningful gradients to the generator, thereby amplifying the mode collapse. [23]

Therefore, previous works mitigate discriminator overfitting with various objectives [24, 25], gradient regularizer [26], normalizing model weights [27], transfer learning [28] and data augmentation methods [24, 25] for regularization of discriminator.

Prior approaches necessitate extensive hyper-parameter tuning and have been predominantly validated on low-resolution images. Additionally, transfer learning approaches rely heavily on large-scale medical datasets and the pre-trained weights. data augmentation methods demand domain-specific customization based on image characteristics and anatomical regions.

While prior studies have demonstrated the efficacy of self-supervised learning-based discriminators in low-data regimes, their application to image-based tasks has been limited. These limitations manifest in two primary ways: the inadequacy of existing self-supervised tasks for image processing, and their restricted application to low-resolution image datasets. This creates a significant gap in the literature, particularly for high-resolution image processing applications  [29, 30].

The self-supervised discriminator [31, 32] enhances GAN training stability and image quality in few-shot scenarios by incorporating an auto-encoding task, which enables comprehensive feature extraction that captures both global composition and local textures. This approach prevents discriminator overfitting by ensuring the extraction of detailed feature maps capable of input image reconstruction, thereby providing meaningful gradients to the generator.

Schrödinger bridge is the process of finding the optimal transport trajectory from an arbitrary source distribution to a target distribution, and is one of the recently studied methodologies in diffusion model based image-to-image translation [33, 34, 35], which is limited by the Gaussian assumption [14]. UNSB [14] solves the SB problem based on entropy regularization and achieves high performance without using additional computation unlike existing SB papers.

The static formulation of intermediate sample, $\{x_{t}\}$ can be described by a Gaussian distribution [14, 23].

The stochastic control formulation [36] found two important properties of $\{x_{t}\}$. First, each element in $\{x_{t}\}$ is a Markov chain. second, $\tau$ controls the randomness for $\{x_{t}\}$ [36] .

Therefore, the distribution of intermediate $\{x_{t}\}$ given $x_{A}$ is as follows :

$Q^{SB}_{B|A}$ denotes the conditional distribution of $x_{B}$ given $x_{A}$ which means optimal trajectory of $x_{A}$ to $x_{B}$. The conditional probability measure $Q^{SB}_{B|A}$ is obtained from $Q^{SB}_{BA}$ by applying Bayes’ rule, where $Q^{SB}_{BA}$ represents an entropy-regularized optimal transport problem.

UNSB shows that SB can be represented composition of adversarial learning and Markov chain.

According to Eq. 7, we can learn the intermediate sampling outputs when we know $x_{A}$. Since we can sample all the samples of each intermediate step, we can finally sample the output predicted $x_{B}$, which we know.

$q_{\phi_{i}}$ is a generator which sample intermediate results to predicted $x_{B}$. $q_{\phi_{i}}$ is parameterized by a neural network. Eq. 10 demonstrates that $L_{adv}$ serves as a crucial learning condition for SB. UNSB proposes the implementation of an enhanced discriminator architecture, which is particularly justified given the constraints of finite sampling and the curse of dimensionality encountered in mid-stage sampling processes. This theoretical framework provides mathematical validation for our discriminator-centric approach in extremely sparse data conditions. The alignment between UNSB’s theoretical foundations and our proposed solution offers compelling evidence for the effectiveness of focusing computational resources on discriminator optimization in low-data scenarios.

UNSB consists of two stages: Generation and Training. In the generation stage for intermediate sample, generation $x_{t}$, the input image, the output of the network, and a Gaussian noise with a predefined variance are used to generate the image of the next stage. The training stage uses adversarial loss and patchNCE loss [37] to guide the output of the network to translate to the target modality while preserving contents details. Unlike traditional diffusion models, The input image is required to perform all time steps.

Despite its innovative approach, direct application of the UNSB framework presents several limitations. While UNSB employs unpaired learning to generate subsequent time step images solely through network input and output without any target modality information, this approach potentially compromises critical lesion information present in the source image. To address this limitation, we incorporate paired target modality information in both generation and training processes. This paired learning paradigm enables the utilization of reconstruction loss [38] and facilitates the integration of additional prior information, such as segmentation masks and intensity-related constraints [39].

For the training of our framework, we sample the time step, $T$. In our experiments, $T$ is [0-4]. To calculate our loss function $L_{FGSB}$, we sample the source image from $dataset_{A}$ and paired target image, $x_{B}$.

We use a Markovian discriminator to determine the domain change from the source image to the target image on the patch-level. The adversarial loss is adopted for generator $G$ and discriminator $D$. $x_{t}$ is calculated using $x_{B}$, $\hat{x}_{B_{t}}$, Gaussian noise with predefined variance.

To preserve anatomic structure information, we use L1 loss and patchNCE loss.

We estimate mutual information loss, $L_{SB}$ using patch-wise mutual information estimator [40]. The mutual information estimator and the generator are jointly optimized through an adversarial training scheme, following a min-max optimization paradigm analogous to the adversarial loss formulation. The mutual information loss functions as a regularization that constrains the intermediate states to adhere to the optimal transport trajectory towards the target domain, ensuring semantic consistency throughout the whole generation stage. mutual information estimator loss is defined as :

Aggregate with $x_{prior}$ as follows:

Mutual information loss is defined as :

Then, the final loss, $L_{FGSB}$ consists of $L_{adv}$, $L_{SB}$, $L_{L1}$ and $L_{patchNCE}$ with weight parameter.

To aggregate additional information such as intensity prior, segmentation map, etc., we use context-preserving loss and weighted patchNCE loss. In our experiment, the binary prior map $x_{prior}$ is defined as 1 if the intensity meets a predefined threshold condition, and 0 otherwise.

The patchNCE loss computation requires a sampling of both positive and negative patches for contrastive learning. Therefore, we first utilize $x_{prior}$ to identify and select critical patches for the sampling procedure.

Except for the IXI experiment, we use combined loss with context-preserving loss and weighted patchNCE loss.

While single-step Generative Adversarial Networks (GANs) have demonstrated promising performance, their performance can be considered suboptimal.

Identity loss functions to translate the suboptimal intermediate results of the refinement process into the target domain as much as possible. Identity loss comprises L1 loss and patchNCE loss computed between the generator output and $x_{B}$

Final loss,$\mathcal{L}_{FGSB}$ as follows :

We describe the sampling procedure, the generation stage for the intermediate generation and final generation. When $t=0$, the source image $x_{A}$ is directly utilized as the initial network input $x_{0}$, as it originates from the source domain.

The predicted target image generated from $x_{0}$ serves as the input for the subsequent time step. During this process, $x_{B}$ is utilized to incorporate information from the original target domain.

During inference, the sampling procedure employs only Gaussian noise, as $x_{B}$ is unavailable and only the refinement process for the predicted target image is required.

Our approach used discriminator to mapping generated image to target image domain. In scenarios where data scarcity is a critical constraint, numerous previous studies have focused on implementing appropriate regularization techniques for the discriminator. It offers a distinct advantage in that it eliminates the need for separate pre-training procedures or elaborate data augmentation strategies, thereby providing a more streamlined approach to model optimization.

The self-supervised discriminator is a novel approach that incorporates auxiliary decoders in conjunction with the discriminator network. The discriminator extracts feature maps at two distinct scales, which are then utilized by two dedicated decoders for image reconstruction tasks. The crop decoder $Dec_{1}$, which processes the first feature map, performs partial image reconstruction utilizing cropped feature representations. Simultaneously, the resize decoder $Dec_{2}$, processing the subsequent feature map, undertakes complete image reconstruction. This self-supervised learning paradigm enables the discriminator to develop robust feature representations that capture both global compositional structures and fine-grained textural details of the input images.

The decoders only consist of four time-conditional conv layers. The training process involves random cropping of the target image at arbitrary positions. To maintain spatial correspondence, the feature map is cropped at positions that align with the target image crops, considering the receptive field. These cropped feature representations are then fed into the decoder network, which is trained to reconstruct the corresponding cropped image regions, establishing a direct mapping between local features and their spatial reconstructions. The secondary decoder receives the extracted feature map as input and is trained to reconstruct the complete target image in its entirety. During this process, the target image is down-sampled to half their original spatial dimensions to optimize computational efficiency while maintaining essential visual information.

The self-supervised discriminator loss is defined as :

We use four datasets (IXI¹¹1https://brain-development.org/ixi-dataset/, BraTS2020 [21], MICCAI2017 WMH [41], private) to evaluate our frameworks. All datasets were randomly split into non-overlapping training and test sets. The source image in all experiments is T1w as the most basic imaging technique. T1-weighted imaging serves as an optimal source sequence for synthesis due to its rapid acquisition time and superior anatomical delineation. The clinical efficacy of our approach is further supported by T1w’s established role as a fundamental sequence in preliminary studies and its widespread use as a reference standard for multimodal imaging protocols. [32, 42, 43] For the BraTS2020 and MICCAI2017 WMH datasets, we performed only skull-stripping [44] since they were already bias-corrected and co-registered. The IXI dataset and the private cohort were implemented additional spatial co-registration using FSL [45]. All images were intensity-normalized to the range [-1, 1] and padded to uniform dimensions (224×224 or 256×256) using minimum intensity values.

We employed multiple evaluation metrics to assess the quality of generated images. The fidelity of image reconstruction was measured using Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), and Normalized Root Mean Square Error (NRMSE).

Furthermore, to evaluate the preservation of clinically relevant features, we computed the Dice similarity coefficient and object-wise recall, specifically focusing on the accurate reconstruction of lesion regions. To evaluate dice and recall, We use U-Net based model [41, 46].

1) Pix2Pix is a paired image-to-image translation method that combines a conditional GAN architecture with pixel-wise reconstruction loss. This approach allows the model to learn direct mappings between source and target domains using paired training data.

2) Syndiff advances this concept by incorporating a DDGAN framework alongside pair-wise reconstruction loss. While the original Syndiff implementation includes CycleGAN components for unpaired learning, we focus on the paired training paradigm for a fair comparison with other methods.

3) MT-Net takes a paired learning approach by using a transformer-based architecture. It first uses masked auto-encoder (MAE) pre-training to build robust feature representations, followed by task-specific fine tuning. This strategy helps to overcome the limitations of traditional paired learning approaches.

The adversarial loss, $L_{\text{adv}}$, and patch-wise mutual information loss, $L_{\text{SB}}$, are adopted with the Markovian-based architecture, where neural mutual information estimator [40] is used for estimating $L_{\text{SB}}$. The detailed network architecture and hyperparameters follow the UNSB frameworks.

Weights for reconstruction loss, $\lambda_{L1}=100.0$ and weighted reconstruction loss, $\lambda_{wL1}=10.0$. All other weight parameters were $1.0$. Weights for identify loss terms were same with above weight parameters. Details of training hyperparameters were : 50 epochs for 25 subjects, 200 epochs for other the number of subjects and $10^{-4}$ learning rate. Starting from the 100th epoch, the learning rate decreases linearly. To minimize the stochastic fluctuations in the output, we constrain the variance of the random noise variable. We applied horizontal flip as a data augmentation.

Both our model and comparison models were implemented in Python and PyTorch framework. All experiments was conducted on one NVIDIA RTX 3090 or A6000.

We conduct FLAIR synthesis task on the private cohort dataset. Our model, FGSB achieves the higher performance than other models. In particular, FGSB shows the best reconstruction performance of whitematter hyperintensities.

We evaluate the performance of T2-weighted MRI synthesis using the IXI dataset of healthy controls. Our model shows superior performance compared to existing baseline methods. For entries marked with ^†, we report the performance values as presented in the original papers.

We evaluate the performance of high resolution FLAIR synthesis using the MICCAI2017 WMH challenge dataset. Given the limited number of axial slices per subject, separate subject-wise reduction was not implemented in this study. Entries denoted by ^† indicate MT-Net pre-trained on T1-weighted and FLAIR sequences using a private dataset comprising 172 subjects.

We evaluate the performance of multi-modal MRI synthesis using the BraTS2020 dataset. For entries marked with ^†, we report the performance values as presented in the original papers.