VariViT: A Vision Transformer for Variable Image Sizes

Authors: aswathi  varma
daniel  scholz
daniel  rueckert
benedikt  wiestler
s  varma
c  shit
d  prabhakar
h b scholz
b  li
d  menze
b wiestler rueckert
Venue: MIDL_io_2024_Conference
Published: 2024-01-30 17:28:40

Abstract: Vision Transformers (ViTs) have emerged as the state-of-the-art architecture in representation learning, leveraging self-attention mechanisms to excel in various tasks. ViTs split images into fixed-size patches, constraining them to a predefined size and necessitating pre-processing steps like resizing, padding, or cropping. This poses challenges in medical imaging, particularly with irregularly shaped structures like tumors. A fixed bounding box crop size produces input images with highly variable foreground-to-background ratios. Resizing medical images can degrade information and introduce artifacts, impacting diagnosis. Hence, tailoring variable-sized crops to regions of interest can enhance feature representation capabilities. Moreover, large images are computationally expensive, and smaller sizes risk information loss, presenting a computation-accuracy tradeoff. We propose VariViT, an improved ViT model crafted to handle variable image sizes while maintaining a consistent patch size. VariViT employs a novel positional embedding resizing scheme for a variable number of patches. We also implement a new batching strategy within VariViT to reduce computational complexity, resulting in faster training and inference times. In our evaluations on two 3D brain MRI datasets, VariViT surpasses vanilla ViTs and ResNet in glioma genotype prediction and brain tumor classification. It achieves F1-scores of 75.5% and 76.3%, respectively, learning more discriminative features. Our proposed batching strategy reduces computation time by up to 30% compared to conventional architectures. These findings underscore the efficacy of VariViT in image representation learning.

Main paper: 
 introduction: 
 Deep neural architectures, notably Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs), have emerged as effective architectures for image representation learning, consistently achieving state-of-the-art performance on real-world data across different domains. ResNet, a CNN-based variant (He et al., 2016) , achieves its efficacy in representation learning by the use of residual connections. ViT (Dosovitskiy et al., 2020) captures long-range dependencies by directly attending to global image information through selfattention mechanisms. ViTs are increasingly attracting attention in medical image analysis (He et al., 2021; Gao et al., 2021a; Chen et al., 2021; Shamshad et al., 2023; Gao et al., 2021b; Jang and Hwang, 2022) . Input images are typically resized to a fixed size before being fed into a ViT model. These models perform well with evenly dispersed signals, allowing operations like interpolation and cropping for fixed-size inputs. However, medical images feature small, irregular regions of interest, where such methods can be detrimental. Fixed-size inputs result in varying foreground-to-background ratios, especially with smaller pathologies, where more background (e.g., healthy brain) is included compared to larger tumors. Moreover, medical images are particularly sensitive to distortions. Resizing them may introduce artificial features or modify existing ones, mimicking or obscuring real abnormalities and interfering with diagnosis. Figure 1 We propose VariViT to handle variable-size images, addressing the limitation of fixedsize inputs. Our method recognizes the heterogeneous nature of real-world medical images where foreground-to-background ratios vary significantly. VariViT retains the favorable properties of ViTs while integrating the capability to handle diverse image sizes. Our contributions are as follows:1. We introduce a novel flexible positional embedding strategy tailored to different image sizes. 2. We propose an alternate batching strategy to improve computational efficiency, leveraging on the inclusion of smaller-sized images. 3. We demonstrate the applicability of VariViT in (i) glioma genotype prediction and (ii) brain tumor classification, two challenging tasks due to the highly variable tumor sizes. Our extensive experiments on two brain MRI datasets highlight the superior performance of VariViT over both vanilla ViT and ResNet architectures. 
