Conditional Generation of 3D Brain Tumor Regions via VQGAN and Temporal-Agnostic Masked Transformer

Authors: meng  zhou
farzad  khalvati
Venue: MIDL_io_2024_Conference
Published: 2024-01-29 22:21:33

Abstract: Neuroradiology studies often suffer from lack of sufficient data to properly train deep learning models. Generative Adversarial Networks (GANs) can mitigate this problem by generating synthetic images to augment training datasets. However, GANs sometimes are unstable and struggle to produce high-resolution, realistic, and diverse images. An alternative solution is Diffusion Probabilistic Models, but these models require extensive computational resources. Additionally, most of the existing generation models are designed to generate the entire image volumes, rather than the regions of interest (ROIs) such as the tumor region. Research on brain tumor classification using magnetic resonance imaging (MRIs) has shown that it is easier to classify the ROIs compared to the entire image volumes. To this end, we present a class-conditioned ROI generation framework that combines a conditional vector-quantization GAN and a class-conditioned masked Transformer to generate high-resolution and diverse 3D brain tumor ROIs. We also propose a temporal-agnostic masking strategy to effectively learn relationships between semantic tokens in the latent space. Our experiments demonstrate that the proposed method can generate high-quality 3D MRIs of brain tumor regions for both low- and high-grade glioma (LGG/HGG) in the BraTS 2019 dataset. Using the generated data, our approach demonstrates superior performance compared to several baselines in a downstream task of brain tumor type classification. Our proposed method has the potential to facilitate accurate diagnosis of rare brain tumors using MRI-based machine learning models.

Main paper: 
 introduction: 
 Gliomas are the most frequent primary adult brain tumor types within the central nervous system (Menze et al., 2014; Bakas et al., 2017) . Among all variations of gliomas, high-grade glioma (HGG) accounts for the majority of cases, and low-grade glioma (LGG) accounts for less common cases. For both variations, a commonly used technique for diagnosing is the multi-parametric Magnetic Resonance Imaging (MRI) equipped with different sequences such as T1-, T2-weighted, and Fluid Attenuated Inversion Recovery (FLAIR) (Menze et al., 2014) . Each modality provides distinct biological information about the tumor, aiding radiologists in determining the tumor type. However, distinguishing between HGG and LGG remains challenging, and misdiagnosis may lead to suboptimal prognoses (Mzoughi et al., 2020) .In recent years, deep learning-based methods have proven to be one of the effective ways for adult brain tumor classification tasks using brain MR images (Ge et al., 2020; Hao et al., 2021; Namdar et al., 2022; Tandel et al., 2020) . However, the requirement for large training datasets poses challenges in medical imaging, especially for rare diseases such as LGG, leading to potential overfitting and poor generalization to unseen datasets. Several works aim to mitigate the imbalanced data problem. One line of work is to use the transfer learning approach by pre-training models on huge datasets (i.e., ImageNet), and then finetuning them on domain-specific datasets (Ghazal et al., 2022; Tak et al., 2023; Ullah et al., 2022) . Another line of work is to synthesize MRIs using Generative Adversarial Network (GAN) (Volokitin et al., 2020; Kwon et al., 2019; Sun et al., 2020; Xia et al., 2020) or Diffusion-based methods (Khader et al., 2022; Peng et al., 2022a; Dorjsembe et al., 2023; Sanchez et al., 2022) to alleviate the need for extensive datasets. However, GANs for image generation can be unstable, produce blurry images, and encounter mode collapse problems (Kwon et al., 2019) . As an alternative approach, Diffusion Probabilistic models have been proposed and demonstrated superior performance over GANs (Müller-Franzes et al., 2022) , but these methods are extremely computationally expensive when synthesizing full-resolution MRIs, thus posing challenges in both the training and inference stage. More recently, autoregressive transformer models have attracted increasing attention in image generation tasks (Esser et al., 2021a,b; Huang et al., 2023) . The key idea behind such models is to obtain discretized feature maps from a Vector-Quantization GAN (VQGAN) model and then use the transformer model to learn the compositions. Autoregressive transformers have been extended to medical images (Pinaya et al., 2023; Tudosiu et al., 2022; Zhou et al., 2023) , to unconditionally generate brain MR images. However, one limitation of these models is the lack of conditioning; separate models need to be trained for different tasks, which becomes time-consuming and resource-intensive. Therefore, a condition-based image generation model is important for practical applications in real clinical settings. Moreover, a group of works (Sajjad et al., 2019; Mzoughi et al., 2020; Srinivasan et al., 2023) have demonstrated the effectiveness of using the tumor region of interest (ROI) for classifying tumor types because ROIs contain less information than the whole-image that may negatively affect the results. Hence, this work aims to explore generating different brain MRI tumor ROIs based on their pathology labels. To this end, we introduce the first class-conditional generation framework for synthesizing 3D brain tumor MRI ROIs. Our model is built upon the previous work (Zhou et al., 2023) and extended to a conditional generation paradigm. Our framework has three modules: a 3D-VQGAN image encoder to extract high-level feature maps while concurrently learning the importance score for each region in the feature maps; an Exponential Moving Averages(EMA) codebook with l 2 -norm lookup for converting feature maps into discrete semantic tokens (Peng et al., 2022b) ; and a temporal-agnostic masked transformer to learn the relationships between discrete tokens. We evaluated our proposed method in the BraTS 2019 dataset and demonstrated superior performance over several baselines on both image generation quality and the downstream HGG vs. LGG classification task. Our contributions are as follows: (1). We propose the first image generation framework for different tumor types based on the given class label.(2). We use a classifier-guidance approach to learn the importance score for each region in the encoded feature maps. (3). We propose a novel temporal-agnostic hybrid masking strategy which uses the importance score to mask tokens to prevent any information leakage.(4). Experiments show our proposed method outperforms several baselines in both image generation quality and the downstream classification task. 
