← 返回 2026-01-16

VQ-Seg:基于向量量化令牌扰动的半监督医学图像分割 VQ-Seg: Vector-Quantized Token Perturbation for Semi-Supervised Medical Image Segmentation

Sicheng Yang, Zhaohu Xing, Lei Zhu 📅 2026-01-15 👍 4 2026-07-13 08:35
一致性学习 医学图像分割 半监督学习 向量量化 基础模型对齐

用向量量化离散空间中的码本索引打乱替代Dropout扰动,实现半监督医学图像分割

前置知识

半监督医学图像分割(Semi-Supervised Medical Image Segmentation)

半监督学习是一种机器学习范式,利用少量带标注数据和大量无标注数据进行模型训练。在医学图像分割场景中,获取高质量的像素级标注极其昂贵且耗时,需要专业放射科医生逐层标注CT或MRI扫描。因此,半监督方法通过引入无标注数据来扩充训练集,利用一致性正则化(consistency regularization)等技术强制模型在不同扰动下产生一致的预测,从而提升泛化能力。

本文的核心目标就是在半监督设置下提升医学图像分割性能,理解这一背景是理解全文方法动机的前提。

一致性学习(Consistency Learning)

一致性学习是半监督学习的核心范式之一,其基本假设是:对输入或中间特征施加微小扰动后,模型的预测不应发生显著变化(平滑性假设)。常用实现方式包括对同一输入施加不同增强、Dropout扰动特征、或使用教师-学生(Teacher-Student)框架。在医学图像分割中,一致性学习通常通过Mean Teacher架构实现:教师网络通过指数移动平均(EMA)更新参数,对学生网络施加扰动后的输出产生伪标签,指导学生学习。

VQ-Seg正是在一致性学习框架下提出新的扰动策略,理解一致性学习的工作机制是理解QPM模块设计动机的关键。

向量量化(Vector Quantization, VQ)

向量量化是一种将连续特征空间离散化的技术,最早在VQ-VAE中被引入深度学习。其核心是维护一个可学习的码本(codebook)$\mathcal{C} = \{c_1, c_2, ..., c_K\}$,包含 $K$ 个码向量(codewords)。对于编码器输出的连续特征 $z$,VQ通过最近邻匹配找到最接近的码向量:$i = \arg\min_{j} \|z - c_j\|$,用该码向量替代原始连续特征。这一离散化过程通常使用直通估计器(Straight-Through Estimator, STE)来实现梯度回传。

VQ-Seg的核心创新就是在VQ离散空间中设计扰动策略,因此理解VQ的基本原理——码本结构、量化过程、STE梯度估计——是读懂QPM模块的必要前提。

Dropout正则化

Dropout是一种经典的正则化技术,在训练过程中以概率 $p$ 随机将神经元输出置零。在一致性学习中,Dropout被用作特征级扰动:对同一输入经过两次前向传播(一次Dropout开、一次Dropout关),要求两次输出一致。Monte Carlo Dropout(MC-Dropout)进一步将此用于不确定性估计。然而,Dropout率 $p$ 是一个敏感超参数——过低则扰动不足、正则化无效,过高则导致模型坍塌、性能急剧下降。

论文正是针对Dropout的这一根本性缺陷(超参数敏感性)提出替代方案,理解Dropout的工作原理和局限性才能理解VQ-Seg的创新价值。

基础模型(Foundation Model)与DINOv2

基础模型是指在大规模数据上预训练、具有强大表征能力的通用模型,可迁移到各种下游任务。DINOv2是Meta提出的视觉基础模型,采用自蒸馏(self-distillation)架构在大规模图像数据上训练,无需标注即可学习到丰富的语义特征和空间结构信息。DINOv2的特征具有很强的语义一致性,即使在自然图像上训练,也能很好地迁移到医学图像领域。在本文中,DINOv2作为外部语义先验,用于指导和对齐量化后的离散特征。

VQ-Seg的PFA模块依赖DINOv2提供语义监督信号,理解DINOv2的特征特性是理解语义对齐策略的基础。

