Re-DiffiNet: Modeling discrepancy in tumor segmentation using diffusion models

Authors: tianyi  ren
abhishek  sharma
juampablo heras rivera
harshitha  rebala
jacob  ruzevick
mehmet  kurt
a  ren
j h sharma
h r rivera
e  honey
a  chopra
Venue: MIDL_io_2024_Conference
Published: 2024-01-31 09:01:12

Abstract: Identification of tumor margins is essential for surgical decision-making for glioblastoma patients and provides reliable assistance for neurosurgeons. Despite improvements in deep learning architectures for tumor segmentation over the years, creating a fully autonomous system suitable for clinical floors remains a formidable challenge because the model predictions have not yet reached the desired level of accuracy and generalizability for clinical applications. Generative modeling techniques have seen significant improvements in recent times. Specifically, Generative Adversarial Networks (GANs) and Denoising diffusion probabilistic models (DDPMs) have been used to generate higher-quality images with fewer artifacts and finer attributes. In this work, we introduce a framework called Re-Diffinet for modeling the discrepancy between the outputs of a segmentation model like U-Net and the ground truth, using DDPMs. By explicitly modeling the discrepancy, the results show an average improvement of 0.55\% in the Dice score and 16.28\% in 95\% Hausdorff Distance from cross-validation over 5-folds, compared to the state-of-the-art U-Net segmentation model. The code is available:

Main paper: 
 introduction: 
 Glioblastoma is the most frequent primary malignant brain tumor in adults, representing approximately 57% of all gliomas and 48% of all primary malignant central nervous system (CNS) tumors (Ostrom et al., 2018; Tan et al., 2020) . This heterogeneous group of tumors is characterized by their resemblance to glia that perform a variety of important functions including support to neurons (Isensee et al., 2021; Ahuja et al., 2020) .The treatment for glioma patients generally consists of surgery, radiotherapy, and chemotherapy and the outcomes of patients with gliomas vary widely according to the glioma type and prognostic factors. Due to the superior soft tissue contrast, multimodal MRI images which allow the complexity and the heterogeneity of the tumor lesion to be better visualized than a CT scan have become the golden standard for surgical decision-making for glioma patients (Hanif et al., 2017; Keunen et al., 2014; van Dijken et al., 2019) . However, visual identification of tumor margins in CT or MRI still remains a challenge for neurosurgeons and researchers (Wang et al., 2019) . Clinically, brain tumor masks are often obtained through Magnetic Resonance Imaging (MRI) scans, which require experienced radiologists to manually segment tumor sub-regions (Baid et al., 2021b) . This is a long process that is unscalable to the needs of all patients. Thus, the recent growth of machine learning technologies holds promise to provide a reliable and automated solution to segmentation to save time and help medical professionals with this process (Luu and Park, 2021) .Deep learning techniques have been widely used in brain tumor segmentation. U-Net is the state of art for tumor segmentation. U-Net and its variants have been used in brain tumor segmentation. such as U-Net++ (Zhou et al., 2018) , 3D U-Net (C ¸içek et al., 2016) , V-Net (Milletari et al., 2016) , Attention-U-Net (Oktay et al., 2018) , Trans-U-Net (Chen et al., 2021) , and Swin-U-Net (Cao et al., 2022) . Transformer architectures has also been applied in brain tumor segmentation. TransU-Net and Swin-U-Net show potential to predict accurate tumor margins. However, the state-of-the-art models in brain tumor segmentation are still based on the encoder-decoder architectures such as U-Net (Isensee et al., 2021) and its variations. For instance, Luu et. al (Luu and Park, 2021) modified the nnU-Net model by adding an axial attention in the decoder. Futrega et. al (Futrega et al., 2021) optimized the U-Net model by adding foreground voxels to the input data, increasing the encoder depth and convolutional filters. Siddiquee et. al (Siddiquee and Myronenko, 2021) applied adaptive ensembling to minimize redundancy under perturbations.While U-Net-based architecture have led to significant improvements in region-based metrics for tumor segmentation e.g. Dice scores, it is also important to improve boundarydistance metrics like HD scores (Karimi and Salcudean, 2019; Yeghiazaryan and Voiculescu, 2018) . Being able to locate boundaries of tumors is crucial for surgical planning. Thus, modeling techniques that are able to capture finer details and high frequency information at the boundaries, are desirable. One of the critical factors that makes predicting tumor boundaries difficult, is the inherent variability in tumor attributes at the boundaries. Thus, the modeling techniques also need to be able to capture the variability in tumor shapes.Generative modeling techniques have seen great improvements in recent times. Specifically, Generative Adversarial Networks and Denoising-Diffusion based models have been used to generate desired images of greater quality. While GANs are able to generate images of high fidelity, they are also prone to mode collapse. Thus, they often fail to capture the variability of the data they seek to model. On the other hand, Denoising-Diffusion based models have been shown to be good at both mode coverage i.e. capturing the variability in the data (Kingma et al., 2021) , as well as at generating high quality images (Dhariwal and Nichol, 2021) . There have been very few instances of Diffusion models being used for brain tumor segmentation, that have shown promising results. (Xing et al., 2023; Wolleb et al., 2022; Wu et al., 2022) .In this work, we introduce a framework called Re-Diffinet, for modeling discrepancy between the outputs of a segmentation model like U-Net and the ground truth, using Denoising Diffusion Probabilistic Models. By explicitly modeling the discrepancy, we intend to build upon previous segmentation models, force diffusion models to focus explicitly on the regions that other models miss, and exploit diffusion models' ability to capture finer details and variability in the data. 