related work: 
 The traditional ViT model uses a fixed input size. It either initializes fixed positional embeddings or learns them during training. These embeddings are linearly interpolated for fine-tuning and evaluation at higher resolutions. The FlexiViT model (Beyer et al., 2023) also adapts to variable image-to-patch size ratios and sequence lengths by resizing the learnable 2D positional embedding grid with bilinear interpolation. While effective in 2D, interpolating the 3D embedding grid for variable-sized images may result in information loss and higher computational requirements due to more complex calculations. The SuperViT model (Lin et al., 2023) patchifies an image at multiple scales and improves computational cost by randomly dropping tokens. However, a random selection of tokens risks information loss and degrades representation quality, particularly in medical images where the region of interest may be small. The Pix2Struct model (Lee et al., 2023) , similar to ours, handles variable image sizes. The vision encoder resizes the input images to extract fixed-size patches fitting a predefined sequence length. Padding is applied to the sequence as needed, allowing it to reach the desired fixed length. The model learns a large grid of 2D absolute positional embeddings, enabling the identification of patch positions based on x and y coordinates. However, this approach can be computationally expensive, especially for 3D images. In contrast, our model efficiently manages size differences without resorting to such computationally expensive operations. Moreover, scaling the images can introduce undesired artefacts.The NaViT model (Dehghani et al., 2023 ) also addresses challenges on computational complexity and variable image sizes. NaViT packs patches from different-sized images into the same sequence. It maintains a fixed sequence length by randomly dropping tokens and padding. The model employs masked attention and pooling to prevent interactions between patches from different images. However, dropping tokens may lead to information loss, and implementing masked attention and pooling introduces more complex architectural changes.Patch Emb.Transformer Encoder Classifier Head Class 0 1 . . . . Sample 3 . . . VariViT Sample 1 Sample 2 L 1 W 1 L 2 W 2 L 3 H 1 H 2 W 3 H 3 N 1 N 2 N 3 
method: 
 Overview. We focus on learning 3D image representation in this work. In the conventional ViT framework, a 3D input image is divided into non-overlapping patches. These patches are flattened and linearly projected to obtain a sequence of patch embeddings. The sequence is fed into the transformer encoder after adding a CLS token. The enriched CLS token is used for the final classification. However, the transformer model lacks intrinsic awareness of the spatial arrangement of patches within the sequence. Therefore, positional embeddings are added to the sequence to explicitly incorporate this information. Both relative and absolute positional encoding can be employed (Vaswani et al., 2017; Wu et al., 2021) . They can either be fixed or learned during training. In its simplest form, absolute fixed sinusoidal embeddings are used for this purpose.Fixed positional embeddings are predefined vectors representing each patch's absolute coordinates within the input sequence, typically generated using sinusoidal functions. The 1D positional encoding for even and odd indices is formulated as (Vaswani et al., 2017) :PE(pos, 2i) = sin pos/10000 2i/d PE(pos, 2i + 1) = cos pos/10000 (2i+1)/d(1) pos represents the position to be encoded. The parameter 1/10000 2i/d governs the wavelength of the sinusoids. Here, d denotes the embedding dimension and i refers to each of the individual dimensions of the embedding.Positional embeddings can be adjusted for various image sizes through interpolation, as suggested in the ViT paper (Dosovitskiy et al., 2020) . However, interpolation increases computational complexity and may introduce approximations, making it suboptimal for resizing 3D embedding grids in variable-sized tumor crops. Pix2Struct suggests learning a large grid of embeddings for a predefined sequence length (Lee et al., 2023) . This, however, can extend training time. Leveraging the consistent center alignment in tumor crops can provide a reliable reference point for resizing positional embeddings without the need for interpolation, thus preserving the original information.Building upon this concept, the VariViT model (Figure 2 ) introduces the center and select method for resizing positional embeddings. Our model adapts the vanillaViT architecture for various input image sizes, particularly tailored to heterogeneous tumor shapes, while maintaining a fixed patch size. This is achieved by extracting 3D bounding box crops categorized into three sizes by tumor volume. Additionally, we exploring batching methods to manage diverse image sizes within batches. The architectural and training modifications to the baseViT are explained in the following sections.Center and Select. Similar to the vanilla ViT, the patch embedding size remains constant in VariViT regardless of the image size. However, with a fixed patch size, variations in image size can lead to a difference in the number of patches or the sequence length. Consequently, positional embeddings must be dynamically resized to accommodate these variations. In tumor classification tasks, 3D crops are often employed to isolate the tumor region by identifying tumor boundaries using segmentation masks. The center coordinates of the 3D crop are aligned with the center of mass of the tumor. Despite variations in size, tumors k , where k = l, h, w. Thus, the original positional information is extracted at a lower computational cost.G l G h G w G' l G' h G' w (C l , C h , C w ) = (⌊G l / 2⌋,⌊G h / 2⌋,⌊G w / 2⌋) CBS GA (a) (b) (c)Batching Methods. Training models with different input dimensions poses challenges due to the variability of bounding box sizes in the dataset. Existing Python frameworks lack seamless solutions to address this issue. To address this challenge, we opt for two specific strategies, namely a custom batch sampler and gradient accumulation.1. Custom Batch Sampler (CBS) -This strategy involves grouping images of the same size into a batch, as shown in Figure 3 (c) while allowing the image size to vary randomly from batch to batch. This maintains consistency within each batch.Including batches with smaller image sizes contributes to significantly faster training. 2. Gradient Accumulation (GA) -In this method, the weight update is performed after accumulating gradients over several mini-batches, resulting in batches with varying image sizes (Figure 3 -(c )). The effective batch size is given by: Batch Size = Mini-Batch Size × Update Interval. Here, we adjust the mini-batch size to 1 and set the update interval to the desired batch size. We make our codes for VariViT and the batching schemes publicly available at https: //github.com/Aswathi-Varma/varivit. 