KL散度(Kullback-Leibler Divergence)

KL散度是衡量两个概率分布之间差异的非对称度量,定义为 $D_{KL}(P \| Q) = \sum_j P(c_j) \log \frac{P(c_j)}{Q(c_j)}$。在分布鲁棒优化(DRO)中,KL散度用于定义不确定性集合的半径,衡量扰动后分布相对于原始分布的偏离程度。在本文中,作者将KL散度作为理论分析工具,比较Dropout和QPM两种扰动策略导致的分布偏移程度,从而论证QPM的优越性。

论文的理论贡献部分使用KL散度分析扰动稳定性,理解这一概念有助于把握QPM相对于Dropout的理论优势。

研究动机

在半监督医学图像分割中,基于Dropout的特征扰动一致性学习是一种广泛采用的策略,但存在一个根本性的超参数敏感问题。Dropout率 $p$ 的选择极为关键:实验表明,当 $p$ 较低(如 $p=0.3$ 或 $p=0.5$)时,扰动效果微乎其微,对分割性能几乎没有提升;而当 $p$ 较高(如 $p \geq 0.7$)时,Dice和Jaccard指标急剧下降,HD95和ASD显著增加,分割结果在视觉上几乎不可用(如Fig. 1所示,$p=0.9$ 时预测完全崩溃)。理论上,作者通过KL散度分析揭示了这一现象的本质原因:Dropout诱导的后验分布 $Q$ 与先验分布 $P$ 之间的KL散度 $D_{KL}(P \| Q) \approx \frac{1}{2}\left(\frac{p}{1-p} + \log(1-p)\right)$ 随 $p$ 增大而急剧上升,导致过度正则化。这意味着Dropout缺乏一个「甜蜜点」——既提供足够正则化又避免模型坍塌的稳定区间极为狭窄,严重限制了其实用性。

本文的目标是本文的核心目标是设计一种替代Dropout的特征扰动策略,使其能够在离散特征空间中提供可控、结构化且稳定的正则化效果,从而摆脱对敏感超参数的依赖。具体而言,作者希望:(1)通过向量量化将连续特征空间离散化,利用码本(codebook)的结构特性设计扰动机制;(2)扰动强度由单一控制项 $\epsilon$ 直接调控,且扰动分布始终有界、数值稳定;(3)在保持扰动有效性的同时,通过双分支架构和基础模型对齐弥补量化带来的信息损失;(4)在自建的大规模肺癌数据集(828例CT扫描)和公开的ACDC数据集上验证方法的有效性。

与已有工作不同的是,现有半监督医学图像分割方法在特征扰动设计上存在一个根本性缺口:扰动空间是连续的、随机的,缺乏结构约束。Dropout在连续特征空间中随机置零,扰动模式完全随机、不可解释,且超参数敏感。VQ-Seg的独特切入角度是将扰动从连续空间转移到离散的VQ空间——利用码本中码向量之间的学习距离定义扰动转移概率,使得扰动是有结构的、可解释的。具体而言,QPM模块通过 $\pi(j|i) = \frac{\epsilon \exp(-d(c_i, c_j))}{\sum_{k \neq i} \exp(-d(c_i, c_k))}$ 定义从码字 $c_i$ 到 $c_j$ 的转移概率,扰动方向由码本学习到的语义距离决定——语义相近的码字之间转移概率更高。这是首次将向量量化引入半监督分割的扰动设计,开创了一个全新的技术路线。

核心方法