model architecture: 
 We adopted and extended the VQGAN (Esser et al., 2021b) and recently proposed 3D-VQGAN (Zhou et al., 2023) with some modifications detailed below for class-conditional generation of brain tumor ROIs. Stage 1. 3D-VQGAN: The first stage is shown at the top of Figure 1 . We train all modules presented to learn efficient data representation through a reconstruction task in this stage. Our encoder, decoder, and discriminator follow the same design as in (Zhou et al., 2023) , except we replace batch normalization with group normalization to stabilize the training process for small batch sizes (Wu and He, 2018) . Importance Score Map: We use a lightweight scoring network f before quantization to assign an importance score for each region in encoded feature maps. Let the encoded map with size z e ∈ R H×W ×Dp×nz where H, W, Dp, n z denote the height, width, depth, and the number of feature maps, respectively. Then, for each region r i ∈ R nz , its score is defined by s i = f (r i ), where i = 1, ..., H × W × Dp. The larger the score s i is, the more important the feature region r i is. To learn f , we use a classifier-guidance approach by using an auxiliary MLP-based classifier f cls after f to classify tumor types (e.g., HGG vs. LGG).We hypothesize that regions with the higher scores (i.e., important regions) are the key to differentiating between two tumor types, and thus by optimizing both f and f cls , the model can identify important feature regions that are specific to either LGG or HGG tumors. Quantization: In the quantization step, the latent feature maps are quantized by replacing each one with its closest corresponding vector in codebook C. Formally, we train a learnable codebook C = {c i } K i=1 that transforms feature vectors z e to K := H × W × Dp discrete tokens c q , q ∈ [1, K] by the nearest neighbor search in C, and each token c q includes an embedding vector c z ∈ R nz . We use the l 2 normalization for codebook lookup, as done in (Yu et al., 2021) . Finally, we stack K quantized feature vectors back to the original latent shape and feed it into the decoder D to produce reconstructed images. Stage 2. Class-conditional Masked Transformer: In this stage, we propose a novel temporal-agnostic hybrid masking strategy based on importance scores computed in Stage 1, which is inspired by DropBlock (Ghiasi et al., 2018) and BERT (Devlin et al., 2018) . The 3D images are initially represented in the latent space, and the encoder E, decoder D, and codebook C are fixed, with only the transformer being trained. The encoded feature map z e of size H×W ×Dp×n z is quantized into a set of L discrete tokens, where L = H×W × Dp. We first set the masking ratio α and randomly sample N = L × α tokens to be masked. These tokens are then divided into two equal subsets: N 1 and N 2 , which denote the number of important tokens and unimportant tokens to be masked based on their importance score. Let Y = {y i } L i=1 be the raster-scan linearized discrete tokens and each of y i associates with their importance score s i . The importance score and the corresponding tokens are sorted in descending order, denoted as Y ′ . For N 1 important tokens, we randomly sample⌈ N 1 2 ⌉ tokens from the top-k = 25% of Y ′ .For each selected token, we also mask along with their spatial or temporal neighborhood tokens. Unlike random masking, this blockwise masking around each selected token prevents information leakage from neighbors, enhancing the model's learning ability on important tokens and preventing short-cut learning. The special [MASK] token is used to mask out these important tokens and their associated blocks. For N 2 unimportant tokens, we randomly sample them from the remaining (1 -k) × Y ′ tokens and replace them with N 2 randomly selected tokens from the codebook C. Importantly, we ensure that N 1 and N 2 are non-overlapping (N 1 ∩ N 2 = ∅). We denote M = {m i } L i=1 be the mask for each of the discrete tokens, where m i = 1 if the token i is unmasked and m i = 0 if the token i is masked out. Finally, we prepend a class label indicating HGG or LGG sample at the start of each indices sequence. During training, the objective is to reconstruct the masked tokens using unmasked ones. The introduced noise from maskedout tokens is hypothesized to enhance the transformer model's ability to learn relationships between semantic tokens, improving overall robustness. The proposed masking strategy and its difference between random masking is depicted in the bottom right of Figure 1 . Classification: For the downstream classification task between LGG and HGG tumor types, we use a standard 3D ResNet-50 model (Hara et al., 2017) that takes 3D tumor ROIs as inputs and outputs two class probabilities for two tumor types.We employ the same loss function as done in (Zhou et al., 2023) . We use a combination of the pixel differences loss (L pixel ), perceptual loss (L perp ) (Johnson et al., 2016) , GAN-based feature matching loss (L match ) (Ge et al., 2022) , 3D image gradient loss (L grad ), codebook loss (L codebook ) (Esser et al., 2021b) , and the discriminator loss (L Dis ) in the first stage. See Equation (1) and Equation (2) for details.L pixel = ∥x -x∥ 1 , L perp = 6 j=1 ∥f i (x j ) -f i ( xj )∥ 2 2 , L match = ∥f i Dis (x) -f i Dis (x)∥ 1 , L grad = ∥∇(A(x)) -∇(A(x))∥ 2 2 + ∥∇(R(x)) -∇(R(x))∥ 2 2 + ∥∇(S(x)) -∇(S(x))∥ 2 2(1)L Dis = E x∼p d [max(0, 1 -D(x))] + E x∼p d [max(0, 1 + D(x)], L codebook = ∥sg[E(x)] -c z ∥ 2 2 + β∥sg[c z ] -E(x)∥ 2 2 (2)Where x is the original image and x is the reconstructed image, ∇(•) computes the x-and y-direction gradients of the image, A(x), R(x), S(x) represents slicing over Axial, Coronal and Sagittal plane, respectively. Additionally, we use the standard cross-entropy loss L ce between class logits and class labels for our auxiliary classifier f cls . Aggregating all the loss terms together yields the loss objective in Equation (3) for the first stage of the framework:EQUATIONWhere , 5] is the weighting factor between different loss terms. We follow previous publications (Ge et al., 2022; Khader et al., 2022) to set λ 1 = λ 3 = 4 and λ 2 = λ 5 = 1. We also set λ 4 = 4 and β in L codebook to be 1. C 1 and C 2 are balancing factors between the main task and auxiliary task, we empirically set C 1 = 0.8 and C 2 = 0.2.λ i , i ∈ [1For the transformer model, we use the cross entropy loss between the reconstructed token sequence and the ground truth token sequence as shown in Equation (4) to optimize the transformer.EQUATION)Where D is the training dataset, Y M denotes the unmasked tokens, thus the masked tokens can conditioned on these unmasked tokens during training.We used the FLAIR sequence data from the BraTS 2019 dataset (Bakas et al., 2017 (Bakas et al., , 2018;; Menze et al., 2014) . The data contains 259 HGG patients and 76 LGG patients. We used this dataset since it is the latest version of the dataset which provides labels for the brain tumor pathology classification, i.e., HGG vs. LGG. We reshaped the data from 240 × 240 × 155 to 128 × 128 × 128 and normalized all pixel intensities in [-1, 1]. To achieve this, we first remove all zero-valued slices in both the brain images and the segmentations, since we are interested in the slices with the brain tumor present. Then, we obtain the ROIs by multiplying the images with masks. Finally, we center-crop the region based on the segmentation mask to a target size of 128 × 128 × 128. 
experiments: 
 For the first stage of the proposed 3D-VQGAN-cond model, we train for 10k epochs with an initial learning rate of 0.0001 and cosine decay to 0 for all sub-modules, a mini-batch size of 3, and with the Adam optimizer (Kingma and Ba, 2014) . We set the codebook size K = 1024. For the second stage, we train the transformer for 5k epochs using a learning rate of 4.5e -06, a mini-batch size of 3, and the AdamW optimizer (Loshchilov and Hutter, 2017) . We set the mask ratio α = 0.5. We randomly held out 25 patients from both HGG and LGG as a standalone test set. The rest of the data is used to train our model.To assess the usability of our generated data, we conducted two sets of classification experiments: Experiment (1). We aimed to determine if models trained on synthetic data are better than those trained without or only using a portion of synthetic data. Following the approach in (Zhou et al., 2023) , we compared our classification model pre-trained with both synthetic LGG and HGG to one pre-trained with real HGG and synthetic/real LGG. We use the same amount of data for pre-training and the same data for fine-tuning. Experiment (2). Our goal was to investigate whether increasing the number of synthetic data for pretraining enhances classification performance. We generated 250 synthetic HGG and LGG from baseline models and our proposed model to pre-train the classification model. The data for fine-tuning remained consistent. More details can be found in Appendix A. The ablations on the top-k ratio can be found in Appendix B due to page limit. Baseline Model & Comparison. For comparison of image generation results, we consider five state-of-the-art methods, 3D-WGAN-GP (Gulrajani et al., 2017) , 3D-αWGAN (Kwon et al., 2019) , 3D-Med-DDPM (Dorjsembe et al., 2023) , Medical Diffusion (Khader et al., 2022) and 3D-VQGAN (Zhou et al., 2023) . We re-implemented the Medical Diffusion to make it class-conditioned (Medical Diffusion-C) which can be jointly trained on LGG and HGG data, as our class-conditioned baseline. Other baselines were rerun separately for two data classes. For classification, we establish a baseline where we only use traditional augmentations, ensuring a fair comparison with other methods. We evaluate the quality of generated MRI ROIs using three commonly used metrics in previous publications (Kwon et al., 2019; Peng et al., 2022a; Dorjsembe et al., 2023; Zhou et al., 2023) : maximum mean discrepancy (MMD) (Gretton et al., 2012) , multi-slice structure similarity (MS-SSIM) (Rosca et al., 2017) , and the Fréchet Inception Distance (FID) (Heusel et al., 2017) . FID score is computed in three views (Axial, Coronal, Sagittal) to reflect the nature of medical images. The classification performance is evaluated using AUC, F1-Score, and Accuracy. Generating Synthetic MRI Data. To generate 3D tumor ROIs, we start with a class token, either 0 for LGG or 1 for HGG, and then have the transformer model predict and complete the rest of the indices. Next, we obtain the embedding vectors of each index from the codebook C and then feed them into decoder D to produce the final images. For other baselines, we follow the exact procedure as stated in their paper. 
results for generated images: 
 In Figure 2 , we compare LGG and HGG ROIs generated by baseline models and our proposed method. The center three slices in the Axial plane are shown for better visual quality.We observed that both 3D-WGAN-GP and 3D-αWGAN produce images that lack detail and exhibit major artifacts. 3D-Med-DDPM and Medical Diffusion/-C introduce noise and checkerboard artifacts. Generated images from both 3D-VQGAN and our proposed method contain the detailed attributes of the tumor and exhibit high image fidelity. Quantitative metrics are computed over 250 generated HGG and LGG samples as shown in Table 1 . It can be seen that our proposed method performs best in terms of preserving diversity based on the MS-SSIM score. For the MMD score on LGG data, our method outperforms all methods except 3D-αWGAN, which we argue that this is not a fair comparison because it exhibits a severe mode collapse problem (99.4 in MS-SSIM). For the MMD score on HGG data, our method is slightly worse than Medical Diffusion/-C, but it still outperforms other baselines. For FID, our method consistently outperforms on FID-A score, and the other two FID-S and FID-C scores are very close to the best performance. It is also worth noting that all baselines except Medical Diffusion-C are trained separately, whereas ours and Medical Diffusion-C are trained jointly on both LGG and HGG data. Our performance 1) outperforms Medical Diffusion-C in most of the quantitative metrics, and 2) exceeds or is on par with other baselines indicating our method can effectively learn and distinguish between two tumor types and significantly reduce the time needed to train separate models. More visualizations can be found in Appendix C. Zoom in for a better view.We trained a classification model to validate the efficacy of the proposed 3D-VQGAN-cond model in distinguishing between HGG and LGG brain tumor types. Table 2 shows the classification results for Experiments (1) and ( 2) from Section 3. For Experiment (1), we showed that the model pre-trained with both synthetic HGG and LGG data outperforms all baselines, including the model trained with traditional augmentations and trained only on synthetic LGG data. This highlights the substantial improvement in classification performance achieved through pre-training with synthetic data, alleviating the demand for extensive real data to train effective classification models. For Experiment (2), we demonstrated that the classification performance improves when the number of synthetic data used for pre-training increases. These results collectively indicate that synthetic ROI data can be effectively used to pre-train deep models and only a small amount is needed for finetuning. In addition, the improvements from both experiments are statistically significant (two-sided t-test p < 0.05) compared to the previous SOTA (3D-VQGAN), which further validate that both LGG and HGG samples generated by our proposed 3D-VQGAN-cond model have better image quality and fidelity compared to other baselines. More results can be found in Appendix D.Table 2 : Results for all experiments as described in Section 3. We run all for three trials and report as mean±standard deviation. Trad. Aug. is the short for traditional augmentations.(a) Experiment (1)  