model architectures: 
 In summary, We first trained a state-of-the-art U-Net model to predict three labels of tumor, then we tested several variations of Diffusion architectures to predict the discrepancy between the ground truth and the previous U-Net labels.We adopted the optimized U-Net (Futrega et al., 2021) as our baseline model architecture (Ren et al., 2024) for comparison purposes. U-Net has a symmetric U-shape that characterizes architecture and can be divided into two parts, i.e., encoder and decoder. The encoder comprises 5 levels of same-resolution convolutional layers with strided convolution downsampling. The decoder follows the same structure with transpose convolution upsampling and convolution operating on concatenated skip features from the encoder branch at the same level. The training dataset is comprised of the pairs {(I, x 0 )} N i=1 , where I ∈ R 4×D×W ×H represents the four 3D-MRI contrast as multi-channel input, x 0 ∈ R 3×D×W ×H represents the associated one-hot encoded segmentation mask, with 3 tumor labels: 1) Whole Tumor, 2) Enhancing Tumor, and 3) Necrotic Tumor Core. The baseline U-Net predicts the tumor labels x0 given the input I:x0 = U (I)(1)This model builds upon the Diff-U-Net (Xing et al., 2023) , which uses conditional DDPMs. DDPMs work by learning to denoise images at various noise levels. Once DDPMs have been trained, they can take a randomly drawn noise (usually from a gaussian) and successively denoise it over several steps to generate a sample from the distribution of images (Ho et al., 2020) . Diff-UNet conditions DDPM on MRIs as shown in equation 2 (Xing et al., 2023) .EQUATIONIn comparison, we condition our diffusion model with predictions (U (I)) from baseline U-Net along with MRIs (I) as shown in figure 1 . We tested 3 variants of this approach (3 different inputs): 1) Conditioning the diffusion model with only the U-Net output U (I), 2) Conditioning the diffusion model with a concatenation of MRI contrasts and baseline U-Net predictions U (I), 3) conditioning the diffusion model with MRIs I masked by U-Net predictions U (I). A mask which has a value 1 for each tumor voxel and 0.2 for non-tumor voxel is applied to each of the 4 MRI contrasts, which are concatenated and used as inputs.The resulting masked input is represented as shown in equation 3 : 2, 3, 4) , where i denotes an MRI contrast (3) Among the 3 variants, we chose the best performing UA-Diffusion approach (Table 1 ) and used it for the remaining experiments. The expression for the best performing variant of UA-Diffusion is shown in equation 4:M ′ (x, y, z) = 1 if U (I)[x, y, z] > 0, 0.2 if U (I)[x, y, z] = 0, where x,y,x are voxel indices mask(I, U (I)) = concat(I i • M ′ |i = 1,EQUATION)where t is the time embedding, x t is the corresponding noise masks, Îf = ξ(cat(U (I), I)) are the multi-scale features extracted using a trainable copy (ξ) of the encoder of the denoising-U-Net (DU ). These multi-scale features are added to the outputs of the corresponding layers in the denoising-U-Net, DU (see Figure 1 ).Our proposed Re-DiffiNet model architecture is similar to the U-Net augmented Diffusion (UA-Diffusion). However, instead of trying to generate ground truth segmentation masks x 0 , we generate the absolute discrepancy between the ground truth segmentation masks and baseline U-Net's predictions i.e. ∆x 0 = abs(U (I) -x 0 ) (Figure 1 ). These discrepancy masks (∆x 0 ) will have a value of 1 for each voxel that is predicted incorrectly by the baseline U-Net (U ), and 0 for voxels where the predictions are correct. Once, we have generated the estimated discrepancies ∆x 0 , the tumor mask predictions can be obtained as shown in equation 5 :EQUATIONSubtracting the estimated discrepancy ∆x 0 from U (I) and taking the absolute, we flip every incorrect voxel (as per our estimate) in U (I) i.e. 1 → 0 and 0 → 1. While, the correct voxels (as per our estimation) in the baseline U-Net, remain the same. Discrepancy U-Net: To test if improvements observed are due to the combination of discrepancy modeling and diffusion model, we also investigated using a second U-Net (Discrepancy U-Net) to predict discrepancies, to correct outputs of the baseline optimized U-Net. We compare the results of discrepancy UNet with ReDiffiNet in Table 1 and 2. Our models were implemented in Pytorch and MONAI, and trained on 2 NVIDIA A40 GPUs. The model was trained on overlapping regions, whole tumor (WT), tumor core (TC), and enhancing tumor(ET). TC entails the ET, as well as the necrotic (NCR) parts of the tumor, and WT describes the complete extent of the disease. The diffusion models was trained using a compound loss function including DICE loss, Binary cross entropy (BCE) loss, and Mean square error(MSE) loss. The model was trained using the AdamW optimizer with a learning rate of 0.0001 and a weight decay equal to 0.0001. The network's performance was evaluated using 5-fold cross-validation. The data were randomly shuffled and equally split into 5 groups for cross-validation. The model will be evaluated on two metrics: Dice similarity coefficient (Dice) measures the similarity between the model prediction and the ground truth; 95% Hausdorff distance (HD95) measures the boundary distance between the model prediction and the ground truth.The training dataset provided for the BraTS 2023 Adult Glioma challenge (Baid et al., 2021a) consists of 1251 brain MRI scans along with segmentation annotations of tumorous regions. The 3D volumes were skull-stripped and resampled to 1 mm 3 isotropic resolution, with dimensions of (240, 240, 155) voxels. For each example, four modalities were given: native (T1), post-contrast T1-weighted (T1Gd), T2-weighted (T2), and T2 Fluid Attenuated Inversion Recovery (T2-FLAIR). Segmentation labels were annotated manually by one to four experts. Annotations consist of three disjoint classes: enhancing tumor (ET), peritumoral edematous tissue (ED), and necrotic tumor core (NCR). For all the MRI contrasts in the BraTS2023 dataset, we rescale the voxel intensity after Z-Score normalization as the preprocessing protocol.  