VQ-Seg的整体思路可以用一个直觉来概括:与其在连续特征空间中随机「挖洞」(Dropout),不如先将特征离散化为有限的码本索引,然后在这些离散索引之间进行有结构的「洗牌」(QPM)。技术路线如下:输入医学图像首先通过编码器提取连续特征,然后通过VQ模块量化为离散的码本索引。在此离散空间中,QPM模块基于码字间的距离定义概率转移矩阵,以概率 $\epsilon$ 将每个码字替换为其他码字(距离越近概率越高),从而引入结构化扰动。为了弥补量化带来的视觉信息损失,框架采用双分支架构:共享的Post-VQ特征同时输入图像重建解码器和分割解码器,通过重建任务作为自监督信号保留结构信息。此外,Post-VQ Feature Adapter(PFA)模块利用冻结的DINOv2基础模型作为语义先验,通过对比学习将量化特征对齐到语义空间,弥补量化过程中丢失的高级语义信息。整体采用Mean Teacher框架,教师网络通过EMA更新,对学生网络的扰动输出生成伪标签进行一致性学习。

VQ-Seg的核心创新在于将一致性学习中的特征扰动从连续空间转移到离散的向量量化空间,这一转换带来了根本性的方法论差异。与Dropout的本质区别体现在三个方面:第一,扰动机制完全不同——Dropout以概率 $p$ 随机置零神经元输出,扰动模式是无结构的随机噪声;而QPM基于码字间的学习距离 $d(c_i, c_j)$ 定义转移概率 $\pi(j|i)$,扰动是在离散语义空间中有方向的迁移。第二,可控性截然不同——Dropout的效果高度依赖 $p$ 的选择,微小变化可能导致性能剧烈波动;而QPM的扰动强度由单一控制项 $\epsilon$ 线性调控,实验表明 $\epsilon \in [0.3, 0.7]$ 范围内性能稳定提升。第三,理论保证不同——作者证明QPM的扰动分布 $Q(c_j|\epsilon)$ 始终有界且数值稳定(详见Appendix B),而Dropout的KL散度在 $p \to 1$ 时趋向无穷。这种从连续随机扰动到离散结构扰动的范式转换,是本文最本质的贡献。

方法步骤详情