conclusions: 
 We propose the first class-conditional generation framework for LGG and HGG brain tumor types based on VQGAN and masked Transformer. The conditional scheme enables generating different types of tumors in a unified framework, rather than in separate models that require a large amount of time and resources to train. Our proposed method performs better or on par with several baseline models in image quality metrics such as MS-SSIM, slicewise FID, and MMD. Using the generated data, our method yields the best classification performance compared to all other baselines.Meng Zhou, Matthias W Wagner, Uri Tabori, Cynthia Hawkins, Birgit B Ertl-Wagner, and Farzad Khalvati. Generating 3d brain tumor regions in mri using vector-quantization generative adversarial networks. arXiv preprint arXiv:2310.01251, 2023. 
 
 Appendix: A.1. Implementation Details of the Scoring NetworkThe scoring network f is learned using an auxiliary MLP, f cls , through the classifierguidance approach we described in Section 2.1. To learn f , we first use the encoder E to get the latent feature map z e ∈ R H×W ×Dp×nz , then for each region r i ∈ R nz , its score is defined by s i = f (r i ),All programs were implemented in Pytorch, and all models were trained on a single TESLA V100 GPU. Additionally, we applied the automatic mixed precision in the PyTorch library during the training process (Subramaniam et al., 2022) . The overall training time for our 3D-VQGAN-cond model takes about 7 GPU days to complete. For our transformer architecture, we used the same one in (Esser et al., 2021b) . Our code is available at here.For classification, we pre-trained for 50 epochs with a batch size of 8 and a learning rate of 0.001 with Adam optimizer for all models using synthetic data. These models are then finetuned using a batch size of 10, and a learning rate of 0.01. For models that do not involve the use of synthetic data (i.e., traditional augmentation), we trained those with a batch size of 8 and a learning rate of 0.01 with Adam optimizer. All classification models are optimized by focal loss (Lin et al., 2017) , as we noticed that it is better than the standard cross entropy loss.Data Split. As discussed in the main text, we randomly hold out 25 HGG and LGG patients (50 in total) as the standalone test data. For the rest of the 234 HGG patients and 51 LGG patients, we design two sets of training data combinations for traditional and non-traditional augmentation methods. For the traditional augmentation baseline, we augment the LGG data by rotating 30 degrees, scaling by 1.5 times larger, left-right flipping, and elastic deformation to form a balanced dataset of 234 cases for both HGG and LGG. For all other non-traditional augmentation models, we pre-trained on 183 real HGGs and 183 synthetic LGGs, and then fine-tuned with 51 real HGGs and 51 real LGGs for Experiment (1) in Section 3 follows the setup in (Zhou et al., 2023) ; we pre-trained on 250 synthetic HGGs and LGGs, and then fine-tuned with 51 real HGGs and 51 real LGGs for Experiment (2). We ensure that the data for fine-tuning is the same across experiments for a fair comparison. Furthermore, during the fine-tuning stage in both experiments, we used 85% of the data for optimizing the model and the remaining 15% of the data for validation, and we repeated this process three times to ensure the robustness of our model. We also ensure that there is no overlap between the validation data in the three runs. We have also conducted ablations on the number of neighbor tokens to be masked around important tokens, as described in Section 2.1. In our work, we selected only one neighborhood token (we denote as single-side) in either spatial or temporal dimension to be masked, given that our latent feature maps are relatively compact (8 × 8 × 8), and masking more tokens around the important one may lose too much feature information for the model to be learned and reconstructed effectively during training. In the ablation, we fix the mask ratio α = 0.5, top-k ratio k = 25%, and mask two neighborhood tokens, i.e., for a given important token, we mask two more tokens to its left and right in the spatial dimension, or successor and predecessor tokens in the temporal dimension (we denote as double-sided).The quantitative results are reported in Table 4 , we observed that masking more tokens led to performance degradation, which validates our claim above. LGG HGG MMD (10 4 ) ↓ MS-SSIM MMD(10 4 ) ↓ MS-SSIM single-side 1.72 87.9 (2.6) 1.57 88.5 (0.1) double-sided 2.07 88.9 (3.6) 1.91 89.9 (1.3) 7 . Our method outperforms MedicalNet by a significant margin in all metrics, demonstrating the effectiveness of our proposed approach. Although there is a difference with the model backbone we used, we hypothesize that this change would not dramatically alter the results. This hypothesis is based on the datasets used to pre-train the model were whole 3D volumes of various organs, including but not limited to the brain. Moreover, the substantial difference between whole 3D volumes and Regions of Interest (ROIs) can also affect the results.  