experiments and results: 
 We trained 3 models 1) Baseline U-Net (section 2.1.1), 2) U-Net augmented diffusion or UA-Diffusion (section 2.1.2), and 3) Re-DiffiNet (section 2.1.3). We first trained the baseline U-Net model. Then, the predictions of the baseline U-Net model were used as inputs in U-Net augmented diffusion (UA-Diffusion), and Re-DiffiNet. In a preliminary study, we trained 3 variants of the U-Net augmented diffusion (UA-Diffusion) on a random train-test split and compared them with the baseline-U-Net (See section2.1.1), 2) U-Net augmented diffusion or UA-Diffusion (section 2.1.2), and 3) Re-DiffiNet (section 2.1.3). We first trained the baseline U-Net model. Then, the predictions of the baseline U-Net model were used as inputs in U-Net augmented diffusion (UA-Diffusion), and Re-DiffiNet. In a preliminary study, we trained 3 variants of the U-Net augmented diffusion (UA-Diffusion) on a random train-test split and compared them with the baseline-U-Net (See section 2.1.2). We found that using diffusion directly to predict tumor masks doesn't lead to any significant performance gains over the baseline U-Net, as shown in Table 1 . On the other hand, using diffusion model to predict discrepancies and using them to correct U-Net's outputs leads to significant performance gains specially in terms of HD95 score. Among the 3 UA-Diffusion approaches concatenating the U-Net prediction and MRI yielded the best performance. Thus, we use the UA-Diffusion with concatenation of MRI and U-Net prediction, for 5-fold cross-validation in Table 2 .The results of 5-fold cross-validation are shown in Table 2 , which reports the Dice Score (DICE) and 95 percentile Hausdorff distance (HD95) and the average scores of all methods on the three overlapping regions whole tumor (WT), tumor core (TC) and Enhancing tumor (ET) for the BraTS2023 dataset (Figure 2 ). We found that while using the diffusion model to directly output the tumor segmentation mask does lead to improvements over the U-Net model, the improvements are modest: 0.12% improvement in Dice, and 5.61% improvement in HD95 score. On the other hand, with Re-DiffiNet we found a 16.28% improvement in HD95 score, indicating the benefits of modeling discrepancy using diffusion models, while simultaneously the Dice score was comparable with the baseline U-Net (0.55%  
discussion and conclusion: 
 In this research, we proposed a tumor segmentation framework Re-DiffiNet, which uses diffusion models to refine and improve predictions of a tumor segmentation model (like optimized U-Net). Most tumor segmentation studies optimize for region-based metrics like Dice scores, and have been able to show high Dice score in the range of 90% or greater. However, boundary-distance metrics like HD scores are also critical, and being able to improve upon these score while not sacrificing performance on Dice score is highly desirable.In this work, we investigated the potential to refine predictions generated by state-of-the-art U-Net models using diffusion models. We found that while using diffusion models to directly generate tumor masks did lead to improvements in performance over the baseline U-Net, it was the use of discrepancy modeling i.e. predicting the differences between ground truth masks and baseline U-Net's outputs, that led to most significant improvements. This was indicated by 16.28% improvement in HD-95 score, highlighting significant improvements on the boundaries of tumors. While, discrepancies can be modeled by any other modeling technique (even a U-Net), effectively acting as a boosting method, we chose diffusion primarily because of its ability to generate high-fidelity visual attributes, as well as capture variability in the data distribution, both of which are exhibited by brain tumors. Another benefit of using diffusion models instead of a U-Net to improve the baseline U-Net's predictions, would be the to learn more robust and diverse representations from the data, due to the inherently different mechanism using which diffusion models are trained.Our work shows the potential of further improving tumor segmentation by combining diffusion models and discrepancy modeling. In this work, we investigated Re-DiffiNet for the segmentation of gliomas. In the future, we intend to test our approach to improve the segmentation of other kinds of tumors like meningioma, and pediatric brain tumors. 
 
 Appendix:  



Reviews: review: paper_summary: paper_summary: This work utilizes denoising diffusion based models (DDPM) to improve tumor segmentation quality. Instead of following standard practice where a tumor itself is modeled and a DL model is trained to predict a mask of a tumor; the proposed idea entails first running a segmentation step using a standard UNet, and then a DDPM predicts the discrepancy between the UNet segmentation's output and the ground truth. As a result a more accurate segmentation is obtained.
main_review: main_review: Thank you for promptly responding to all the reviewers comments, I appreciate it. I would like to improve my rating. I think this paper is really relevant to the conference and the field could find value in it.
strength_weakness: strength_weakness: It is not clear whether a second UNet could have been trained to model the predictive error of the baseline UNet without needing to use a diffusion model. Perhaps the task could have been solved in a easier way - has this been considered?
questions: questions: I would like the authors to answer my questions in the weakness and detail comments section.
paper_summary: paper_summary: The authors propose ReDiffiNet for improved segmentation of brain tumours that predicts the error from an initial prediction using UNet to improve final predictions based on the structure of a UNet-Diffusion architecture. Comparisons of individual tumour tissue predictions are performed between UNet, UA-Diffusion and ReDiffiNet are performed in terms of DSC and HD95 showing numerical improvement for the proposed discrepancy based learning.
main_review: main_review: The authors have made a strong case for their solution in their rebuttal and I am happy to modify my rating to borderline accept given that the organisation of the paper is much clearer and that the novelty and contributions are better stated.
strength_weakness: strength_weakness: - The organisation of the paper does hinder clarity as the reader is referred to section of explanations happening very late in the manuscript
- Limited description of the Diffusion setting at the basis of the proposed architecture
- No indication of statistical evaluation across the different method and no comment on what seems to be a numerically small difference in DSC (0.10 to 0.94% difference across the different folds). No indication of variability within each fold. 
- Novelty is not entirely clear - is the focus entirely on ReDiffiNet or also on the use of diffusion models in general?
questions: questions: Statistical testing should be included as the variability may mask the observed numerical differences.
The authors should really justify their added contribution (is it simply the Discrepancy training) or is the aim also to assess the benefit of diffusion models as a whole
paper_summary: paper_summary: Conventional segmentation models like UNet tend not to capture the variability in tumor segmentation well enough. The authors propose Re-DiffiNet, which uses Denoising Diffusion Probabilistic Models (DDPM) to adjust the segmentation predictions of a UNet network. It is based on Diff-UNet (https://arxiv.org/pdf/2303.10326.pdf), which conditions the diffusion model with the input image (e.g. MRI). The authors, on the other hand, condition it on both the MRI and the UNet mask predictions and call that variation U-Net Augmented Diffusion (UA-Diffusion). Then, they re-define the framework’s task to generate an “absolute discrepancy” mask that is used to adjust the previously generated UNet mask instead of generating the whole adjusted mask as Diff-UNet does.

The authors employed the BraTS 2023 dataset and demonstrated that while the Dice score remains similar, the 95th percentile Haussdorf Distance (HD95) is significantly improved. The experiments included a baseline UNet, Diff-UNet, several versions of the proposed UA-Diffusion, and the proposed Re-DiffiNet.
main_review: main_review: The proposed addition of the prediction mask to the diffusion model conditioning and the re-definition of the prediction task as absolute discrepancy are both straightforward improvements that yield better performance. The discrepancy between the Diff-U-Net in this versus the original paper is the reason why I decided to give a weak instead of a strong accept. I hope the authors address that as the paper itself is useful and interesting.
strength_weakness: strength_weakness: - The baseline Diff-U-Net performance seems lower than what is reported in the original paper. I would like to see that addressed at least in the Discussion.

- It would be useful to see several examples of masks where HD95 was increased but Dice score was not.
questions: questions: - It would be interesting to include a figure where we can see the generated masks that achieved similar Dice but much better HD95 against the baseline UNet and Diff-UNet.

- I was initially confused by the introduction of "UA-Diffusion" and "Re-DiffiNet" names. Is there a way to simplify that? Also, why not keep it similar to Diff-UNet as it is based on it, e.g. "Re-Diff-UNet"

- Could you please add the code to the linked repository?

- Diff-U-Net performs significantly worse than the baseline UNet experiment in your experiments. However, the Diff-U-Net paper reports improved performance over their baselines. They used a very similar dataset - BraTS 2020. Could you please address that? Are their results incorrect, or was there something off with your implementation?

- Could you please report the time and memory needed during training and inference for all methods, as well as the batch size?

Comments: comment: Thank you very much for these updates and clarifications that would be helping to consider a changed rating to borderline.
* "The organization on of the paper..."

Thank you for the suggestions. We have revised the paper structure by incorporating a summary paragraph at the beginning of the methods section so the reader will have a better understanding of the model. 

* "Limited description of the Diffusion…"

Thank you for the suggestion. We have added a paragraph in section 2.1.2 to describe the diffusion model.

* "No indication of statistical evaluation across the different method …"

We conducted a 5-fold cross-validation in order to account for variability that might arise due to training. While, this doesn’t establish statistical significance, it is a standard evaluation methodology for models in Machine Learning.

In addition, tumor masks generated by Diff-UNet (which Re-DiffiNet builds upon) are obtained by averaging over multiple denoising processes [Xing et al., 2023]. This implicitly minimizes the variability that might arise from randomly drawn noise sample.

One of the primary goals of this manuscript was to improve boundary distance metrics like HD score. Therefore, we included discrepancy modeling as an important component of our methods. Tumor segmentation models like UNet are typically prone to making errors on the boundary. Thus, modeling discrepancy explicitly forced diffusion to improve the boundaries. 

We agree with the reviewer’s point that Dice score improvement is marginal, and we have highlighted this point in the discussion session. The main contribution of Re-DiffiNet is that it improves the boundary distance metric (HD score) while not causing any decrease in the Dice Score.	We add the statistical analysis in the appendix that shows that our proposed Re-DiffiNet can improve the HD95 score when compared to the baseline U-Net for all five folds.


* "Novelty is not entirely clear…?"

The reviewer raises an excellent point. This paper contributes in two different ways: 1) It shows that diffusion can help improve predictions generated by a segmentation model like U-Net. This is the first study to the best of our knowledge where diffusion has been able to improve over a state-of-the-art method in brain tumor segmentation. 2) It also shows that while diffusion does help improve, its effectiveness is increased when it is used to model discrepancy. To the best of our knowledge, this is the first study that uses discrepancy modeling for brain tumor segmentation.

