?CVPR 2024 | 遷移預(yù)訓(xùn)練生成模型中的知識到邊緣小模型

    

     

      

       

        

         
本文約2900字，建議閱讀6分鐘
         

          本文介紹的是我們的一篇收錄于 CVPR 2024 的論文。
         
        
       
      
     
    

本文介紹的是我們的一篇收錄于 CVPR 2024 的論文。為了解決某個機構(gòu)或邊緣設(shè)備上數(shù)據(jù)缺乏的問題，一方面我們通過聯(lián)邦學(xué)習(xí)獲得來自所有參與設(shè)備的任務(wù)相關(guān)的共享知識；另一方面我們將該共享知識作為預(yù)訓(xùn)練生成模型（如 Stable Diffusion）的輸入來生成任務(wù)相關(guān)的全局?jǐn)?shù)據(jù)；最后我們將生成的數(shù)據(jù)傳輸?shù)綑C構(gòu)或邊緣設(shè)備上并通過額外的有監(jiān)督學(xué)習(xí)任務(wù)，將該全局?jǐn)?shù)據(jù)注入到邊緣模型中。

我們稱這一過程為“知識遷移鏈路（KTL）”，并在每一輪聯(lián)邦學(xué)習(xí)過程中進行一次知識遷移。此過程中生成模型只作推理不做訓(xùn)練，這種利用預(yù)訓(xùn)練模型的方式對資源的需求較少。

運行實驗所需倉庫-個性化聯(lián)邦學(xué)習(xí)算法庫：

https://github.com/TsingZ0/PFLlib

運行實驗所需倉庫-異構(gòu)聯(lián)邦學(xué)習(xí)算法庫：

https://github.com/TsingZ0/HtFLlib

1、動機

隨著新一輪 AI 時代的到來，模型的量級越來越大，對數(shù)據(jù)的需求也越來越大。不論是哪個領(lǐng)域，有效且高質(zhì)量的數(shù)據(jù)一直是一種稀缺存在，甚至成為了一種數(shù)據(jù)資產(chǎn)。與此同時，互聯(lián)網(wǎng)上也廣泛存在著能力強勁的開源生成模型。如果能夠利用這些生成模型中存儲的知識，來生成具體任務(wù)所需的數(shù)據(jù)，便可以讓小公司和邊緣智能設(shè)備都能享受到大模型帶來的豐富成果。

為了解決數(shù)據(jù)稀缺問題，有以下四種常見途徑：

1. 利用來自公開數(shù)據(jù)集的數(shù)據(jù)，但這類數(shù)據(jù)很難做到與具體任務(wù)相關(guān)，且任務(wù)無關(guān)數(shù)據(jù)甚至?xí)a(chǎn)生負(fù)面影響；

2. 由數(shù)據(jù)請求方上傳數(shù)據(jù)生成需求（比如分類任務(wù)中的標(biāo)簽語義）到云端的生成模型來生成數(shù)據(jù)，但這種文本很容易導(dǎo)致隱私問題；

3. 利用云端大模型隨機生成無標(biāo)簽數(shù)據(jù)，但這樣做依舊存在與利用公開數(shù)據(jù)集同樣的問題，甚至無標(biāo)簽數(shù)據(jù)的引入增加了模型訓(xùn)練的難度；

4. 利用預(yù)訓(xùn)練模型引入額外知識，但適用于具體任務(wù)的預(yù)訓(xùn)練模型稀少且其中的額外知識不一定匹配當(dāng)前任務(wù)。

換言之，不管是引入額外數(shù)據(jù)還是引入額外知識，都盡量需要與當(dāng)前任務(wù)相關(guān)，才能最大限度地起到正面作用。

考慮到聯(lián)邦學(xué)習(xí)技術(shù)可以將需要實現(xiàn)相似任務(wù)的設(shè)備或機構(gòu)聯(lián)合起來，實現(xiàn)知識遷移和共享；我們便將任務(wù)相關(guān)知識的獲取通過聯(lián)邦學(xué)習(xí)實現(xiàn)。考慮到傳統(tǒng)聯(lián)邦在模型架構(gòu)上的局限性，我們引入異構(gòu)聯(lián)邦學(xué)習(xí)技術(shù)，允許各個參與方采用不同架構(gòu)的模型。

之后，我們將該任務(wù)相關(guān)知識作為輸入，傳遞給預(yù)訓(xùn)練生成模型，并針對當(dāng)前任務(wù)做了域?qū)R，從而生成任務(wù)相關(guān)的數(shù)據(jù)。為了有效利用該數(shù)據(jù)，我們將其傳輸?shù)铰?lián)邦學(xué)習(xí)參與方后，運行一個額外的有監(jiān)督任務(wù)實現(xiàn)知識遷移。

2、異構(gòu)聯(lián)邦學(xué)習(xí)技術(shù)

傳統(tǒng)聯(lián)邦學(xué)習(xí)考慮了數(shù)據(jù)異質(zhì)性，但依舊要求所有參與方訓(xùn)練同一個架構(gòu)的模型，增加了尋找相似任務(wù)參與方的難度。于是我們考慮取消這一點要求，允許參與方采用各自的模型進行知識共享。然而，這樣一來，傳統(tǒng)聯(lián)邦學(xué)習(xí)中基于模型參數(shù)共享的范式不再可用，對新型的知識共享機制提出了要求。其中包括：1）保護隱私，2）保護知識產(chǎn)權(quán)，3）輕量化，4）易于獲得等。

