跳转至

Simple yet Effective Semi-supervised Knowledge Distillation from Vision-Language Models via Dual-Head Optimization

会议: NeurIPS 2026
arXiv: 2505.07675
代码: https://github.com/erjui/DHO
领域: 多模态 VLM
关键词: 知识蒸馏, 半监督学习, 视觉语言模型, 双头优化, 梯度冲突

一句话总结

针对传统单头蒸馏在少样本/低标注半监督场景下监督损失与蒸馏损失间的严重梯度冲突,本文提出双头优化(DHO),解耦监督分类头与蒸馏头并在推理阶段轻量线性插值,以极低开销实现全免重训调参且全面超越单头蒸馏及教师模型。

研究背景与动机

在大规模图文对上预训练的视觉语言模型(如 CLIP)展现出非凡的零样本与少样本泛化能力,而传统半监督学习(SSL)虽能利用海量未标注数据,却往往受限于有限标注而难以习得如此丰富的多模态表示。直觉上,借助知识蒸馏(KD)将教师 VLM 的通用表征迁移至紧凑的特定任务学生模型中,同时结合少量标注样本的有监督学习,是解决数据稀缺场景的理想途径。

然而,传统的单阶段单头蒸馏(SHO)在半监督设定下效果并不理想。根本矛盾在于:来自稀缺标注样本的有监督交叉熵损失与来自未标注数据上教师概率输出的蒸馏损失之间存在严重的梯度冲突。教师模型的蒸馏信号覆盖面广且置信度稳定,而少量标注样本的梯度方向则针对特定 hard label,二者在共享分类头及特征提取器深层反向传播时,梯度余弦相似度频繁跌破零值,导致监督信号被过度淹没或产生破坏性干扰,严重损害了紧凑学生模型的特征表征学习。

现有文献多采用繁琐的多阶段自监督预蒸馏后微调,或者在多任务中使用代价高昂的梯度投影手术(如 PCGrad),这不仅增加了显存与计算开销,也丧失了推断期快速平衡两者权重的灵活性。本文从解耦表征与分类决策的极简视角切入:核心 idea:通过解耦构建独立的有监督头与蒸馏头(Dual-Head Optimization, DHO),从根源上消除分类头及特征提取器上的梯度冲突,并在推理阶段通过线性插值在无需重训的前提下动态调节任务特异性与通用泛化先验的平衡。

方法详解

整体框架

DHO 面向少样本或低比例标注(如 1% / 10%)半监督分类任务。给定带有少量标注的有标注集 \(\mathcal{D}^{(l)}\) 和规模较大的未标注集 \(\mathcal{D}^{(u)}\),学生模型包含共享特征提取器 \(g(x)\) 以及解耦并行的两个独立线性分类头:有监督分类头 \(h_{\text{CE}}(z)\) 和蒸馏预测头 \(h_{\text{KD}}(z)\)。在训练阶段,有标注样本同时传入两头计算损失,未标注样本经由教师 VLM 提取软标签后仅传给蒸馏头计算 KL 散度,两头各自回传梯度并在共享特征骨干中实现协同更新;在推理阶段,两头预测概率经由轻量超参数直接线性插值合并输出。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入图像 x<br/>(有标注样本 / 未标注样本)"] --> Feat["特征提取器 g(x)<br/>输出共享特征表示 z"]
    Feat --> H1["独立双预测头解耦<br/>h_CE(z) 与 h_KD(z) 并行计算"]
    In --> VLM["教师 VLM 编码器<br/>计算类别文本提示余弦相似度"]
    VLM --> H1
    H1 --> Loss["协同目标反向传播<br/>L = λ L_CE + (1-λ) L_KD (消解梯度冲突)"]
    Loss --> Infer["推断期双头概率插值<br/>p_DHO = α σ(h_CE) + (1-α) σ(h_KD / β)"]
    Infer --> Out["最终预测类别 argmax(p_DHO)"]

关键设计

1. 独立双预测头解耦:从源头消解单头蒸馏的梯度冲突