We have also added a new result in the revised manuscript i.e., Discrepancy UNet- We have shown that if we try to predict discrepancy using a standard UNet instead of DDPM, the performance doesn’t improve compared to the Baseline UNet. This shows that it is the combination of discrepancy modeling and diffusion model, that leads to performance improvements in the case of Re-DiffiNet.

* "A paragraph clearly stating how a diffusion.." 

We have added a paragraph in section 2.1.2 to describe the diffusion model.

* "Instead of referring the reader to section 3…."

Thank you for the suggestion. We have moved the explanations of the 3 variants of UA-Diffusion from section 3 to section 2.1.2. Due to space limitations, we have not used figures. Instead, we have used equations and text for the explanation.

We have also updated Figure 1 to illustrate how the absolute discrepancy map is obtained. To summarize, we obtain it by subtracting each ground truth mask with the corresponding UNet prediction in an element-wise manner and then taking the absolute value of each element.  

* "Statistical testing should be included…"

Thanks for the suggestion. We conducted a 5-fold cross-validation to account for variability that might arise due to training.  We included statistical testing in the appendix.

This paper contributes in two different ways: 1) It shows that diffusion can help improve predictions generated by a segmentation model like U-Net. This is the first study to the best of our knowledge where diffusion has been able to improve over a state-of-the-art method in brain tumor segmentation.
2) It also shows that while diffusion does help improve, its effectiveness is increased when it is used to model discrepancy. To the best of our knowledge, this is the first study that uses discrepancy modeling for brain tumor segmentation.
We have also added a new result in the revised manuscript i.e., Discrepancy U-Net We have shown that if we try to predict discrepancy using a standard UNet instead of DDPM, the performance doesn’t improve compared to the Baseline UNet. This shows that it is the combination of discrepancy modeling and diffusion model that leads to performance improvements in the case of Re-DiffiNet.
Thanks for your comments.