experimental setup: 
 Datasets. To highlight the effectiveness of VariViT, we perform experiments on two distinct 3D brain MRI datasets (Appendix A): The glioma dataset comprising 1856 MRI scans sourced from various studies (van der Voort et al., 2021; Sayah et al., 2022; Calabrese et al., 2022; Bakas et al., 2022 Bakas et al., , 2017)) ; and the brain tumor dataset containing 1699 MRIs from publicly available datasets (Baid et al., 2021; Suter et al., 2022; Moawad et al., 2023) . Both datasets contain FLAIR, T2w, T1w, and T1w+contrast MR images, all registered to the SRI24 atlas and resampled to a uniform voxel size of 1x1x1 mm³, forming a four-channel multi-modal input for our experiments. These datasets are chosen to evaluate the model's performance on two binary classification tasks: (i) identifying the isocitrate dehydrogenase (IDH) mutation status, a key biomarker that separates two adult-type diffuse gliomas groups (Louis et al., 2021) . (ii) distinguishing between primary brain tumors (gliomas) and secondary brain tumors (metastases). We use the glioma dataset for multi-class classification task, targeting three glioma subtypes: glioblastoma, astrocytoma, and oligodendroglioma.We extract the largest tumor in each patient for our baseline models by cropping a 96×96×96 mm³ bounding box guided by segmentation masks provided in the datasets. To address different tumor sizes in VariViT, we categorize the datasets into three size bins with equal sample distribution (Appendix B). These bins correspond to crop sizes of 64×64×64 mm³, 80×80×80 mm³, and 96×96×96 mm³ for the largest tumors. Additionally, we rescale all image intensities to the range [0, 1].Training and Evaluation. In our training configuration, we opt for the 3DViT model (Prabhakar et al., 2023) . We employ the ViT -S/16 configuration with a patch size of 16. This setup consists of 12 encoder blocks, each having an embedding dimension of 384 and 6 attention heads. A linear layer is used as the classification head. For both the VariViT -S/16 models (GA and CBS), we use the same vanilla ViT base, differing only in the positional embedding. The model comprises approximately 28 million trainable parameters. We utilize ResNet-18, a CNN with 33 million parameters to compare our results with a convolutional model of roughly the same number of parameters. To benchmark against recent stateof-the-art variable image size models, we incorporate the Pix2Struct vision encoder into a 3D framework. This model serves as one of our baselines, with 34.8 million trainable parameters.We fix the sequence length of Pix2Struct at 216, corresponding to our dataset's maximum image size. For all the models, we utilize absolute positional encoding. We explore relative positional embedding with our batching methods in Appendix C but do not observe any significant advantage over the absolute approach.All models undergo training for 100 epochs, utilizing the AdamW optimizer (Loshchilov and Hutter, 2017) with a weight decay of 0.05. The base learning rate is set to 1e-3, and cosine decay (Loshchilov and Hutter, 2016) is applied for learning rate decay. A batch size of 8 is used, and the experiments are conducted on Nvidia RTX A6000. A warm-up schedule (Goyal et al., 2017) of 40 epochs is applied. Data augmentations such as random affine, random noise, random gamma, random blur, and random flips are incorporated during training. Cross-entropy loss with class weights is employed to mitigate class imbalance in the dataset. For the evaluation metrics, we utilize the Area Under the Curve (AUC), F1-score, and Matthews Correlation Coefficient (MCC) score. We train our proposed model on the glioma dataset for IDH status classification task. Then, we compare its performance against three baseline models: fixed-size 3D ResNet-18, 3D vanilla ViT, and variable-size 3D Pix2Struct. We employ a k-fold cross-validation strategy with k=5 for all the models and report the mean metrics and standard deviation obtained from the test sets. The VariViT-GA model consistently outperforms the baseline models in terms of AUC, F1-score, and MCC (Table 1 ). Our model also exhibits notably faster training times compared to its counterparts. VariViT-CBS not only outperforms ResNet-18 and vanilla ViT on various performance metrics, including AUC, but also further reduces the training time. The t-SNE plot (Van der Maaten and Hinton, 2008) in Figure 4 illustrates the improved cluster separation achieved by our model. Brain Tumor Type. To further highlight the effectiveness of our proposed model, we apply it to the brain tumor dataset for the classification of primary versus metastatic tumors. Comparing the VariViT-GA model with the baselines, it distinctly outperforms Pix2Struct and ResNet-18. Our model shows superior performance in MCC and F1-scores compared to VanillaViT (Table 2 ). The VariViT-CBS surpasses Pix2Struct in performance and achieves similar results to ResNet-18, with faster training times. This underscores its efficiency and suitability for practical applications.Ablation Study -Positional Embedding. Here, we compare the effectiveness of different positional embedding methods within the VariViT-GA architecture. We analyze three position-embedding strategies alongside the center and select method: (i) Indepen- 3 ), but the center and select approach produces better results for the IDH status classification task.Ablation Study -Multi-Class Classification. We extend the glioma dataset to classify three glioma subtypes, thereby evaluating the model's performance in this more complex task. Our model achieves comparable MCC scores to both ResNet-18 and Pix2Struct (Table 4 ), while demonstrating faster training times. This study underscores the effectiveness of our model across diverse classification tasks, especially in scenarios involving multiple classes.  
