最近的研究强调了扩散模型与表征学习之间的相互作用。扩散模型的中间表征可用于下游视觉任务，同时视觉模型表征能够提升扩散模型的收敛速度和生成质量。然而，由于输入不匹配和 VAE 潜在空间的使用，将视觉模型的预训练权重迁移到扩散模型中仍然具有挑战性。
为了解决这些问题，来自高德地图的研究者提出了统一自监督预训练（USP, Unified Self-Supervised Pretraining），该方法通过在变分自编码器（VAE）的潜在空间中进行潜在掩码建模（Masked Latent Modeling）预训练，预训练得到的 ViT 编码器等权重可以无缝迁移到下游任务，包括图像分类、语义分割以及基于扩散模型的图像生成。
论文名称：USP: Unified Self-Supervised Pretraining for Image Generation and Understanding
论文地址：pdf/2503.06132
代码地址：cxxgtxy/USP
USP 在理解任务上实现有竞争力的表现；在生成任务中，USP 显著加速 DiT 和 SiT 模型收敛速度，比从头训练快 11.7 倍（DiT-XL）和 46.6 倍（SiT-XL）。
研究背景
在过去十年中，预训练 - 微调（pretraining-finetuning）范式在图像识别任务中取得了显著成功。但在图像生成领域，该范式的探索仍然有限。
DDAE 等近期的研究表明，生成模型不仅可以合成高质量图像，还可以学习到优越的视觉表征，图像理解和图像生成之间存在着深层次的联系。
例如，iGPT（Image GPT）探索了基于像素空间的自回归预训练，但该方法在大规模数据集和模型上扩展时面临着严重的计算成本问题。此外，这种方法与扩散模型不兼容。
为了弥合这一差距，REPA 提出通过对齐扩散模型与预训练的视觉模型（如 DINOv2）的表征，可以让扩散模型更高效地学习判别特征，从而提高训练效率和生成质量。然而，REPA 存在几个主要问题：
1. 高昂的计算成本：REPA 依赖于 DINOv2 这样的大规模预训练视觉模型，其预训练需要超过 22,000 GPU 小时（A100），计算资源需求极高。
2. 额外的教师网络（Teacher Network）：使用 DINOv2 作为教师模型会增加 GPU 显存消耗，同时降低扩散模型的训练速度。
尽管已有研究揭示了生成和理解任务间的联系，但仍然有一些关键问题尚未解决：
1. 预训练是否对扩散模型的训练是可行且必要的？
2. 是否可以找到一种同时适用于生成和理解任务的预训练方法？
3. 现有的 “预训练 - 微调” 范式是否能成功应用于生成模型？
该论文旨在提出一种简单而有效的方法来解决这些问题。
方法设计
实现能够同时适用于图像理解和生成的统一预训练 - 微调范式面临诸多挑战：
C1: 输入不匹配：图像理解模型通常接收干净的图像作为输入，而扩散模型接受的是添加了噪声的图像。
C2: 结构不匹配：生成模型多为基于 VAE 的潜空间扩散模型，而大多数图像理解任务并不使用 VAE。此外， ViT 结构在图像生成任务中通常会进行修改。
C3: 损失函数和标签格式不同：图像理解任务和图像生成任务通常采用不同的优化目标，这使得直接共享预训练模型变得困难。
尽管面临这些挑战，研究者也观察到了一些有利的现象：
P1: 神经网络对噪声具有鲁棒性：预训练的视觉模型在噪声或数据增强下仍然可以保持较高的分类精度，例如在 ImageNet-C 数据集上测试。说明即使扩散模型处理的是加噪图像，预训练模型仍可以学习到有效的特征。
P2: 扩散模型可以学习到判别性特征：扩散模型能学习到用于图像分类等任务的判别性特征。如果能够有效地对齐这些表征，扩散模型的收敛速度和最终性能都可以得到显著提升。
P3: ViT 结构具有较强的适应性：尽管 ViT 在应用到扩散模型时经历了一定修改（如 AdaLN-Zero 层归一化和额外的条件输入）。但如果设计得当，这些修改仍然可以与 ViT 的预训练权重兼容。
P4: VAE 具有强大的压缩和重建能力：扩散模型中使用的 VAE（如 SD-VAE）能够有效地保留原始图像的重要信息。即使在 VAE 的潜空间中进行训练，仍然可以获得高质量的视觉特征。
基于以上观察，本文的研究者提出了统一的自监督预训练架构，见下图 1：
USP 架构基于一个简单的自编码器（Autoencoder），但在 VAE 潜空间中进行操作，而非像素空间。输入图像首先经过 VAE 编码到潜空间，并通过 PatchConv 进行图片分块。部分块按照设定 mask 比例被随机掩码，未掩码的块输入到 ViT 编码器，而解码器负责重建掩码块，损失函数仅使用简单的 MSE loss。在预训练阶段，VAE 参数被冻结，仅训练 ViT 编码器。预训练完成后，ViT 编码器的权重可用于初始化下游任务，如分类、分割和生成。
在将预训练模型权重适配到下游理解和生成模型时，针对图像分类任务，ViT 编码器的预训练权重可直接继承，无需额外调整，且仍然使用 Class Token 作为最终表征。
对于生成任务，由于 DiT 和 SiT 的结构与 ViT 略有不同，对初始化策略进行了优化。首先，在 AdaLN-Zero 层归一化中，恢复可训练的偏置（β）和缩放因子（γ），使其与预训练的 ViT 权重对齐。其次，由于预训练是在 224×224 进行，而 ImageNet 生成任务通常在 256×256 进行，因此本文采用 Bicubic Interpolation 扩展位置编码。最后，由于生成任务不需要 class token，在 DiT/SiT 中直接将其移除。这种初始化策略确保了 ViT 预训练权重能够无缝适配到下游分类和生成任务，而不引入额外计算开销或存储需求。
实验设置
本论文涵盖三种模型规模，见表 1。预训练阶段所有模型共享相同的解码器 —— 由 8 个 Transformer 块组成。
在额外实验中，将预训练时长扩展到 1600 轮，以证明 USP 在更长预训练时间上的可扩展性。为了与 MAE 进行公平比较，本文在 224×224 分辨率上进行预训练，尽管消融实验表明更高的分辨率可以带来更好的性能。
图像生成实验
本文在两种基于 Transformer 的扩散模型，DiT 和 SiT 上验证了 USP。评估均在 ImageNet 256×256 上进行，使用 50,000 个样本，不使用 CFG。
表 2 显示了在不同规模的 DiT 模型上的对比结果。USP 在所有模型规模上均显著提升了生成质量，且随着训练时间延长，生成质量不断提高。相比最近的 DiT 变体在 2.5M 步的 FID，USP 仅在 400K 内就能达到更好的效果。
表 3 显示了在不同规模的 SiT 模型上的对比结果。USP 表现出和 DiT 一致的提升效果。同时，表 3 与近期利用表征对齐来提升 DiT/SiT 性能的方法进行了比较，USP 在所有模型设置下均优于其他方法。
图像理解实验
论文在 ImageNet-1k 数据集上进行了线性探测（Linear Probe）和微调（Fine-tuning）的图像分类评估。在线性探测（LP）任务上，USP 的性能优于 MAE；在微调（SFT）任务上，USP 表现与 MAE 相当，表 5 总结了分类结果：
进一步，论文在 ADE20 数据集上进行了分割性能评估。表 6 显示了 USP 在单尺度 mIoU 指标上的表现，相比 MAE 提升了 0.5%。
消融实验
研究者进行了全面的消融实验，以探讨 USP 设计中不同组件的影响。例如 VAE、输入分辨率、掩蔽率等。更多的消融实验说明见原论文。
讨论
VAE 在图像理解任务中的作用
在图像分类任务中应用 VAE 并不是最理想的选择。因为 VAE 的核心目标是在保证重建能力的同时尽可能压缩信息，而原始图像本身是无损的，因此直接在原始图像上进行分类可能更高效。然而，我们的实验表明，如果使用高质量的 VAE 进行编码，图像分类任务的性能至少可以达到与标准方法相当的水平。
研究者认为，这种现象的主要原因是：
VAE 的潜空间编码仍然能保留足够的判别信息，即使经过压缩，仍能支持良好的分类表现。
VAE 提供了一种对抗噪声的方式，通过潜空间中的信息提取，模型可能学习到更鲁棒的特征。
工作机制（对比 REPA）
为了更深入地理解 USP，研究者在 DiT-XL/2 训练过程中，对不同层的线性探测性能进行了分析（见图 4）。比较了以下几种情况：
1.DiT-XL/2 预训练后的线性探测性能（“Pre.”）
2.DiT-XL/2 生成微调后的线性探测性能（“Ft.”）
3.SiT-XL/2 模型的线性探测性能
4.SiT-XL/2 在应用 REPA 方法后的线性探测性能
主要发现：
与 REPA 不同，USP 不依赖额外的对齐损失，而是通过精心设计的初始化策略，让模型自动找到最适合线性分类的层。
经过 40 万步训练后，USP 的第 20 层成为最优的线性分类层，这表明 USP 能够自适应地优化表征学习，在分类和生成任务之间找到平衡点。
REPA 通过人为设计的表征对齐方式来增强判别能力，但这种方法可能会限制生成模型的潜力。
这些实验表明，USP 的初始化策略比基于表征对齐的方法更高效，并且更适用于同时提升分类和生成任务的统一框架。
对基础范式的另一种视角
论文中使用经过 800 轮预训练且掩码比例为 0.75 的 ViT-Large 模型研究了图像修复任务。如图 5 所示，USP 在图像修复方面显著优于 MAE，凸显了强表征能力对有效修复的重要性。
这一结果与下面基于扩散的生成框架相契合，表明过度使用监督标签微调编码器以增强判别能力，并不会显著提升图像生成效果。
为了实证这一假设，论文中采用了一个监督微调模型，模型在 ImageNet 验证集上的准确率达 82.6%，并使用其初始化 DiT-B/2 进行 400K 步训练。表 8 总结了实验结果。其性能明显低于预训练，进一步验证了对该范式的分析。
0 条