基本信息

项目 内容
论文标题 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使分布更“尖锐”

网络架构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
原始图像 

├── 弱增强(Resize + CenterCrop) → 教师输入

└── 强增强(RandomResizedCrop + ColorJitter + ...) → 学生输入
│ │
↓ ↓
教师网络(ViT) 学生网络(ViT)
│ │
↓(无梯度) ↓(计算梯度)
输出:z_teacher 输出:z_student
│ │
↓(中心化 + 锐化) ↓
目标分布: 预测分布:
p = softmax((z_t - mean)/τ) q = softmax(z_s / τ)
│ │
└─────────── 交叉熵损失 ─────┘

L = KL(q || p)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
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 + 自监督预训练