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

AI summary

Overview

Research area: Semi-supervised medical image segmentation, specifically consistency learning with feature-level perturbation, combined with vector quantization (VQ) and visual foundation models.

Technical level: Intermediate. Readers benefit from familiarity with semi-supervised learning, teacher-student consistency training, and vector-quantized autoencoders.

Scope: The paper proposes VQ-Seg, a framework that replaces dropout-based feature perturbation with a controllable, codebook-based perturbation inside a discrete VQ space, and evaluates it on a newly collected Lung Cancer CT dataset and the public ACDC dataset. The provided content is truncated — the ACDC results appear in Appendix C (Table 7), which is not included in the text supplied here.

What This Paper Is About

Semi-supervised medical image segmentation uses a small set of labeled scans plus many unlabeled scans, and a common trick is to add random noise (dropout) to internal features and force the model to make consistent predictions anyway. The problem is that dropout is governed by a single dropout rate hyperparameter that must be tuned by hand — too low and the noise does nothing, too high and the model collapses. This paper's goal is to replace that fragile dropout knob with a structured perturbation applied to discrete vector-quantized codebook indices, so that regularization strength is controllable and stable.

Key Contributions

1. A new large-scale dataset. The authors collected a multi-center Lung Cancer (LC) dataset of 828 chest CT scans annotated for central-type lung carcinoma, with one segmentation target per volume and the dominant lesion annotated precisely for each case.

2. Quantized Perturbation Module (QPM). The paper states it is the first approach to employ vector quantization to discretize the feature space for this purpose. QPM perturbs discrete representations by shuffling the spatial locations of codebook indices, using distances between codewords to define the transition probabilities, which the authors describe as more interpretable and stable than dropout.

3. Dual-branch architecture with a shared Post-VQ space. The post-quantization feature space is used simultaneously for image reconstruction and segmentation, so reconstruction acts as a self-supervisory signal that encourages the VQ encoder to preserve structural information.

4. Post-VQ Feature Adapter (PFA) with foundation model guidance. A frozen foundation model (DINOv2) supplies an external semantic prior, and the PFA — a resize operation plus a 1×1 convolution — transforms quantized codebook embeddings into a semantically aligned space via patch-wise contrastive learning.

Main Findings

• Dropout rate is a hard-to-tune failure point. In fully supervised experiments on the LC dataset, low dropout rates (DR = 0.3, DR = 0.5) had negligible impact on segmentation performance, while high rates (DR ≥ 0.7) caused Dice and Jaccard to drop sharply and HD95 and ASD to increase substantially. DR = 0.9 produced unusable predictions.

• Theoretical argument against dropout. Using a moment-matching approximation, the authors derive the KL divergence between prior and dropout-perturbed posterior as approximately ½(p/(1−p) + log(1−p)) for dropout rate p in (0,1), which grows sharply as p increases. In contrast, the QPM perturbed distribution Q(c_j|ε) is described as always well-defined and bounded, ensuring numerical stability.

• State-of-the-art results on the LC dataset at 5% labeled data. VQ-Seg reached Dice 0.6643, Jaccard 0.5257, HD95 12.2525, and ASD 4.2276. The paper reports this outperforms second-best Unimatch (Dice 0.6493, Jaccard 0.5071) by 1.5% and 1.86%, improves on second-best ABD for HD95 (12.5608) by 0.3083, and on second-best MCNet for ASD (4.9231) by 0.6955.

• State-of-the-art results on the LC dataset at 10% labeled data. VQ-Seg reached Dice 0.7852, Jaccard 0.6731, HD95 11.6179, and ASD 4.2094, surpassing second-best MCNet (Dice 0.7555, Jaccard 0.6414) by 2.97% and 3.17%, improving on second-best UA-MT for HD95 (11.6724) by 0.0545 and on second-best ARCO for ASD (4.3660) by 0.1566.

