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

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

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

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

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

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

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

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

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

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

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

解決方案核心思路

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

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

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

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

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

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

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

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

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

三階段過(guò)程如下圖所示：

性能及計(jì)算能效比

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

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

實(shí)驗(yàn)結(jié)果

實(shí)驗(yàn)設(shè)置

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

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

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

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

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

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

蒸餾過(guò)程圖像生成動(dòng)畫

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

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

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

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