新智元報道

編輯：LRST

【新智元導讀】數(shù)據(jù)集蒸餾是一種用少量合成數(shù)據(jù)替代全量數(shù)據(jù)訓練模型的技術，能讓模型高效又節(jié)能。WMDD和GUARD兩項研究分別解決了如何保留原始數(shù)據(jù)特性并提升模型對抗擾動能力的問題，使模型在少量數(shù)據(jù)上訓練時既準確又可靠。

數(shù)據(jù)集蒸餾（Dataset Distillation，DD）正在成為大模型時代一條務實的「降本增效」路徑：用每類極少的合成樣本訓練出接近使用全量數(shù)據(jù)的模型，從而顯著降低時間、算力與能耗。

真正的挑戰(zhàn)包括兩點：

其一，如何在極少的樣本中盡可能保留原始數(shù)據(jù)分布的幾何特性；

其二，如何在不顯著增加成本的情況下，使在DD上訓練的模型在對抗擾動下仍保持可靠。

加州大學圣地亞哥分校、伊利諾伊大學厄巴納-香檳分校等機構的研究人員在「ICCV 2025上發(fā)表的WMDD」和「AAAI 2025上提出的GUARD」分別針對這兩個問題提出了解決方案。

WMDD

用Wasserstein度量保留數(shù)據(jù)分布的幾何特性

數(shù)據(jù)蒸餾要做的，歸根到底是「少而不失真」。

WMDD（Wasserstein Metric-based Dataset Distillation）的出發(fā)點很直接：既然我們關心分布的幾何特性，就用最能保留形狀和結構的距離來對齊它——Wasserstein 距離，并以其在幾何意義上的「重心」（barycenter）作為蒸餾目標。

論文鏈接：https://arxiv.org/abs/2311.18531

代碼鏈接：https://github.com/Liu-Hy/WMDD

最優(yōu)傳輸（Optimal Transport, OT）的思想可追溯至法國數(shù)學家Monge，后由Kantorovich在二十世紀鑄成現(xiàn)代優(yōu)化理論的重要分支。

與之緊密相關的Wasserstein距離，刻畫的是把一個分布搬運成另一個分布所需的最小代價，因此能夠保留分布的幾何特性。

直觀地說，如果把數(shù)據(jù)分布想象成一團可以流動的「質量」，Wasserstein衡量的是把一團搬運成另一團的最小「功」：

帶著幾何去對齊每一份質量，而不是只盯住均值或協(xié)方差。相應的 Wasserstein 重心，也不是把若干分布簡單「糊」在一起，而是找到能同時接近它們的中心形態(tài)——在形狀上保留各自結構特征。這種「保形」的能力，正是數(shù)據(jù)蒸餾最需要的。

圖1 二維玩具分布的「重心」對比：原始分布由圓環(huán)（紅）與交叉（藍）兩類形狀組成，色深表示概率密度、色相區(qū)分源分布類型。KL與MMD的重心更像是僵硬的混合；Wasserstein重心以最小搬運代價在幾何上做自然插值，較好保留結構特征。

WMDD把合成數(shù)據(jù)的學習搬到特征空間中：先用在原始數(shù)據(jù)上訓練好的分類器，把每類樣本映射到最后一層線性分類器之前的特征空間，再在該空間分別計算各類的Wasserstein重心。

這里的重心不是一個點，而是一組帶權代表點，用盡量少的「錨」概括類內多樣性與幾何結構。

隨后，研究人員用梯度下降優(yōu)化每類少量合成圖像，使其特征貼近對應的重心代表點。

同時引入一個簡單但關鍵的正則：按類統(tǒng)計的BatchNorm約束（Per-Class BN，PCBN）。

不同于以往全局BN對齊容易讓類別間的梯度相互牽扯，PCBN在類內對齊各層BN的均值與方差，既利用了預訓練網(wǎng)絡的先驗，又不打亂類內結構的學習。

為保證可擴展性，研究人員采用高效的最優(yōu)傳輸求解，對重心的「位置—權重」交替優(yōu)化，計算與存儲開銷與主流的高效分布匹配式蒸餾方法相當。

