Vector Quantization: A Case Study on Robust Medical Semantic Segmentation

Abstract

Pixel-wise medical image segmentation is essential for accurate disease diagnosis and analysis. Although these models exhibit high accuracy and impressive performance, they often face issues of trustworthiness due to their sensitivity to minor perturbations in the original images, which can cause significant deviations in segmentation tasks. To address these challenges, we focus on enhancing both the accuracy and the robustness of these systems against adversarial perturbations. We introduce a robust projection-head that utilizes a non-parametric method, Vector Quantization Momentum (VQM), to increase the hypothesis margin in the input space and reduce disturbances in the feature space. Additionally, we propose a new loss function, Anchor-based Vector Quantization (AVQ), designed to achieve denser pixel-prototype embedding. Our evaluation, using two white-box adversarial attack algorithms across two datasets, demonstrates a defense improvement of up to 14.46%.

Keywords: robustness, medical segmentation, dversarial attack, momentum