论文阅读_DINO

基本信息

项目	内容
论文标题	Emerging Properties in Self-Supervised Vision Transformers
作者
作者单位	Meta AI
时间	2021
发表会议/期刊
全称	Self-Distillation with No Labels

方法概览

特点	文章性质
输入	单张 RGB 图像
输出
所属领域	分类，语义分割

在 DINO 之前，主流的自监督方法有：

• 对比学习（如 MoCo、SimCLR）：需要负样本，显存消耗大
• 生成式模型（如 MAE）：重建像素，可能关注低级细节

创新点

1. 通过“学生-教师”自蒸馏框架，无需任何标签，让 Vision Transformer（ViT）自动学习强大的视觉表征，并且在训练过程中涌现出语义分割、定位等能力。
2. 自蒸馏框架（Student-Teacher Distillation）
   1. 学生网络（Student）：接收强增强图像，实时更新
   2. 教师网络（Teacher）：接收弱增强图像，不计算梯度
   3. 教师网络是学生的指数移动平均 （EMA）, 教师参数不是独立训练的而是学生参数的平滑版本,随训练teacher部分权重逐渐增大。- 避免模型崩溃（collapse）
   4. 教师输出经过中心化和锐化（Centering & Sharpening）
      1. 中心化：减去所有样本输出的平均值，防止所有输出趋同（collapse）
      2. 锐化（Sharpening）：对教师输出进行 温度缩放 + softmax使分布更“尖锐”

网络架构

原始图像 
│ 
├── 弱增强（Resize + CenterCrop） → 教师输入 
│ 
└── 强增强（RandomResizedCrop + ColorJitter + ...）          → 学生输入  
                                    │                         │
                                    ↓                         ↓
                                教师网络（ViT）             学生网络（ViT） 
                                    │                         │ 
                                    ↓（无梯度）                 ↓（计算梯度）
                                输出：z_teacher             输出：z_student 
                                    │                          │
                                    ↓（中心化 + 锐化）           ↓ 
                                  目标分布：                  预测分布：
                              p = softmax((z_t - mean)/τ)   q = softmax(z_s / τ) 
                                    │                          │ 
                                    └─────────── 交叉熵损失 ─────┘ 
                                                      ↓ 
                                                      L = KL(q || p)

import torch
import torch.nn as nn

class DINO(nn.Module):
    def __init__(self, student, teacher, embed_dim=256, warmup_teacher_temp=0.04,
                 teacher_temp=0.04, warmup_teacher_temp_epochs=30):
        super().__init__()
        self.student = student
        self.teacher = teacher
        self.student_head = MLP(student.embed_dim, embed_dim)
        self.teacher_head = MLP(teacher.embed_dim, embed_dim)
        
        # 教师温度（可变）
        self.teacher_temp = teacher_temp
        self.warmup_teacher_temp = warmup_teacher_temp
        self.warmup_teacher_temp_epochs = warmup_teacher_temp_epochs
        
        # 教师参数初始化为学生
        for param in teacher.parameters():
            param.requires_grad = False  # 不计算梯度

    @torch.no_grad()
    def update_teacher(self, m):
        # m: 动量，如 0.996
        for param_q, param_k in zip(self.student.parameters(), self.teacher.parameters()):
            param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)

    def forward(self, global_view, local_view):
        # global_view: 弱增强（教师输入）
        # local_view: 强增强（学生输入）
        
        # 学生网络（强增强）
        student_feats = self.student(local_view)
        student_logits = self.student_head(student_feats)  # (B, D)
        
        # 教师网络（弱增强，无梯度）
        with torch.no_grad():
            teacher_feats = self.teacher(global_view)
            teacher_logits = self.teacher_head(teacher_feats)  # (B, D)
            
            # 中心化
            teacher_logits = teacher_logits - teacher_logits.mean(dim=0, keepdim=True)
            
            # 温度缩放（可带 warmup）
            curr_temp = self.get_teacher_temp(epoch)
            teacher_probs = torch.softmax(teacher_logits / curr_temp, dim=1)
        
        # 损失：KL 散度
        loss = -torch.sum(teacher_probs * torch.log_softmax(student_logits, dim=1), dim=1).mean()
        
        return loss

    def get_teacher_temp(self, epoch):
        if epoch < self.warmup_teacher_temp_epochs:
            return self.warmup_teacher_temp
        else:
            return self.teacher_temp

ViT 的局限

问题	说明
❌ 依赖 ViT	在 CNN 上效果不如 ViT
❌ 对增强敏感	强/弱增强策略需精心设计
❌ 训练不稳定	初期可能崩溃，需 warm-up 和参数调整

后续发展：

方法	改进点
DINOv2	使用更大数据集（LAION），支持多任务（检测、分割）
iBOT	在 patch 级别进行自蒸馏，增强局部语义
MoBY	结合对比学习与 DINO
CAE	结合 MAE 的掩码重建与 DINO 的蒸馏
SAM（Segment Anything）	受 DINO 启发，使用 ViT + 自监督预训练