* "HD95 abbreviation not introduced at the first occurrence."
We have fixed the abbreviation in the Abstract and 2.2. Training details section.

* "The authors cited a BraTS 2021 challenge for BraTS 2023. Please verify the citation."
Thanks for pointing this out. The glioma dataset from the BraTS 2021 challenge is the same as BraTS 2023 Adult Glioma dataset.

* "Part of the Section 2.1.1. can be read as if the authors performed the annotation. Perhaps add a sentence just to clarify that.
Additionally, please state that the data comes from the “Adult Glioma” sub-challenge of the BraTS 2023 challenge."
We agree this could bring confusion to the readers, we removed the sentences related to the ground truth label annotation in the previous section Section 2.1.1., since the details of ground truth annotation can be found in the literature with more details. Also, we state that the data is from Adult Glioma sub-challenge of the BraTS 2023 challenge in the new section 2.3. Dataset.

* "The abstract refers to Denoising-diffusion-based models as DDPM. Should it not be “Denoising diffusion probabilistic models”?"
Thanks for pointing out the error, We fixed the full name of DDPM to “Denoising diffusion probabilistic models” in abstract:
"Specifically, Generative Adversarial Networks (GANs) and Denoising diffusion probabilistic models (DDPMs) have been used to generate higher-quality images with fewer artifacts and finer attributes."