Decision: After a busy and productive discussion between authors and reviewers (including new results), the paper is very much improved and all authors agree that the MIDL community will see the value on their method.

Reviews: review: paper_summary: paper_summary: The authors presented a generative GAN-based method for generating synthesized brain tumor ROIs, and verified its added benefit for improving the tumor classification performance.
main_review: main_review: Heartfelt thanks to the authors for carefully addressing my questions and concerns. The authors present an innovative generative GAN-based methodology and assess its added value on performance with experiment results; this could offer valuable insight to the research community. I would still strongly recommend the authors make the code publicly available, with good documentation.
strength_weakness: strength_weakness: The selection of modules and methods is less motivated and verified. The code seems not to be available, making it hard for other researchers to reproduce and compare. It is strongly recommended to release the code for the convenience of the research community.
questions: questions: see detailed comments.
paper_summary: paper_summary: This paper proposes to use 3D VQ-GAN to generate 3D brain tumour ROIs, instead of the entire brain, which can stabilises the training of GANs by reduce data dimension. Evaluation is performed by evaluating the quality of synthetic brain tumour ROIs and also by downstream task training on synthetic data.
main_review: main_review: This paper proposes to use 3D VQ-GAN to generate 3D brain tumour ROIs conditioned on labels of LGG and HGG. I found the technical novelty is limited, and I do not follow the advantages of generating ROIs rather than whole brains. The experiments only focus on comparing different generative models for ROIs, but would be better to see experiments comparing generating ROIs and whole brains. The claimed benefit of focusing on ROIs is because GANs are unstable, but the size for ROIs is set as 128128128, which is large and I am not sure how this is going to help stabilise GAN training.

