🤖 machine learning

VariViT: A Vision Transformer for Variable Image Sizes

The paper proposes VariViT, a Vision Transformer architecture that handles variable image sizes through a novel positional embedding resizing scheme and an efficient batching strategy, achieving superior performance in brain tumor classification and genotype prediction while reducing computational costs by up to 30%.

Original authors: Aswathi Varma, Suprosanna Shit, Chinmay Prabhakar, Daniel Scholz, Hongwei Bran Li, Bjoern Menze, Daniel Rueckert, Benedikt Wiestler

Published 2026-02-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Aswathi Varma, Suprosanna Shit, Chinmay Prabhakar, Daniel Scholz, Hongwei Bran Li, Bjoern Menze, Daniel Rueckert, Benedikt Wiestler

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery by looking at a series of photographs. In the world of Artificial Intelligence (AI), these "detectives" are called Vision Transformers (ViTs). They are incredibly smart at looking at pictures to find patterns, like spotting a tumor in a brain scan.

However, there's a big problem with how these detectives usually work: They only accept photos of one specific size.

The Problem: The "One-Size-Fits-All" Suit

Imagine you have a collection of photos of tumors. Some are tiny specks, and some are huge masses.

The Old Way (Vanilla ViT): To feed these photos into the AI, you have to force them all into a standard frame.
- If the tumor is small, you have to stretch the photo out to fit the frame. This makes the tumor look blurry and introduces fake details (artifacts), like stretching a rubber band until it snaps.
- If the tumor is huge, you have to chop off the edges or shrink the whole image down. This might cut off important clues or make the tumor look too small to see clearly.
- The Result: The AI gets confused. It spends a lot of time looking at the empty background (the healthy brain) because the tumor got distorted or shrunken, wasting its energy.

The Solution: VariViT (The Flexible Detective)

The authors of this paper created a new AI model called VariViT. Think of VariViT as a detective who doesn't need a rigid frame. They can look at a photo exactly as it is, whether it's a tiny square or a giant rectangle.

Here is how VariViT works, using some simple analogies:

1. The "Center and Select" Trick (The Positional Map)

To understand a picture, an AI needs a map to know where things are (e.g., "the tumor is in the top-left corner").

The Old Way: If the photo size changes, the AI tries to stretch or squish its map to fit. This is like trying to stretch a rubber map; the streets get distorted, and the AI gets lost.
The VariViT Way: Imagine the AI has a giant, perfect map of the largest possible tumor. When it looks at a smaller tumor, instead of stretching the map, it simply zooms in on the center of that giant map.
- Since tumors are usually centered in medical scans, the "center" of the big map and the "center" of the small tumor are the same.
- VariViT just grabs the relevant chunk of the map from the middle. No stretching, no distortion, just a clean, sharp view of the tumor.

2. The "Smart Grouping" Strategy (Batching)

Training an AI is like teaching a class of students. Usually, you put students of the same height in the same row so they can all see the board.

The Old Way: If you have students of different heights (different image sizes), you either have to make the short ones stand on boxes (padding) or the tall ones sit on the floor (cropping). This is slow and messy.
The VariViT Way: The authors invented two clever ways to group the students:
1. Group by Size: Put all the "small tumor" photos in one batch and all the "large tumor" photos in another. The AI learns them separately but efficiently.
2. The "Accumulate" Method: If you must mix sizes, the AI takes a few small photos, learns from them, then takes a few large ones, learns from them, and then updates its brain. This saves a massive amount of computing power.

Why Does This Matter?

The researchers tested VariViT on brain scans to do two things:

Identify the type of brain tumor (Primary vs. Metastatic).
Predict genetic mutations (IDH status) which helps doctors choose the right treatment.

The Results:

Better Accuracy: VariViT was better at spotting the tumors and predicting their genetics than the old "stretch-and-shrink" models. It achieved higher scores (F1-scores around 75-76%) because it wasn't distracted by blurry backgrounds.
Faster Speed: Because it didn't waste time stretching images or processing empty background space, it trained 30% faster.

The Bottom Line

VariViT is like giving the AI detective a pair of smart glasses that automatically adjust to the size of the object they are looking at. Instead of forcing the world to fit the AI's rigid rules, VariViT adapts to the messy, irregular reality of human anatomy. This means doctors can get faster, more accurate diagnoses without the AI getting confused by distorted images.

1. Problem Statement

Vision Transformers (ViTs) have achieved state-of-the-art performance in image representation learning but face significant limitations when applied to medical imaging, particularly 3D MRI scans:

Fixed Input Constraints: Standard ViTs require fixed-size inputs, necessitating pre-processing steps like resizing, padding, or cropping.
Information Loss & Artifacts: Resizing medical images can distort anatomical structures, introduce artificial features, or obscure real abnormalities, negatively impacting diagnosis.
Variable Foreground-to-Background Ratios: Medical lesions (e.g., tumors) vary greatly in size. A fixed bounding box crop results in highly variable ratios of tumor (foreground) to healthy tissue (background). Standard ViTs often waste computational resources attending to irrelevant background data or fail to capture small, irregular pathologies effectively.
Computational Trade-offs: Large images are computationally expensive, while aggressive cropping/resizing risks losing critical diagnostic information.

2. Methodology: VariViT

The authors propose VariViT, a modified ViT architecture designed to handle variable image sizes while maintaining a consistent patch size. The core innovations are:

A. Center and Select Positional Embedding Strategy

Unlike standard ViTs that use interpolation to resize positional embeddings (which can cause information loss), VariViT employs a "Center and Select" mechanism:

Concept: In tumor classification, 3D crops are typically centered on the tumor's center of mass. This provides a consistent spatial reference point across images of different sizes.
Implementation:
1. Initialize a fixed 3D sinusoidal positional embedding grid for the largest possible image size ( $G_l \times G_h \times G_w$ ).
2. For smaller input images, the model dynamically selects a subset of embeddings centered around the middle of the large grid.
3. This avoids interpolation, preserving the original positional information and reducing computational complexity associated with resizing 3D grids.

B. Novel Batching Strategies

To address the challenge of training on batches containing images of different dimensions (which standard frameworks struggle with), VariViT implements two strategies:

Custom Batch Sampler (CBS): Groups images of the same size into a single batch. The batch size remains constant, but the image dimensions vary between batches. This ensures consistency within a batch and accelerates training by utilizing smaller images efficiently.
Gradient Accumulation (GA): Allows for batches with varying image sizes within the same iteration. Gradients are accumulated over several mini-batches (effectively setting the mini-batch size to 1) before performing a weight update. This mimics a larger batch size while handling variable inputs.

3. Key Contributions

Flexible Positional Embedding: Introduction of the "Center and Select" method, enabling ViTs to process variable-sized 3D images without interpolation artifacts.
Efficient Batching: Development of CBS and GA strategies to manage heterogeneous image sizes within training batches, significantly reducing training time.
Empirical Validation: Demonstration of superior performance in two challenging medical tasks: Glioma genotype prediction (IDH status) and Brain Tumor classification (Primary vs. Metastatic).

4. Experimental Results

The model was evaluated on two 3D brain MRI datasets (Glioma dataset: 1,856 scans; Brain Tumor dataset: 1,699 scans) against baselines including 3D ResNet-18, Vanilla ViT, and Pix2Struct.

Task 1: IDH Mutation Status Prediction (Glioma Dataset)

Performance: VariViT-GA achieved the highest F1-score (75.5%) and MCC (0.709), outperforming Vanilla ViT (F1: 74.4%) and ResNet-18 (F1: 71.6%).
Efficiency: VariViT-CBS and GA models demonstrated significantly faster training times compared to fixed-size baselines.
Visualization: t-SNE plots showed clearer cluster separation for the VariViT model, indicating better feature discrimination.

Task 2: Primary vs. Secondary Brain Tumor Classification

Performance: VariViT-GA achieved an F1-score of 76.3% and MCC of 0.706, outperforming all baselines including Pix2Struct.
Efficiency: The model maintained high accuracy while reducing computation time by up to 30% compared to conventional architectures.

Ablation Studies:

Positional Embeddings: The "Center and Select" method outperformed independent fixed embeddings, trilinear interpolation (fixed), and learned interpolation methods.
Multi-class Classification: In a 3-class glioma subtype task, VariViT-GA achieved comparable MCC scores to ResNet-18 and Pix2Struct but with faster training.

5. Significance

Medical Imaging Adaptation: VariViT directly addresses the "fixed-size" bottleneck of ViTs in medical imaging, allowing models to focus on regions of interest without distorting anatomical data through resizing.
Computational Efficiency: By leveraging variable-sized crops and efficient batching, the method reduces the computational burden of training large 3D models, making high-performance ViTs more accessible for resource-constrained medical applications.
Feature Learning: The ability to process variable-sized inputs without interpolation leads to more discriminative feature learning, as evidenced by improved F1-scores and MCC in tumor classification tasks.
Future Potential: The approach is modality-agnostic and can be extended to other 3D medical imaging tasks or high-resolution datasets where anatomical variability is high.

In conclusion, VariViT represents a significant architectural advancement for applying Vision Transformers to real-world medical data, balancing the need for high-resolution, distortion-free input with computational efficiency.