NeurIPS 2024 Spotlight 轻微的预训练数据扰动如何造就更优秀的扩散模型

在人工智能领域，基础模型(Foundation Models)的预训练数据质量一直被认为至关重要。我们最新的研究《Slight Corruption in Pre-training Data Makes Better Diffusion Models》（轻微的预训练数据扰动如何造就更优秀的扩散模型），第一次全面系统地探讨了这些扰动对模型性能的影响，并提出了一种创新的方法来将数据的“瑕疵”转化为模型性能提升的“助力”。本文延续我们之前的工作[1]，发现预训练数据中适度的噪声也可以提升模型预训练和下游性能。这项研究成果又卡耐基梅隆大学，香港大学，RIKEN等共同发表于NeurIPS 2024，并获得了spotlight论文荣誉。

• 论文标题：Slight Corruption Makes Better Diffusion Models
• 论文链接：https://arxiv.org/pdf/2405.20494
• 论文模型及代码：https://huggingface.co/DiffusionNoise

研究背景

扩散模型（Diffusion Models）近年来在生成高质量图像、音频和视频方面取得了令人瞩目的成功，其背后的一大推动力是大规模的预训练数据。例如，著名的Stable Diffusion模型依赖于从互联网上爬取的数十亿对图文数据（如LAION-2B），这些数据不仅为模型提供了丰富的训练语料，还帮助其实现了高度多样化和条件化的生成能力。这种方法显著提升了扩散模型的生成效果，使其在艺术创作、内容生成等任务中表现出色。

然而，大规模数据的收集和处理并非毫无挑战。尽管这些数据集通常会经过严格的筛选和过滤，但误配对样本（condition corruption）依然不可避免，例如错误的图像标签、文本描述不匹配、甚至无意义的内容。这些误配对样本被认为是大规模数据不可避免的“副产品”，传统观点通常认为它们会削弱模型性能，增加生成的噪声和不一致性。

然而，我们的研究挑战了这一传统认知。通过系统性实验，我们发现，轻微的预训练数据扰动（例如2.5%-7.5%的条件误配对）不仅不会对模型性能造成显著负面影响，反而能够帮助扩散模型生成更高质量、更具多样性的内容。令人惊讶的是，这一发现与我们此前在CLIP和ImageNet分类器的研究结论高度一致：轻微的标签噪声在特定情况下可以帮助模型在域内任务（In-Domain Tasks）中提升性能[1]。

这一观察背后的潜在原因可能是，轻微的噪声引导模型利用更多的数据维度来学习特征分布，从而增加生成数据的多样性和质量。同时，噪声还能避免模型过度拟合数据中的偏差或冗余模式，使生成的分布更贴近真实数据分布。这一发现不仅为扩散模型的优化提供了新的视角，也对理解大规模数据预训练过程中的数据偏差提出了重要见解。

随着生成模型的广泛应用，我们的研究不仅为扩散模型的设计和训练带来了新思路，还为开发更健壮、更通用的生成系统提供了理论基础和实践方法。

实验发现

为了全面研究预训练数据扰动对扩散模型的影响，我们设计了系统性实验，涵盖预训练和下游任务两个部分。以下是实验的具体设置以及主要发现。

实验设置

模型：我们采用了多种扩散模型，包括Latent Diffusion Models (LDMs)、Diffusion Transformers (DiT-XL/2) 和 Latent Consistency Models (LCMs)，分别对类条件（class-conditional）和文本条件（text-conditional）的生成任务进行实验：

• LDM-4: Latent Diffusion Model，基于VQ-VAE编码的潜空间（down-sampling factor为4），图像分辨率为256×256。
• DiT-XL/2: 使用Transformer作为主干网络的扩散模型。
• LCM-v1.5: 针对快速推理优化的Latent Consistency Models，基于Stable Diffusion v1.5

预训练数据：我们在两个广泛使用的数据集上引入了合成扰动：

• ImageNet-1K (IN-1K): 为类条件生成任务引入标签翻转扰动（随机将标签替换为其他类别标签）。
• CC3M: 为文本条件生成任务引入文本替换扰动（随机交换图文对中的文本）。

扰动比例（corruption ratio, η）分别设置为0%、2.5%、5%、7.5%、10%、15%、20%。

预训练评价指标：

• 生成质量和保真度：使用Fréchet Inception Distance (FID) 和 Inception Score (IS)。
• 多样性和覆盖率：使用Precision-Recall (PR) 曲线和样本熵值（Sample Entropy）。
• 我们在Appendix里提供了其他metric的测试结果，如sFID, Memorization Ratio, RMD (sample complexity), L2 distance.

下游任务实验设置：预训练之后，我们研究了扩散模型在两个下游个性化任务中的表现：