* "Introduction, page 2 - TransUNet and Swin-U-Net should be cited."
Thanks for the suggestion, these two papers are important in this field. We added these two papers in the introduction section as suggested, in Introduction paragraph 3: 
" U-Net and its variants have been used in brain tumor segmentation. such as U-Net++ (Zhou et al., 2018), 3D U-Net (Cicek et al., 2016), V-Net (Milletari et al., 2016), Attention-U-Net (Oktay et al., 2018), Trans-U-Net (Chen et al., 2021), and Swin U-Net (Cao et al., 2022)."

* "It would be interesting to include a figure where we can see the generated masks that achieved similar Dice but much better HD95 against the baseline UNet and Diff-UNet."
Thanks for the suggestion, we included two extra cases in the supplementary file: The first example shows that Re-DiffiNet produces a segmentation that has a slightly worse Dice score but a higher HD95 score, and the second example shows the Re-DiffiNet produces a segmentation that has a much better Dice and HD95 score.

* "I was initially confused by the introduction of "UA-Diffusion" and "Re-DiffiNet" names. Is there a way to simplify that? Also, why not keep it similar to Diff-UNet as it is based on it, e.g. "Re-Diff-UNet"."
We appreciate the reviewer's suggestion. We think that the most important improvement in Re-DiffiNet could be attributed to predicting the difference between baseline UNet and ground truth, effectively “redefining” the outputs of UNet. And thus we chose Re-DiffiNet as an informative title about the method. The results of UA-Diffusion are more attributable to the method in Diff-UNet.  

