劉夏鳴
(華中光電技術(shù)研究所-武漢光電國家研究中心,湖北武漢 430223)
近年來深度遷移學(xué)習(xí)在包括計算機視覺、自然語言處理等人工智能相關(guān)領(lǐng)域都取得了很大進展。例如基于海量語料的預(yù)訓(xùn)練的語言表征模型(Bidirectional Encoder Representation from Transformers,BERT)模型[1]已經(jīng)被證明可以顯著提升幾乎所有的自然語言處理任務(wù)。在計算機視覺領(lǐng)域,基于上百萬圖像數(shù)據(jù)集[2]標簽數(shù)據(jù)訓(xùn)練的分類模型可以遷移到語義分割、目標檢測等任務(wù)上,大幅提升這些任務(wù)的準確率等指標。本文研究的問題是如何在一定的約束條件下用遷移學(xué)習(xí)的方法同時提升多個任務(wù)的指標。我們主要考慮兩個在實際應(yīng)用中廣泛存在的約束條件:一,計算資源和存儲資源受限。此約束條件在某些硬件例如低功耗設(shè)備上格外突出。二,任務(wù)之間的耦合度應(yīng)當盡可能,使得每個任務(wù)可以在不影響其它任務(wù)的情況下得到迭代更新。
在傳統(tǒng)的遷移學(xué)習(xí)方案中,下游任務(wù)模型是在預(yù)訓(xùn)練模型的基礎(chǔ)上通過“微調(diào)”的方式獲得的[3]。在微調(diào)的過程中,任務(wù)模型以預(yù)訓(xùn)練模型的網(wǎng)絡(luò)權(quán)重作為初始權(quán)重,在任務(wù)標注數(shù)據(jù)上以梯度下降算法進行訓(xùn)練,逐步改變網(wǎng)絡(luò)權(quán)重,最終得到了一個新的模型。當有多個任務(wù)存在時,人們需要為每個任務(wù)微調(diào)得到一個獨一無二的任務(wù)模型,且這些任務(wù)模型中的權(quán)重和計算無法共用。顯而易見,傳統(tǒng)遷移學(xué)習(xí)方案計算復(fù)雜度和存儲復(fù)雜度都為O(n),其中n 是任務(wù)的數(shù)量。這樣的復(fù)雜度在大多數(shù)應(yīng)用中是不可接受的。
本文考慮擁有12 個Transformer 模塊的ViT-Base 模型。我們提出的方案流程由三個步驟構(gòu)成:(1)單任務(wù)微調(diào)訓(xùn)練;(2)單任務(wù)模型蒸餾;(3)多任務(wù)模型合并。模型方案的細節(jié)如下:
2.1.1 Transformer 編碼器
本文使用的ViT-Base 模型擁有12 層Transformer 編碼器網(wǎng)絡(luò)。Transformer 模型在2017 年由A. Vaswani 等人提出,原用于機器翻譯任務(wù)。目前基于Transformer 結(jié)構(gòu)的模型已經(jīng)在許多自然語言處理和計算機視覺任務(wù)中取得了遠超傳統(tǒng)卷積神經(jīng)網(wǎng)絡(luò)(CNN)或長短時期記憶網(wǎng)絡(luò)(LSTM)的表現(xiàn)。Transformer 網(wǎng)絡(luò)有兩種類型,分別是編碼器和解碼器。一個Transformer 編碼器模塊由若干個編碼器和解碼器堆疊形成。如下圖的左側(cè)部分為編碼器,由多頭注意(Multi-Head Attention)和一個全連接組成,用于將輸入語料轉(zhuǎn)化成特征向量。右側(cè)部分是解碼器,其輸入為編碼器的輸出以及已經(jīng)預(yù)測的結(jié)果,保證訓(xùn)練與推理的一致性,用于輸出最后結(jié)果的條件概率,圖1。
圖1 Transformer 結(jié)構(gòu)圖
2.1.2 ViT 模型結(jié)構(gòu)
在計算機視覺有關(guān)的任務(wù)中,要讓Transformer 模型能夠接受圖像輸入,我們必須設(shè)法將一副二維圖像轉(zhuǎn)化成一個一維的向量序列。在ViT 模型中,我們把一副高度為H 寬度為W 的圖像拆分成N 個邊長為P 的正方形圖像小塊,其中N=HW/P2。若圖像中的每個像素包含C 個通道,則每個正方形小塊對應(yīng)一個維度為CP2 的向量,該向量記錄了圖像小塊的每個像素每個通道的原始數(shù)值。ViT 模型用一個可學(xué)習(xí)的全連接網(wǎng)絡(luò)將該向量非線性映射成了一個D 維的特征向量。通過這種方法,原始圖像變成了一個包含N 個D 維向量的序列,后者可以作為Transformer層的輸入。為了區(qū)分每個D 維向量對應(yīng)的正方形小塊在原始圖像中的位置,ViT 為每個位置的小塊引入了一個可學(xué)習(xí)的位置向量。該位置向量會加到特征向量中,作為后續(xù)Transformer 層的輸入。最后,ViT 引入了一個類似BERT 的[CLS]向量作為零號特征向量,該向量作為整幅圖像的特征表示。
2.1.3 ViT 模型預(yù)訓(xùn)練
我們使用的ViT-Base 模型在圖像網(wǎng)絡(luò)數(shù)據(jù)集上進行預(yù)訓(xùn)練。圖像網(wǎng)絡(luò)數(shù)據(jù)集是一個圖像分類數(shù)據(jù)集,其訓(xùn)練集、驗證集和測試集分別包含128 萬、5 萬和10 萬張標注圖片,分屬于1000 個類別。我們使用的ViT-Base模型在預(yù)訓(xùn)練任務(wù)上經(jīng)過300 次循環(huán)后最終可以達到84.15%的準確率。
2.2.1 單任務(wù)模型微調(diào)
第一個步驟,是針對每個任務(wù)微調(diào)預(yù)訓(xùn)練ViT 模型深層(靠近輸出層)的L 個模塊,凍結(jié)剩余位于淺層的12-L 個模塊。具體的L 數(shù)值是一個超參數(shù),可以隨著任務(wù)的不同而變化。一個任務(wù)擁有較大的L 值意味著該任務(wù)被分配了更多的可訓(xùn)練參數(shù)以及計算資源。在極端情況下L=12,此時我們的微調(diào)方案和傳統(tǒng)的全微調(diào)等價。為了確保最低程度的參數(shù)和計算共享,我們在一個預(yù)先設(shè)定好的較小的區(qū)間N_min <= L <= N_max 內(nèi)搜尋最佳的L 值。在我們后續(xù)的實驗中,我們?nèi)_min=4 以及N_max=10。
第一個步驟最終為每一個任務(wù)輸出一個任務(wù)模型,這些任務(wù)模型的底層和預(yù)訓(xùn)練模型共享參數(shù),頂層使用任務(wù)獨占的網(wǎng)絡(luò)參數(shù)。我們將這些模型稱為教師模型(見圖2 左),用于后續(xù)的模型蒸餾。
圖2 多任務(wù)ViT 模型設(shè)計圖
2.2.2 單任務(wù)模型蒸餾
在上一個步驟中,我們用微調(diào)部分網(wǎng)絡(luò)參數(shù)的辦法針對每個任務(wù)微調(diào)得到了一個任務(wù)模型。這些任務(wù)模型沒有直接的聯(lián)系,因此可以用知識蒸餾的方法壓縮這些模型的非共享網(wǎng)絡(luò)。具體的說,對一個第一步中得到的有著L_t 層任務(wù)獨占網(wǎng)絡(luò)以及12-L_t 層凍結(jié)網(wǎng)絡(luò)任務(wù)模型,我們可以用任意一種模型蒸餾算法將前者壓縮成一個僅有l(wèi)_t 層的更小的模塊,而讓后者保持不變(見圖2 中)。在實驗中,我們發(fā)現(xiàn)讓學(xué)生模型的網(wǎng)絡(luò)參數(shù)從教師模型的對應(yīng)網(wǎng)絡(luò)層中初始化可以達到最佳的學(xué)習(xí)效果。
2.2.3 多任務(wù)模型合并
在最后一步里,我們把所有的單任務(wù)學(xué)生模型合并成一個多任務(wù)模型,這樣凍結(jié)網(wǎng)絡(luò)層中的參數(shù)和所有計算可以被各個任務(wù)所共享,從而提升效率(見圖2 右)。在實際操作中,我們只需要在創(chuàng)建多任務(wù)模型計算圖后將各任務(wù)模型中的參數(shù)載入該計算圖中的對應(yīng)位置即可。
在此我們將對比在一系列任務(wù)上我們的模型與主流基線模型的表現(xiàn)。
3.1.1 單任務(wù)模型(無模型蒸餾):此模型由預(yù)訓(xùn)練模型ViT 在任務(wù)數(shù)據(jù)上以全參數(shù)微調(diào)的方式得到。此為預(yù)訓(xùn)練模型進行遷移學(xué)習(xí)的最佳實踐。如圖3 所示。
圖3 預(yù)訓(xùn)練ViT 模型結(jié)構(gòu)圖
3.1.2 單任務(wù)模型(有模型蒸餾)我們把擁有12 個Transformer 層的單任務(wù)模型通過模型蒸餾的方式壓縮至6 個模塊,從而降低計算復(fù)雜度。
3.1.3 多任務(wù)模型:采用傳統(tǒng)的多任務(wù)學(xué)習(xí)的方式訓(xùn)練,除了各任務(wù)輸出層以外共享全部的12 層傳遞參數(shù),具體方案為每個任務(wù)分配任務(wù)專屬的1~3 個傳遞層用于學(xué)習(xí)任務(wù)相關(guān)的特征。所有任務(wù)共享預(yù)訓(xùn)練模型的前6 層網(wǎng)絡(luò)作為特征提取層。為了更加清晰的展示ViT 模型結(jié)構(gòu)和訓(xùn)練過程中的向量變化,圖4 給出了ViT的向量維度變化圖。
圖4 ViT 的向量維度變化圖
3.2.1 分類任務(wù):分類樣本包含100 個類別,每個類別共600 張圖片作為訓(xùn)練樣本,每張圖片的分辨率為32x32,如圖5 示例。
圖5 分類樣本示例
3.2.2 目標檢測:目標樣本包含20 萬張圖片,50 萬標注目標,共80 個類別。
3.2.3 語義分割:本樣本包含5000 張細粒度標注的圖片,2 萬張粗粒度標注的圖片,標注共分屬于30 個類別。
在接下來的實驗中使用Adam 的優(yōu)化器,采用學(xué)習(xí)率為2e-5 的32 批次訓(xùn)練超參數(shù),從實驗結(jié)果中可以得出,我們的模型在圖像分類、語義分割和目標檢測這三個經(jīng)典的視覺任務(wù)上都超過了單任務(wù)模型蒸餾和多任務(wù)聯(lián)合訓(xùn)練的基線結(jié)果,僅僅只是略微遜于單任務(wù)模型全微調(diào)時的表現(xiàn)。僅使用3 層傳遞網(wǎng)絡(luò),我們在分類任務(wù)、目標檢測和語義分割這三個數(shù)據(jù)集上任務(wù)上分別達到了12 層基線網(wǎng)絡(luò)指標的99.4%,98.3%,98.5%,如表1 所示。
表1 模型在三個視覺任務(wù)上的表現(xiàn)對比
本文提出的模型方案的主要優(yōu)點是它的靈活性。首先,不同的視覺任務(wù)可以從不同深度的網(wǎng)絡(luò)層獲取不同抽象層級的表征。其次,我們的方案可以允許給不同的任務(wù)分配不同的計算資源,例如更重要的任務(wù)可以分配更多的計算資源,或者降低資源投入邊際效費比較低的任務(wù)資源,從而使得模型在多任務(wù)環(huán)境下的平均表現(xiàn)打到更優(yōu)。第三,通過對任務(wù)之間的依賴進行解耦,我們的模型有效回避了傳統(tǒng)多任務(wù)模型的任務(wù)之間相互沖突的可能性,也不再需要針對任務(wù)權(quán)重超參數(shù)進行細致的調(diào)整。我們的多任務(wù)模型可以非常直觀的對單個任務(wù)子網(wǎng)絡(luò)模型的參數(shù)進行更新。
適配器模型是另一類在資源受限條件下應(yīng)用較為廣泛的模型。適配器模型通過在一個模型網(wǎng)絡(luò)的層與層之間添加一個小型可訓(xùn)練模塊(即適配器)來實現(xiàn)預(yù)訓(xùn)練模型對特定任務(wù)的適配。在遷移學(xué)習(xí)的過程中,原預(yù)訓(xùn)練模型中既有參數(shù)不發(fā)生變化,僅有適配器中的參數(shù)得到更新。這類方法的優(yōu)勢是只需要引入極少量的參數(shù)即可實現(xiàn)較優(yōu)秀的任務(wù)表現(xiàn),但其劣勢在于計算開銷不會降低,因為任務(wù)網(wǎng)絡(luò)中的數(shù)據(jù)流被適配器完全改變了,無法跨任務(wù)復(fù)用。也就是說,在多任務(wù)場景下適配器模型的計算復(fù)雜度仍為O(n),因此無法與本文中提出的方法進行直接的比較。
本文提出了一種新型多任務(wù)模型框架,可以在保證靈活性、低耦合性的同時極大降低模型對內(nèi)存的開銷和對計算資源的開銷。我們在多個視覺任務(wù)上證明了我們可以在保證模型預(yù)測精度的同時降低至多75%的資源開銷,因此具有較強的應(yīng)用推廣價值。