discussion and conclusion: 
 ViTs excel at image feature learning, but limitations exist in medical image analysis due to computational burden and anatomical variability. We address this with VariViT, a method that efficiently scales to various 3D image sizes while effectively learning for improved classification. Our approach centers around maintaining focus on the region of interest, a strategy that demonstrably improves feature learning. By adapting to variable image sizes, VariViT concentrates on critical areas despite inherent anatomical and/or pathological variability. This targeted approach, however, necessitates an initial bounding box or segmentation.While we tested our framework on brain MRIs, its versatility allows for adaptation to other modalities and regions, and we encourage to adapt our method to individual needs. A significant avenue for future research is the exploration of our batching strategy and positional embedding technique in the context of extremely large datasets or high-resolution images.In such scenarios, we anticipate that the efficiency gains of our method will be particularly advantageous when compared to traditional ViT implementations. 
 
 Appendix:  


Decision: This paper proposes strategies for handling variable image sizes in a ViT model.

Strengths: 
+ Handling variable size images is an interesting problem in medical imaging and useful for the community
+ Proposed strategies (flexible positional embedding, alternate batching strategy) could be applied to other ViT architectures
+ Paper is clearly written

Weaknesses:
- Experimental comparisons are largely made to simple baselines (added 1 sota comparison with rebuttal)
- The added multi-class experiment shows slightly weaker or comparable performance to baseline 
- Model cannot handle test image larger than that seen in training

