跳转至

CleanDIFT: Diffusion Features without Noise

会议: CVPR 2025
arXiv: 2412.03439
代码: https://compvis.github.io/cleandift
领域: 扩散模型 / 自监督学习
关键词: 扩散特征, 无噪声特征提取, 语义对应, 知识蒸馏, 时间步无关

一句话总结

提出 CleanDIFT,通过轻量级无监督微调(单卡 A100 仅 30 分钟),使扩散模型直接在干净图像上提取高质量语义特征,消除了传统方法需要加噪和调时间步的限制,在语义对应、深度估计、分割等多任务上显著超越标准扩散特征。

研究背景与动机

领域现状:预训练扩散模型的内部特征(diffusion features)已被证明是强大的语义描述子,在语义对应匹配、分割、检测等任务上表现出色。然而,提取这些特征需要先对图像加噪声再输入模型,因为扩散模型在训练时就是处理带噪图像的。

现有痛点:加噪声带来三个问题。(1)信息损失:噪声破坏了图像的高频细节,即使在 t=261 的低噪声下,重建图像也已损失大量信息;(2)特征编码噪声:实验证明即使在低 t 值下,扩散特征中相当一部分方差来自噪声而非图像信息;(3)时间步依赖:不同下游任务需要不同的最优 t 值(语义对应用 t=261,分割用 t=100),需要逐任务调参。一些方法通过多噪声集成缓解(如 DIFT 平均 8 次),但计算开销成倍增加。

核心矛盾:扩散模型在不同噪声水平下学到了不同语义层次的信息(高噪声→粗结构,低噪声→精细细节),这提供了多尺度特征但也造成了时间步依赖。需要将所有时间步的信息整合到单次无噪声前向传播中。

本文目标 如何从扩散模型中提取无噪声、时间步无关的通用语义特征,同时保留或超越有噪声多时间步特征的性能。

切入角度:将扩散模型视为 T 个不同的特征提取器家族(每个时间步一个)。训练一个学生模型在干净图像上输出特征,通过时间步条件投影头对齐到所有 T 个教师特征提取器,将多时间步信息整合到单一特征空间。

核心 idea:用一个在干净图像上运行的特征提取模型,通过时间步条件投影头将其特征与冻结扩散模型在所有噪声水平下的特征对齐,实现无噪声时间步无关的统一特征提取。

方法详解

整体框架

以预训练扩散模型(SD 1.5/2.1)为教师,初始化一个可训练副本作为学生特征提取器。学生接收干净图像(无噪声、无时间步输入),教师接收对应的加噪图像和时间步。在 K=11 个特征提取位置(U-Net 中间块 + 解码器块),通过时间步条件投影头将学生特征投影到教师特征空间,最大化余弦相似度。训练仅需 400 步/30 分钟。推理时丢弃投影头,直接使用学生模型的内部特征。

关键设计

  1. 时间步条件投影头(Timestep-Conditioned Projection Heads)

    • 功能:将学生模型的时间步无关特征对齐到教师在不同噪声水平下的时间步特定特征
    • 核心思路:每个特征提取位置 k 有一个独立的投影头 \(\text{proj}^{(k)}(\text{feat}_c^{(k)}(\mathbf{x}_0); t)\),由 3 层零初始化的 FFN 组成(因残差连接初始为恒等映射)。投影头以时间步 t 作为额外输入,学习将统一特征映射到对应时间步的特征空间。损失为 \(\mathcal{L} = -\sum_{k=1}^{K} \text{sim}(\text{proj}^{(k)}(\text{feat}_c^{(k)}(\mathbf{x}_0); t), \text{feat}^{(k)}(\mathbf{x}_t; t))\)
    • 设计动机:不同时间步的扩散特征具有不同语义特性,投影头将学生的统一特征分别映射到各时间步空间,迫使学生特征包含所有时间步的信息
  2. 分层时间步采样(Stratified Timestep Sampling)

    • 功能:每张训练图像高效覆盖完整噪声频谱
    • 核心思路:将 [0, T] 分为 I=3 个均匀区间,每区间内均匀采样一个时间步 \(t_i \sim \mathcal{U}(\frac{i}{I}T, \frac{i+1}{I}T)\)。每张图像计算 3 个不同噪声水平的对齐损失,共享同一次学生前向传播
    • 设计动机:确保训练覆盖低/中/高三个噪声范围,避免偏向某一噪声区域
  3. 噪声-特征关系分析

    • 功能:提供理论动机——证明标准扩散特征中存在大量噪声编码
    • 核心思路:通过线性回归分析扩散特征中纯噪声 \(\text{feat}(\epsilon; T)\) 对完整特征 \(\text{feat}(\mathbf{x}_t; t)\) 的方差解释比例。发现即使在 t=261,噪声已解释了相当大比例的特征方差。进一步发现 t=0 的干净图像特征也无法完全解释剩余方差——说明中间时间步有独特的语义信息(不在 t=0 或 t=T 中)
    • 设计动机:直接给扩散模型输入干净图像(t=0)不好用,因为模型没在干净图像上训练过。需要专门的微调来让模型在无噪声输入下也能输出有用特征

