DiNO (Knowledge Distillation with No Labels)（一）

DiNO模型是由facebook团队研发，使用自监督训练方式训练Transformer模型。

• 论文地址：Emerging Properties in Self-Supervised Vision Transformers
• Github地址：DiNo

本文在ChatGPT的帮助下完成

基础知识（教师-学生模型）

教师学生模型（Teacher-Student Model），又称为知识蒸馏（Knowledge Distillation），是一种广泛使用的机器学习方法，尤其在深度学习领域中得到了广泛应用。这种模型架构涉及两个关键部分：一个教师模型（Teacher model）和一个学生模型（Student model）。教师模型通常是一个预训练好的、性能较高的大模型，而学生模型则是一个较小的模型。

教师学生模型的核心思想是通过教师模型来指导学生模型，使学生模型能学习到教师模型的知识，但以更小、更高效的形式表达。具体来说，学生模型学习模仿教师模型在相同输入下的输出或特征行为。

知识蒸馏的步骤

• 教师模型训练：首先训练一个大型的教师模型，这个模型通常具有较深的网络结构，能够在特定任务上达到高精度。

• 学生模型训练：训练一个结构更简单的学生模型。在训练过程中，学生模型不仅要学习数据的标签，还要尽可能地模仿教师模型的输出。这通常通过最小化教师输出和学生输出之间的差异来实现，比如使用软标签或输出概率分布的温度缩放（Temperature Scaling）。

下面给出一个知识蒸馏的Example帮助理解：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torch.nn.functional import softmax, log_softmax

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义教师模型
class TeacherModel(nn.Module):
    def __init__(self):
        super(TeacherModel, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(784, 1200),
            nn.ReLU(),
            nn.Linear(1200, 1200),
            nn.ReLU(),
            nn.Linear(1200, 10)
        )
    
    def forward(self, x):
        return self.network(x)

# 定义学生模型
class StudentModel(nn.Module):
    def __init__(self):
        super(StudentModel, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(784, 100),
            nn.ReLU(),
            nn.Linear(100, 10)
        )
    
    def forward(self, x):
        return self.network(x)

# 加载数据
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)

# 实例化模型
teacher = TeacherModel().to(device)
student = StudentModel().to(device)

# 教师模型训练
def train_teacher(model, loader):
    model.train()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    for epoch in range(5):
        for images, labels in loader:
            images = images.view(images.shape[0], -1).to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
        print(f'Epoch {epoch+1}, Loss: {loss.item()}')

# 知识蒸馏
def distill_student(teacher, student, loader, T=2.0):
    teacher.eval()
    student.train()
    optimizer = optim.Adam(student.parameters(), lr=0.001)
    criterion = nn.KLDivLoss(reduction='batchmean')
    for epoch in range(5):
        for images, labels in loader:
            images = images.view(images.shape[0], -1).to(device)
            with torch.no_grad():
                teacher_outputs = softmax(teacher(images) / T, dim=1)
            student_outputs = log_softmax(student(images) / T, dim=1)
            optimizer.zero_grad()
            loss = criterion(student_outputs, teacher_outputs)
            loss.backward()
            optimizer.step()
        print(f'Epoch {epoch+1}, Distillation Loss: {loss.item()}')

# 训练教师模型
train_teacher(teacher, train_loader)

# 执行知识蒸馏
distill_student(teacher, student, train_loader)

几个有帮助的关于知识蒸馏的链接：

• Pytorch: Knowledge Distillation Tutorial
• Knowledge Distillation, aka Teacher-Student Model

自监督学习中教师-学生模型中的应用

在自监督学习中，教师学生模型的应用称为自监督知识蒸馏，这种策略主要利用教师模型来生成无标签数据的软目标（soft targets），而学生模型则学习这些软目标，以此来学习数据的有用表示。这种方法特别适用于数据标注稀缺或昂贵的情况，比如医疗图像分析、大规模视频处理或任何需要处理大量未标记数据的领域。

自监督教师学生模型的工作机制

教师模型先自监督训练，之后将知识传递给学生

• 教师模型初始化：
  在自监督设置中，教师模型通常是通过一些先前定义好的自监督任务预训练的。例如，可以使用随机数据增强（如随机裁剪、颜色变换等）生成图像的多个视图，然后让教师模型学习区分或预测这些变换。

• 无标签数据的利用：
  教师模型处理无标签数据，生成输出，这些输出反映了数据中的复杂模式和结构。这些输出可能是特征表示、类别概率分布或其他形式的编码信息。

• 学生模型的训练：
  学生模型不直接从真实的标签学习，而是试图复制或逼近教师模型的输出。通过这种方式，学生模型可以从教师模型学习如何处理数据，尽管没有明确的监督信号（即没有真实标签）。

• 损失函数：
  在自监督知识蒸馏中，通常使用如交叉熵或KL散度的损失函数来最小化学生输出和教师输出之间的差异。这种损失帮助学生模型在特征空间中模仿教师模型的行为。

• 教师模型的更新：
  与传统知识蒸馏不同，自监督学习中的教师模型有时也可以进行更新，例如使用指数移动平均（EMA）更新教师模型的参数，以保持教师的知识是最新的，同时也是平滑的。