Following rebuttal the paper, the authors included significant enhancements (new comparisons, clarifications/discussion) and the reviewers then all tended toward weak/strong accept. While there is some question regarding the experimental settings and performance, I think the the paper topic and presented strategies would be of high interest to the MIDL community, and therefore recommend accept.

Reviews: review: paper_summary: paper_summary: This paper proposed a novel strategy of "center and select" for resizing positional embedding in Vision Transformer, handling variations in object size. The authors also proposed and tested two strategies of batching methods to address the challenge of different input dimensions. On two 3D brain MRI datasets, VariViT surpasses vanilla ViTs in genotype prediction and tumor classification tasks, with a clearer separation of clusters in tSNE visualization of embedding layer output.
main_review: main_review: I appreciate that the authors addressed the concerns and made substantial changes to the paper. I am keeping the weak acceptance decision since the requirement of segmentation masks is still a limitation and further optimizations are in future directions
strength_weakness: strength_weakness: The proposed method seems to require segmentation masks as part of the dataset since it computes the largest size of tumor bounding boxes the first, which could be a potential limitation.

The two tasks involved in this paper are only binary classification tasks, i.e., outputs 0 and 1. It will be great to also extend the evaluation for multi-class tasks.

The authors only compared their proposed method with vanilla ViTs and Res-Net, and through ablation studies with fixed/learned and independent/interpolated positional embedding strategies, but it will also be of interest to include a comparison with other related works, especially the two similar works Pix2Struct and NaViT.
questions: questions: Following the previous thoughts, it would be great if the authors would discuss the need for segmentation masks and whether this could be optimized or automated, and a comparison with other closely-related works.
paper_summary: paper_summary: The key idea behind this paper is the development of VariViT, "an improved ViT model" designed to handle variable sized medical images. This paper includes:

1. "A flexible positional embedding strategy tailored to different image sizes"
2. "An alternate batching strategy" 
3. Demonstration of VariViT on a couple of different medical imaging tasks.
main_review: main_review: The authors have addressed most of my concerns clearly and have provided proof of additional experiments. Authors have provided sufficient work for me to increase from my initial rating of weak accept.
strength_weakness: strength_weakness: The paper focuses exclusively on brain MRI datasets for glioma genotype prediction and brain tumor classification. The generalizability of VariViT to other medical imaging modalities, body parts, and disease types remains unexplored. Extending the evaluation to include diverse datasets could provide a more comprehensive understanding of the model's applicability and robustness across different medical imaging challenges.

The paper presents an ablation study on positional embedding methods, which is valuable. However, a more comprehensive set of ablation studies, including the impact of different patch sizes, the effect of the novel batching strategy on different types of medical images, and variations in the model architecture, could offer deeper insights into the model's critical components and their contributions to its performance.

Although VariViT shows improved computational efficiency over traditional architectures, the scalability of the proposed batching strategy and positional embedding approach in extremely large datasets or very high-resolution images is not thoroughly discussed.
questions: questions: While the paper presents its strengths, a more explicit discussion on the limitations of VariViT, including potential biases, dataset limitations, and areas where the model might not perform as expected, would provide a balanced view and suggest avenues for future research.

Minor revisions for language and clarity could improve the paper's accessibility. This includes simplifying complex sentences and ensuring consistency in terminology.
paper_summary: paper_summary: Most of the current Vision Transformer (ViT) architectures rely on fixed-size images which implies preprocessing steps on images like resizing, cropping, and padding. For example, a fixed crop size produces input images with highly variable foreground-to-background ratios while resizing medical images may degrade information and introduce artifacts. This may impact the robustness of the model. To handle variable-sized images, this paper adapts the original ViT architecture proposing more flexible absolute positional embeddings and a batching strategy that groups images of the same size into the same batch. The method is evaluated on two 3D brain MRI datasets.
main_review: main_review: The authors have adressed most of my concerns and answered some questions I had.
I particularly appreciate the effort of providing comparisons with previous work while extending their work to the 3D data case. Therefore, I am increasing my grade to weak accept.
strength_weakness: strength_weakness: I like the fact that the method is simple and effective but I have concerns about the comparisons with related works that also propose ViT architecture that can handle variable-sized images. How does the proposed VariViT perform compared to these architectures like Pix2Struct or NaViT? Authors said that because these architectures are dropping tokens it may lead to information loss and hence worse representation quality. However, no quantitative evaluations allow them to conclude about this and the superiority of VariViT compared to Pix2Struct or NaViT. That is why, in addition to the baseline (3D ResNet-18, 3D Vanilla ViT), it would be interesting to see how the method compares to these concurrent works.
questions: questions: I would like the authors to address my concerns in the Weaknesses section but also discuss the points raised in the Detailed Comments section.

