VLM学习-DINOv2三大损失DINO / iBOT / KoLeo 解析

三个损失是 DINOv2 自监督训练的核心，分工是：DINO 损失管全局语义、iBOT 损失管局部细节、KoLeo 损失管特征空间不坍塌。逐个结合源码解析。

DINO / iBOT / KoLeo 三大损失解析

〇、共同背景：无标签的"自蒸馏"

三者都建立在同一个范式上（回顾之前解析的 train/ssl_meta_arch.py）：

复制代码

同一张图裁剪出多个视角（2 global 224 + 8 local 96）
   ├─ student 网络: 吃全部视角，输出可训练
   └─ teacher 网络: EMA 滑动平均自 student，no_grad，输出当"软标签"

损失 = 让 student 的预测分布去拟合 teacher 的预测分布（交叉熵）

没有真实标签------teacher 的输出本身就是标签，这就是"自蒸馏"。三个损失决定"在什么粒度上蒸馏、以及怎么防退化"。

一、DINO 损失：CLS token 的交叉蒸馏（全局语义）

源码：loss/dino_clstoken_loss.py

作用对象 ：每个裁剪的 CLS token（经 DINOHead 投影到 65536 维原型空间后的输出）。

核心计算：交叉熵

python 复制代码

def forward(self, student_output_list, teacher_out_softmaxed_centered_list):
    total_loss = 0
    for s in student_output_list:                       # student 的每个视角
        lsm = F.log_softmax(s / self.student_temp, dim=-1)   # student 温度 0.1，分布平滑
        for t in teacher_out_softmaxed_centered_list:   # teacher 的每个视角
            loss = torch.sum(t * lsm, dim=-1)           # CE = -Σ t·log p
            total_loss -= loss.mean()
    return total_loss

L DINO = − ∑ c t c log ⁡ p c , p = softmax ( s / 0.1 ) \mathcal{L}{\text{DINO}} = -\sum{c} t_c \log p_c, \qquad p = \text{softmax}(s/0.1) LDINO=−c∑tclogpc,p=softmax(s/0.1)

关键点：student × teacher 的全配对交叉------student 的 local 小裁剪要去预测 teacher 的 global 大裁剪的分布（跨分辨率、跨视角的互预测），且排除"自己预测自己"。这迫使"局部看到猫耳朵 → 推断整图是猫"的全局语义对齐。

非对称温度：锐化 teacher、平滑 student

teacher 侧除以很小的 teacher_temp（0.04→0.07 预热）→ softmax 输出接近 one-hot，标签尖锐；
student 侧除以 0.1 → 分布较平，预测宽容 。
这种不对称避免 student 走捷径输出常数分布。

防坍塌机制一：centering

python 复制代码

def softmax_center_teacher(self, teacher_output, teacher_temp):
    self.apply_center_update()
    return F.softmax((teacher_output - self.center) / teacher_temp, dim=-1)

center 是 teacher 输出的指数滑动均值（动量 0.9），并经过 all_reduce 全卡同步（异步 all_reduce，L84-L86）。减去均值后再 softmax，防止所有样本的 teacher 输出挤到同一个原型上（输出坍塌）。

防坍塌机制二（替代方案）：Sinkhorn-Knopp

DINOv2 实际默认用这个更硬的方案：对 exp(teacher_output/T) 构成的 K×B 矩阵交替做行归一化和列归一化，迭代 3 次：

python 复制代码

for it in range(n_iterations):      # 3 次迭代
    Q /= sum_of_rows; Q /= K        # 每行归一: 每个原型分到的总权重 = 1/K（原型均匀使用）
    Q /= sum_of_cols; Q /= B        # 每列归一: 每个样本分到的总权重 = 1/B

效果：在 batch 内强制 K 个原型被均匀分配 ------不允许某个原型独占所有样本，从分配层面杜绝坍塌。在 ssl_meta_arch.py 中通过 cfg.train.sinkhorn_knopp_teacher 开关二选一。

二、iBOT 损失：mask 位置的 patch 级交叉蒸馏（局部细节）

源码：loss/ibot_patch_loss.py

作用对象 ：patch tokens ，且只在被 mask 的位置上计算。

机制：自监督版 BERT

回顾 prepare_tokens_with_masks：student 的输入中，一部分 patch 被替换为可学习的 mask_token。于是：

复制代码

student 看到被遮的图 → 必须根据上下文推断被遮 patch 的内容
teacher 看到完整的图 → 其 patch token 输出就是"标准答案"

损失 = 仅在被 mask 的位置上，student patch 分布拟合 teacher patch 分布

核心计算：带 mask 权重的交叉熵

python 复制代码