====update======
The authors’ responses have addressed most of my concerns. I would like to update rating accordingly.
strength_weakness: strength_weakness: - Evaluation is only on one dataset with 259 HGG patients and 76 LGG patients. 
- Technical novelty/contribution is limited. Despite it focused on ROIs, the differences between this paper and previous works on 3D brains are minimal.
- Brain tumour can not only affect the ROI, but also affect the other parts of brains by causing deformations. Maybe have some discussion on it?
- ROIs are center-cropped to "a target size of 128 × 128 × 128”. As the authors mentioned, they focused on ROIs because “GANs can be unstable, produce blurry images, and encounter mode collapse problems”. However, 128*128*128 is still large, and since authors used GANs in this paper, how did they ensure that the unstable training did not occur? 
- What if we reshaped a whole 3d brain to 128*128*128, and generate the downsampled version brains?
questions: questions: - Brain tumour can not only affect the ROI, but also affect the other parts of brains by causing deformations. Maybe have some discussion on it?
- ROIs are center-cropped to "a target size of 128 × 128 × 128”. As the authors mentioned, they focused on ROIs because “GANs can be unstable, produce blurry images, and encounter mode collapse problems”. However, 128*128*128 is still large, and since authors used GANs in this paper, how did they ensure that the unstable training did not occur? 
- What if we reshaped a whole 3d brain to 128*128*128, and generate the downsampled version brains?
- Improve Figure 2 and 3, I think might be worth to have similar real images by side to better demonstrate image quality.
paper_summary: paper_summary: This work aimed at 3D brain tumor ROI generation task as well as downstream classification task. This work adopts VQVAE and VQGAN structure to learn a discretized latent space for brain tumor ROI generation. The main contribution and novelty of this work could be summarized as 1. converting the latent feature into importance map by introducing an auxilary classification module and loss; 2. using a class-conditioned transformer and a proposed temporal-agnostic hybrid token masking method to improve the performance of the transformer. The proposed approach is demonstrated on a single dataset with LGG and HGG data. The results show the success of the proposed method in both quality of generated data and downstream classification accuracy.
main_review: main_review: I think the authors have successfully addressed my concerns and I'm satisfied with the additional experiments. They added more experiments on the ablation studies of different ratio of masking. Thank the authors for their efforts and this great work.
strength_weakness: strength_weakness: 1. The size of the dataset is not that large. Also the dataset is not balanced for this two classes. I wonder if the authors have done anything to address that issue.
2. I didn't see an ablation study in terms of the masking strategy. I wonder if the authors have done that or not.
3. The comparisons with diffusion model is not that fair. It's very easy and straight forward to add class condition to diffusion models. In this work, the diffusion model is trained separately.
questions: questions: Similar to detailed comments.