• ControlNet：基于空间条件（如Canny边缘图和语义分割掩码）控制生成。
• T2I-Adapter：在文本条件生成中结合额外的输入模态以提高可控性。数据集：

我们选用ImageNet-100（IN-100）作为个性化任务的数据集，并使用OpenCV Canny检测器和SegmentAnything（SAM）生成边缘和分割掩码作为输入条件。

下游任务评价指标：与预训练一致，包括FID、IS、Precision 和 Recall。

Pre-training结果

• 轻微扰动提升生成质量和多样性

• 更高的质量：通过在不同程度扰动的条件下训练模型，我们发现2.5%-7.5%的扰动比例显著提高了生成图像的质量。例如，在IN-1K数据集上的类条件生成任务中，2.5%扰动的LDM-4模型，其FID从未扰动模型的9.44降至7.00。这表明轻微扰动帮助模型更好地捕捉特征分布，生成更接近真实分布的样本。
• 更丰富的多样性：轻微扰动同样显著提升了生成样本的多样性。以文本条件生成任务（CC3M数据集）为例，7.5%扰动的模型在CLIP Score上从19.85提升至17.93，同时熵值（Sample Entropy）和相对马氏距离（Relative Mahalanobis Distance, RMD）也有所增加，反映出生成分布的覆盖范围更广。

• 扰动程度的非线性效应

• 提升-退化：随着扰动比例从0%逐步增加到20%，生成质量和多样性呈现出“先提升后下降”的非线性趋势。例如，在IN-1K数据集中，FID在扰动比例达到7.5%时达到最优（7.00），而当扰动比例增加到15%时，FID略有上升至9.18。这表明，过高的扰动可能会干扰模型捕捉有用的特征。
• 偏向于低噪声的设置：与以往研究的高噪声标签学习不同，我们的结果表明低比例（2.5%-7.5%）的扰动不仅是可行的，还可以成为优化生成模型性能的有效策略。

• 跨模型一致性

• 轻微扰动的效应在多种扩散模型上得到了验证，包括LDM-4、DiT-XL/2和LCM-v1.5。这表明这种方法具有广泛适用性。即使在Transformer架构，如DiT-XL/2中，轻微扰动也能提高生成能力。

• 视觉表现的显著改进

• 我们进一步通过生成样本的视觉评估验证了轻微扰动的优越性。例如，在文本条件生成任务中，7.5%扰动的模型生成的图像更加逼真且多样化，展示了更丰富的细节和场景变化。

Downstream结果

• 轻微扰动增强个性化生成性能

• ControlNet实验：在基于Canny边缘图和SAM分割掩码的个性化任务中，轻微扰动的预训练模型表现优异。这表明，预训练中的轻微扰动为下游任务提供了更通用和更具适应性的特征表示。
• T2I-Adapter实验：文本条件生成任务中，预训练扰动的优势也更加明显。

• 增强模型的多样性与保真度平衡

• 在下游任务中，我们观察到轻微扰动的模型不仅在FID指标上表现更优，还在Precision-Recall曲线上表现出更好的平衡性。例如，在IN-100的ControlNet实验中，2.5%扰动的模型在Precision和Recall上分别达到0.84和0.58，显著优于未扰动模型的0.82和0.61。这表明轻微扰动有效改善了模型生成的覆盖率，同时维持了较高的精确性。

• 下游任务的视觉评估改进

• 我们通过对生成图像的视觉质量进行对比，发现轻微扰动的模型能够生成更加自然且多样化的图像。

• 高效适配的能力

• 我们发现轻微扰动的预训练模型在进行下游任务的微调时具有更高的适应性。例如，与未扰动模型相比，轻微扰动的模型在少量标注数据的情况下，仍能实现显著的性能提升。这说明轻微扰动不仅能提升预训练阶段的表现，还能为下游任务提供更具泛化性的初始权重。

理论分析

我们通过理论分析揭示了轻微数据扰动如何在生成质量和多样性上影响扩散模型的性能。本部分分为两部分：生成多样性的提升和生成质量的改善，并结合我们的数学建模和分析，提供直观的解释和数学依据。

生成多样性：轻微扰动扩展生成分布

核心观察：
轻微扰动能有效提升生成样本的多样性，使生成分布覆盖更多的数据空间。我们通过计算生成分布的熵值（Entropy）作为多样性的度量，高熵表示生成分布更广泛地覆盖数据空间，而低熵则意味着生成分布过于集中。

理论证明：
在生成多样性方面，我们提出并证明了以下定理：

定理 1：
对于任意类别 ，在生成长度 足够大的情况下，若训练数据的经验协方差矩阵为满秩，且类别的期望范数 为常数，则有：

• 和 分别为带有扰动和无扰动条件的生成分布的熵值；
• 为扰动控制参数；
• 为数据维度。