• Every module contributes. On the LC dataset with 10% labeled data, the VQ-embedded Unimatch baseline scored Dice 0.7443 / Jaccard 0.6238 / HD95 14.2153 / ASD 5.2301. Adding QPM raised Dice to 0.7701 (Jaccard 0.6559, HD95 13.0246, ASD 4.9378); adding the dual-branch architecture further raised Dice to 0.7784 (Jaccard 0.6620, HD95 12.4728, ASD 4.6013); using PFA yielded Dice 0.7761 (Jaccard 0.6597, HD95 12.7381, ASD 4.7005). All three combined gave the best results across all metrics: Dice 0.7852, Jaccard 0.6731, HD95 11.6179, ASD 4.2094.

• Perturbation strength peaks at ε = 0.7. Dice moved from 0.7741 at ε = 0.3, to 0.7803 at ε = 0.5, to 0.7852 at ε = 0.7, then fell to 0.7418 at ε = 0.9.

• Loss weights. Best overall results came from λ_a = 5 and λ_u = 1. For λ_a = 1 the Dice was 0.7720 (HD95 11.5080, ASD 4.2672); for λ_a = 10 it was 0.7765 (HD95 11.9235, ASD 4.1926). For λ_u = 5 Dice was 0.7670 (ASD 4.1652); for λ_u = 10 Dice was 0.7843 (ASD 4.1013).

• DINOv2 beats medical-domain foundation models as the semantic prior. Under both regimes DINOv2 scored highest (0.6643 at 5%, 0.7852 at 10%), ahead of CLIP (0.6421 / 0.7483), BiomedCLIP (0.6507 / 0.7629), MAE (0.6386 / 0.7541), and Rad-DINO (0.6535 / 0.7793).

• Codebook size of 16,384 is the sweet spot. Dice rose from 0.6531 at size 1,024 through 0.6582 at 2,048, 0.6627 at 4,096, to 0.6643 at 16,384 (5% labeled), then fell to 0.6595 at 32,768 and 0.6415 at 65,536. Codebook utilization fell from 100% at sizes 1,024 and 2,048 to 98% at 16,384 and 92% at 65,536.

• Gains persist as labels increase. VQ-Seg scored Dice 0.6643 (5%), 0.7852 (10%), 0.8100 (20%), 0.8507 (50%), and 0.9102 (100%), compared with Unimatch at 0.6493 / 0.7511 / 0.7855 / 0.8279 / 0.8871, ABD at 0.6414 / 0.7468 / 0.7780 / 0.8235 / 0.8824, MCNet at 0.6378 / 0.7555 / 0.7812 / 0.8203 / 0.8751, and UNet-S at 0.4343 / 0.6490 / 0.7205 / 0.7880 / 0.8345.

Methodology in Plain English

The framework starts with an encoder that turns an image into continuous features. Those features are then snapped to the nearest entry in a learned codebook of discrete codewords, where the entry index is chosen by minimum distance to the continuous feature.

Perturbation without a dropout knob. Instead of zeroing out features at random, VQ-Seg swaps a codeword for another codeword. The swap follows a probability rule: with probability 1 − ε the codeword stays the same, and with probability ε it is replaced by a different codeword, weighted by an exponential of the negative distance between the two codewords and normalized by a factor Z_i. Intuitively, nearby codewords are more likely substitutes than distant ones. A worked example in the paper uses K = 4 and ε = 0.7, where the transition from codeword c₁ to c₂ has 49% probability. Because the codebook is uniform at 1/K and the resulting perturbed distribution is well-behaved, the authors argue the perturbation radius is bounded and controllable through a single ε, unlike dropout's rate.

Two decoders on one shared discretized space. The student and teacher networks both feed the quantized features into an image decoder and a segmentation decoder. For labeled data, the loss combines an L1 reconstruction term and a cross-entropy segmentation term. For unlabeled data, the teacher produces a segmentation map, an argmax over it becomes a pseudo-label, and the student is trained on the QPM-perturbed quantized representation of the same unlabeled input against that pseudo-label. The total dual-branch loss adds the labeled loss to the unlabeled loss scaled by λ_u.

