💻 computer science

Geometry-Aware Dataset Condensation for Diffusion Model Training

本文提出了几何感知数据集压缩（GADC），该方法通过使用单侧部分最优传输和语义正则化，将真实子集选择重新表述为一个几何感知分布对齐问题，旨在构建能够保留有效扩散模型训练所需的几何结构和分布保真度的紧凑数据集。

原作者： Xiao Cui, Yulei Qin, Mo Zhu, Wengang Zhou, Hongsheng Li, Houqiang Li

发布于 2026-06-19

📖 1 分钟阅读☕ 轻松阅读

原作者： Xiao Cui, Yulei Qin, Mo Zhu, Wengang Zhou, Hongsheng Li, Houqiang Li

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你想教一位大师级厨师（一个扩散模型）如何烹饪出一顿完美的佳肴。传统上，你会给他们提供一个包含数百万份食谱和食材的海量图书馆（全量数据集）。虽然这种方法可行，但阅读所有内容需要耗费极长时间，成本高昂，且需要一个巨大的厨房来进行处理。

数据集压缩（Dataset Condensation） 的理念是创建一个仅包含几十份完美食谱的微型“速查表”，在没有冗余负担的情况下，教会厨师所需的一切知识。

然而，以往尝试制作这些速查表的尝试都面临两个大问题：

合成伪造： 有些人试图从零开始发明新的食谱。结果是？这些食谱看起来很怪异，吃起来感觉很假，让厨师感到困惑。
糟糕的选择： 另一些人则试图直接挑选现有的“最佳”食谱。但他们是基于单一、简单的评分（比如“这个食谱有多难？”）来进行挑选的。这忽略了大局，导致遗漏了重要的变化并产生了偏差。

本论文提出了一种构建这种速查表的新方法，称为几何感知数据集压缩（Geometry-Aware Dataset Condensation）。以下是其工作原理，使用简单的类比进行说明：

1. 问题所在：“一维化”错误

想象一下，完整的数据集是一个巨大的、复杂的城市，拥有街区、公园和摩天大楼（数据分布）。

旧方法 (D2C)： 这种方法试图通过基于“难度”的一条直线来对最好的房子进行排名。这就像是试图通过仅仅观察高度来挑选一个三维城市中的最佳房屋。你可能会选出一座高耸却空旷的摩天大楼，却错过了一间温馨且必不可少的农舍。你丢失了城市的形状。
目标： 我们需要挑选一小组能够完美代表整个城市整体形状的房屋，保持公园、街道和街区的完整性。

2. 解决方案：“单侧部分传输”

作者使用了一个名为**最优传输（Optimal Transport）**的数学工具，它就像一家物流公司，试图将货物从一个仓库（全量数据集）转移到一个较小的仓库（压缩后的子集）。

旧方式（平衡传输）： 旧规则规定：“你必须将大仓库里的每一粒沙子都搬运到小仓库里，并完美匹配重量。”
- 缺陷： 由于小仓库非常小，这迫使物流公司不得不把城市边缘（低密度区域）那些沉重且无用的沙子也拖到小仓库里，以填满重量配额。这扭曲了地图。
新方式（单侧部分传输）： 作者说：“我们只需要搬运重要的货物。我们不需要搬运那些空旷、低密度郊区的沙子。”
- 益处： 这使得小仓库能够完全专注于城市的“核心”——繁忙的街道和受欢迎的街区。它确保了小规模子集能够捕捉到原始数据的真实几何结构（形状和结构），而不会被噪声所拖累。

3. 安全网：“统计正则化”

仅仅搬运货物是不够的；我们需要确保新的仓库仍然感觉像原来的城市。作者添加了两个“安全网”：

均值-方差检查： 他们确保小仓库中建筑物的平均“高度”和“分布范围”与大城市相匹配。如果大城市有高有矮的建筑组合，那么小仓库也必须具备同样的组合。
置信度检查： 他们确保所选的房屋是清晰可辨的。如果一座房子看起来像个模糊的团块，分不清是谷仓还是车库，他们就会将其剔除。这确保了“厨师”不会被模棱两可的例子所迷惑。

