VQ-Seg: Vector-Quantized Token Perturbation for Semi-Supervised Medical Image Segmentation

TL;DR

This work tackles the instability of dropout-based perturbations in semi-supervised medical image segmentation by introducing VQ-Seg, which performs controlled perturbations in a discrete vector-quantized (VQ) space via a Quantized Perturbation Module (QPM). A dual-branch architecture shares a Post-VQ feature space for simultaneous image reconstruction and segmentation, while a Post-VQ Feature Adapter (PFA) aligns these quantized features with semantic priors from a frozen foundation model (DINOv2) through patch-wise contrastive learning. The method is validated on a large Lung Cancer CT dataset (828 scans) and the ACDC MRI dataset, achieving state-of-the-art performance under low-label regimes and demonstrating robustness through extensive ablations on codebook size, perturbation strength, and foundation-model influence. The approach reduces information loss from quantization and provides a stable, interpretable regularization mechanism with practical impact for semi-supervised medical image analysis, albeit with additional computational overhead from FM integration and discrete-space constraints.

Abstract

Consistency learning with feature perturbation is a widely used strategy in semi-supervised medical image segmentation. However, many existing perturbation methods rely on dropout, and thus require a careful manual tuning of the dropout rate, which is a sensitive hyperparameter and often difficult to optimize and may lead to suboptimal regularization. To overcome this limitation, we propose VQ-Seg, the first approach to employ vector quantization (VQ) to discretize the feature space and introduce a novel and controllable Quantized Perturbation Module (QPM) that replaces dropout. Our QPM perturbs discrete representations by shuffling the spatial locations of codebook indices, enabling effective and controllable regularization. To mitigate potential information loss caused by quantization, we design a dual-branch architecture where the post-quantization feature space is shared by both image reconstruction and segmentation tasks. Moreover, we introduce a Post-VQ Feature Adapter (PFA) to incorporate guidance from a foundation model (FM), supplementing the high-level semantic information lost during quantization. Furthermore, we collect a large-scale Lung Cancer (LC) dataset comprising 828 CT scans annotated for central-type lung carcinoma. Extensive experiments on the LC dataset and other public benchmarks demonstrate the effectiveness of our method, which outperforms state-of-the-art approaches. Code available at: https://github.com/script-Yang/VQ-Seg.

Figures (7)

• Figure 1: Effect of dropout rate on segmentation performance in a fully supervised setting on the LC dataset. Low dropout rates show negligible impact, whereas a high dropout rate (DR $\geq$ 0.7) severely degrades both quantitative metrics and visual outputs. Notably, DR = 0.9 leads to unusable predictions, highlighting the challenge of selecting an optimal dropout rate.
• Figure 2: Overview of the VQ-Seg framework. The input image $x$ is encoded into continuous features $z$, which are then quantized into a discrete codebook space via vector quantization (VQ). Quantized Perturbation Module (QPM) introduces controllable perturbations for consistency learning. The dual-branch architecture jointly optimizes image reconstruction and segmentation using the shared Post-VQ features. Additionally, a Post-VQ Feature Adapter (PFA) aligns the quantized features with semantic embeddings from a foundation model (FM).
• Figure 3: A concrete example of the Quantized Perturbation Mechanism (QPM) with a codebook size of $K=4$ and a perturbation strength $\epsilon = 0.7$. It illustrates the probabilistic transitions from the original codeword $c_1$ (index 1) to itself and other codewords ($c_2, c_3, c_4$) with their respective probabilities $\pi(j|1)$, where the transition to $c_2$ (49%) exhibits the highest probability of replacement.
• Figure 4: Architecture of the Post-Quantization Feature Adapter (PFA) designed for aligning post-quantization features with a frozen Foundation Model (FM) via a patch-wise contrastive loss, $\mathcal{L}_{\text{align}}$. The PFA initially employs a resizing operation followed by a $1\times1$ convolution to match the spatial resolution and channel dimensionality of the FM features, thereby facilitating subsequent semantic alignment.
• Figure 5: Visual results on LC with 5% and 10% labeled data show that VQ-Seg consistently yields more accurate predictions of anatomical structures and boundaries than all other compared methods.