Aligning the discretized features with a foundation model. Because quantization can discard fine detail and shift semantics, the PFA resizes the VQ features and applies a 1×1 convolution so their spatial resolution and channel count match a frozen DINOv2 feature map. A patch-wise contrastive loss using cosine similarity and a temperature τ pulls each adapted VQ patch toward its matching foundation model patch and away from other patches, providing localized semantic supervision. The overall objective is the dual-branch loss plus λ_a times the alignment loss.

Training setup. Implemented in PyTorch 2.4.1. Slices resized to 224×224. Each iteration uses a mini-batch of 8 labeled and 4 unlabeled samples. 100,000 iterations with AdamW and a polynomial learning rate schedule; initial learning rate 5×10⁻⁶ for the backbone, scaled by a factor of 40 for task-specific heads, decayed by (1 − t/T)^0.9. Augmentations include random rotation, color jittering, and CutMix-based strong perturbations for unlabeled data. The Straight-Through Estimator is used for VQ training, the codebook size is 16,384, and an entropy-based regularization loss accelerates codebook learning. The teacher is updated by exponential moving average of the student with decay α = 0.996. Training ran on four NVIDIA GeForce RTX 4090 GPUs. The LC dataset uses a 70–10–20 training/validation/testing split, and all experiments are on 2D slices. VQ-Seg uses the same encoder and decoder architecture as Unimatch, with the quantization and alignment process added after the encoder output.

Why This Matters

Impact on research. The paper argues that a widely used regularization mechanism — feature dropout for consistency learning — has an inherent stability problem that can be shown both empirically and via a KL divergence argument, and that moving the perturbation into a discrete latent space offers a principled alternative. It also demonstrates that a foundation model trained on natural images (DINOv2) transfers better than medical-specific alternatives (BiomedCLIP, Rad-DINO) for this alignment task, which is a useful data point for how medical imaging pipelines should source semantic priors.

Real-world applications:

• Disease diagnosis, where segmentation is a fundamental upstream step.
• Anatomical structure delineation of organs and tissues.
• Lesion localization, including the central-type lung carcinoma targeted by the new LC dataset.
• Surgical planning, where accurate boundaries inform the procedure.

Industry relevance. The core practical value is reducing annotation cost: VQ-Seg leads the comparison at 5%, 10%, and 20% labeled data on the LC dataset, and the paper reports the largest relative gains in low-label regimes. That directly affects how much expert radiologist time a clinical AI product needs. The requirement for four RTX 4090 GPUs plus the additional computational overhead of the foundation model are real deployment considerations, and the paper acknowledges the overhead explicitly in its limitations.

Future Directions

1. Extending controllable perturbation to continuous feature spaces. The paper states the current perturbation operates solely in the discrete VQ space, making it difficult to extend to the continuous feature representations commonly used in existing semi-supervised frameworks.

2. Reducing foundation model overhead. The authors identify the added computation from the frozen foundation model as a limitation and list more efficient foundation model integration as future work.

3. Exploring the untested regions of the design space. Codebook utilization falls and performance degrades at sizes 32,768 and 65,536, and Dice falls at perturbation strength ε = 0.9, so the mechanisms governing codebook redundancy and over-regularization are not fully explained.

4. Confirming generalization beyond the two evaluated datasets. The paper reports a similar trend on ACDC in Appendix C (Table 7), but those specific ACDC numbers are not present in the provided content, and the ACDC collection is described only as 100 short-axis cine-MRI scans acquired on 3T and 1.5T scanners.

Target Audience

Researchers and graduate students working on semi-supervised medical image segmentation, consistency regularization, or vector-quantized representation learning. Also relevant to applied machine learning engineers in medical imaging who need to reduce annotation requirements, and to anyone evaluating which visual foundation model to use as a semantic prior for domain-specific segmentation. Readers without background in teacher-student consistency training or codebook-based quantization will find the theoretical derivations and architectural details harder to follow.

Code: https://github.com/script-Yang/VQ-Seg

Funding acknowledged: Guangdong Science and Technology Department (2024ZDZX2004) and the Guangzhou Industrial Information and Intelligent Key Laboratory Project (No. 2024A03J0628).

Authors’ 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.