過去幾年，數(shù)據(jù)壓縮或蒸餾任務(wù)引起了人們的廣泛關(guān)注。通過將大規(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ān)鍵作用。通過采用先進(jìn)的算法，數(shù)據(jù)壓縮取得了顯著的進(jìn)展。然而，現(xiàn)有解決方案主要擅長壓縮低分辨率的小數(shù)據(jù)集，這種局限性是因?yàn)樵陔p層優(yōu)化過程中執(zhí)行大量未展開的迭代會(huì)導(dǎo)致計(jì)算開銷巨大。

MBZUAI 和 CMU 團(tuán)隊(duì)的最新工作 SRe2L 致力于解決這一問題。該工作是目前唯一實(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ù)，蒸餾過程和訓(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è)生成算法來高效可行地生成需要的樣本，生成的樣本需要包含 / 保留原始數(shù)據(jù)集中核心的信息。目前比較常用的方法包括梯度匹配、特征匹配、軌跡匹配等等，但是這些方法的一個(gè)共同缺點(diǎn)就是沒法 scale-up 到大規(guī)模數(shù)據(jù)集上。比如，由于計(jì)算量和 GPU 顯存的限制，無法蒸餾標(biāo)準(zhǔn)的 ImageNet-1K 或者更大的數(shù)據(jù)集。計(jì)算量和 GPU 顯存需要過大的主要原因在于這些方法生成過程需要匹配和保存的信息過多，目前很多 GPU 顯存沒法容納所有需要匹配的數(shù)據(jù)信息，導(dǎo)致這些方法大多數(shù)只適用于較小的數(shù)據(jù)集。

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

解決方案核心思路

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

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

具體來說，本文提出了一種新的數(shù)據(jù)集壓縮框架，稱為擠壓、恢復(fù)和重新標(biāo)記 (SRe2L)，如下圖所示，該框架在訓(xùn)練過程中解耦模型和合成數(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è)模型之中，通過模型參數(shù)來存儲(chǔ)原始數(shù)據(jù)集中的信息，類似于我們通常進(jìn)行的模型訓(xùn)練；
• 第二步是將這些高度抽象化的信息從訓(xùn)練好的模型參數(shù)中恢復(fù)出來，本文討論了多種不同損失和正則函數(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ǔ)起來。

三階段過程如下圖所示：

圖片

性能及計(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ù)合成過程中內(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è)試階段，跟之前工作相同，文本通過從頭開始訓(xùn)練模型來評(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ì)量、清晰度還是語義信息，都明顯更高。

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

此外，包含 50、200 個(gè) IPC（具有 4K 恢復(fù)預(yù)算）的壓縮數(shù)據(jù)集文件可從以下鏈接獲?。篽ttps://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è)類別。可以看到本文的結(jié)果明顯優(yōu)于基線（baseline）性能。

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