1. I would like to see some explanations from the authors about the proposed masking strategy with ablation studies. Is it working? How does it work? Also whether the choice of the number of neighbors matters?

2. I would like to see some comparisons with diffusion models with condition.

3. I wonder if the author could explore their proposed method on more complicated data, e.g., larger ROI, more classes instead of 2.

Comments: title: Thank you
Thank you
Looking forward to further feedback
Looking forward to further feedback
Thank you
Reply to Clarification Regarding the Updated Rating
Clarification Regarding the Updated Rating
Update of comments
Author Response to Reviewer h3no
References
Author Response to Reviewer jbth
Author Response to Reviewer Y4ZB
Summary of Changes
comment: Thank you for your feedback! We are happy to hear that we were able to address all of your concerns.
Heartfelt thanks to the authors for addressing my concerns. I have no more questions.
Thank you for your feedback! We are happy to hear that we were able to address all of your concerns.
Congrats on your successful revisions! My concerns are addressed and thank you for your efforts on the paper and revision.
Dear Reviewer Y4ZB,

Thank you again for your valuable comments and suggestions, which are very helpful to us. We have posted responses above to your concerns.

We understand that this is quite a busy period, so we sincerely appreciate it if you could take some time to return further feedback on whether our responses address your concerns. If there are any other comments, we will try our best to answer them.

Best,

The Authors
Dear Reviewer h3no,

Thank you again for your valuable comments and suggestions, which are very helpful to us. We have posted responses above to your concerns.

We understand that this is quite a busy period, so we sincerely appreciate it if you could take some time to return further feedback on whether our responses address your concerns. If there are any other comments, we will try our best to answer them.

Best,

The Authors
Thank you for letting us know, we appreciate it.
Thanks. I have double checked it has been successfully updated. Somehow it is only visible to ACs, Reviewers and PCs. I tried to set it visible to Authors but couldn’t make it. Anyway, Rest assured that the rating is updated.
We would like to express our gratitude again for your insightful comments and we are happy to see that our responses address most of your concerns and thanks for updating the rating.

However, upon checking the author console, we could not see the change in the rating. Could you please kindly double-check that the rating has been successfully updated? Thank you once again for your review!
Thank the authors for the response. It addresses most of my concerns and I update my rating accordingly.
We thank the reviewer for the constructive comments.

**Q1: Ablations on the proposed masking strategy**

We have conducted two ablation studies. We focused on the top-k ratio we defined in Section 2.1 (ratio of important tokens of all discrete semantic tokens obtained from encoded feature maps), and also the number of neighboring tokens you have suggested. We claimed that the top-k ratio controls the trade-off between the diversity of generated images (measured by MS-SSIM score) and their quality (measured by MMD score). We have conducted experiments on k=0, 25, and 50% and reported two scores mentioned above on the generated samples. Notice that when k=0, we treat every token as unimportant, thus the proposed masking strategy degrades to random masking. Results show that when k=0%, we have a slightly lower MMD but much higher MS-SSIM, when k=25%, we have a slightly higher MMD than k=0% but much lower MS-SSIM score, when k=50%, we have a much higher MMD and almost the same MS-SSIM with k=25%. Hence, to balance this trade-off, we select k=25% in this study. Next, we conducted ablations on the number of neighboring tokens to be masked. We masked only one neighbor token in either spatial or temporal dimension in our work, so the ablation is done by masking two neighboring tokens around the important one. We observed a dramatic increase in the MMD score for double-sided masking. Please refer to Appendix B for full details.