在常规单头优化(SHO)中,分类权重矩阵 \(W\) 同时接收监督交叉熵梯度 \(\nabla_W \mathcal{L}_{\text{CE}} = (\hat{p} - y) z^\top\) 与蒸馏梯度 \(\nabla_W \mathcal{L}_{\text{KD}} = (\hat{p} - p) z^\top\)。其内积正比于 \((\hat{p} - y)^\top (\hat{p} - p) \|z\|^2\),只要学生预测与硬标签/软标签的残差方向不一致,该内积便转为负值,导致梯度方向严重冲突,并通过链式法则进一步传递至特征提取器参数 \(\theta\): $\(\mathtt{CosSim}(\nabla_\theta \mathcal{L}_{\text{CE}}, \nabla_\theta \mathcal{L}_{\text{KD}}) \propto (\hat{p} - y)^\top W \left(\frac{\partial z}{\partial \theta}\right)^\top \frac{\partial z}{\partial \theta} W^\top (\hat{p} - p) < 0\)$ 针对这一痛点,DHO 构建相互独立的双预测头 \(h_{\text{CE}}(z) = W_{\text{CE}} z + b_{\text{CE}}\) 和 \(h_{\text{KD}}(z) = W_{\text{KD}} z + b_{\text{KD}}\)。由于分类头参数完全正交,头层面的梯度冲突彻底归零;而在共享骨干网络 \(\theta\) 上,由于两头各自逼近目标分布避免了分类层残差的恶性抵消,其实际反向梯度的余弦相似度在整个训练周期内稳定维持正值(\(\mathtt{CosSim} > 0\)),从而在最轻量的架构改动下恢复了高质量的特征聚类与表示学习。

2. 推断期双头概率插值:免重训模拟最优超参数平衡

训练完成后,若仅保留单头做推断会丢失互补性。DHO 引入推断期概率加权插值规则,结合温度缩放控制预测平滑度: $\(\hat{p}_{\text{DHO}} = \alpha \, \sigma\big(h_{\text{CE}}(z)\big) + (1 - \alpha) \, \sigma\big(h_{\text{KD}}(z) / \beta\big)\)$ 其中 \(\alpha \in [0, 1]\) 负责平衡监督头与蒸馏头的置信度,\(\beta > 0\) 为蒸馏头的推断温度参数。理论上,论文证明了在两头分别收敛至 \(\epsilon\)-界限时,设置 \(\alpha = \lambda\) 且 \(\beta = 1\) 的 DHO 推断输出在 \(\ell_1\) 范数下 \(\epsilon\)-逼近最优单头蒸馏解 \(\hat{p}^* = \lambda y + (1 - \lambda) p_\tau\)。更重要的是,在传统单头蒸馏中调整平衡权重 \(\lambda\) 必须全量重新训练模型,而 DHO 将该权衡解耦为推断期超参 \(\alpha\) 和 \(\beta\),可在验证集上做网格搜索快速调优,以零重训成本获取最佳下游性能。

3. 语言感知权重初始化与余弦对齐:提升 VLM 跨模态蒸馏兼容性

针对学生模型本身亦为轻量 VLM 或需向 VLM 架构对齐的场景,随机初始化分类头易在训练初期破坏预训练对齐空间。DHO 提出语言感知初始化策略:利用教师文本编码器 \(f_{\mathcal{T}}\) 生成的全部类别提示词嵌入矩阵 \([f_{\mathcal{T}}(t_1), \dots, f_{\mathcal{T}}(t_C)]^\top\) 直接赋初值给 \(W_{\text{CE}}\) 与 \(W_{\text{KD}}\)。同时,将蒸馏头 \(h_{\text{KD}}\) 的输出逻辑重构为与教师 VLM 相同的余弦相似度计算范式: $\(h_{\text{KD}} = \frac{1}{\zeta} \big[\mathtt{CosSim}(g(x), w_1), \dots, \mathtt{CosSim}(g(x), w_C)\big]^\top\)$ 该设计拉齐了学生蒸馏头与教师预测空间的几何几何流形,有效防止了初期的蒸馏震荡,进一步提升收敛速度与最终准确率。

损失函数 / 训练策略

