CVPR2024|遷移預訓練生成模型中的知識到邊緣小模型

©PaperWeekly 原創 · 作者 | 張劍清

單位 | 上海交通大學、清華大學（AIR）

研究方向 | 聯邦學習

本文介紹的是我們的一篇收錄於 CVPR 2024 的論文。為了解決某個機構或邊緣裝置上資料缺乏的問題，一方面我們透過聯邦學習獲得來自所有參與裝置的任務相關的共享知識；另一方面我們將該共享知識作為預訓練生成模型（如 Stable Diffusion）的輸入來生成任務相關的全域性資料；最後我們將生成的資料傳輸到機構或邊緣裝置上並透過額外的有監督學習任務，將該全域性資料注入到邊緣模型中。

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

論文標題：

An Upload-Efficient Scheme for Transferring Knowledge From a Server-Side Pre-trained Generator to Clients in Heterogeneous Federated Learning

論文連結：

https://arxiv.org/abs/2403.15760

程式碼連結：

https://github.com/TsingZ0/FedKTL（含有PPT和Poster）

執行實驗所需倉庫-個性化聯邦學習演算法庫：

https://github.com/TsingZ0/PFLlib

執行實驗所需倉庫-異構聯邦學習演算法庫：

https://github.com/TsingZ0/HtFLlib

動機

隨著新一輪 AI 時代的到來，模型的量級越來越大，對資料的需求也越來越大。不論是哪個領域，有效且高質量的資料一直是一種稀缺存在，甚至成為了一種資料資產。與此同時，網際網路上也廣泛存在著能力強勁的開源生成模型。如果能夠利用這些生成模型中儲存的知識，來生成具體任務所需的資料，便可以讓小公司和邊緣智慧裝置都能享受到大模型帶來的豐富成果。

為了解決資料稀缺問題，有以下四種常見途徑：

1. 利用來自公開資料集的資料，但這類資料很難做到與具體任務相關，且任務無關資料甚至會產生負面影響；

2. 由資料請求方上傳資料生成需求（比如分類任務中的標籤語義）到雲端的生成模型來生成資料，但這種文字很容易導致隱私問題；

3. 利用雲端大模型隨機生成無標籤資料，但這樣做依舊存在與利用公開資料集同樣的問題，甚至無標籤資料的引入增加了模型訓練的難度；

4. 利用預訓練模型引入額外知識，但適用於具體任務的預訓練模型稀少且其中的額外知識不一定匹配當前任務。

換言之，不管是引入額外資料還是引入額外知識，都儘量需要與當前任務相關，才能最大限度地起到正面作用。

考慮到聯邦學習技術可以將需要實現相似任務的裝置或機構聯合起來，實現知識遷移和共享；我們便將任務相關知識的獲取透過聯邦學習實現。考慮到傳統聯邦在模型架構上的侷限性，我們引入異構聯邦學習技術，允許各個參與方採用不同架構的模型。

之後，我們將該任務相關知識作為輸入，傳遞給預訓練生成模型，並針對當前任務做了域對齊，從而生成任務相關的資料。為了有效利用該資料，我們將其傳輸到聯邦學習參與方後，執行一個額外的有監督任務實現知識遷移。

異構聯邦學習技術

傳統聯邦學習考慮了資料異質性，但依舊要求所有參與方訓練同一個架構的模型，增加了尋找相似任務參與方的難度。於是我們考慮取消這一點要求，允許參與方採用各自的模型進行知識共享。然而，這樣一來，傳統聯邦學習中基於模型引數共享的正規化不再可用，對新型的知識共享機制提出了要求。其中包括：1）保護隱私，2）保護智慧財產權，3）輕量化，4）易於獲得等。

▲ 圖1：異構聯邦學習技術

目前異構聯邦學習技術還未形成統一的知識共享機制，我們考慮一種輕量化且不需要額外資料的知識共享機制：共享 prototype。本文考慮的是面向影像的多分類任務，其 prototype 的定義就是每個類別的代表性特徵向量，可透過平均該類所有的特徵向量獲得。我們將 prototype 當作共享知識，輸入到生成模型後得到相應圖片資料，並將圖片-向量對（image-vector pairs）傳回參與者，如下圖。

