只要一塊6年前的2080Ti，就能做大模型數據蒸餾？

來自上交大EPIC實驗室等機構的一項最新研究，提出了一種新的數據集蒸餾方法——NFCM。

與前SOTA相比，新方法的顯存占用只有1/300，并且速度提升了20倍，相關論文獲得了CVPR滿分。

NCFM引入了一個輔助的神經網絡，將數據集蒸餾重新表述為一個極小化極大（minmax）優化問題。

在多個基準數據集上，NCFM都取得了顯著的性能提升，并展現出可擴展性。

在CIFAR數據集上，NCFM只需2GB左右的GPU內存就能實現無損的數據集蒸餾，用2080Ti即可實現。

并且，NCFM在連續學習、神經架構搜索等下游任務上也展現了優異的性能。

將數據蒸餾轉化為minmax優化

NCFM的核心是引入了一個新的分布差異度量NCFD，并將數據集蒸餾問題轉化為一個minmax優化問題。

通過交替優化合成數據以最小化NCFD，以及優化采樣網絡以最大化NCFD，NCFM在提升合成數據質量的同時，不斷增強分布差異度量的敏感性和有效性。

特征提取與頻率參數采樣

NCFM的第一步，是進行特征提取，也就是從真實數據集和合成數據集中分別采樣一批數據，并將其輸入到特征提取網絡中。

特征提取網絡將原始數據從像素空間映射到一個特征空間，得到對應的特征表示，目的是提取數據的高層語義特征，為后續的分布匹配做準備。

特征提取網絡可以是一個預訓練的模型，也可以是一個隨機初始化的模型，這里NCFM采用了一種混合方式。

接下來，NCFM引入了一個輕量級的神經網絡作為采樣網絡，它接受一個隨機噪聲作為輸入，輸出一組頻率參數。

這些頻率參數將用于對特征函數（Characteristic Function，CF）進行采樣。

特征函數計算與分布差異度量

對于每一個頻率參數，將其與特征表示進行內積運算，然后取復指數，就得到了對應的CF值。

這兩個CF值都是復數，其中實部刻畫了數據在該頻率上的分布范圍，捕捉分布的散度或多樣性；虛部則反映了數據在該頻率上的分布中心，捕捉分布的典型性或真實性。

通過比較真實數據和合成數據的CF值，就可以全面地度量它們在特征空間上的分布差異。

為了定量地度量真實數據和合成數據之間的分布差異，NCFM引入了一個稱為神經特征函數差異（Neural Characteristic Function Discrepancy，NCFD）的度量。

NCFD綜合考慮了所有采樣頻率上的CF差異，將其匯總為一個標量值。NCFD越小，說明兩個分布越接近；NCFD越大，說明兩個分布差異越大。

minmax優化

有了NCFD這個分布差異度量，NCFM的優化目標就很清晰了——

最小化NCFD，使得合成數據和真實數據的分布盡可能接近；同時，望最大化NCFD對合成數據的敏感度，使之能夠準確反映合成數據的變化。

為了同時實現這兩個目標，NCFM引入了一個minmax優化框架：

• 在極小化階段，固定采樣網絡的參數，調整合成數據，目標是最小化NCFD。這一步使得合成數據向真實數據分布不斷靠攏。
• 在極大化階段，固定合成數據，調整采樣網絡的參數，目標是最大化NCFD。這一步使得NCFD對合成數據的差異更加敏感，提升其作為差異度量的有效性。

通過交替進行極小化階段和極大化階段的優化，NCFM不斷改進合成數據的質量，同時也不斷強化NCFD度量的敏感性和準確性。

模型微調與標簽生成

為了進一步提升合成數據的質量，NCFM在優化過程中還引入了兩個額外的步驟——模型微調和標簽生成。

• 在模型微調階段，NCFM用合成數據微調特征提取網絡，使其更加適應合成數據的特征分布，從而進一步縮小合成數據和真實數據之間的特征差異，提高合成數據的真實性；
• 在標簽生成階段，用一個預訓練的教師模型來為合成數據生成軟標簽。軟標簽提供了更加豐富和細粒度的監督信息，可以指導合成數據更好地模仿真實數據的類別分布，提高合成數據的多樣性。

一塊2080Ti搞定CIFAR實驗

相比于此前方法，NCFM在多個數據集上實現了顯著的性能提升。

在CIFAR-10、CIFAR-100、等數據集中上，NCFM在每類1/10/50張圖片的情況下的測試精度均超過了所有baseline方法。

在ImageNet的各個子集上，NCFM也展現了卓越的性能。

例如在ImageNette上，每類10張圖片時，NCFM達到了77.6%的測試精度，比現有最佳方法（RDED）高出14.4個百分點；

在ImageSquawk上，每類10張圖片時，NCFM達到了72.8%的測試精度，比現有最佳方法（MTT）高出20.5個百分點。

在性能提升的同時，NCFM還實現了大量的速度提升和資源節約。

在CIFAR-100上，NCFM每輪迭代的平均訓練時間比TESLA快了29.4倍，GPU內存消耗僅為TESLA的1/23.3（每類50張圖片）；

在Tiny ImageNet上，NCFM每輪迭代的平均訓練時間比TESLA快了12.8倍，GPU內存消耗僅為TESLA的1/10.7（每類10張圖片）。

并且，NCFM在CIFAR-10和CIFAR-100上實現了無損的數據集蒸餾，僅使用了約2GB的GPU內存，使得CIFAR上的所有實驗都可以在一塊2080Ti上進行。

此外，NCFM生成的合成數據在跨模型泛化能力上超過了現有方法。

例如在CIFAR-10上，用NCFM生成的合成數據訓練AlexNet、VGG和ResNet，都取得了比現有方法更高的測試精度。

來自上交大“最年輕博導”課題組

本文第一作者，是上交大人工智能學院EPIC實驗室博士生王少博。

王少博本科就讀于哈工大軟件工程專業，專業排名第一名；然后在上交大讀研，導師是嚴駿馳教授，研究方向為深度學習理論和可解釋性機器學習，其間專業排名第二。

現在王少博正在張林峰助理教授負責的EPIC實驗室讀博，研究方向為“高效、可解釋的深度學習和”大模型。

王少博現在的導師張林峰，是本文的通訊作者。

張林峰出生于1997年，去年6月在清華叉院取得博士學位，然后到上交大人工智能學院擔任博導并負責EPIC實驗室，時年僅有27歲。

同時，張林峰還在NeurIPS、ICML、ICLR、CVPR等頂級學術會議當中擔任審稿人。

張林峰還曾到香港科技大學（廣州）擔任訪問助理教授，他的邀請人胡旭明同樣是一名年輕博導，并且也參與了本項目。

此外還有EPIC實驗室的其他成員，以及來自上海AI實驗室的學者，亦參與了NFCM的研究。

論文地址：https://github.com/gszfwsb/NCFM/blob/main/asset/paper.pdf
GitHub倉庫：https://github.com/gszfwsb/NCFM