ImageNet-1K壓縮20倍，Top-1精度首超60%：大規(guī)模數(shù)據(jù)集蒸餾轉(zhuǎn)折點

該工作是目前唯一實現(xiàn)了大規(guī)模高分辨率數(shù)據(jù)集蒸餾的框架

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

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

該工作目前已完全開源，包括蒸餾后的數(shù)據(jù)，蒸餾過程和訓(xùn)練代碼。

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

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

數(shù)據(jù)集蒸餾 / 壓縮任務(wù)的定義和難點

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

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

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

解決方案核心思路

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

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

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

本文提出的方法展示了在不同數(shù)據(jù)集規(guī)模的靈活性，并在多個方面表現(xiàn)出多種優(yōu)勢：1）合成圖像的任意分辨率，2）高分辨率下的低訓(xùn)練成本和內(nèi)存消耗，以及 3）擴展到任意評估網(wǎng)絡(luò)結(jié)構(gòu)的能力。本文在 Tiny-ImageNet 和 ImageNet-1K 數(shù)據(jù)集上進行了大量實驗，并展示出非常優(yōu)異的性能。

三階段數(shù)據(jù)集蒸餾框架

本文提出一個三階段數(shù)據(jù)集蒸餾的框架：

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

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

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

三階段過程如下圖所示：

性能及計算能效比

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

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

實驗結(jié)果

實驗設(shè)置

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

在 full ImageNet-1K 數(shù)據(jù)集上的結(jié)果

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

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

壓縮后的數(shù)據(jù)可視化

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

蒸餾過程圖像生成動畫

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

將該方法擴展到持續(xù)學(xué)習(xí)任務(wù)上的結(jié)果

上圖展示了 5 步和 10 步的增量學(xué)習(xí)策略，將 200 個類別（Tiny-ImageNet）分為 5 個或 10 個學(xué)習(xí)步驟，每步分別容納 40 個和 20 個類別?？梢钥吹奖疚牡慕Y(jié)果明顯優(yōu)于基線（baseline）性能。

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