**Q2: Diffusion model with conditions**

We have re-implemented a class-conditioned version of Medical Diffusion (termed Medical Diffusion-C), since it is the most similar architecture compared to ours (VQGAN+Diffusion vs. VQGAN+Transformer) and is computationally friendly. Although Medical Diffusion-C outperforms in FID score based on generated samples, ours is better in other quantitative metrics as well as classification results. Please see Table 1 and 2, Figure 2,3, and 4, and Appendix D for details. 

**Q3: Hard to tell which generated ROI is better**

We agree with the reviewer. As a non-radiologist, it is challenging to tell the difference between generated ROIs from a clinical perspective. But here, we provide some of our own thoughts when we look at these ROIs in Figure 3 and 4. All GAN-based baselines produce images with noise and blurry edges, and the image quality is low. For the 3D-Med-DDPM, the intensity range in the generated samples seems to mismatch the real samples, and it looks unreal compared with the real ROIs. For Medical Diffusion/-C, the generated images suffer from minor checkerboard artifacts (visible when zoomed in). 3D-VQGAN sometimes produce blurry images (sample 2 in HGG), but overall, the generated images are smooth and do not have any checkerboard artifacts or noises. Finally, for our method, the images exhibit high-resolution with no noise, no blurry edges, and no checkerboard artifacts, the contrast inside the generated ROIs looks similar to the real ROIs. Therefore, we think our proposed method produces the best ROIs in terms of visual quality. We have also included this in Appendix C, highlighted in red.

**Q4: Proposed method on more complicated datasets**

We agree with the reviewer that it would make our proposed method stronger if we could evaluate it on different datasets. However, considering the costs and efforts to get more complex datasets, we think this would be a great line of future work to work on. We envision that our method can adopt to any other tumor types as well.
[1]. Sajjad, M., Khan, S., Muhammad, K., Wu, W., Ullah, A., & Baik, S. W. (2019). Multi-grade brain tumor classification using deep CNN with extensive data augmentation. Journal of computational science, 30, 174-182.

[2]. Mzoughi, H., Njeh, I., Wali, A., Slima, M. B., BenHamida, A., Mhiri, C., & Mahfoudhe, K. B. (2020). Deep multi-scale 3D convolutional neural network (CNN) for MRI gliomas brain tumor classification. Journal of Digital Imaging, 33, 903-915.

[3]. Srinivasan, S., Bai, P. S. M., Mathivanan, S. K., Muthukumaran, V., Babu, J. C., & Vilcekova, L. (2023). Grade Classification of Tumors from Brain Magnetic Resonance Images Using a Deep Learning Technique. Diagnostics, 13(6), 1153.

[4]. Esser, P., Rombach, R., & Ommer, B. (2021). Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 12873-12883).
Thank you for your review and insightful comments.

**Q1: Discussion on why ROI-based classification for the brain tumor, and missing literature**

This is a great question. In fact, the brain tumor types diagnosis/classification based on MRI are ROI-based in recent studies [1,2,3]. We are classifying tumor types in this work (HGG vs. LGG), rather than classifying whether the tumor is present in the brain or not. In this case, the key difference between **tumor pathologies (types)** would be based on the tumor area (ROI), and hence, the information outside the ROI may negatively affect the results. We have added this to Section 1 (page 2), highlighted in red. We have also added you suggested literature to Section 1 (page 2).

**Q2: GAN stability for 128x128x128 images**

We would like to re-emphasize on the VQGAN architecture that we used in this work. The goal of VQGAN is to learn a CNN-based autoencoder, and a codebook to quantize the latent feature maps to discrete semantic tokens. The overall workflow of VQGAN is first to encode the image into latent features, perform quantization and obtain quantized features from the codebook, and decode them back to the original image. To supervise the training of VQGAN, we follow the original paper [4] to use an adversarial training approach (GAN) to differentiate between real and reconstructed images, in addition to conventional reconstruction objectives such as L1 loss, perceptual loss, etc. Therefore, in our case, the GAN is only used to force the reconstructed image from the entire workflow to be the same as the original image in an adversarial manner. Unlike other conventional GAN models, **we are not using GAN to generate any new images**. In Section 1, we refer GANs that are unstable and produce blurry images as those models that generate images from random noise.  In our work, the image generation happens in our proposed class-conditional masked transformer model.

**Q3: Re-emphasize novelty**

We want to reemphasize our novelty here; we agree with the review that this paper is based on some existing previous works on image generation. However, there are substantial differences with previous publications. 

1. To the best of our knowledge, this is the first attempt to use VQGAN and transformer to generate high-resolution brain tumor types (HGG and LGG) in a class-conditioned manner.