* "Could you please add the code to the linked repository?"
Our code is open-source now!

* "Diff-U-Net performs significantly worse than the baseline UNet experiment in your experiments. However, the Diff-U-Net paper reports improved performance over their baselines. They used a very similar dataset - BraTS 2020. Could you please address that? Are their results incorrect, or was there something off with your implementation?"
We appreciate the reviewer’s observation. The Diff-UNet paper uses several variants of UNet but does not use NVIDIA’s optimized U-Net variant that we used in our paper. The results for Diff-UNet reported in our paper are similar to those reported in the original paper, and any minor differences could be attributed to the different Train-Test split used in our paper, as we used 5-fold cross-validation. Another reason for the differences could also be attributed to differences in the dataset used as pointed out by the reviewer- We used BraTS 2023 (n=1251) while the original paper used BraTS 2020 (n=371).

* "Could you please report the time and memory needed during training and inference for all methods, as well as the batch size?"
Thank you for pointing this out. The Batch size for all the methods was 1. Training time for each session of ReDiffiNet on average was around 84h when trained on a single cluster node with 2 GPUs (48 GB RAM each). The inference time was around 10 hours.
Thanks for your comments.

* "It is not clear whether a second UNet could have been trained to model the predictive error of the baseline UNet without needing to use a diffusion model. Perhaps the task could have been solved in a easier way - has this been considered?"

Thanks for highlighting this point. We have trained a second UNet to predict the discrepancy between the optimized UNet output and the ground truth labels. The results show that UNet is not able to further improve the results of a current UNet prediction by predicting discrepancies. The results are indicated in Table 1 and Table 2.

We have added the following paragraph in section 2.1.3 (Re-DiffiNet) highlighting this point:
”To test if any improvements observed are due to the combination of discrepancy modeling and the diffusion model, we also investigated using a second U-Net (Discrepancy U-Net) to predict the discrepancy maps and correct the outputs of the baseline optimized U-Net. We compare the results of discrepancy UNet with ReDiffiNet in Table 1 and 2.”

We selected the state-of-the-art model from the 2021 Brain Tumor Segmentation Challenge (Optimized U-Net (Futrega et. al. 2021)) as our baseline. We also tested an alternative U-Net model with slightly inferior performance compared to the optimized UNet, however, our conclusion regarding diffusion for discrepancy modeling remains unchanged. 

* "The optimized U-Net was utilized as baseline and later to condition the DDPM model - does the architectural choice has a strong effect on the ultimate result?"

The reviewer raises a great point. Diffusion models are capable of learning distributions with high diversity [Kingma et. al. 2023]. We expected the discrepancies between UNet’s predictions and ground truth to have a high diversity. Thus, we used diffusion models to learn them. We expect that as long as the discrepancies do not follow a completely random pattern, diffusion models should be able to improve the predictions of any model using discrepancy modeling. Whether the improvement will be as good as Re-DiffiNet or greater can only be empirically determined and is the subject of future work.