总体训练目标定义为双头联合损失加权和: $\(\mathcal{L} = \lambda \mathcal{L}_{\text{CE}} + (1 - \lambda) \mathcal{L}_{\text{KD}}\)$ 其中 \(\mathcal{L}_{\text{CE}}\) 仅在有标注集 \(\mathcal{D}^{(l)}\)(大小为 \(N\))上计算交叉熵损失: $\(\mathcal{L}_{\text{CE}} = \frac{1}{N} \sum_{n} \ell\big(\sigma(h_{\text{CE}}(z_n^{(l)})), y_n\big)\)$ 而 \(\mathcal{L}_{\text{KD}}\) 在有标注集与未标注集 \(\mathcal{D}^{(u)}\)(大小为 \(M\))上统一计算与教师概率 \(p\) 的 KL 散度: $\(\mathcal{L}_{\text{KD}} = \frac{1}{N} \sum_{n} D_{\text{KL}}\big[p_n^{(l)} \,\|\, \sigma(h_{\text{KD}}(z_n^{(l)}))\big] + \frac{1}{M} \sum_{m} D_{\text{KL}}\big[p_m^{(u)} \,\|\, \sigma(h_{\text{KD}}(z_m^{(u)}))\big]\)$ 默认超参数设定为 \(\lambda = 0.5\)、缩放因子 \(\eta = 2\)、温度 \(\zeta = 0.01\)。优化器采用 AdamW,使用余弦退火学习率策略调度。在少样本实验中,当缺乏验证集时直接固定 \(\beta = 0.5\),零样本教师时设置 \(\alpha = 0.4\),少样本教师时设置 \(\alpha = 0.2\)。

实验关键数据

主实验

论文在 ImageNet 少样本半监督基准(ResNet-18 从头训练 / 自监督 ResNet-50)以及 1% 与 10% 标注的低样本场景下进行了充分验证,对比单头基准与先前 SOTA 方法。

表 1:ImageNet 上少样本半监督分类准确率对比(Top-1 %)

方法类型 骨干模型与方法 0-shot 1-shot 2-shot 4-shot 8-shot 16-shot
单头基线 ResNet-18 (CE) - 0.7 1.1 1.8 3.4 8.2
ResNet-18 (CE+KD feature) - 17.1 23.5 28.0 32.2 33.8
ResNet-18 (CE+KD logit) - 50.5 50.6 50.6 51.0 51.2
双头基线 ResNet-18 (SSKD) - 42.5 46.2 48.0 50.6 52.0
ResNet-18 (DHKD) - 19.7 23.5 23.4 23.7 26.8
本文方法 ResNet-18 (DHO, 零样本教师) - 51.8 52.4 52.6 53.3 54.5
ResNet-18 (DHO-F, 少样本教师) - 53.7 54.2 54.8 56.2 57.7
教师模型 CLIP ResNet-50 (零样本教师) 60.3 - - - - -
Tip-Adapter-F ResNet-50 (少样本教师) - 61.0 61.6 62.5 63.8 65.4
自监督骨干 ResNet-50 (CE+KD logit) - 60.4 60.8 61.2 61.6 62.3
ResNet-50 (DHO, 零样本教师) - 61.0 62.1 62.5 63.7 64.7
ResNet-50 (DHO-F, 少样本教师) - 62.3 63.1 63.9 65.5 66.8

在 ImageNet 低样本(1% 与 10% 标注)设定下,DHO 配合 ViT-L/14 分别达到 84.6% 与 85.9% Top-1 准确率,大幅超越此前半监督 SOTA(如 Semi-ViT ViT-H/14 632M 参数下的 80.0% / 84.3%),以不足一半的参数量(304M)实现最高 +3.0% 的绝对提升。

消融实验

表 2:推断期双头插值策略及语言感知对齐消融分析

实验项目 设定 / 配置 Top-1 准确率 (%) 说明 / 相对变化
推断头选择
(16-shot ImageNet, ResNet-18)
仅 CE 头预测 (\(\alpha=1.0\)) 60.64 缺失蒸馏头的平滑泛化先验
仅 KD 头预测 (\(\alpha=0.0\)) 61.55 缺失任务监督头的尖锐特异性
DHO 双头插值 (\(\alpha=0.4, \beta=0.5\)) 65.37 插值融合显著提升 +3.82% / +4.73%
初始化与对齐
(1% ImageNet, ViT-B/16 学生)
无文本初始化,无余弦对齐 78.3 纯随机线性头基线
仅 CE 头采用文本初始化 78.5 +0.2% 提升
双头均采用文本初始化 78.6 +0.3% 提升
双头文本初始化 + KD 头余弦对齐 78.7 累积提升 +0.4%
特征表征线性评估
(16-shot ImageNet 冻结特征提取器)
冻结特征 + 线性评估 (CE+KD feature) 62.3 (Top-5: 85.0) 特征表征受梯度干扰
冻结特征 + 线性评估 (CE+KD logit) 66.2 (Top-5: 88.8) 单头 logit 蒸馏基线
冻结特征 + 线性评估 (DHO) 67.1 (Top-5: 89.3) Top-1 +0.9%,证明骨干表征质量实质提升