Comments: comment: The authors have adressed most of my concerns and answered some questions I had.
I particularly appreciate the effort of providing comparisons with previous work while extending their work to the 3D data case. Therefore, I am increasing my grade to weak accept.
We thank the reviewer for their thoughtful evaluation of our paper on *VariViT*. We appreciate your recognition of the strengths of our work, particularly in handling variable-sized images with *ViT*, and the clarity of our problem statement and proposed solution. We value your feedback regarding potential weaknesses and suggestions for improvement. We have taken your input into consideration and made the necessary adjustments accordingly.

### 1. Comparison with Related Works

Regarding the comparison with related works such as *Pix2Struct* and *NaViT*, we agree that it is essential to provide a comprehensive evaluation to demonstrate the superiority of *VariViT*. In response to your comment, we conducted additional experiments to compare *VariViT* not only with vanilla *ViT* and ResNet but also with *Pix2Struct*. Despite *Pix2Struct's* initial operation in a 2D domain, we extended its application to a 3D framework. Our model demonstrates better performance over *Pix2Struct*, achieving a **1.3%** improvement in the F1-score for *IDH* status classification and a **4.3%** improvement for tumor type classification. Notably, this is achieved with up to an **8%** reduction in training time. To ensure full reproducibility, we also release our 3D implementation of *Pix2Struct* along with the *VarViT* code publicly at [https://github.com/Aswathi-Varma/varivit](https://github.com/Aswathi-Varma/varivit). We were unable to include *NaViT* in our comparison due to two factors. First, the authors did not release their code publicly. Second, certain aspects of their paper lacked clarity, making it difficult to implement *NaViT* within the limited time available for rebuttal.

### 2. Experimentation with Relative Positional Encoding

Regarding the reviewer's suggestion on relative positional encoding, we appreciate your insight into the potential relevance and usefulness of relative position information in certain medical imaging tasks. We experimented with relative positional embeddings alongside our custom batch sampling methods. Although we did not observe significant advantages over absolute positional embeddings in our experiments, we acknowledge the importance of relative position information within certain contexts. As such, we have included a comparison in the *Appendix. C* of the revised paper.

### 3. Handling Test Images Larger Than the Largest Training Image Size

Indeed, our method relies on the assumption that all positions have been encountered during training, including those for the largest image size present in the dataset. We achieve this by utilizing segmentation masks to determine tumor sizes and dividing the dataset into predefined bins. Consequently, our train-test split ensures that both sets contain similar bin sizes, guaranteeing that all sizes in the test set have been encountered during training. If a test image exceeds the size of the largest one in the training dataset, our current model does not support such scenarios. While resizing by interpolation could address this issue, we've prioritized lower training and inference times without interpolation for efficiency.
Thank you to the reviewer for his/her insightful feedback on our paper. We are pleased to see that the reviewer acknowledges the significance of our contribution in introducing *VariViT*, an improved *ViT* model tailored to address the challenges posed by variable-sized medical images. We agree that traditional *ViTs* and *CNNs* often struggle with the variability and sensitivity inherent in medical imaging, and our innovation aims to bridge this gap.

### 1. Limited Dataset Diversity:

Our research primarily focused on brain MRI datasets, influenced by the expertise of our **neuroradiology** research group. However, we recognize the significance of assessing *VariViT* across various medical imaging modalities. We expect that the adaptability and flexibility of *VariViT* make it easily extendable to pathologies of other body regions. Therefore, we plan to extend our evaluation to include a diverse range of medical imaging tasks involving different anatomical sites to provide a more comprehensive understanding of *VariViT's* applicability and robustness in various medical imaging challenges. We now include this aspect (as well as scalability, see below) in our discussion.

### 2. Scope of Ablation Studies:

The reviewer pointed out that our ablation studies could be expanded to include a more comprehensive set of experiments. We agree that exploring aspects like patch sizes, the new batching strategy on various medical images, and variations in the model architecture would be valuable. We plan to conduct additional experiments to gain deeper insights into these critical components and their contributions to *VariViT's* performance. Additionally, we incorporate an analysis of the model's performance in multi-class classification tasks within our revised manuscript. We also explore *relative* positional encoding because of its relevance in certain medical imaging tasks. While integration with our batch sampling methods showed feasibility, no significant advantages over absolute positional embeddings were observed. We include a comparison of both methods in Appendix C of the revised paper.

### 3. Scalability Concerns:
 
We appreciate the reviewers raising concerns about the scalability of our proposed batching strategy and positional embedding approach in extremely large datasets or high-resolution images. To maintain a focused initial study, we limited the scope of our experiments. However, we plan to conduct further investigations specifically in the context of handling large datasets and high-resolution images.
Thank you for the detailed feedback on our paper. We appreciate the acknowledgement of the strengths of our proposed _VariViT_ model in addressing challenges in medical image computing, particularly in the context of irregularly shaped structures like tumors. Additionally, we thank the reviewer for recognizing the intuitiveness and effectiveness of our _center and select_ method for resizing positional embeddings and the efficiency of our proposed _CBS_ and _GA_ batching strategies.

### 1. Segmentation Mask:

We recognize the importance of segmentation masks for tasks involving tumor bounding boxes, as they assist in delineating tumors and calculating their volume. This data is vital for organizing images into suitable size categories. The process necessitates prior knowledge of the Region of Interest's (ROI) size. While we understand the concern about the requirement for segmentation masks, it's important to note that these masks are essential for accurate classification and cannot be omitted. We will explore potential optimizations or automation of tumor delineation in future iterations of our paper. Further, we have added this to the discussion section of the revised report.

### 2. Multi-Class Classification: 

Regarding the limitation of our paper to binary classification tasks, we acknowledge the importance of addressing the evaluation of our model to multi-class tasks. To address this concern, we have extended the glioma dataset to include multi-class labels (tumor type as per _WHO2021_ classification of brain tumors [1]). We demonstrate that our proposed method maintains comparable performance to the baseline models, even in the context of multi-class classification, while exhibiting faster training times. This extension not only broadens the applicability of our approach but also highlights its versatility across various classification scenarios. 

### 3. Comparison to Other SOTA Methods:

The reviewer raised the issue of limited comparative analysis with other related works. We recognize the significance of such comparisons in providing a comprehensive understanding of state-of-the-art methodologies. To rectify this, we conducted additional experiments comparing our method not only with vanilla _ViT_ and _ResNet_ but also with _Pix2Struct_. While _Pix2Struct_ initially operates in a 2D domain, we extended its application to a 3D framework. Our model demonstrates better performance over _Pix2Struct_, achieving a **1.3%** improvement in the F1-score for _IDH_ status classification and a **4.3%** improvement for tumor type classification. Notably, this is achieved with up to an **8%** reduction in training time. To ensure full reproducibility, we also release our 3D implementation of _Pix2Struct_ along with the _VarViT_ code publicly at https://github.com/Aswathi-Varma/varivit. We were unable to include _NaViT_ in our comparison due to two factors. First, the authors did not release their code publicly. Second, certain aspects of their paper lacked clarity, making it difficult to implement _NaViT_ within the limited time available for rebuttal.

[1] Louis, David N., et al. "The 2021 WHO classification of tumors of the central nervous system: a summary." Neuro-oncology 23.8 (2021): 1231-1251.