2. The proposed masking strategy introduces a novel approach by considering the importance of each token during the masking process, contrasting with conventional random masking methods. The proposed strategy prioritizes tokens based on their importance scores and we perform block-wise masking around them to improve the model learnability and robustness. We have shown the effectiveness of this approach by conducting ablation studies in Appendix B.

3. To learn the importance score, we proposed a classifier-guidance approach by using an auxiliary classifier to guide the learning process. Although classifier-guidance is not new, it is the first time this approach applied to the medical image data to effectively learn the importance of each region feature.

**Q4: Reshape whole 3D brain to 128x128x128 and generate the downsampled version brains**

As we stated in the previous question, the current effort focuses on ROI-based classification of brain tumor pathology. Hence in this work we continue this line of work and only use ROI instead of the whole brain. We agree with the reviewer that it is possible to work with the whole brain. However, in that case, we need either to generate synthetic whole-brain images with the corresponding segmentations simultaneously, or we need to use a subsequent segmentation model to get the ROI, this is somehow out of the scope of our study. We think the proposed approach by the reviewer is a great line of future work to continue.

**Q5: Include real ROIs in figures**

It is quite hard to enlarge figures in the main text due to the page limit, but we have included the real ROIs in Figure 3 and 4 and enlarged them in Appendix C for better visualization and comparison.
We thank the reviewer for the detailed feedback and insightful comments.

**Q1: Class-conditioned vs. Class-guided**

When we train the transformer, we prepend a class token to the discrete semantic representation for each sample we obtain from the codebook. The transformer will learn the relationship between the semantic tokens given the information of the class label during training. When we generate the synthetic image, we only provide a class label (0 or 1) to the transformer, and let it autoregressively generate the next token to obtain a full sequence of tokens, which are then decoded back to the image. Hence, since we use a single label describing the overall image class and use it as a condition to generate the synthetic image, we believe it should be class-conditioned, following the terminology from previous publications [1,2]. We have changed ‘conditional 3D-VQGAN’ to ‘3D-VQGAN’ and ‘Conditional Masked Transformer’ to ‘class-conditional masked transformer’ to avoid any confusion (Figure 1, page 3).

**Q2: Misclassification in importance score map**

The goal of our proposed auxiliary MLP classifier is to learn which region in the latent feature maps is important and assign a weight (score) with respect to the tumor type (HGG/LGG). If the region can best reflect the tumor or can be used to differentiate tumor types, it should receive a higher weight. In the second stage, when we train our transformer model, we use the importance score map to get the important/unimportant regions in the feature maps and perform masking for each sample. The importance score map is not used in either ROI generation or the downstream classification task. We agree that there is a chance that the misclassification problem would occur, but this only affects the weight of feature regions, and whether or not they are important, and the way they are masked. Since our masking strategy is hybrid, we argue that this hybrid masking is relatively robust to this problem.

**Q3: Compare with transfer learning and Add p-value to the results**

We have conducted the independent-sample two-sided t-test on the AUC and F1-score of our proposed method and compared them with the previous SOTA (3D-VQGAN). Both p-values for AUC and F1-score are statistically significant (p < 0.05 in all cases). For the transfer learning baseline, we used MedicalNet [3] pre-trained weights with their ResNet-34 backbone instead of ResNet-50 due to our computational limitation. We froze the feature extractor and only trained the classification head. Results and details can be found in Appendix D.2.

References:

[1]. Esser, P., Rombach, R., & Ommer, B. (2021). Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 12873-12883).

[2]. Yu, J., Li, X., Koh, J. Y., Zhang, H., Pang, R., Qin, J., ... & Wu, Y. (2021). Vector-quantized image modeling with improved vqgan. arXiv preprint arXiv:2110.04627.

[3]. Chen, S., Ma, K., & Zheng, Y. (2019). Med3d: Transfer learning for 3d medical image analysis. arXiv preprint arXiv:1904.00625.
We would like to thank all the reviewers for their helpful and constructive feedback, and insightful suggestions. We have addressed these comments and below, we summarize the major changes to the paper.

1.	We added the p-value of our proposed method compared with the previous SOTA (3D-VQGAN) in Table 1, and included the results for class-conditional Medical Diffusion as our ‘class-conditional’ baseline (Reviewer Y4ZB, h3no).

2.	We compared our proposed method with a transfer learning-based approach, (MedicalNet [1]), and reported in classification results in Appendix D.2. (Reviewer Y4ZB).

3.	We included some discussions in ROI-based vs. whole-brain-based classification, and added more references in Section 1, page 2 (Reviewer jbth).

4.	We conducted extensive ablation studies and reported the results in Appendix B (Reviewer h3no).

5.	We included the generated samples of class-conditional Medical Diffusion in Figure 2,3, and 4. We also enlarged Figure 3 and 4 in Appendix C and included real ROIs for better visualization and comparison. We provided a brief analysis of the image quality of the generated images for all methods in Appendix C, highlighted in red (Reviewer jbth, h3no).

6.	Due to the page limit, we decided to move some implementation details regarding the importance score map from the main text to Appendix A.1.

We are working on the code and will release them in the near future.

Reference:

[1]. Chen, S., Ma, K., & Zheng, Y. (2019). Med3d: Transfer learning for 3d medical image analysis. arXiv preprint arXiv:1904.00625.