AG-VIM

Context & Problem

Accurate segmentation of brain tumors in 3D MRI scans is one of the most computationally demanding tasks in medical imaging. The BraTS benchmark — the standard evaluation dataset for this problem — requires models to simultaneously segment three nested tumor sub-regions (enhancing tumor, tumor core, and whole tumor) across four MRI modalities (T1, T1ce, T2, FLAIR), each at a resolution of 240×240×155 voxels. Standard 3D CNNs struggle here because their convolutional kernels capture only local spatial context; the limited receptive field means distant anatomical relationships — critical for identifying diffuse infiltration boundaries — are not modelled. 3D Vision Transformers solve the receptive field problem through self-attention across all voxel patches, but this comes at quadratic computational cost relative to sequence length. For a 3D volume, the sequence length explodes. Training a 3D ViT on full-resolution BraTS volumes requires enormous GPU memory budgets and extremely long training times, making iterative research impractical without access to large-scale cluster infrastructure.

Solution & Approach

AG-VIM proposes a hybrid architecture that resolves this trade-off. The backbone replaces standard transformer blocks with Vision Mamba layers — implementations of the Mamba selective state space model (SSM) adapted for 2D and 3D visual data. Unlike attention, SSMs process sequences with linear time and memory complexity, making full-resolution 3D volumes tractable on a single consumer-grade GPU. The critical insight of AG-VIM is that raw Mamba alone can lose spatial precision at skip connections in the decoder — it models long-range dependencies efficiently but without explicit spatial weighting. To address this, the architecture introduces attention gates at every encoder-decoder skip connection. These gates learn to suppress activations from irrelevant healthy tissue regions and amplify signals from tumour-boundary voxels before feature concatenation. This combination — linear-complexity sequence modelling in the encoder with spatially selective skip connections in the decoder — allows the model to simultaneously capture global structural context and fine-grained boundary detail. Training used MONAI's medical imaging framework for robust pre-processing pipelines including intensity normalisation per modality, random spatial augmentations, and sliding window inference.

Key Features & Highlights

• Vision Mamba encoder-decoder backbone: replaces quadratic self-attention with linear-complexity selective SSM layers, enabling full-resolution 3D MRI volumes (240×240×155) to be processed on a single GPU without patch-tiling artefacts.
• Attention Gates at skip connections: learned spatial weighting mechanisms that suppress activations from irrelevant healthy tissue and amplify tumour-boundary signals before decoder feature concatenation, recovering spatial precision without reintroducing quadratic cost.
• MONAI-based training pipeline: standardised intensity normalisation per MRI modality, random flipping, rotation, and elastic augmentation, combined with Dice + cross-entropy compound loss and sliding window inference for evaluation on full volumes.
• Memory-efficient training footprint: the linear SSM complexity allows significantly larger effective batch sizes and higher input resolutions compared to 3D ViT baselines of equivalent parameter count, enabling faster ablation cycles during research.

Results & Impact

AG-VIM achieved competitive Dice Similarity Coefficients across all three BraTS tumour sub-regions while operating within the memory budget of a single consumer-grade GPU — a configuration that comparable 3D Vision Transformer baselines cannot match without patch-based approximations. The attention gate ablation confirmed statistically significant improvements in boundary-region Dice scores versus a plain Mamba backbone without gating, validating the core architectural hypothesis. The project produced a clean, reproducible PyTorch codebase on GitHub and deepened applied expertise in state-space models, medical image pre-processing with MONAI, and the practical trade-offs between architectural expressivity and computational cost in volumetric deep learning.