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

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

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

NCFM引入了一個輔助的神經(jīng)網(wǎng)絡，將數(shù)據(jù)集蒸餾重新表述為一個極小化極大（minmax）優(yōu)化問題。

在多個基準數(shù)據(jù)集上，NCFM都取得了顯著的性能提升，并展現(xiàn)出可擴展性。

在CIFAR數(shù)據(jù)集上，NCFM只需2GB左右的GPU內(nèi)存就能實現(xiàn)無損的數(shù)據(jù)集蒸餾，用2080Ti即可實現(xiàn)。

并且，NCFM在連續(xù)學習、神經(jīng)架構搜索等下游任務上也展現(xiàn)了優(yōu)異的性能。

將數(shù)據(jù)蒸餾轉化為minmax優(yōu)化

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

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

特征提取與頻率參數(shù)采樣

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

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

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

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

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

特征函數(shù)計算與分布差異度量

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

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

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

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

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

minmax優(yōu)化

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

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

為了同時實現(xiàn)這兩個目標，NCFM引入了一個minmax優(yōu)化框架：

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

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

模型微調(diào)與標簽生成

為了進一步提升合成數(shù)據(jù)的質量，NCFM在優(yōu)化過程中還引入了兩個額外的步驟——模型微調(diào)和標簽生成。

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

一塊2080Ti搞定CIFAR實驗

相比于此前方法，NCFM在多個數(shù)據(jù)集上實現(xiàn)了顯著的性能提升。

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

在ImageNet的各個子集上，NCFM也展現(xiàn)了卓越的性能。

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

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

在性能提升的同時，NCFM還實現(xiàn)了大量的速度提升和資源節(jié)約。

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

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

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

此外，NCFM生成的合成數(shù)據(jù)在跨模型泛化能力上超過了現(xiàn)有方法。

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

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

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

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

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

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

張林峰出生于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