圖2 WMDD方法流程：真實數(shù)據(jù)集T與合成集S先經(jīng)特征網(wǎng)絡f獲得特征；以真實特征計算按類的Wasserstein重心（含位置與權重）；在蒸餾過程中，以「特征匹配損失+按類BN統(tǒng)計（PCBN）正則」共同優(yōu)化合成樣本，使其對齊重心并生成更具代表性的合成數(shù)據(jù)。

不同于傳統(tǒng)的基于MMD的分布匹配式蒸餾，為什么這套方案能在高分辨率數(shù)據(jù)集上既可算、又好用？

從Wasserstein更能捕捉分布幾何的角度看，它帶來的「搬運」視角會讓優(yōu)化的梯度更有指向性：不是簡單拉近某個統(tǒng)計量，而是明確地把合成樣本朝真實數(shù)據(jù)的相應區(qū)域推近；而MMD的意義嚴重依賴核的選擇與帶寬設定，未必能穩(wěn)定地刻畫我們關心的幾何關系。

另外，從DD算法實際效率的角度，研究人員在附錄中給出了另一種解釋：完整的核化MMD計算開銷隨原始數(shù)據(jù)集的樣本量平方增長，很難在ImageNet-1K量級上落地。

許多基于MMD的蒸餾方法為了規(guī)?；?，退而只對齊特征均值（等同于使用線性核），這在統(tǒng)計上并不能唯一標定分布，往往會忽略高階差異與幾何結構，從而無法給優(yōu)化提供足夠準確的信號。

Wasserstein一方面具有理論保障，即在合理的Lipschitz假設下，真實分布與合成分布在期望風險上的差異可以被W1（Wasserstein-1）距離上界；

而另一方面，固定合成樣本量下單次OT求解的復雜度只隨原始數(shù)據(jù)樣本量線性增長，同時可以借助特征空間降維與高效OT求解保持可計算性。

實驗上，WMDD在ImageNette、Tiny-ImageNet與ImageNet-1K三個高分辨率數(shù)據(jù)集上，1/10/50/100 IPC的多種預算下都給出了強競爭甚至領先的結果。

以100 IPC為例，研究人員在三套數(shù)據(jù)上的top-1分別達到約87.1%、61.0%和60.7%，逼近用全量數(shù)據(jù)訓練的同架構模型（約89.9%、63.5%、63.1%）。

更值得一提的是跨架構泛化：研究人員用ResNet-18蒸餾得到的合成數(shù)據(jù)，在50IPC下訓練ResNet-50/101還能持續(xù)增益，遷移到ViT-Tiny/Small也有可觀表現(xiàn)，這說明WMDD學到的不是「過擬合某個骨干」的捷徑，而是真正在特征幾何上貼近了真實分布。

效率方面，得益于重心計算的可解性，時間與顯存開銷與當前最為高效的分布匹配方法處于同一數(shù)量級，卻在大多數(shù)設置下取得SOTA表現(xiàn)。

圖3 WMDD的合成圖像可視化：在ImageNet-1K的10IPC設置下（蒸餾與可視化使用ResNet-18），從若干選定類別中隨機各采樣1張，呈現(xiàn)貼近人類感知的關鍵類別特征。

研究人員還做了一系列消融來拆解設計取舍。首先，把特征匹配從交叉熵替換為Wasserstein重心回歸，在三個數(shù)據(jù)集上都帶來穩(wěn)定增益；其次，把全局BN正則改為PCBN后，類內多樣性保持得更好，特征分布不再「塌縮」，與重心匹配目標形成共振。

研究人員也嘗試了Sliced Wasserstein作為更快的替代，得到的精度幾乎不損，速度略有提升，說明WMDD的核心收益確實來自Wasserstein幾何，而非某個工程技巧或脆弱的超參。

從更宏觀的視角看，WMDD把數(shù)據(jù)蒸餾拉回「分布幾何」的基本圖景：用最小搬運代價貼近真實數(shù)據(jù)的形，既不丟掉類內的多樣性，也不破壞類間的相對關系。這種方式天然適合與預訓練表征協(xié)同，也為與生成式模型的結合打下基礎。

研究人員也在思考兩個自然的延伸：

