PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sicheng Yang, Zhaohu Xing, Lei Zhu

arXiv:2601.10124v1cs.CV

TL;DR

Semi-supervised medical image segmentation needs stable feature perturbations because dropout-based consistency learning depends on sensitive rate selection. VQ-Seg discretizes features and replaces dropout with controllable codebook-index perturbations, while sharing post-VQ features across reconstruction and segmentation and adding foundation-model guidance. Experiments on LC and ACDC report state-of-the-art performance, while the authors identify discrete-space perturbation and foundation-model overhead as limitations.

  • Problem

    Dropout-based feature perturbation requires careful manual tuning of a sensitive dropout rate, making stable regularization difficult.

  • Method

    VQ-Seg uses QPM to shuffle codebook-index locations in discrete VQ space, with shared reconstruction-segmentation features and a foundation-model-guided PFA.

  • Results

    VQ-Seg achieves state-of-the-art performance on the LC and ACDC datasets under limited supervision.

  • Takeaways & Limitations

    Structured perturbation in quantized feature space provides the framework's controllable alternative to dropout for semi-supervised segmentation.

  • Takeaways & Limitations

    Perturbation currently operates only in discrete VQ space, and foundation-model integration adds computational overhead.

Abstract

from arXiv · show

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.

1 Introduction

Semi-supervised medical image segmentation uses consistency learning with feature perturbations, but dropout requires difficult rate selection that can either provide negligible regularization or severely degrade predictions. VQ-Seg addresses this with structured perturbation in a discrete VQ space, dual-branch reconstruction and segmentation, foundation-model guidance, and a new LC dataset.

  • Consistency learning uses feature-level dropout, but its effectiveness depends critically on selecting the dropout rate.
  • DR=0.3 and DR=0.5 have negligible impact, whereas DR ≥0.7 sharply reduces Dice and Jaccard and increases HD95 and ASD.High dropout also produces frequently unusable segmentation outputs.
  • The LC dataset contains 828 chest CT scans annotated for central-type lung carcinoma, and VQ-Seg achieves state-of-the-art performance across the LC and open-source datasets.
  • VQ-Seg introduces QPM to perturb encoded features by shuffling codebook-index locations in discrete VQ space.The module is designed as a structured and controllable alternative to traditional dropout.
  • A dual-branch architecture shares the post-quantized space between image reconstruction and segmentation, using reconstruction as a self-supervisory signal.This design aims to preserve essential visual information while improving learned representations.
  • A foundation-model-guided PFA aligns quantized codebook embeddings with semantic information through contrastive learning.The strategy is intended to mitigate loss of high-level semantic features after quantization.

2 Related work

Related work emphasizes semi-supervised segmentation under limited annotation settings, with pseudo-labeling and feature-level dropout among widely used strategies. Dropout-based consistency methods use random feature perturbations to encourage stable predictions and model uncertainty.

  • Semi-supervised learning addresses the expense and labor of obtaining large-scale, high-quality medical-image annotations.
  • Pseudo-labeling methods train an initial model on labeled data and generate pseudo-labels for semi-supervised learning.
  • Feature-level dropout is widely used for consistency learning in semi-supervised medical image segmentation.
  • Monte Carlo Dropout models uncertainty and enforces prediction consistency through random perturbations of intermediate features.These perturbations are used to encourage stable and robust predictions despite data uncertainty.

3 Methodology