4. 策略：“贪婪构建 + 交换优化”

你如何实际挑选出这些特定的房屋？你不能检查所有可能的组合（组合数量太多了！）。因此，他们使用了两步走的策略：

贪婪构建： 从一片空地开始，逐一添加房屋，每次都选择当前对地图提升最大的那一个。这就像是一块一块地拼凑拼图。
交换： 一旦拼图完成，他们就开始寻找错误。“嘿，这个角落里的房子不太行，让我们把它和外面的那个换一下。”他们不断进行交换，直到地图变得尽可能完美。

实验结果

当他们在 ImageNet（一个包含 140 万张图像的庞大数据库）上测试这种方法以训练 AI 图像生成器时：

更好的质量： 与之前的方法相比，他们生成的图像看起来更清晰、更多样化（具有更低的“FID”分数）。
高效性： 他们仅使用 0.8% 的原始数据（10,000 张图像而非 140 万张）即可进行训练，并且效果甚至优于使用原始数据的随机切片。
速度： 选择这 10,000 张图像的过程比之前的方法要快得多。

总结：
这篇论文告诉我们，要用一个小数据集来训练强大的 AI，你不应该仅仅挑选“最难”或“最容易”的样本。相反，你应该通过数学手段挑选出一组微小的图像，它们能完美地保留原始庞大数据集的形状、结构和多样性，同时忽略掉那些空旷、嘈杂的边缘。这就像是在策划一场博物馆展览，仅用一个房间就捕捉到了整个艺术史收藏的灵魂。

技术摘要：面向扩散模型训练的几何感知数据集凝聚

1. 问题陈述

数据集凝聚旨在通过合成或选择，从真实数据中构建紧凑的数据集，以降低与深度学习模型训练相关的存储和计算成本。虽然现有的凝聚方法在判别式任务（如分类）中取得了成功，但它们并不适用于训练扩散模型。

合成局限性： 通过连续像素优化生成合成数据的方法往往会产生低保真度的样本。扩散模型对噪声和结构失真高度敏感，这使得这些合成样本不适合进行真实的建模。
选择局限性： 现有的真实子集选择方法通常依赖于固定标准或启发式排名（例如，D2C 根据一维轴上的扩散难度对图像进行排名）。这些方法无法保留扩散似然目标所需的复杂分布几何结构。它们通常缺乏与扩散训练对齐的原则性优化目标，导致选出的子集与全量数据分布对齐不良，且无法捕捉必要的流形结构。

核心挑战在于如何选择一个紧凑的真实子集，使其能够保留完整数据集的几何结构和分布保真度，特别是针对扩散模型的似然训练目标进行优化。

2. 方法论

作者提出了几何感知数据集凝聚（GADC），将真实子集选择重新表述为一个几何感知的分布对齐问题。该框架由三个主要部分组成：

A. 单侧部分最优传输 (One-Sided Partial Optimal Transport, POT)
为了解决小规模子集与全量数据集之间的容量不匹配问题，作者采用了单侧部分最优传输。

概念： 与强制执行全质量对齐的平衡 OT 不同，单侧 POT 允许源质量（选定的子集）被完全传输，同时在容量约束下放宽对目标质量（全量数据集）的要求。
机制： 这允许在低密度边缘区域存在未匹配的质量，从而将传输集中在高密度、几何稳定的数据流形主导区域。
实现： 该问题通过“虚拟源（dummy-source）”重构法和熵正则化进行高效求解，并通过 Sinkhorn 迭代进行优化。这实现了灵活的质量分配，使子集能够专注于代表性区域，同时保持必要的覆盖范围。

B. 统计与语义正则化
为了补充几何对齐并确保分布保真度，引入了两个轻量级正则化项：

