隻要一塊 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 實驗室讀博，研究方向爲 " 高效、可解釋的深度學習和 " 大模型。
王少博現在的導師張林峰，是本文的通訊作者。
同時，張林峰還在 NeurIPS、ICML、ICLR、CVPR 等頂級學術會議當中擔任審稿人。
張林峰還曾到香港科技大學（廣州）擔任訪問助理教授，他的邀請人胡旭明同樣是一名年輕博導，并且也參與了本項目。
此外還有 EPIC 實驗室的其他成員，以及來自上海 AI 實驗室的學者，亦參與了 NFCM 的研究。
論文地址：
https://github.com/gszfwsb/NCFM/blob/main/asset/paper.pdf
GitHub 倉庫：
https://github.com/gszfwsb/NCFM