关键发现

  • 双头协同优于单头独立:消融表明即便使用相同的训练骨干,单独采用 CE 头(60.64%)或 KD 头(61.55%)推断均显著低于双头线性加权推断(65.37%),说明两头学到了互补的分类流形。错误率归约分析显示,结合后 KD 头的单头失误率从 11.5% 骤降至 2.5%。
  • 骨干表征确实得到净化:在线性评估协议下,冻结特征提取器仅微调顶层线性分类器,DHO 的 Top-1 准确率仍比单头 logit 蒸馏高出 0.9%(67.1% vs 66.2%),直接证明消除梯度冲突使共享骨干学到了聚类更紧凑、区分度更高的通用特征。
  • 计算开销极低:在 ImageNet 1000 类下,双头结构仅引入极少的参数增加(ResNet-18 增加 4.4% 参数,FLOPs +0.0%,吞吐量仅下降 0.20%),在类别数较少的数据集上额外开销可忽略不计。

亮点与洞察

  • 从梯度根源定位蒸馏失效机制:论文通过严谨的公式推导阐明了半监督蒸馏中监督残差与蒸馏残差的内积失配现象,解释了为何传统单头蒸馏在少样本下难以释放教师 VLM 的潜能。
  • 推断期插值等价于免重训调参:通过理论证明 DHO 推断阶段的凸组合 \(\alpha \hat{p}_{\text{CE}} + (1 - \alpha) \hat{p}_{\text{KD}}\) 逼近最优单头联合损失的鞍点解,将高成本的训练期超参数 \(\lambda\) 搜索转化为推断期零成本的网格调优。
  • 无缝嵌入多模态下游适配:DHO 可即插即用集成到 VPT、CoOp、PromptSRC 等提示词微调方法中,不仅在分布内 ImageNet 持续涨点,在 OOD 变体(ImageNet-V2/Sketch/R/A)上也建立了新的 SOTA。

局限与展望

  • 复杂密集预测任务尚未验证:当前验证主要集中在图像分类基准,对于目标检测、语义分割等需要密集像素级对齐和结构化标签的多任务场景,双头解耦与推断期聚合策略仍需进一步探索。
  • 基于熵的自适应加权机制未达预期:作者尝试用预测熵自适应计算动态权重 \(\alpha\),但由于现代神经网络预测普遍过度自信,未经校准的熵无法真实反映不确定性,导致自适应融合效果未能超越固定网格搜索权重。
  • 模型跨架构分布对齐约束:当教师与学生来自不同训练分布时(如 DFN 预训练的 ViT-H 与 CLIP 预训练的 ViT-B),在剧烈分布漂移任务(如 ImageNet-A)上仍会出现一定程度的迁移瓶颈。

相关工作与启发

  • vs 单头 Logit 蒸馏 (Hinton, 2015; Chen et al., 2020):传统单头方法强行让单一分类头承受稀疏监督硬标签与平滑教师软标签,造成梯度相互打架;DHO 通过结构解耦消除了梯度冲突,并在推断期保留了双重知识集成优势。
  • vs 多任务梯度手术 (PCGrad, Yu et al., 2020b):PCGrad 采用反向传播期间的实时梯度正交投影,带来 \(\mathcal{O}(|\theta|)\) 的显存与计算时间开销;DHO 仅增加 \(\mathcal{O}(d \times C)\) 的分类头参数,从优化流形源头规避冲突且效果更优(8 个数据集平均提升 79.1% vs 71.6%)。
  • vs 既有双头蒸馏 (SSKD / DHKD):SSKD 仅在推断时使用 CE 单头且服务于无监督域迁移,DHKD 旨在防止二值分类崩溃;DHO 首次针对基础多模态模型到紧凑学生的少样本迁移,并利用推断期双头概率插值实现最优权衡。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 核心想法大道至简,从梯度冲突角度解耦双头并给出推断等价性证明,视角清晰有力。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 15 个基准数据集、多种模型架构(ResNet/ViT/MobileNet)及少样本/低样本/OOD 全面评估,证据极其详实。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机陈述直接流畅,理论分析与实证观察呼应紧密,图表清晰。
  • 价值: ⭐⭐⭐⭐⭐ 开销极小、即插即用,极具工业落地与后续学术扩展价值。