VQ-Seg的方法步骤如下:(1)编码阶段:输入图像 $x$ 通过编码器 $f_{enc}$ 提取连续特征 $z = f_{enc}(x)$。(2)向量量化:将连续特征 $z$ 通过最近邻匹配量化到码本 $\mathcal{C} = \{c_1, ..., c_K\}$ 中的码字,得到量化索引 $i = \arg\min_j \|z - c_j\|$,使用直通估计器(STE)回传梯度。(3)QPM扰动:对量化后的码字索引应用扰动策略——以概率 $1-\epsilon$ 保留原始码字,以概率 $\epsilon$ 按照距离权重 $\exp(-d(c_i, c_j))$ 替换为其他码字,得到扰动后的量化表示 $q'(z)$。(4)双分支解码:扰动后的量化特征同时输入图像重建解码器 $D_i$ 和分割解码器 $D_s$,分别产生重建图像 $\hat{x} = D_i(q'(z))$ 和分割预测 $\hat{y} = D_s(q'(z))$。(5)PFA语义对齐:通过Post-VQ Feature Adapter将量化特征调整尺寸后,与DINOv2提取的语义特征进行patch-wise对比学习,损失为 $\mathcal{L}_{align} = -\frac{1}{HW}\sum_i \log \frac{\exp(\text{sim}(f_i^{pfa}, f_i^{fm})/\tau)}{\sum_j \exp(\text{sim}(f_i^{pfa}, f_j^{fm})/\tau)}$。(6)教师-学生一致性学习:教师网络(EMA更新)对无标注数据生成伪标签,指导学生网络在QPM扰动后的量化特征上学习一致的分割预测。(7)联合优化:总损失 $\mathcal{L} = \mathcal{L}_{db} + \lambda_a \mathcal{L}_{align}$,其中 $\mathcal{L}_{db}$ 包含重建损失(L1)和分割损失(Cross-Entropy),$\lambda_a$ 平衡两项贡献。

技术新颖性

VQ-Seg在技术层面具有多方面的新颖性。首先,这是首个将向量量化引入半监督医学图像分割扰动设计的工作,开辟了「离散空间扰动」这一全新技术方向。QPM模块的设计尤为精巧:扰动转移概率 $\pi(j|i) = \frac{\epsilon \exp(-d(c_i, c_j))}{\sum_{k \neq i} \exp(-d(c_i, c_k))}$ 将码本的拓扑结构(码字间距离)转化为扰动方向的先验,使得扰动不是均匀随机的,而是倾向于向语义相近的码字迁移——这比Dropout的无差别随机置零更具物理意义。其次,双分支架构的设计巧妙地将信息保留(重建)和任务学习(分割)统一在同一离散空间中,重建任务为VQ编码器提供了自监督信号,鼓励码本学习到更具结构信息的表示。第三,PFA模块引入基础模型对齐的思路新颖——将冻结的DINOv2作为「外部语义锚点」,通过对比学习拉近量化特征与语义特征的距离,这是一种在离散空间中注入连续语义先验的有效方式。理论上,作者首次从KL散度角度分析了Dropout扰动的不稳定性根源,并证明了QPM扰动分布的有界性,为方法设计提供了坚实的理论支撑。

VQ-Seg框架整体架构图
Figure 2: VQ-Seg框架整体架构图
QPM量化扰动机制的具体示例
Figure 3: QPM量化扰动机制的具体示例
Post-VQ Feature Adapter(PFA)架构
Figure 4: Post-VQ Feature Adapter(PFA)架构

实验结果

VQ-Seg在多个数据集和标注比例设置下均取得了最优性能。在自建的Lung Cancer(LC)数据集(828例CT扫描)上:5%标注设置下,VQ-Seg达到Dice 0.6643、Jaccard 0.5257,分别超过次优方法Unimatch(Dice 0.6493、Jaccard 0.5071)1.50%和1.86%;同时HD95为12.2525(优于ABD的12.5608),ASD为4.2276(优于MCNet的4.9231)。10%标注设置下优势更加明显:Dice达0.7852(超过MCNet的0.7555达2.97%),Jaccard达0.6731(超过MCNet的0.6414达3.17%),HD95和ASD均为最优。消融实验(Table 2)验证了各模块的协同效应:基线模型(VQ-embedded Unimatch)Dice为0.7443,加入QPM后提升至0.7701,再加入双分支架构提升至0.7784,加入PFA单独使用达0.7761,三者联合达最优0.7852。超参数实验(Table 3)表明扰动强度 $\epsilon=0.7$ 时性能最优,$\lambda_a=5$、$\lambda_u=1$ 为最佳损失权重。基础模型消融(Table 4)显示DINOv2在5%和10%标注下均优于CLIP、BiomedCLIP、MAE和Rad-DINO,包括在医学数据上预训练的模型。码本大小实验(Table 5)表明K=16384为最优,利用率达98%,过大码本(65536)因利用率下降(92%)导致性能下降。

LC数据集上的定量比较(5%和10%标注设置)
Table 1: LC数据集上的定量比较(5%和10%标注设置)
不同标注比例下的性能比较
Table 6: 不同标注比例下的性能比较
LC数据集上的可视化比较结果
Figure 5: LC数据集上的可视化比较结果
查看结构化数据
任务指标本文基线提升
半监督肺癌CT分割(5%标注) Dice 0.6643 0.6493 (Unimatch) +1.50%
半监督肺癌CT分割(5%标注) Jaccard 0.5257 0.5071 (Unimatch) +1.86%
半监督肺癌CT分割(5%标注) HD95 12.2525 12.5608 (ABD) -0.3083(越低越好)
半监督肺癌CT分割(5%标注) ASD 4.2276 4.9231 (MCNet) -0.6955(越低越好)
半监督肺癌CT分割(10%标注) Dice 0.7852 0.7555 (MCNet) +2.97%
半监督肺癌CT分割(10%标注) Jaccard 0.6731 0.6414 (MCNet) +3.17%
半监督肺癌CT分割(10%标注) HD95 11.6179 11.6724 (UA-MT) -0.0545(越低越好)
半监督肺癌CT分割(10%标注) ASD 4.2094 4.3660 (ARCO) -0.1566(越低越好)

局限与改进

论文存在多方面局限性。作者明确承认的局限包括:(1)QPM的扰动仅在离散VQ空间中操作,难以直接扩展到现有半监督框架中普遍使用的连续特征表示,限制了方法的通用适用性;(2)引入基础模型(DINOv2)虽然带来了丰富的语义先验,但也引入了额外的计算开销,增加了推理和训练的资源需求。从独立分析角度,还存在以下局限:(3)实验仅在两个数据集上验证(LC和ACDC),未在更广泛的医学图像分割基准(如BraTS脑肿瘤分割、ISIC皮肤病变分割等)上测试,泛化性证据不够充分;(4)LC数据集为自建数据集,虽然规模达828例,但缺乏独立第三方验证,可能存在数据选择偏差;(5)ACDC数据集仅有100例,规模较小,且为心脏MRI而非CT,与LC数据集的模态差异较大,难以判断方法在同模态不同器官上的表现;(6)消融实验仅在10%标注设置下进行,未验证各模块在更低标注比例下的贡献差异;(7)论文未报告训练时间和计算成本的详细对比,无法评估方法的实际部署可行性。

独立分析的弱点

VQ-Seg存在几个值得关注的弱点。首先,QPM模块的扰动机制依赖于码本中码字间距离的准确性,但码本在训练初期是随机初始化的,此时距离度量不可靠,可能导致早期训练阶段的扰动方向不合理。改进方向:可以在训练初期使用较小的扰动强度 $\epsilon$ 并逐步增大(curriculum learning策略),或引入预热阶段先稳定码本再启用QPM。其次,双分支架构同时优化重建和分割两个任务,虽然通过 $\lambda_u$ 控制权重,但两个任务的梯度方向可能存在冲突,尤其在训练后期重建任务可能主导优化方向。改进方向:可以设计自适应权重调节机制或梯度投影策略来解耦任务冲突。第三,PFA模块使用DINOv2作为语义先验,但DINOv2在自然图像上训练,对医学图像的域差距(domain gap)可能影响对齐质量,尤其对于形态学独特的病变(如中央型肺癌)。改进方向:可以探索在医学图像上微调DINOv2或使用混合域适应策略。第四,扰动强度 $\epsilon$ 仍然是一个超参数,虽然比Dropout率更稳定,但最优值(0.7)仍需通过实验确定,自动化选择机制值得探索。

未来方向

论文作者提出两个明确的未来方向:(1)开发连续空间中的可控扰动机制,使方法能兼容现有半监督框架中广泛使用的连续特征表示,扩大适用范围;(2)探索更高效的基础模型集成方式,降低引入DINOv2带来的计算开销。基于VQ-Seg的成果,还可以延伸以下研究方向:(3)将QPM扩展到3D医学图像分割(如体积CT或MRI),探索在时空维度上的码字扰动策略;(4)研究自适应码本大小——根据数据复杂度动态调整码本容量,在表征能力和计算效率间取得平衡;(5)探索多尺度VQ——在编码器的不同层级使用不同分辨率的码本,捕捉从局部纹理到全局结构的多层次语义;(6)将QPM思想扩展到其他半监督任务,如目标检测、图像分类和语义分割;(7)结合主动学习策略,智能选择最有价值的标注样本,在标注预算有限时最大化性能提升。

复现评估

VQ-Seg的复现条件较为理想。代码已开源(https://github.com/script-Yang/VQ-Seg),为复现提供了直接基础。ACDC数据集为公开基准,可免费下载;但LC数据集为自建数据集(828例中心型肺癌CT扫描),论文未明确说明是否公开发布,这可能成为复现的一个障碍。算力方面,论文使用4块NVIDIA GeForce RTX 4090 GPU进行训练,训练100,000次迭代,batch size为8个标注样本+4个无标注样本,输入尺寸224×224——这对大多数实验室而言是可承受的计算资源。实现基于PyTorch 2.4.1,采用标准的AdamW优化器和多项式学习率调度。总体而言,代码开源+公开数据集+合理算力需求的组合使得复现难度较低,但LC数据集的可获取性是关键瓶颈。