ImageNet-1K壓縮20倍，Top-1精度首超60%：大規模數據集蒸餾轉折點

該工作是目前唯一實現了大規模高分辨率數據集蒸餾的框架

過去幾年，數據壓縮或蒸餾任務引起了人們的廣泛關注。通過將大規模數據集壓縮成具有代表性的緊湊子集，數據壓縮方法有助于實現模型的快速訓練和數據的高效存儲，同時保留原始數據集中的重要信息。數據壓縮在研究和應用中的重要性不可低估，因為它在處理大量數據的過程中起著關鍵作用。通過采用先進的算法，數據壓縮取得了顯著的進展。然而，現有解決方案主要擅長壓縮低分辨率的小數據集，這種局限性是因為在雙層優化過程中執行大量未展開的迭代會導致計算開銷巨大。

MBZUAI 和 CMU 團隊的最新工作 SRe2L 致力于解決這一問題。該工作是目前唯一實現了大規模高分辨率數據集蒸餾的框架，可以將 Imagenet-1K 原始的 1.2M 數據樣本壓縮到 0.05M (壓縮比 1:20)，使用常用的 224x224 分辨率進行蒸餾，在 ImageNet-1K 標準驗證集（val set）上取得了目前最高的 60.8% Top-1 精度，遠超之前所有 SOTA 方法，如 TESLA (ICML’23) 的 27.9% 的精度。

該工作目前已完全開源，包括蒸餾后的數據，蒸餾過程和訓練代碼。

論文：https://arxiv.org/abs/2306.13092

代碼：https://github.com/VILA-Lab/SRe2L

數據集蒸餾 / 壓縮任務的定義和難點

傳統的模型蒸餾是為了得到一個更加緊湊的模型，同時保證模型性能盡可能得高。與之不同，數據集蒸餾任務關注于如何得到一個更緊湊同時更具表達能力的壓縮后的數據集，數據樣本相比原始數據集會少很多（節省從頭訓練模型的計算開銷），同時模型在該壓縮后的數據集上訓練，在原始數據驗證集上測試依然可以得到較好的精度。

數據集蒸餾任務的主要難點在于如何設計一個生成算法來高效可行地生成需要的樣本，生成的樣本需要包含 / 保留原始數據集中核心的信息。目前比較常用的方法包括梯度匹配、特征匹配、軌跡匹配等等，但是這些方法的一個共同缺點就是沒法 scale-up 到大規模數據集上。比如，由于計算量和 GPU 顯存的限制，無法蒸餾標準的 ImageNet-1K 或者更大的數據集。計算量和 GPU 顯存需要過大的主要原因在于這些方法生成過程需要匹配和保存的信息過多，目前很多 GPU 顯存沒法容納所有需要匹配的數據信息，導致這些方法大多數只適用于較小的數據集。

針對這些問題，新論文通過解耦數據生成和模型訓練兩個步驟，提出了一個三階段數據集蒸餾算法，蒸餾生成新數據過程只依賴于在原始數據集上預訓練好的模型，極大地降低了計算量和顯存需求。

解決方案核心思路

之前很多數據集蒸餾方法都是圍繞樣本生成和模型訓練的雙層優化 (bi-level optimization) 來展開，或者根據模型參數軌跡匹配 (trajectory matching) 來生成壓縮后的數據。這些方法最大的局限在于可擴展性不是很強，需要的顯存消耗和計算量都很大，沒法很好地擴展到完整的 ImageNet-1K 或者更大的數據集上。

針對這些問題，本文作者提出了解耦數據生成和模型訓練的方法，讓原始數據信息提取過程和生成數據過程相互獨立，這樣既避開了更多的內存需求，同時也避免了如果同時處理原始數據和生成數據導致原始數據中的噪聲對生成數據造成偏差 (bias)。