目前異構(gòu)聯(lián)邦學(xué)習(xí)技術(shù)還未形成統(tǒng)一的知識共享機制，我們考慮一種輕量化且不需要額外數(shù)據(jù)的知識共享機制：共享 prototype。本文考慮的是面向圖像的多分類任務(wù)，其中prototype 的定義就是每個類別的代表性特征向量，可通過平均該類所有的特征向量獲得。我們將 prototype 當(dāng)作共享知識，輸入到生成模型后得到相應(yīng)圖片數(shù)據(jù)，并將圖片-向量對（image-vector pairs）傳回參與者，如下圖。

3、知識遷移鏈路（KTL）

上一節(jié)的最后已經(jīng)簡單描述了我們提出的知識遷移鏈路（KTL），但省略了很多細(xì)節(jié)，這里我們對重點步驟進行展開（其他步驟及細(xì)節(jié)詳見論文）。下圖是我們的整體框架，其中最重要的是步驟 3 和步驟 6。

步驟3：當(dāng)我們在生成模型的特征空間采樣時，可以生成清晰圖像，但這樣的圖像并非任務(wù)相關(guān)。根據(jù)我們的實驗觀察，如果直接將參與方上傳的 prototype 輸入到預(yù)訓(xùn)練生成模型，由于參與方模型的特征空間和生成模型的特征空間不匹配（通常連維度都不一致），導(dǎo)致生成的圖像跟隨機輸入一樣模糊不清。

所以我們需要先將 prototype 映射到高維的生成模型特征空間，并保證這些 prototype 依舊是任務(wù)相關(guān)的。因為我們考慮的是分類問題，任務(wù)相關(guān)指的就是 prototype 映射后得到的特征向量依舊保持類別可分離特性。我們稱這一過程為域?qū)R（domain alignment），如下圖可見，對齊后的特征向量可以使生成模型產(chǎn)生清晰圖片。

為了實現(xiàn)域?qū)R，我們在服務(wù)器端額外訓(xùn)練了一個輕量化的特征轉(zhuǎn)換器（F），并定義其訓(xùn)練目標(biāo)為對齊特征空間（使用 MMD 損失）和保證類別可分離性（使用 MSE 損失）。

步驟6：在進行域?qū)R后，將映射后的向量輸入到生成模型即可得到清晰的圖片。我們將圖片-向量對下載到各個參與方，而后通過一個額外的有監(jiān)督任務(wù)實現(xiàn)知識遷移。圖 3 中的   指的是原本的分類任務(wù)目標(biāo)函數(shù)，  則指的是我們提出的額外有監(jiān)督任務(wù)的目標(biāo)函數(shù)。

由于我們只對參與方模型的特征提取器進行知識遷移，而特征提取器需要的正好是特征提取相關(guān)的通用知識，且預(yù)訓(xùn)練生成模型含有大量通用知識。所以只要是生成模型生成的且由任務(wù)相關(guān)的 prototype 誘導(dǎo)出的圖片，都可以在 KTL 中發(fā)揮作用。

進一步地，我們不需要生成模型在特定數(shù)據(jù)集上進行預(yù)訓(xùn)練，在任意圖像數(shù)據(jù)集上預(yù)訓(xùn)練的生成模型都可以在我們的框架中發(fā)揮作用。我們也通過實驗證實了我們方法（FedKTL）的這一個能力。如下圖表所示，在四種不同數(shù)據(jù)集上預(yù)訓(xùn)練的生成模型帶來的表現(xiàn)幾乎是類似的，且這些數(shù)據(jù)集并不與參與方本地的數(shù)據(jù)集重合。

▲ 表1：使用在不同預(yù)訓(xùn)練數(shù)據(jù)集上預(yù)訓(xùn)練的生成模型的效果。其中 是與服務(wù)器端數(shù)據(jù)生成相關(guān)的超參數(shù)。

4、部分實驗

由于篇幅有限，主體實驗請查閱論文。這里僅討論與我們方法特性相關(guān)的實驗結(jié)果，包括“上傳的通訊成本低”、“可擴展到 Stable Diffusion 等不同架構(gòu)的生成模型模型”、“適用 cloud-edge 框架（不使用聯(lián)邦學(xué)習(xí)）”。

上傳的通訊成本低：由于我們只上傳 prototype，且對特征空間進行了降維（詳見論文），我們方法所需的上傳通訊成本極低。在現(xiàn)實環(huán)境中，上傳鏈路帶寬往往比下載鏈路低好幾個數(shù)量級。我們方法的這一特性使得它更能適應(yīng)現(xiàn)實環(huán)境，且充分利用上傳下載的帶寬。

可擴展到 Stable Diffusion 等不同架構(gòu)的生成模型模型：論文展示了我們的知識遷移框架不但可以使用 StyleGAN 等較小的生成模型，也可以在不加修改的情況下使用當(dāng)下最熱門的基于 diffusion 技術(shù)的大型生成模型 Stable Diffusion。但由于 Stable Diffusion 的特征空間維度非常高，不加修改直接使用無法完全挖掘其生成能力，如何匹配高維特征空間是未來有待探索的方向。

適用 cloud-edge 框架（不使用聯(lián)邦學(xué)習(xí)）：在我們的知識遷移框架（KTL）中，我們引入了兩種外來信息給到參與方，一種是來自其他參與方的任務(wù)相關(guān)知識（來自聯(lián)邦學(xué)習(xí)），另一種是預(yù)訓(xùn)練生成模型中存儲的知識。

在某些情況下，尋找其他運行相似任務(wù)的設(shè)備或機構(gòu)參與到聯(lián)邦學(xué)習(xí)中較為困難。這時候我們僅用第二種預(yù)訓(xùn)練生成模型中存儲的知識，也可以實現(xiàn)知識遷移。而且這種無聯(lián)邦學(xué)習(xí)的知識遷移范式，可以依托現(xiàn)有的 cloud-edge 框架實現(xiàn)。如下表，在只有一個參與者的情況下，參與者擁有的數(shù)據(jù)越少，我們的 KTL 方法效果越好。