▲ 圖2：異構聯邦學習中使用prototype進行知識遷移

知識遷移鏈路（KTL）

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

▲ 圖3：Federated Knowledge-Transfer-Loop (FedKTL)

步驟3：

當我們在生成模型的特徵空間取樣時，可以生成清晰影像，但這樣的影像並非任務相關。根據我們的實驗觀察，如果直接將參與方上傳的 prototype 輸入到預訓練生成模型，由於參與方模型的特徵空間和生成模型的特徵空間不匹配（通常連維度都不一致），導致生成的影像跟隨機輸入一樣模糊不清。

所以我們需要先將 prototype 對映到高維的生成模型特徵空間，並保證這些 prototype 依舊是任務相關的。因為我們考慮的是分類問題，任務相關指的就是 prototype 對映後得到的特徵向量依舊保持類別可分離特性。我們稱這一過程為域對齊（domain alignment），如下圖可見，對齊後的特徵向量可以使生成模型產生清晰圖片。

▲ 圖4：生成模型在不同輸入下得到的圖片

為了實現域對齊，我們在伺服器端額外訓練了一個輕量化的特徵轉換器（F），並定義其訓練目標為對齊特徵空間（使用 MMD 損失）和保證類別可分離性（使用 MSE 損失）。

▲ 圖5：域對齊實現方案

▲ 圖6：域對齊的一個例子。這是一個三分類任務，其中參與方模型特徵空間維度為 2，生成模型特徵空間維度為 3，W 是生成模型的有效特徵空間。

步驟6：在進行域對齊後，將對映後的向量輸入到生成模型即可得到清晰的圖片。我們將圖片-向量對下載到各個參與方，而後透過一個額外的有監督任務實現知識遷移。圖 3 中的指的是原本的分類任務目標函式，則指的是我們提出的額外有監督任務的目標函式。

由於我們只對參與方模型的特徵提取器進行知識遷移，而特徵提取器需要的正好是特徵提取相關的通用知識，且預訓練生成模型含有大量通用知識。所以只要是生成模型生成的且由任務相關的 prototype 誘匯出的圖片，都可以在 KTL 中發揮作用。

進一步地，我們不需要生成模型在特定資料集上進行預訓練，在任意影像資料集上預訓練的生成模型都可以在我們的框架中發揮作用。我們也透過實驗證實了我們方法（FedKTL）的這一個能力。如下圖表所示，在四種不同資料集上預訓練的生成模型帶來的表現幾乎是類似的，且這些資料集並不與參與方本地的資料集重合。

▲ 圖6：參與方本地資料樣本和四種用於生成模型預訓練的資料集樣本。

▲ 表1：使用在不同預訓練資料集上預訓練的生成模型的效果。其中是與伺服器端資料生成相關的超引數。

部分實驗

由於篇幅有限，主體實驗請查閱論文。這裡僅討論與我們方法特性相關的實驗結果，包括“上傳的通訊成本低”、“可擴充套件到 Stable Diffusion 等不同架構的生成模型模型”、“適用 cloud-edge 框架（不使用聯邦學習）”。

上傳的通訊成本低：由於我們只上傳 prototype，且對特徵空間進行了降維（詳見論文），我們方法所需的上傳通訊成本極低。在現實環境中，上傳鏈路頻寬往往比下載鏈路低好幾個數量級。我們方法的這一特性使得它更能適應現實環境，且充分利用上傳下載的頻寬。

▲ 表2：每輪的通訊成本

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

▲ 表3：使用 StyleGAN 和 Stable Diffusion 的表現

適用 cloud-edge 框架（不使用聯邦學習）：在我們的知識遷移框架（KTL）中，我們引入了兩種外來資訊給到參與方，一種是來自其他參與方的任務相關知識（來自聯邦學習），另一種是預訓練生成模型中儲存的知識。

在某些情況下，尋找其他執行相似任務的裝置或機構參與到聯邦學習中較為困難。這時候我們僅用第二種預訓練生成模型中儲存的知識，也可以實現知識遷移。而且這種無聯邦學習的知識遷移正規化，可以依託現有的 cloud-edge 框架實現。如下表，在只有一個參與者的情況下，參與者擁有的資料越少，我們的 KTL 方法效果越好。