其一，直接在聯(lián)合空間蒸餾P(X, Y)，將標簽嵌入表示空間并在聯(lián)合度量下對齊，有望同時捕捉標簽邊際分布與標簽間關系結構（例如標簽嵌入的相似性），從而改進跨類遷移與長尾類別的刻畫；

其二，「最優(yōu)」的合成數(shù)據(jù)未必就是全局最接近原始分布的那一組，決策邊界附近的代表性樣本可能更關鍵，而一種折中方案是讓部分樣本的位置介于「重心代表點」和「靠近判別邊界的支持點」之間，把最優(yōu)傳輸?shù)娜謱R與邊界敏感的判別信號結合起來。

在「保形」之外，可靠性同樣關鍵：當模型面對對抗擾動時，訓練在蒸餾數(shù)據(jù)上的模型還能扛得住嗎？這正是第二項工作所聚焦的方向。

GUARD

在蒸餾過程中平滑損失景觀以獲得對抗魯棒性

讓少量蒸餾數(shù)據(jù)上正常訓練的模型具有對抗魯棒性，最直接的想法，是把對抗訓練嵌入蒸餾流程；

但實踐顯示，這會顯著拉低干凈集性能，魯棒收益也并不穩(wěn)定。

研究人員因此轉向更本質的做法：不在外層堆疊昂貴的對抗環(huán)路，而是從損失景觀出發(fā)，直接在真實數(shù)據(jù)鄰域平滑模型的局部曲率。這就是GUARD（Geometric Regularization for Adversarially Robust Dataset）的核心。

論文鏈接：https://arxiv.org/abs/2403.10045

代碼鏈接：https://github.com/yumozi/GUARD

理論上，把樣本在小球擾動內的對抗損失近似展開，可以得到一個上界，其中主導項來自損失對輸入的最大曲率（Hessian最大特征值）。當蒸餾數(shù)據(jù)與真實分布在特征空間足夠接近時，真實數(shù)據(jù)與蒸餾數(shù)據(jù)上的對抗損失上界只相差一個與偏差σ成正比的常數(shù)項。

因此，「在蒸餾數(shù)據(jù)上優(yōu)化魯棒性」能夠可靠地遷移到「在真實數(shù)據(jù)上評測魯棒性」。

實現(xiàn)上，研究人員使用一個高效近似去降低最大曲率：借助一個經(jīng)驗事實，即神經(jīng)網(wǎng)絡輸入空間的梯度方向與最大曲率方向往往高度相似。

按照常用的近似手段，如果在單位梯度方向做小步長擾動，最小化兩處梯度的差異范數(shù)，就能降低最大曲率的近似，從而使損失景觀在該方向上更接近局部線性。

研究人員把該正則嵌入SRe2L的squeeze階段，將標準訓練損失替換為「原損失+曲率正則」，每步僅多一次前向與梯度計算，無需內層對抗環(huán)路，代價極低。

在ImageNette、Tiny-ImageNet與ImageNet-1K上，覆蓋10/50/100images per class（IPC）的多種設置，GUARD在多種白盒/黑盒攻擊（如PGD100、AutoAttack等）下普遍提升魯棒指標，并常?！疙槑А固嵘蓛艟?。

例如，在ImageNette 10 IPC上，干凈準確率從42.42%提至57.93%，AutoAttack從4.99%提至19.69%；在Tiny-ImageNet 50 IPC中，PGD100從0.27%提升到15.63%，AutoAttack從0.16%提升到13.84%；

當壓縮比例放寬到ImageNette 100 IPC，PGD100由31.65%提升到57.50%，AutoAttack由17.93%提升到64.84%

在更具挑戰(zhàn)的ImageNet-1K 10 IPC上，干凈準確率與魯棒性也都獲得顯著提升。

更重要的是，GUARD作為一種「損失景觀正則」的做法，具有良好的普適性：將其加入DC、CDA等不同蒸餾范式，同樣能同時提升干凈與魯棒性能。

圖4 GUARD的合成圖像可視化：在ImageNet-1K的1 IPC設置下，展示若干類別的蒸餾樣本（每類1張）。圖像呈現(xiàn)出類別相關的關鍵元素與混合外觀，體現(xiàn)蒸餾過程對主要語義特征的捕捉。

背后的直覺并不復雜：