具體來說，本文提出了一種新的數據集壓縮框架，稱為擠壓、恢復和重新標記 (SRe2L)，如下圖所示，該框架在訓練過程中解耦模型和合成數據雙層優化為兩個獨立的操作，從而可以處理不同規模的數據集、不同模型架構和高圖像分辨率，以實現有效的數據集壓縮目的。

本文提出的方法展示了在不同數據集規模的靈活性，并在多個方面表現出多種優勢：1）合成圖像的任意分辨率，2）高分辨率下的低訓練成本和內存消耗，以及 3）擴展到任意評估網絡結構的能力。本文在 Tiny-ImageNet 和 ImageNet-1K 數據集上進行了大量實驗，并展示出非常優異的性能。

三階段數據集蒸餾框架

本文提出一個三階段數據集蒸餾的框架：

• 第一步是將整個數據集的核心信息壓縮進一個模型之中，通過模型參數來存儲原始數據集中的信息，類似于我們通常進行的模型訓練；

• 第二步是將這些高度抽象化的信息從訓練好的模型參數中恢復出來，本文討論了多種不同損失和正則函數對于恢復后圖像的質量以及對數據集蒸餾任務的影響；

• 第三步也是提升最大的一步：對生成的數據進行類別標簽重新校準。此處作者采用了 FKD 的方式，生成每個 crop 對應的 soft label，并作為數據集新的標簽存儲起來。

三階段過程如下圖所示：

性能及計算能效比

在 50 IPC 下 (每個類 50 張圖)，本文提出的方法在 Tiny-ImageNet 和 ImageNet-1K 上實現了目前最高的 42.5% 和 60.8% 的 Top-1 準確率，分別比之前最好方法高出 14.5% 和 32.9%。

此外，本文提出的方法在速度上也比 MTT 快大約 52 倍 (ConvNet-4) 和 16 倍 (ResNet-18)，并且在數據合成過程中內存需求更少，相比 MTT 方法分別減少了 11.6 倍 (ConvNet-4) 和 6.4 倍 (ResNet-18)，具體比較如下表所示:

實驗結果

實驗設置

該工作主要聚焦于大規模數據集蒸餾，因此選用了 ImageNet-Tiny 和 ImageNet-1K 兩個相對較大的數據集進行實驗。對于骨干網絡，本文采用 ResNet-{18, 50, 101} 、ViT-Tiny 和自己構建的 BN-ViT-Tiny 作為目標模型結構。對于測試階段，跟之前工作相同，文本通過從頭開始訓練模型來評估壓縮后數據集的質量，并報告 ImageNet-Tiny 和 ImageNet-1K 原始驗證集上的測試準確性。

在 full ImageNet-1K 數據集上的結果

可以看到，在相同 IPC 情況下，本文實驗結果遠超之前方法 TESLA。同時，對于該方法蒸餾得到的數據集，當模型結構越大，訓練得到的精度越高，體現了很好的一致性和擴展能力。

下圖是性能對比的可視化結果，可以看到：對于之前方法 TESLA 蒸餾得到的數據集，當模型越大，性能反而越低，這對于大規模數據集蒸餾是一個不好的情況。與之相反，本文提出的方法，模型越大，精度越高，更符合常理和實際應用需求。

壓縮后的數據可視化

從上圖可以看到，相比于 MTT 生成的數據（第一和第三行），本文生成的數據（第二和第四行）不管是質量、清晰度還是語義信息，都明顯更高。

蒸餾過程圖像生成動畫

此外，包含 50、200 個 IPC（具有 4K 恢復預算）的壓縮數據集文件可從以下鏈接獲取：https://zeyuanyin.github.io/projects/SRe2L/

將該方法擴展到持續學習任務上的結果

上圖展示了 5 步和 10 步的增量學習策略，將 200 個類別（Tiny-ImageNet）分為 5 個或 10 個學習步驟，每步分別容納 40 個和 20 個類別。可以看到本文的結果明顯優于基線（baseline）性能。

更多細節歡迎閱讀其論文原文和代碼。