VQ-Seg quantizes encoded features into a discrete codebook space, then applies structured perturbation, shared reconstruction-segmentation learning, and foundation-model semantic alignment. Its components target controllable regularization while mitigating information and semantic loss from quantization.

  • Framework overview: VQ-Seg encodes images into continuous features and quantizes them into a discrete codebook space for subsequent processing.The encoder output is projected into a learnable codebook by selecting nearest codewords.
  • Quantized Perturbation Module: QPM perturbs quantized representations by probabilistically replacing codewords according to learned codeword distances and a control term ϵ.The perturbed codeword is sampled from π(j | i), with transitions determined by codeword distances and perturbation strength.
  • Quantized Perturbation Module: QPM provides a bounded, numerically stable perturbation controlled by a single ϵ while remaining within the learned discrete latent space.The method describes these transitions as more interpretable and controllable than dropout-based stochastic perturbations.
  • Dual-branch architecture: A dual-branch architecture uses shared Post-VQ features for image reconstruction and segmentation to reduce information loss from quantization.The shared representation is optimized using reconstruction and segmentation objectives, with labeled and pseudo-labeled supervision.
  • Post-VQ Feature Adapter: PFA aligns quantized features with semantic embeddings from a frozen foundation model through resizing, channel matching, and patch-wise contrastive learning.This alignment is intended to mitigate high-level semantic loss and representation drift introduced by quantization.
  • Optimization objective: The total objective combines the dual-branch loss Ldb with the alignment loss Lalign, weighted by λa.The balance parameter λa controls the contribution of semantic alignment relative to the dual-branch objective.

4 Experiments

Experiments on the LC and ACDC datasets evaluate VQ-Seg under limited-label settings, compare its components and hyperparameters, and show strong quantitative and visual performance. Ablations examine QPM, dual-branch reconstruction, PFA, foundation models, codebook size, and labeled-data ratios.

  • Datasets: The LC dataset contains 828 chest CT scans of central-type lung carcinoma, with one precisely annotated dominant lesion per volume.
  • Experimental setup: Experiments use 5% and 10% labeled-data settings, with LC data split 70–10–20 for training, validation, and testing and all experiments conducted on 2D slices.
  • Quantitative comparisons: At 5% labeled LC data, VQ-Seg achieves the highest Dice (0.6643), Jaccard (0.5257), HD95 (12.2525), and ASD (4.2276).It surpasses the second-best methods by 1.5% and 1.86% on Dice and Jaccard, and improves HD95 and ASD by 0.3083 and 0.6955.
  • Quantitative comparisons: At 10% labeled LC data, VQ-Seg maintains the lead with Dice 0.7852, Jaccard 0.6731, HD95 11.6179, and ASD 4.2094.Results on ACDC exhibit a similar trend across datasets.
  • Visual comparisons: Visual comparisons show more consistent segmentation, clearer boundaries, and better preservation of cancer-region structure than other state-of-the-art techniques.
  • Ablation studies: Ablations show Dice increasing from 0.7443 to 0.7701 with QPM, then to 0.7784 with the dual-branch architecture, while PFA alone reaches 0.7761.The full model with QPM, dual-branch architecture, and PFA achieves the best performance across all metrics.
  • Ablation studies: DINOv2 outperforms alternative foundation models under both 5% and 10% labeled regimes, including models specialized for medical domains.
  • Ablation studies: Dice increases as codebook size grows from 1,024 to 16,384, but slightly degrades at 32,768 or 65,536 because of decreased code utilization and overfitting.

5 Conclusion and Limitations

VQ-Seg achieves state-of-the-art performance on the LC and ACDC datasets under limited supervision, while its perturbation remains confined to the discrete VQ space and foundation-model guidance adds computational overhead.

  • VQ-Seg achieves state-of-the-art performance on the LC and ACDC datasets, substantially improving segmentation accuracy under limited supervision.
  • The current perturbation operates solely in the discrete VQ space, making extension to continuous feature representations difficult.
  • Foundation-model guidance introduces additional computational overhead.

A KL Divergence Approximation under Dropout Perturbation

This section models dropout perturbation by treating features as Gaussian before dropout and as a mixture of zeroed and retained activations afterward.

  • The analysis approximates a feature activation h with a Gaussian prior distribution P(h).
  • Dropout with rate p zeroes the feature with probability p and retains it with probability 1 − p.
  • The resulting perturbed distribution is represented as a mixture containing a Dirac delta at zero.