均值-方差正则化 ( $L_{sta}$ )： 对齐所选子集与全量数据集之间特征表示的一阶和二阶统计量（均值和方差），以保持全局分布保真度。
置信度正则化 ( $L_{conf}$ )： 通过惩罚预测类别概率较低的样本，确保语义一致性。这可以防止引入可能损害对齐效果的不可靠几何锚点。

C. 两阶段离散优化策略
由于固定大小子集的选择是一个组合优化问题，作者提出了一个高效的两阶段求解器：

第一阶段（贪婪几何引导选择）： 一个增量构建过程，通过选择样本来最小化复合目标函数（POT 损失 + 正则化项），建立广泛的流形覆盖。
第二阶段（基于交换的精炼）： 一个精炼阶段，通过在已选样本和未选样本之间进行成对交换，来纠正早期由于局部最优导致的决策，从而进一步优化全局几何对齐和一致性。

3. 核心贡献

问题重构： 本文将扩散模型的数据集凝聚重新定义为分布对齐问题，而非依赖于启发式排名或基于标量的选择。它引入了带有统计正则化的单侧 POT 目标，以选择具有代表性的真实子集。
优化框架： 提出了一种两阶段离散优化框架（贪婪构建结合基于交换的精炼），能够高效解决对齐问题，避免了固定准则或基于排名的采样所带来的局限性。
实验性能： 大量实验表明，在各种扩散变体（DiT, SiT）、不同子集规模、图像分辨率和训练轮数下，该方法均取得了优于以往方法的保真度和分布覆盖度。

4. 实验结果

该方法在 ImageNet-1K 上使用 DiT-L/2 和 SiT-L/2 架构进行了评估，数据预算范围从 0.8% (10K 图像) 到 8% (100K 图像)。

性能指标： 与包括 Random Sampling, K-Center, Herding, CCS, DQ, D2C 在内的基线方法相比，该方法始终实现了最低的 Fréchet Inception Distance (FID) 以及最高的 Inception Score (IS)、Precision 和 Recall。
- 示例： 在 ImageNet 256×256 分辨率、10K 数据预算 (0.8%) 下，该方法达到了 3.43 的 FID，优于 D2C (4.20) 和 Random (35.86)。
- 示例： 在 512×512 分辨率、10K 预算下，该方法达到了 6.17 的 FID，显著优于 D2C (14.8)。
鲁棒性： 这种改进在不同的扩散变体（DiT 和 SiT）、不同的评估协议（10K vs. 50K 生成样本）以及更长的训练迭代次数（高达 300K）中均保持稳定。
效率： 选择过程在计算上是高效的。在单块 GPU 上，该方法选择 ImageNet 中 10K 子集仅需 5.5 小时，而 D2C 需要 41.9 小时，DQ 需要 30.4 小时。

消融实验：

移除 POT 损失 ( $L_{OT}$ ) 会导致 FID 退化并降低 Precision/Recall，证实了几何对齐的必要性。
用平衡 OT 替换单侧 POT 会导致性能下降，验证了在边缘区域允许未匹配质量的益处。
移除统计 ( $L_{sta}$ ) 或置信度 ( $L_{conf}$ ) 正则化会导致 FID 或 IS 分别变差，强调了全局统计特性和语义清晰度的重要性。

5. 意义与主张

本文声称，通过将数据集凝聚视为几何感知的分布对齐问题，成功弥合了数据选择与扩散模型训练特定需求之间的差距。

数据效率： 该方法能够在凝聚的数据集上实现高质量的生成建模，显著降低了与扩散模型训练相关的存储和计算成本。
保真度： 与合成生成方法不同，该方法在确保子集与全量数据分布几何结构对齐的同时，保留了真实样本的高保真结构。
实用性： 所提出的两阶段优化策略具有可扩展性且计算高效，是资源受限环境下的实用解决方案。

作者总结道，其方法为扩散式生成训练提供了一个可靠的数据基质，为提升效率提供了一条与模型端优化相辅相成的、以数据为中心的路径。