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

源码 code github
三个损失是 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 维原型空间后的输出）。

核心计算：交叉熵

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

$${\mathcal{L}}_{\text{DINO}}=-\sum _c{t}_c\log p_c,p=\text{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

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 次：

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 权重的交叉熵

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]

两个细节：

1. xFormers 融合算子：优先用 xformers.ops.cross_entropy(..., bw_inplace=True)——65536 维 softmax 的中间张量巨大，融合实现把 softmax/log/乘加合成一个 kernel 并原地反传，大幅省显存；无 xFormers 时回退手写 Σ t·log_softmax；
2. 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 归一化后）。

核心计算

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

$${\mathcal{L}}_{\text{KoLeo}}=-\frac{1}{B}\sum _{i=1}^B\log d(x_i,x_{NN(i)})$$

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

两个实现细节

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

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

四、三者协同总览

在 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 评测都能高分的根源。