Path-Guided Flow Matching for Dataset Distillation¶
会议: NeurIPS 2026
arXiv: 2602.05616
OpenReview: https://openreview.net/forum?id=mZhtTOUDIu
领域: 优化/理论
关键词: 数据集蒸馏, Flow Matching, 原型引导, 连续时间路径控制, 生成式蒸馏
一句话总结¶
本文提出首个基于流匹配的生成式数据集蒸馏框架 PGFM,通过在 VAE 隐空间引入带信任域约束和早停门控的轻量级原型路径引导,仅需少量 ODE 求解步数即实现了高效、高多样性且保真的紧凑数据集合成。
研究背景与动机¶
数据集蒸馏旨在将海量训练集压缩为每个类别仅含少量样本(低 IPC)的合成集合,使下游模型在此紧凑数据上训练能达到接近全量数据集的泛化性能。然而,传统基于双层优化或轨迹匹配的方法存在显存开销极大、难以扩展到高分辨率等瓶颈。近年来,以 MinimaxDiffusion、MGD3 为代表的扩散生成模型成为高分辨率数据集蒸馏的主流范式,通过在隐空间采样多样化代表性实例来拟合真实数据分布。
然而,扩散模型严重依赖复杂的离散噪声调度方案、分类器自由引导(CFG)以及启发式原型约束,不仅采样耗时严重,而且反向去噪轨迹天然具有随机波动与不稳定性。在严苛的低 IPC 预算下,过度依赖启发式引导极易引发轨迹扰动,导致合成样本过平滑或偏离真实流形;同时数十至上百步的迭代推理使得数据合成成本居高不下,严重制约了生成式蒸馏的实用扩展性。
流匹配(Flow Matching)作为新兴的连续时间生成范式,通过学习从先验分布到数据分布的连续速度场,将生成过程转化为确定性常微分方程(ODE)数值积分。相比扩散模型,流匹配具有直线化传输路径、采样步数极少(仅需数步至数十步)且动力学平滑稳定的天然优势。然而,纯流匹配采样缺乏针对下游分类任务的多样性与模态覆盖控制机制。本文的核心 idea 是:将确定性流匹配引入数据集蒸馏,在预训练流模型速度场中叠加早停门控与动态信任域约束的轻量级隐空间原型路径引导,实现高速、高模态覆盖且细节保真的紧凑数据蒸馏。
方法详解¶
整体框架¶
PGFM 的整体流水线包含三个紧密衔接的阶段:隐空间预处理与原型发现、带路径引导的流匹配 ODE 采样、以及面向解码的稳定隐变量选择。首先,利用预训练冻结的 VAE 编码器将类别图像映射至紧凑隐空间,并通过 \(K\)-means++ 聚类提取出能够覆盖类内关键模态的类别原型;随后,从标准高斯噪声出发,利用预训练的高斯混合流匹配(GMFlow)速度场进行 ODE 积分,在前中期采样阶段叠加指向指定原型的归一化引导速度场,并配合早停门控与自适应信任域约束;最后,在 ODE 终点与单步去噪预测隐变量之间挑选批次方差最低的稳定表征,送入冻结 VAE 解码器重构出高分辨率合成图像。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入数据与类内潜变量编码<br/>VAE 隐空间预处理"] --> B["类别隐空间模态发现与原型构建"]
B --> C["原型定向偏置的流匹配热启动"]
C --> D["自适应信任域与早停的路径引导流采样"]
D --> E["批次离散度最小化的稳定隐变量选择"]
E --> F["冻结 VAE 解码生成紧凑蒸馏数据集"]
关键设计¶
1. 类别隐空间模态发现与原型构建:紧凑捕获类内分布核心模式 针对直接在像素空间优化计算量庞大且难以捕捉语义模式的问题,PGFM 采用预训练冻结的 VAE 编码器 \(E\) 将训练集图像映射至隐空间均值表示 \(\bar{z}(x) = \mathbb{E}_{q_E(z|x)}[z]\),并利用全局均值和方差对潜变量进行标准化。为了在低 IPC 预算(例如 \(\text{IPC}=10, 20\))下充分代表类内多峰分布,避免样本坍缩到单一的类别平均中心,系统在每个类别 \(y\) 的隐样本集合 \(\mathcal{Z}_y\) 上运行 \(K\)-means++ 聚类(聚类簇数 \(K = \text{IPC}\)),将得到的聚类中心 \(\{\mu_{y,k}\}_{k=1}^K\) 作为后续引导的确定性目标原型。实验表明,直接选用均值质心而非离质心最近的个别样本,能够有效滤除单样本的实例级噪声与伪影,为流线积分提供平滑而稳健的宏观模态着陆点。
2. 原型定向偏置的流匹配热启动:加速轨迹早期对齐与模态分配 流匹配从先验高斯噪声出发积分时,若初始噪声完全随机,可能会导致部分采样轨迹与预分配的目标原型在相空间中距离过远,加剧后续控制的纠偏负担。为此,PGFM 设计了轻量级的热启动(Warm Start)初始化策略: $\(z_0 \leftarrow (1 - \eta_{\mathrm{init}})\,z_0 + \eta_{\mathrm{init}}\,\mu_{y,k}\)$ 其中 \(\eta_{\mathrm{init}}\) 为适度偏置超参数(在 ImageNette 上通常设为 \(0.09\))。该机制在初始阶段将采样点轻微拉向指定原型的吸引盆地(basin of attraction),确保各采样实例在进入 ODE 积分时便具备明确的模态归属感。消融实验证实,热启动能够将原型命中率从 33% 跃升至 78%,直接解决了生成样本在类内各模态之间分配失衡的核心痛点。
3. 自适应信任域与早停的路径引导流采样:平衡宏观拓扑收敛与微观细节保真 在连续时间流积分过程中,若施加过强或全程存在的外加引导力,会破坏预训练速度场原有的光滑流形结构,导致生成的潜在表征脱离解码器有效流形,产生严重的模糊或灰暗伪影。PGFM 提出了带有自适应信任域缩放与早停门控的受控速度场方程: $\(\frac{dz(t)}{dt} = u_{\phi}(t, z(t), y) + g(t)\,\tilde{u}_{\mathrm{proto}}(t, z(t), y, k)\)$ 其中原型控制速度项定义为当前时刻模型预测的最终纯净隐变量 \(\hat{z}_1(t)\) 与目标原型之间的位移:\(u_{\mathrm{proto}} = \lambda (\mu_{y,k} - \hat{z}_1(t))\)。为了严格限制外加控制相对于模型内生动力学的比重,设计了信任域缩放因子 \(\alpha(t)\): $\(\alpha(t) = \min\left(1,\; \rho(t) \frac{\|u_{\phi}(t, z(t), y)\|_2}{\|u_{\mathrm{proto}}(t, z(t), y, k)\|_2 + c}\right)\)$ 该约束确保了外加控制范数始终被限制在内生流场范数的 \(\rho(t)\) 倍之内(\(\rho(t)\) 随积分进度自 \(\rho_0=0.5\) 线性衰减至 \(\rho_{\min}=0.1\))。同时,引入时间门控 \(g(t) = \mathbb{I}[s(t) \le s_{\mathrm{end}}]\)(默认 \(s_{\mathrm{end}}=0.6\)),仅在积分前半程执行宏观模态牵引,后半程关闭引导、完全交由未扰动的骨干流场进行细粒度高频纹理的确定性恢复。
4. 批次离散度最小化的稳定隐变量选择:消除采样终点微小数值震荡 在有限数值积分步长(如 48 步 Heun 解法器)下,数值求解器在末步的状态 \(z_1\) 与网络接口在倒数步通过单步去噪直接预测的终点 \(\hat{z}_1\) 之间可能存在细微的数值残差。为了防止异常样本产生伪影,PGFM 在最终解码前对比两种候选表征在批次维度上的范数标准差,动态挑选分散度更低、数值分布更稳定的表征进行最终重构: $\(\ell^{\star} = \arg\min_{\ell \in \{z_1, \hat{z}_1\}} \mathrm{Std}_{b \in [B]}\big(\|\ell_b\|_2\big)\)$ 选出最优表征 \(\ell^{\star}\) 后,经 VAE 缩放系数调整输入冻结解码器 \(D\),恢复为无伪影的清晰图像。
实验关键数据¶
主实验¶
在 ImageNette、ImageIDC(分辨率 \(256 \times 256\))和更具挑战性的 ImageNet-100、ImageNet-1K 上,将 PGFM 生成的蒸馏数据集用于从头训练 ConvNet-6、ResNet-18 和 ResNet-AP(ResNet with Average Pooling),与无引导流匹配(FM)、经典分布匹配(DM)、扩散基线(DiT、MiniMaxDiffusion、MGD3)进行了深入对比。
| 数据集 | 评估网络 | IPC 预算 | PGFM (本文) | 前序SOTA (MGD3) | 基础流匹配 (FM) | 相对前序SOTA提升 |
|---|---|---|---|---|---|---|
| ImageNette | ConvNet-6 | 10 | 63.0 ± 1.1 | 58.5 ± 2.1 | 60.5 ± 1.1 | +4.5% |
| ImageNette | ConvNet-6 | 20 | 67.6 ± 1.1 | 64.3 ± 1.1 | 65.3 ± 2.1 | +3.3% |
| ImageNette | ConvNet-6 | 50 | 80.4 ± 1.0 | 78.1 ± 2.1 | 77.7 ± 1.1 | +2.3% |
| ImageNette | ResNet-18 | 10 | 64.0 ± 1.1 | 62.8 ± 2.1 | 61.9 ± 0.7 | +1.2% |
| ImageNette | ResNet-18 | 20 | 70.2 ± 0.7 | 68.1 ± 2.1 | 67.5 ± 1.1 | +2.1% |
| ImageNette | ResNet-AP | 10 | 66.0 ± 1.1 | 65.8 ± 2.4 | 62.4 ± 1.3 | +0.2% |
| ImageIDC | ConvNet-6 | 10 | 53.6 ± 1.1 | 49.7 ± 1.5 | 51.5 ± 1.2 | +3.9% |
| ImageIDC | ResNet-18 | 20 | 61.3 ± 1.0 | 59.3 ± 1.5 | 58.8 ± 1.4 | +2.0% |
| ImageNet-100 | ConvNet-6 | 10 | 25.7 ± 0.4 | 23.4 ± 0.9 | 24.9 ± 0.4 | +2.3% |
| ImageNet-100 | ResNet-18 | 10 | 26.5 ± 0.4 | 23.8 ± 0.5 | 24.8 ± 0.7 | +2.7% |
| ImageNet-100 | ResNetAP-10 | 10 | 27.2 ± 0.3 | 25.7 ± 0.7 | 25.5 ± 0.4 | +1.5% |
| ImageNet-1K | ResNet-18 | 10 | 22.8 | 20.5 | 19.8 | +2.3% |
| ImageNet-1K | ResNet-50 | 10 | 23.8 | 22.5 | 21.4 | +1.3% |
在计算开销方面,PGFM 展现出显著的轻量化优势:单图合成仅需 2.27 TFLOPs 与 32 步 NFE(模型参数 542.77M),相较 MiniMaxDiff(23.73 TFLOPs,50 步,674.93M)实现了 10.5× 的效率提升,相较 MGD3(17.30 TFLOPs,50 步)实现了 7.6× 的合成速度提升。
消融实验¶
在 ImageNette(\(\text{IPC}=10\))上对 PGFM 各核心组件及初始化偏置进行的详细消融结果如下:
| 配置 / 组合变体 | 路径控制 (Path C.) | 早停门控 (Early S.) | 信任域约束 (Trust R.) | 热启动 (Warm S.) | ConvNet-6 准确率 (%) | ResNet-18 准确率 (%) | 原型命中率 (Hit Rate %) |
|---|---|---|---|---|---|---|---|
| 纯无引导流匹配 (FM) | – | – | – | – | 60.5 ± 2.1 | 62.4 ± 2.3 | 10% |
| 仅路径控制 | ✓ | – | – | – | 61.6 ± 1.2 | 63.0 ± 1.5 | 18% |
| 路径控制 + 早停 | ✓ | ✓ | – | – | 62.1 ± 1.3 | 64.8 ± 2.0 | 27% |
| 完整受控采样 (无热启动) | ✓ | ✓ | ✓ | – | 62.1 ± 1.3 | 63.5 ± 1.6 | 33% |
| PGFM 全要素模型 | ✓ | ✓ | ✓ | ✓ | 63.0 ± 1.1 | 66.0 ± 1.1 | 78% |
| 聚类选用最邻近样本点替代质心 | ✓ | ✓ | ✓ | ✓ | 61.1 ± 1.2 | 61.3 ± 1.4 | - |
针对热启动偏置系数 \(\eta_{\mathrm{init}}\) 的独立敏感性实验表明:当 \(\eta_{\mathrm{init}}=0\) 时,Hit Rate 仅为 33%(ResNet-AP 准确率 64.8%);适度提升至 \(\eta_{\mathrm{init}}=0.09\) 时达到最佳平衡,Hit Rate 跃升至 78%,ResNet-AP 达到最高的 66.0%;若进一步增大至 \(0.15\),虽然 Hit Rate 升至 86%,但样本过度集中于质心导致多样性下降,ResNet-AP 准确率反而回落至 64.2%。
关键发现¶
- 纯流匹配已是极强基线:在 ImageNet-100 上,未经任何引导的纯 FM 在 ConvNet-6(24.9%)和 ResNet-18(24.8%)上的表现已全面超越精调扩散基线 MGD3(23.4%、23.8%),证明流匹配积分产生的确定性轨迹天生更贴近真实数据流形。
- 早停与信任域是抗伪影生命线:若不加信任域限制或全程开启路径控制,外力会彻底扭曲速度场,导致合成图像严重泛灰泛白(参见原文图 5);引入 \(\alpha(t)\) 范数截断与 \(s_{\mathrm{end}}=0.6\) 早停后,模型在确保拓扑着陆的同时让基础流场恢复清晰边缘与纹理。
- 热启动提供进入吸引盆地的关键动量:原型命中率从 33% 飞跃至 78% 的核心使能因素在于 \(\eta_{\mathrm{init}}\);仅靠过程中的微分方程导向力不足以扭转远离原型的初值轨迹,初值上的轻度先验注能至关重要。
亮点与洞察¶
- 将流匹配范式成功拓展至数据集蒸馏:跳出了既往扩散模型堆叠反向去噪启发式技巧的窠臼,充分利用 ODE 确定性采样的平滑与高速特性,将合成效率提高近一个数量级。
- 动态信任域设计精妙:将外加引导速度强制绑定在基础速度场强度的比例上限内(\(\rho(t)\) 线性衰减),避免了传统分类器引导或梯度引导经常出现的“梯度爆炸”及流形漂移失真现象。
- 可复用的可控 ODE 引导机制:该“热启动初始化 + 前中期信任域速度叠加 + 后期纯净自由弛豫”的设计模式,可直接迁移推广到条件生成、可控流匹配修复及带硬约束的科学计算反问题中。
局限与展望¶
- 依赖预训练骨干网络与隐空间表达:PGFM 采用冻结的 VAE 和 GMFlow 权重,若预训练模型在特定长尾类别或特殊专业领域表征能力受限,蒸馏质量将面临刚性上限。
- 超参数需按数据集规模微调:热启动系数 \(\eta_{\mathrm{init}}\)、引导强度 \(\lambda\) 和信任域衰减阈值在 ImageNette 与 ImageNet-100 之间仍需人工微调,尚未实现完全自适应的数据驱动自调节。
- 未来方向:探索结合类别长尾特性的分层/密度敏感聚类原型构建方法;将下游分类网络的梯度反馈直接作为流匹配的代理强化信号,进一步实现端到端任务驱动的数据集紧凑蒸馏。
相关工作与启发¶
- vs MGD3 (Chan-Santiago et al., 2025):MGD3 采用扩散模型与反向去噪过程中的模式引导,虽然提升了低 IPC 下的多样性,但采样步数多、FLOPs 极高且命中率仅 30%;PGFM 换用流匹配 ODE 积分,步数更少、能耗降低 7.6×,并将原型命中率提升至 78%。
- vs MiniMaxDiffusion (Gu et al., 2024):MiniMaxDiffusion 需要在蒸馏阶段显式微调扩散模型的参数以最小化最大均值差异,训练和采样成本极其昂贵;PGFM 无需微调任何生成网络权重,纯粹在推理阶段通过零样本控制流场即可取得更优泛化性能。
- vs GMFlow (Chen et al., 2025a):GMFlow 仅作为通用的无监督/条件生成模型;PGFM 首次将其与数据集蒸馏的紧凑代表性需求相结合,提出了针对离散聚类原型的轨迹控制理论与算法实现。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ [首次将流匹配 ODE 理论与路径控制引入生成式数据集蒸馏,打破扩散垄断]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 ImageNette、ImageIDC、ImageNet-100/1K,包含跨模型迁移、消融及细致的命中率评测]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密,图文对照自洽,数学推导清晰且动机阐述透彻]
- 价值: ⭐⭐⭐⭐⭐ [极大降低了高分辨率高表现数据集蒸馏的生成算力门槛,实用价值显著]