损失函数 / 训练策略

余弦相似度对齐损失,K=11 个特征位置求和。全量微调学生模型 + 投影头(45M 参数)。Adam 优化器,学习率 2e-6,线性预热,批量 8。在 COYO-700M 随机子集上训练(与 SD 训练数据同分布,确保提升来自方法而非数据)。仅 400 步即收敛。

实验关键数据

主实验

任务/方法 标准扩散特征 CleanDIFT 提升
语义对应 DIFT PCK_img 66.53 68.32 +1.79
语义对应 TaleOfTwo PCK_img 72.31 73.35 +1.04
语义对应 TellLeftRight PCK_img 77.07 78.40 +1.33
深度估计 NYUv2 RMSE↓ 0.469 0.444 -0.025
语义分割 VOC mIOU 最优 t 值结果 超越所有 t
分类 ImageNet kNN 最优 t 值结果 略超

消融实验

配置 PCK_img 说明
DIFT (8次噪声集成) 66.53 8× 计算开销
DIFT (单次噪声) 65.51 基线
CleanDIFT (无噪声单次) 68.32 1× 开销超越 8× 集成
监督微调 DHF (50步反演) 72.75 50× 开销
CleanDIFT + 监督微调 72.48 1× 开销接近 50×

关键发现

  • 无噪声特征在所有任务上一致优于最优时间步的有噪声特征,证明噪声确实损害特征质量
  • 不能简单地给扩散模型输入干净图像 + 非零 t(性能更差),需要专门微调
  • 提升在高级方法上仍然成立(+DINOv2 融合、+姿态对齐,都有 1+ 点提升)
  • 30 分钟微调、无需集成的 1× 计算开销,效率提升 8-50×
  • 可作为标准扩散特征的直接替代品(drop-in replacement)

亮点与洞察

  • 噪声是特征质量的瓶颈这一洞察虽然事后看显而易见,但之前没有被系统验证。方差分解实验清楚展示了噪声在特征中的占比
  • 整合 T 个特征提取器到一个的思路非常优雅——通过时间步条件投影头实现"一次前向传播,所有时间步的知识"
  • 极低训练成本(400 步/30 分钟/单卡)使方法高度实用,任何使用扩散特征的工作都可以直接受益

局限与展望

  • 仅在 SD 1.5/2.1 (U-Net) 上验证,对 DiT 架构验证较少(虽有 Tab. 7 初步结果)
  • 对 SDXL、Flux 等更大模型的适用性未充分探索
  • 投影头增加了 45M 参数(虽推理时丢弃)
  • 分类任务提升较小,说明池化后噪声影响减弱,主要优势在密集预测任务

相关工作与启发

  • vs DIFT: DIFT 需加噪 + 8 次集成。CleanDIFT 无噪声单次前向即超越,8× 加速 + 更好性能
  • vs A Tale of Two Features / Telling Left Right: 这些方法融合 DINOv2 和扩散特征。替换为 CleanDIFT 后持续提升,说明改进源于特征本身质量提升
  • vs DHF (监督微调): DHF 需 50 步 DDIM 反演。CleanDIFT 单次前向即接近其性能,50× 加速

评分

  • 新颖性: ⭐⭐⭐⭐ 洞察深刻(噪声损害特征),方法优雅(多时间步整合),但技术上是标准蒸馏
  • 实验充分度: ⭐⭐⭐⭐⭐ 四种下游任务+多种方法集成+详细消融+方差分析,非常全面
  • 写作质量: ⭐⭐⭐⭐⭐ 动机分析(方差分解图)极具说服力,实验组织清晰
  • 价值: ⭐⭐⭐⭐⭐ 所有使用扩散特征的工作都可直接受益,低训练成本使其高度可部署

相关论文