Attention-Guided Vector Quantized Variational Autoencoder for Brain Tumor Segmentation
摘要
Precise brain tumor segmentation is critical for effective treatment planning and radiotherapy. Existing methods rely on voxel-level supervision and often struggle to accurately delineate tumor boundaries, increasing potential surgical risks. We propose an Attention-Guided Vector Quantized Variational Autoencoder (AG-VQ-VAE)—a two-stage network specifically designed for boundary-focused tumor segmentation. Stage 1 comprises a VQ-VAE which learns a compact, discrete latent representation of segmentation masks. In stage 2, a conditional network extracts contextual features from MRI scans and aligns them with discrete mask embeddings to facilitate precise structural correspondence and improved segmentation fidelity. Additionally, we propose an attention scaling module to reinforce discriminative feature learning and a soft masking module to refine attention in uncertain tumor regions. Comprehensive evaluations on BraTS 2021 demonstrate that our AG-VQ-VAE sets a new benchmark, improving the HD95 metric by 4.83 mm (Whole Tumor), 2.14 mm (Tumor Core), and 2.39 mm (Enhancing Tumor), compared to state-of-the-art methods, while achieving a 0.23% improvement in Dice score for whole tumor. Furthermore, our qualitative results and ablation study demonstrate that feature-level supervision significantly enhances boundary delineation compared to voxel-level approaches. The code is available at https://github.com/danishali6421/AG-VQVAE-MICCAI .