DINO自监督学习论文复现ViT特征提取能力的实验验证一、DINO的自监督训练为何能产生语义分割级别的特征图DINODIstillation with NO labels在2021年引起关注的核心现象是通过自监督训练的ViT其自注意力图能够自动分割出图像中的前景物体——无需任何像素级标注。这一现象在监督训练的ViT上完全观测不到。理解这个现象背后的机制对于设计更好的自监督学习目标具有指导意义。DINO的核心设计是两个关键组件动量教师网络momentum teacherEMA更新和多裁剪策略multi-crop教师看全局视图学生看局部视图。这两个设计共同创造了一个约束教师必须从完整图像中提取语义一致的表示学生则从局部裁剪中学习匹配——这使得模型被迫学习物体的整体性概念。flowchart TB A[输入图像] -- B{数据增强} B --|全局视图 ×2| C1[教师输入: 224×224] B --|局部视图 ×8| C2[学生输入: 96×96] C1 -- D[教师网络: ViT] D -- D1[EMA 更新参数] C2 -- E[学生网络: ViT] D -- F[教师输出: softmax 概率] E -- G[学生输出: softmax 概率] F -- H[交叉熵损失: 学生 → 教师] G -- H H -- I[梯度仅更新学生] D -- J[自注意力图可视化] J -- K[语义分割级别的前景物体检测] style J fill:#fff3e0 style K fill:#e8f5e9二、动量教师的设计与崩溃预防DINO中最容易出错的机制是动量教师的更新。使用EMA指数移动平均更新的教师网络其参数是学生参数的指数滑动平均。这个设计有两个层面的作用提供稳定的学习目标教师更新慢于学生momentum通常从0.996逐渐升温到1.0避免了目标表示的频繁变化导致学生无法收敛。防止表示坍缩DINO 使用了 centered softmax 来防止教师输出坍缩到均匀分布。具体做法是在 softmax 前减去一个批次统计的均值使教师输出保持尖锐sharp而非均匀。import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms from typing import Tuple import copy class DINOStudent(nn.Module): DINO 学生网络。 为什么学生和教师共享相同的架构但不共享参数 EMA 更新使教师成为学生的慢动作版本 这种时间延迟使得教师输出相对于当前学生参数 更加稳定——类似于强化学习中的 target network。 def __init__(self, vit_model: nn.Module, output_dim: int 65536): super().__init__() self.backbone vit_model # 投影头将 ViT 的 [CLS] token 映射到高维空间 # 高维输出65536增加了表示的判别性 self.head nn.Sequential( nn.Linear(vit_model.embed_dim, 2048), nn.GELU(), nn.Linear(2048, 2048), nn.GELU(), nn.Linear(2048, output_dim) ) def forward(self, x: torch.Tensor) - torch.Tensor: features self.backbone(x) # [B, embed_dim] return self.head(features) # [B, output_dim] class DINOTeacher(nn.Module): DINO 动量教师网络。 EMA 更新策略 momentum 从 0.996 按照 cosine schedule 升温到 1.0。 训练初期快速跟随学生momentum0.996 训练后期几乎不变momentum≈1.0。 为什么不用 PyTorch 的 moving_average 直接操作 需要精确控制每次 step 的 momentum 值 而不是使用固定的 decay 参数。 def __init__(self, student: DINOStudent): super().__init__() self.backbone copy.deepcopy(student.backbone) self.head copy.deepcopy(student.head) # 教师参数不参与梯度计算 for param in self.parameters(): param.requires_grad False torch.no_grad() def update_ema(self, student: DINOStudent, momentum: float): EMA 更新θ_teacher m × θ_teacher (1-m) × θ_student。 for teacher_param, student_param in zip( self.parameters(), student.parameters() ): teacher_param.data.mul_(momentum).add_( student_param.data, alpha1 - momentum ) class DINOLoss(nn.Module): DINO 的交叉熵损失包含 centered softmax 和 temperature。 centered softmax 的作用 在 softmax 之前减去 teacher 输出的批次均值centering 防止 teacher 输出坍缩到某个单一的万能表示。 没有 centering 时teacher 可能对所有输入都输出相似的表示 导致学生无法从教师学到有判别性的特征。 def __init__( self, output_dim: int 65536, teacher_temp: float 0.04, student_temp: float 0.1, center_momentum: float 0.9 ): super().__init__() self.teacher_temp teacher_temp self.student_temp student_temp self.center_momentum center_momentum # 可学习的中心向量在训练过程中通过 EMA 更新 self.register_buffer(center, torch.zeros(1, output_dim)) def forward( self, student_output: torch.Tensor, # [B, D] teacher_output: torch.Tensor, # [B, D] ) - torch.Tensor: 计算 DINO 损失。 学生和教师使用不同的温度参数τ_s τ_t 教师的温度更低 → 输出更尖锐 → 更像硬标签 学生的温度更高 → 输出更平滑 → 更容易匹配 # 教师侧centered softmax teacher_out teacher_output.detach() # 不传梯度到教师 # 更新中心向量 batch_center teacher_out.mean(dim0, keepdimTrue) self.center self.center * self.center_momentum batch_center * (1 - self.center_momentum) # centered softmax teacher_out teacher_out - self.center teacher_probs F.softmax(teacher_out / self.teacher_temp, dim-1) # 学生侧标准 softmax student_log_probs F.log_softmax( student_output / self.student_temp, dim-1 ) # 交叉熵-sum(teacher_probs * log(student_probs)) loss -torch.sum(teacher_probs * student_log_probs, dim-1).mean() return loss三、特征图质量评估的方法论评估自监督学习特征质量的最直接方式是线性探测linear probe——在冻结的预训练特征上训练一个线性分类器。但DINO的独特贡献是揭示了线性探测之外的特征质量维度自注意力图的前景分割能力无需任何微调直接从最后一层的[CLS]token 对其他token的注意力权重中提取前景物体区域。评估方式将注意力图与真实分割标注计算 IoU。k-NN分类性能在DINO特征空间中直接使用k-NN进行图像检索和分类。这个评估不需要任何额外的训练连线性层都不要直接反映表示空间的语义结构质量。跨数据集的迁移能力在ImageNet上预训练的DINO特征在CUB-200鸟类细粒度分类、Stanford Cars等数据集上的线性探测性能反映特征的泛化性。四、复现中的关键训练稳定性问题多裁剪策略的内存管理2个全局视图 8个局部视图意味着每个样本需要10次前向传播。batch size需要相应调小如从256降到64否则显存不够。centering 的预热阶段中心向量的初始化来自第一批数据的教师输出均值。如果第一批数据类别不均衡中心向量会有偏差影响早期训练的稳定性。建议在中心向量稳定前使用较高的温度。ViT 的 patch_size 选择DINO 原始使用 ViT-S/16patch_size16。减小 patch_size如8可以获得更精细的自注意力图但计算量增加4倍。五、总结DINO 自监督学习的核心贡献是揭示了 ViT 在无标签训练下涌现出的语义分割能力动量教师EMA更新提供稳定的学习目标centered softmax 防止表示坍缩。多裁剪策略强制模型从局部推断全局间接学习了物体的完整性概念。特征质量不仅通过线性探测评估自注意力图的前景分割能力是一个同等重要的评估维度。训练稳定性受中心向量初始化和温度参数选择的显著影响。