def forward_masked(self, student_patch_tokens_masked, teacher_patch_tokens_masked,
                   student_masks_flat, n_masked_patches=None, masks_weight=None):
    loss = lossfunc(t, s, self.student_temp)       # 逐 patch 的交叉熵
    if masks_weight is None:
        masks_weight = (1 / student_masks_flat.sum(-1).clamp(min=1.0))...  # 按每张图 mask 数归一
    loss = loss * masks_weight
    return -loss.sum() / student_masks_flat.shape[0]

两个细节：

xFormers 融合算子 ：优先用 xformers.ops.cross_entropy(..., bw_inplace=True)------65536 维 softmax 的中间张量巨大，融合实现把 softmax/log/乘加合成一个 kernel 并原地反传，大幅省显存 ；无 xFormers 时回退手写 Σ t·log_softmax；
sinkhorn_knopp_teacher与 DINOLoss 版几乎相同，唯一差异是 B = 全部被 mask 的 patch 总数（跨卡 all_reduce），即在所有 masked patch 上做原型均匀分配。

分工意义 ：DINO 损失只蒸馏 CLS（全局），patch token 得不到直接监督------而密集任务（分割/深度）恰恰依赖 patch tokens。iBOT 损失补上了这一块：逐 patch 的上下文预测让每个 token 都编码有意义的局部语义，这是 DINOv2 密集特征质量远超 DINO v1 的关键原因。

三、KoLeo 损失：最近邻距离的熵正则（防堆叠）

源码：loss/koleo_loss.py

作用对象：batch 内 student 的 CLS 特征（L2 归一化后）。

核心计算

python 复制代码

def forward(self, student_output, eps=1e-8):
    student_output = F.normalize(student_output, eps=eps, p=2, dim=-1)  # 投影到单位球面
    I = self.pairwise_NNs_inner(student_output)        # 每个样本的最近邻索引
    distances = self.pdist(student_output, student_output[I])   # 到最近邻的欧氏距离
    loss = -torch.log(distances + eps).mean()          # 负对数距离
    return loss

L KoLeo = − 1 B ∑ i = 1 B log ⁡ d ( x i , x N N ( i ) ) \mathcal{L}{\text{KoLeo}} = -\frac{1}{B}\sum{i=1}^{B} \log d(x_i, x_{NN(i)}) LKoLeo=−B1i=1∑Blogd(xi,xNN(i))

直觉：最小化该损失 ⇔ 最大化每个点到其最近邻的距离 ⇔ 特征互相"排斥"，均匀铺满整个超球面。名字来自 Kozachenko-Leonenko 熵估计器（引自 Sablayrolles et al. 2018）。

两个实现细节

球面最近邻用点积找 ：L2 归一化后，欧氏距离最小 ⟺ 点积最大，所以直接 torch.mm(x, x.t()) 取每行最大值的索引（对角线先填 -1 排除自己），全程 GPU 上完成、不用 Faiss；
强制 fp32 ：autocast(enabled=False)------距离是对数里的微小量，半精度会失真。

分工意义 ：centering/Sinkhorn 防的是"所有样本挤成一个点"的输出坍塌，KoLeo 防的是另一种退化------不同类别的特征在球面上成团堆叠 。它保证特征空间的均匀性，这正是 kNN 评测能工作的几何前提（上一轮讲过：同类紧聚 + 均匀铺展 → 余弦相似度可判别）。注意 KoLeo 只加在 student 的 global 裁剪 CLS 上（ssl_meta_arch.py 中 do_koleo 分支）。

四、三者协同总览

损失	蒸馏对象	监督信号	防止的问题	贡献的能力
DINO	CLS token	teacher CLS 分布（cross-view）	输出坍塌（centering / Sinkhorn）	全局语义：kNN/线性探针高分
iBOT	被 mask 的 patch tokens	teacher 完整图的 patch 分布	同上（patch 版 Sinkhorn）	局部细节：分割/深度等密集任务
KoLeo	student CLS（无蒸馏）	无------纯几何正则	特征成团堆叠	特征均匀分布：相似度检索可用

在 train/ssl_meta_arch.py 的 forward_backward 中，三者由 do_dino / do_ibot / do_koleo 开关控制、按配置权重加权求和后反传：

复制代码

total = w_dino·L_DINO + w_ibot·L_iBOT + w_koleo·L_KoLeo

DINO 教模型"看图识整体"，iBOT 教模型"补全局部细节"，KoLeo 保证学到的特征空间不塌不堆------三者合力塑造出 DINOv2 "同类紧聚、异类分离、均匀铺展"的特征几何，这正是前面 kNN/LogReg/Linear 评测都能高分的根源。