这一结论表明，条件扰动能够增加生成分布的熵值，提升生成样本的多样性。此外，熵值的提升程度与扰动大小 和数据维度 成正比。这也与我们实验中发现的结果一致，即轻微扰动（ 小且适中）能够提升多样性，而过大的扰动会导致分布失衡。

我们在不同扰动比例下观察到样本熵值的变化，如IN-1K和CC3M数据集中，2.5%-7.5%扰动的模型的熵值显著高于未扰动的模型，进一步验证了理论结论。

生成质量：轻微扰动优化分布对齐

核心观察：
轻微扰动能够改善生成分布与真实数据分布的对齐程度，进而提升生成样本的质量。我们采用2-Wasserstein距离作为衡量分布对齐程度的指标，距离越小，生成分布越接近真实数据分布，生成样本质量越高。

理论证明：
在生成质量方面，我们提出并证明了以下定理：

定理 2：
对于任意类别 ，在生成长度 足够大的情况下，若类别的期望范数 为常数，假设 ，则有：

其中：

• 为真实分布， 和 分别为无扰动和有扰动条件下的生成分布；
• 为2－Wasserstein距离；
• 为类别 的样本数量；
• 为数据维度。

这一结果表明，轻微的条件扰动能够使生成分布更接近真实分布，即生成质量有所提升。此外，当取适当的值时（如实验中的2.5%-7.5%扰动），生成分布的对齐效果最为显著，进一步验证了我们的实验结论。

在IN-1K和CC3M数据集中，我们观察到，随着轻微扰动的引入，FID显著下降，表明生成分布与真实分布的对齐程度得到了改善。

Conditional Embedding Pertubation

轻微扰动在预训练数据中展现出提升模型性能的潜力，但直接操作预训练数据可能并不总是现实可行。为此，我们提出了一种创新的扰动方法，称为Conditional Embedding Perturbation (CEP) ，通过对条件嵌入层（conditional embedding layer）直接添加随机扰动，模拟数据中的轻微噪声，从而提升模型的生成质量和多样性。

核心思想

CEP的核心思想是，在预训练和下游任务的训练过程中，对模型的条件嵌入（例如文本嵌入或类别嵌入）添加轻微扰动，而无需修改底层的数据或模型结构。具体而言，假设条件嵌入为，我们将其调整为：

其中：

• 是服从正态分布或均匀分布的随机噪声；
• 是控制扰动幅度的超参数。。

通过在条件嵌入中加入轻微扰动，CEP可以有效模拟预训练数据中的轻微噪声，同时避免了直接修改数据可能带来的额外开销或不一致性。

实验验证

我们在多个生成任务和数据集上评估了CEP的有效性，包括IN-1K的类条件生成任务和CC3M的文本条件生成任务。以下是CEP的具体表现：

1. 生成质量和生成多样性的提升: 在IN-1K数据集上，使用CEP的模型（扰动幅度 (\gamma=0.05)）的FID从未加扰动模型的9.44显著降低，显示了更高的生成质量。Precision-Recall曲线显示，CEP不仅提升了生成的多样性（Recall），还维持了高精确度（Precision）。

2. 对不同模型的通用性: CEP在LDM、DiT和LCM等多种扩散模型中均表现出一致的性能提升，证明了该方法的通用性。

• 例如，在LDM-4模型上，CEP将FID从9.44降至7.01，而在DiT-XL/2模型上，CEP将FID从6.76降至5.51。

3. 对下游任务的适用性: 在下游任务中，CEP对生成性能的提升尤为显著。例如，在基于ControlNet的个性化生成任务中，使用CEP预训练的模型，其FID从11.59进一步降低至11.14，同时多样性指标（熵值和Recall）也显著提高。类似地，在T2I-Adapter任务中，CEP的模型生成的图像更加逼真，且对额外的条件输入（如边缘图或分割掩码）的响应更精确。

未来展望

随着生成模型和预训练数据规模的不断增长，理解数据扰动对模型性能的影响变得尤为重要。我们的理论分析不仅揭示了轻微扰动的潜在优势，也为未来研究更复杂的数据偏差和噪声类型奠定了基础。此外，我们相信，这种分析方法还可以推广到其他生成模型（如GAN和自回归模型），从而为生成任务的优化提供更广泛的指导。CEP提供了一种高效且灵活的方式来模拟和利用数据中的轻微噪声。未来，我们计划将CEP推广到更多类型的生成任务中，例如对话生成、语言生成，以及多模态生成任务。同时，我们还将探索如何动态调整使CEP能够适应不同数据分布和任务需求，进一步提升模型性能和鲁棒性。

References

[1] Chen, H., Wang, J., Shah, A., Tao, R., Wei, H., Xie, X., Sugiyama, M. and Raj, B., 2023. Understanding and mitigating the label noise in pre-training on downstream tasks. ICLR 2024.