教师模型和学生模型同时进行自监督训练

这种情况下，教师和学生模型将会同时训练，但是教师模型的参数更新会先学生模型一步，使得教师模型可以提供更稳定，更平滑的目标给学生模型。

对于每个输入图像，通过应用一系列随机数据增强（如裁剪、颜色扰动等）来生成多个不同的视图。这些不同的视图被用来训练学生模型，并通过教师模型来生成目标输出。教师模型对一个视图生成输出，而学生模型则尝试预测另一个增强视图的相同输出。学生模型的目标是最小化其输出与教师模型输出之间的差异。这种差异通常通过KL散度或其他相似性度量来评估。而DiNO模型中正是采用了这种训练方式

DiNO 模型详解

Facebook除了论文外，也在其博客上介绍了DiNO模型的工作原理，可以参考阅读。

DiNO模型结构如上图。$x_1$, $x_2$ 是$x$的随机变化，它们分别被输入到学生模型和教师模型中。学生模型和教师模型具有相同的结构但是不同的参数数量。教师网络的输出通过减去批次内计算出的平均值进行中心化（图中centering模块）。每个网络输出一个经过温度 softmax 归一化的 K 维特征。然后通过交叉熵损失来测量它们的相似性。在教师上应用停止梯度（sg）操作符，只通过学生来传播梯度。教师的参数通过学生参数的指数移动平均（ema）进行更新。

对于上述的学生模型和教师模型，它们都会给一个$K$维度的概率分布输出$P_s$和$P_t$。这种概率分布是通过一个softmax得到的:
$$P_{{\theta }_s}(x^{(i)})=\frac{\exp (g_{{\theta }_s}(x^{(i)})/{\tau }_s)}{{\sum }_{k=1}^K\exp (g_{{\theta }_s}(x^{(k)})/{\tau }_s)},$$

教师-学生模型的目的是最小化学生模型和教师模型的概率分布：
$$mi{n}_{{\theta }_s}H(P_t(x),P_s(x))$$
其中， $H(a,b)=-a\log b$

更具体的，假设我们有一个图片，那么它可以通过数据增强产生一个集合$V$。在构建这个集合时，使用了multi-crop策略，即$V$中包含两个global views： $x_1^g$, $x_2^g$，以及几个smaller resolution 的local views。所有的 local views输入到学生模型中，两个global views输入到教师模型中，然后最小化教师模型和学生模型的损失：

min ⁡ θ s ∑ x ∈ { x 1 g , x 2 g } ∑ x ′ ∈ V , x ′ ≠ x H ( P t ( x ) , P s ( x ′ ) ) \min_{\theta_s} \sum_{x \in \{x_1^g, x_2^g\}} \sum_{x' \in V, x' \neq x} H(P_t(x), P_s(x')) θs​min​x∈{x1g​,x2g​}∑​x′∈V,x′=x∑​H(Pt​(x),Ps​(x′))
论文中把global views的大小设置为$224^2$以覆盖大于原图50%的区域， 把local views的大小设置为$96^2$以覆盖小于原图50%的区域。

DiNO中的教师模型

在DiNO中，教师模型不会预先训练，而是通过学生网络在过去迭代中的状态构建而成。这种方法依赖于学生网络的历史学习成果来定义教师网络的行为和参数。教师网络的参数更新采用指数移动平均（EMA）策略，
$${\theta }_t\leftarrow \lambda {\theta }_t+(1-\lambda ){\theta }_s$$
λ 控制了教师网络旧参数的保留量。当 λ 接近 1 时，教师网络的更新较慢，参数变化平滑，更多地保留历史信息。1−λ 表示新信息（即学生网络的当前参数）在更新中的权重。当 λ 较小，即更接近于 0 时，新信息在更新中的比例增大，教师网络的参数更新更快，反应更为敏感。
在实际应用中，
𝜆 的值通常不是固定的，而是随训练过程动态调整，这种调整方法称为余弦调度。余弦调度的主要思想是模拟余弦函数的起伏变化，使 𝜆的值在训练过程中按照一定的周期性模式变化：

• 初始阶段：λ 可以从一个较高的值开始，确保教师网络在训练初期稳定，不会由于学生网络的初期不稳定性而进行剧烈的参数调整。
• 训练中后期：随着学生网络逐渐稳定并表现出较好的性能，可逐渐减小，使教师网络更快地吸收学生网络的新学习成果。

论文为了防止模型崩塌，还对网络输出进行中心化和锐化处理。中心化操作通过减去一项偏置
𝑐 来调整教师网络的输出，使输出不倾向于任何特定的方向（防止某一维度主导），从而防止模型输出过于一致而导致的崩溃。
$$c\leftarrow mc+(1-m)\frac{1}{B}\sum _{i=1}^B{g}_{{\theta }_t}(x_i)$$

其中 ，𝑚 是衰减系数，𝐵是批次大小。

另外，在使用Vision Transformer时，作者使用了$N=16,8$，把图像划分为$N\times N$个patch，进而转换为序列。

以上便是DiNO的基本工作流程，更多细节，请参考原文。