DINO

v2

DINOv2 的目标是

  • 构建一个在 约 142M 张精心“自动筛选” 并去重的图像上预训练的自监督 ViT。
  • 结合图像级与 Patch 级对比目标，利用大模型蒸馏与高效并行优化，训练出可直接用于分类、分割、深度、检索、视频等多种视觉任务的“随取即用”特征。

架构

DINOv2 是一种基于 Vision Transformer (ViT) 的自监督学习框架，通过教师–学生（teacher–student）机制，结合图像级（DINO）和补丁级（iBOT）对比损失，配合 Sinkhorn-Knopp 中心化与 KoLeo 正则化，实现对大规模、多样化图像的无监督预训练

  • Backbone 采用 ViT 变体（ViT-S/14、ViT-B/14、ViT-L/14、ViT-g/14），支持不同参数规模与计算预算的模型选择
  • 教师网络由学生网络参数的指数移动平均（EMA）生成，教师输出的“原型分布”作为学生目标，以保持稳定的学习信号
  • 学生网络从同一图像的不同视角（多尺度全局剪裁 + 小视窗随机剪裁）中提取特征，分别计算图像级和补丁级对比损失

ViT Backbone 的核心流程与组件：

  • 输入图像切分为固定大小补丁（patch），每个补丁经线性映射后附加位置编码。
  • 添加可学习的分类标记（[CLS]）作为全局图像表示。
  • 多层 Transformer Block（自注意力 + 前馈网络）堆叠进行信息交互
  1. Patch Embedding
    • 将输入图像切分为 patch_size×patch_size（默认为 14×14） 的小块
    • 通过线性投影映射到 embed_dim 维度，产生序列化的 patch 特征
  2. Class Token 与 Register Tokens
    • cls_token：添加可学习的分类标记（[CLS]）作为全局图像表示。
    • register_tokens（可选）：额外的 R 个可学习标记，用于增强不同任务的表达能力
  3. Positional Embedding 与插值
    • 固定长度位置编码，长度等于 num_patches + num_tokens
    • 支持对不同输入分辨率进行双三次插值，兼容可变图像大小
  4. Transformer 块
    • Memory-Efficient Attention：基于 xFormers 的高效实现，减少显存占用
    • Feed-Forward Network：可选 MLP 或 SwiGLU 结构
    • LayerScale：缩放残差输出
    • Stochastic Depth（Drop Path）：按线性衰减概率随机跳过层，实现正则化
  5. Mask Token 与 MIM 支持
    • mask_token：用于掩码图像建模时替换被遮挡的 patch
    • prepare_tokens_with_masks() 方法在输入序列中注入 mask token，配合 iBOT 预训练目标使用
  6. 最终归一化与 Head
    • 全局输出先经 LayerNorm，再送入 head（通常为空或 Identity），以便下游任务接入不同头部结构
flowchart LR
  A[输入图像] --> B(PatchEmbed)
  B --> C[加 cls_token / register_tokens]
  C --> D[加 Positional Embedding]
  D --> E{Transformer Blocks × depth}
  E --> F[LayerNorm]
  F --> G[输出序列: cls, registers, patches]

为了同时服务于图像级和补丁级对比学习，DINOv2 在学生网络末端设计了两个独立的 MLP 头（原型数量通常设置为 128k）：

  • DINO Head：针对 [CLS] token 输出一组“原型分数”，通过 Softmax + Sinkhorn-Knopp 中心化后，与教师端原型分布计算交叉熵损失，强化全局对比学习。
  • iBOT Head：针对被随机掩码（mask）的补丁 token，输出对应补丁的原型分数，采用相同 Softmax+中心化策略，与教师端可见补丁一一匹配，强化局部特征学习。

优化组件

  • Sinkhorn-Knopp 中心化：替换原有教师 Softmax 平滑策略，通过迭代保证原型分布在批次内的均衡，提升聚类一致性。
  • KoLeo 正则化：基于最近邻熵估计，鼓励批次内特征在空间中均匀分布，提升检索与分割等任务的表现。
  • LayerScale + Stochastic Depth：在高容量模型训练中有效缓解梯度不稳定和过拟合风险，可达 40% 的残差层丢弃率。
  • SwiGLU 前馈网络：相较于标准 MLP 提升非线性表达能力，配合 FlashAttention 自研加速版，进一步节省显存并提升 throughput。

实施细节

  • FlashAttention 精简版：在 ViT 的自注意力模块中使用高效、低显存的 Attention 实现，结合分片量化优化，适配 64 维/头硬件特性，确保 1.1B 参数模型能在 A100 40GB 上运行。
  • Sequence Packing：将多尺度剪裁的不同 token 序列拼接为单一长序列，使用块对角 Attention mask 保持跨序列隔离，减少多次前向所需计算量。
  • FSDP 混合精度：利用 PyTorch FSDP 将模型、优化器一阶/二阶矩储存在不同 GPU 上，并在通信时自动下采样梯度为 float16，节省约 50% 通信带宽，助力多节点线性伸缩。

数据

  • 精选“种子”集合：ImageNet-22K、Google Landmarks、若干细粒度分类、室内／室外分割、深度估计基准等
  • 未标注候选：1.3B+ 网络爬取图像，经过安全过滤、NSFW/人脸模糊、PCA 哈希自去重 → 744M 自去重后图像
  • 基于检索的“半监督”扩充
    • 将 curated 集合中每张图像用 ViT-H/16（自监督初步预训练）提取特征
    • 对于大集合（>1M）取 4 个最近邻、小集合则按簇采样，最终汇集 142M 张多样图像