A.2 Intractability of Exact KL Divergence

The exact KL divergence for the dropout-induced mixture is intractable because of the Dirac delta singularity, so the analysis uses moment matching to obtain an analytical approximation.

  • The KL divergence between the original distribution P(h) and dropout-induced distribution Q(h) is formulated but lacks a closed-form solution.
  • The intractability arises from the singular nature of δ(h) at h = 0.
  • A Gaussian Qapprox(h) matching the first two moments of the dropout-perturbed activations approximates Q(h).
  • The approximation yields a clean analytical KL-divergence expression interpreted as the dropout perturbation radius.
  • As p → 1, the approximated divergence grows rapidly and diverges, indicating increasingly unstable dropout perturbations and motivating QPM.

B Proof of Numerical Stability of QPM

The QPM perturbed codeword distribution is shown to be well-defined, positive, normalized, and numerically stable, including at the extreme perturbation value ϵ = 1.

  • The appendix establishes that Q(cj | ϵ) is well-defined and bounded, ensuring numerical stability of its KL divergence even when ϵ = 1.
  • A uniform prior distribution over codewords is assumed for the stability analysis.
  • The perturbation mechanism π(j | i) defines transitions between codewords, which determine the perturbed distribution Q(cj | ϵ).
  • Non-negativity and mild positivity conditions ensure Q(cj | ϵ) is nonnegative and positive for every codeword.
  • The analysis verifies that Q is a valid probability distribution through normalization.
  • The KL divergence remains numerically stable when Q(cj | ϵ) is positive and bounded away from zero.
  • With bounded pairwise codebook distances, transition terms admit finite lower and upper bounds.
  • At ϵ = 1, log(KQ(cj)) remains finite, supporting numerical stability for all valid ϵ ∈ [0, 1].

C Performance Analysis on the ACDC Dataset

On ACDC, VQ-Seg outperforms competing methods under both 5% and 10% labeled-data settings, with gains reported in Dice and Jaccard metrics.

  • VQ-Seg consistently outperforms leading baselines in Dice and Jaccard under both 5% and 10% labeled-data settings.The comparison includes Unimatch, ABD, and ARCO.
  • 0.9057 Dice under 5% labeled data exceeds Unimatch (0.8915), ABD (0.8874), and ARCO (0.8879), approaching nnUNet-F (0.9185).The result is reported as evidence of strong representation ability under limited supervision.
  • 0.9103 Dice and 0.8327 Jaccard under 10% labeled data are the best reported values, surpassing ABD (0.8992) and ARCO (0.9026).The passage characterizes these gains as consistent with robustness and scalability as labeled data increase.

D Statistical Analysis

Repeated trials on the 5% labeled LC dataset indicate that VQ-Seg has lower performance variance and statistically significant improvements over the compared baselines.

  • Ten runs with different random seeds compare MCNet, Unimatch, and VQ-Seg on the 5% labeled LC dataset.The averaged Dice scores and standard deviations are visualized in Fig. 6.
  • VQ-Seg has a noticeably smaller error bar than the competing methods, indicating lower performance variance across random seeds.The error bars represent standard deviations across repeated runs.
  • p < 0.05 in paired two-tailed t-tests indicates statistically significant differences between VQ-Seg and each baseline over ten trials.

E Codebook Evolution

The t-SNE visualization shows codebook activation expanding from sparse initial use to near-complete, more uniform utilization by convergence.

  • The t-SNE visualization tracks all codebook vectors across initial, intermediate, and converged training stages.Orange marks activated codewords and blue marks inactive codewords.
  • 3.4% of codewords are active initially, with activation increasing to 67% during training and 98% by convergence.The distribution also becomes more uniform as training proceeds.
  • The reported progression demonstrates increased codebook utilization and representation diversity through quantization.
Loading 2601.10124v1…