如果教師模型在真實數(shù)據(jù)鄰域的損失景觀崎嶇、曲率大，那么最陡上升方向會頻繁改變，蒸餾過程就難以從教師的梯度場穩(wěn)定地恢復帶有類別核心特征的少量數(shù)據(jù)；

相反，降低最大曲率，能減少對抗擾動最敏感方向的「尖銳性」，有利于合成數(shù)據(jù)免受高頻噪聲的影響，學到真實穩(wěn)定的類別特征，從而同時提升魯棒性與干凈泛化。

研究人員在實驗中發(fā)現(xiàn)，在更小的IPC（例如1或10）下，曲率正則往往能給干凈準確率和魯棒性同時帶來顯著提升；

當IPC較大時，若正則過強則可能抑制必要的細粒度判別，此時適當減弱正則強度便能在保持魯棒收益的同時避免精度回落。

這似乎說明，合成數(shù)據(jù)量越小，能承載的模型復雜度就應越低。如何從理論角度闡釋這一現(xiàn)象，以及如何控制模型的復雜度（如曲率）使合成數(shù)據(jù)達到最優(yōu)泛化和魯棒性，有待進一步研究。

需要說明的是，研究人員并不把GUARD當作形式化的全局魯棒保證；若干假設（如局部凸性與特征映射的Lipschitz）是合理但理想化的。

然而，它以極低的額外開銷，在大規(guī)模數(shù)據(jù)與強攻擊設置下給出了兼顧精度與魯棒的均衡解。

收束與展望

這兩項工作分別回答了DD的兩個關鍵問題：如何在極低樣本預算下盡可能保留原始數(shù)據(jù)分布的重要幾何特性，以及如何在不顯著增加訓練成本的情況下獲得可遷移的對抗魯棒性。

更凝練地說，通過讓神經(jīng)網(wǎng)絡在蒸餾數(shù)據(jù)上學到更穩(wěn)健的表示，并在表示空間對數(shù)據(jù)分布的重要方面做更精準的匹配，DD有望在準確性與穩(wěn)健性之間取得兼顧。

向前看，值得探索的方向包括：在聯(lián)合空間直接蒸餾P(X, Y)，在復雜場景中將生成建模與重心表達結合以提升類內多樣性，對曲率正則進行更細粒度的自適應以適配不同的IPC，以及將這些思路推廣到分布外與多模態(tài)設置中。

希望「小而強」的數(shù)據(jù)集能成為走向可信與高效AI的一條可持續(xù)道路。

作者信息

劉昊洋是UIUC信息學博士生。主要研究興趣在于可信機器學習，數(shù)據(jù)集蒸餾，并探索大模型多智能體方法用于自動化科研與數(shù)據(jù)分析。

黎羿江在約翰斯?霍普金斯大學獲得計算機科學碩士學位，現(xiàn)在是加州大學圣地亞哥分校二年級博士生，主要研究方向聚焦于實現(xiàn)多模態(tài)、交互式以及三維具身環(huán)境下高效和魯棒的AI學習。

王沛然現(xiàn)在是加州大學洛杉磯分校一年級博士生，主要研究方向聚焦于AI安全、程序分析、基于AI的安全研究以及系統(tǒng)安全方法在大模型安全上的應用

邢天騁本科畢業(yè)于南京大學，現(xiàn)就讀于新加坡國立大學計算機碩士（人工智能方向）。目前主要方向集中在LLM可靠性，特別關注其知識邊界的識別與刻畫，從而減少并控制可能產(chǎn)生的有害或錯誤行為。

沈易凡在伊利諾伊大學香檳分校獲得計算機科學碩士學位，現(xiàn)在是伊利諾伊大學香檳分校的計算機系一年級博士生，主要研究多模態(tài)，強化學習，交互式視頻生成。

薛宇亮 (Eric) 本科畢業(yè)于多倫多大學，現(xiàn)就讀于哥倫比亞大學計算機科學碩士。主要研究方向聚焦于數(shù)據(jù)集蒸餾與大模型智能體，致力于探索在資源受限環(huán)境下實現(xiàn)樣本高效的學習，以及用智能體自動化機器學習的創(chuàng)新方法。

參考資料：

https://arxiv.org/abs/2311.18531

https://arxiv.org/abs/2403.10045