IJCAI 2021 | 中科院计算所:自监督增强的知识蒸馏方法
知識蒸餾作為模型壓縮里的一項重要技術,在學術界和工業界都得到了廣泛的研究。
近日,AI Drive 邀請中科院計算所在讀博士生楊傳廣,分享其在 IJCAI 2021 上發表的最新工作:自監督增強的知識蒸餾。
論文標題:
Hierarchical Self-supervised Augmented Knowledge Distillation
論文鏈接:
https://arxiv.org/abs/2107.13715
代碼鏈接:
https://github.com/winycg/HSAKD
背景介紹:為什么要對模型進行知識蒸餾
今天分享的自監督增強知識蒸餾方法,是我們在 IJCAI 2021 中發表的用于通用圖像分類任務的一種方法,該方法可以提升網絡的性能。
首先介紹,為什么要對模型進行知識蒸餾。
第一,資源困境。深度學習模型,可以完成復雜的任務,但是往往具有龐大的參數量,計算開銷和內存需求大的問題。因此,在計算能力受限平臺上部署這些龐大的模型,通常是一個挑戰,在不影響模型性能的前提下進行模型加速與壓縮,也是一個研究熱點。通常,我們使用的方法有參數剪枝、參數共享、模型量化、知識蒸餾等。文中我們重點研究知識蒸餾相關的模型壓縮。
知識蒸餾的基本思想,是從教師網絡中提取知識,然后將知識作為監督信號來訓練學生網絡,使得學生網絡能夠有更好的泛化能力。
上圖展示了知識蒸餾的整體示意圖,其中,Big Model 通常代表教師網絡,而 Small Model 通常代表學生網絡。知識蒸餾的核心思想,就是將這個教師網絡所產生的信息作為監督信號,來訓練學更小的學生模型,使得這個小模型能夠在完成任務需求的情況下,復雜度要遠遠小于原有的教師模型。
相關工作:知識蒸餾技術的發展脈絡
上圖重點介紹了知識蒸餾技術的發展脈絡,以及一些代表性的相關工作。
現有的知識蒸餾主要分為兩個層面的研究內容:
1. 怎樣在教師網絡中挖掘有益的知識來引導學生網絡學習。
通常可以概括為三個層面:第一個層面是決策概率分布,response-based;第二個是特征信息,feature-based;第三個是高階的跨樣本相似度關系,relation-based。
2. 有了教師網絡中的知識形式,以何種方式來遷移給學生?
第一種最常見的方式是采用 KL divergence 和 L2 誤差的方式,分別對概率分布和特征圖信息作為一種逼近手段;第二種是從模型的角度,在教師和學生之間引入一個中間模型。第三種,可以利用對抗生成的思想來依次進行知識遷移。
對于最傳統的概率層面知識蒸餾方法,首先預訓練一個 Teacher Model,輸入 x 計算網絡可以得到一個概率分布輸出,這里稱為 soft labels,同樣的給定輸入x學生網絡也可以得到對應的概率分布信息,這里稱為 soft predictions。
然后,將這兩個概率分布經過一個溫度T平滑,再進行一個 KL divergence 的逼近,使得這個學生網絡的分布逼近教師網絡產生的概率預測分布。
這種方法在直覺上是 work 的,因為教師網絡具有更好的表現,通常是因為它具有更高質量的一個預測的概率分布,然后使得學生網絡去逼近更高質量的概率分布,通常可以提升學生網絡的泛化能力。
概率層面中,CNN 在做出推斷的時候,通常會產生很多特征圖,這些多層次特征圖通常代表 CNN 對輸入的學習過程,因此,傳統的方法 Feature map 采用遷移特征圖的方式,這種方式通常是采用 L2 誤差的方式,對教師網絡和學生網絡之間不同階段的特征圖進行對應,使得學生網絡在決策過程中能夠逼近教師網絡的決策過程。
同樣,我們可以在特征圖中凝練出更高更好的信息,最常見的可能就是注意力圖(Attention map)。在這里,這個方法是給將給定的 Feature map,從參照維度通過聚合生成 Spatial Feature map,作為一種知識遷移的形式。
另一種比較經典的是 FSP 方法,該方法利用 CNN 不同階段的特征圖,求取相關度——Gram 矩陣,以它作為一種知識表達形式。
第三種比較經典的方法是,可以基于給定特征圖中的決策邊界。
對于高階的框架的相似度層面的知識蒸餾,代表性的工作有兩個,一個是基于高斯和距離的相關性,稱為 Correlation,該方法是將 Embedding 映射到一個高維的 Embedding 空間,然后在高維空間中作為相似度的求取。
第二種方法是 relationship 的 R KD,該方法是將這個網絡不同樣本所產生的 Embedding 求距離相關或者是余弦值角度相關的相似度關系,作為一種知識表達形式。
有了這個知識之后,怎樣進行知識遷移?
有一種比較好的方法就是:基于一種觀測,即采用過大的教師網絡直接來遷移知識給過小的學生網絡,可能會存在一種支持 gap,于是引入網絡協助器來對蒸餾的知識進行抽象。這個網絡協助器,通常是介于教師網絡和學生網絡之間規模的模型,然后該模型可以對教師知識進行一定程度的抽象,再傳遞給學生網絡。
另一個思想用對抗生成(GAN)。可以將學生網絡看成 GAN 框架中的生成器,來模仿教師網絡的特征向量,同時需要引入一個額外的判別器,來對特征向量進行二分類,然后通過對抗生成的框架,可以使得學生網絡學習到教師網絡的概率分布——這就是特征分布。
下文是對自監督表示學習引導的知識蒸餾的介紹。
研究動機
自監督學習的主要思想,是利用人工構造的偽標簽和數據增強,使得網絡學習更加泛化性的特征表示。
現有的自監督學習表示通常可以分為兩大模式,第一個是基于分類的自監督學習,通常,是先提出一種基于分類的 pretext task,該任務通常施加在輸入的圖像上,例如比較經典的四種旋轉角度,顏色通道打亂等等,可以引導網絡去正確預測其中施加的是哪種數據增強,從而使得網絡學到更好的特征表示。
自監督學習特征表示,通常可以用于到下游的一些語義識別任務,比如圖像分類和目標檢測。
第二個層面是基于對比學習的自監督學習,該種方法近年來比較流行,通常的思想是對輸入一張圖像增加不同的數據增強,然后使用對比誤差來逼近一張圖像的不同視角。
在這里,一張圖像的不同視角通常認為是正樣本,然后其他圖像所產生的 Embedding 特征向量是負樣本。
第一個基于分類的自監督學習模式的代表工作 RotNet。該方法是對輸入的 Image X,施加四種轉換,分別是零度(零度代表對圖像沒有變化)、90 度、180 度、270 度,最后進一步訓練網絡對輸入的圖像做何種旋轉做出判斷。
例如上圖由于是 4 種旋轉角度,相當于是一個自監督的自分類問題。對一個圖像旋轉 90 度,然后要求卷積神經網絡能在 4 種角度中,正確判斷出旋轉的是 90 度,最后分類的 loss 也是基于傳統的交叉分類 loss。
第二比較有代表性的對比自監督學習工作,是 SimCLR。該工作的核心思想就是采用對比學習,在不同的數據增強之間最大化一致性。有關論文中研究了很多種不同的數據增強,上圖左邊部分展示的,便是 SimCLR 的整體過程:針對這個輸入 x,通常可以隨機采樣兩種不同的數據增強,然后在這里施加了兩種不同數據增強之后,分別得到了 xi 和 xj 增強之后的數據,并認為這是正樣本。
對于這兩種增強之后的數據,我們輸入到同一個卷積神經網絡f中,可以得到對應的特征表示分別為 xi 和 xj。然后經過一個特征向量轉換之后,得到表示 zi 和 zj,最后最大化 zi 和 zj 的一致性,同時將其他的圖像作為負樣本,做一個對比學習。
另一個比較有代表性的對比自監督學習工作是 PIRL,其對比學習的框架跟 SimCLR 是一致的,它的主要思想也是對一張圖像施加轉換,將一個圖像的不同轉換的視角作為正樣本,其他圖像作為負樣本,文中引入了一個 Jigsaw Puzzle 轉換。
如上圖左邊部分,給定一個輸入圖像 I 施加一個 Transform t,然后得到一 個It,即拼圖之后的一個圖像。標準的自監督學習,判斷出給定的這個 It,并判斷出這個拼圖是基于怎樣的順序,而這里采用的也是鼓勵一致性的做法,針對卷積神經網絡,輸入的兩個視角是原始的圖像I和轉換之后的圖像It,然后分別得到它的表示之后,采用對比學習來鼓勵一致性,從而進行自監督學習。
先前的自監督學習的場景,主要是用在沒有標簽的情況下。一種可行的思想是,利用自監督學習來進一步增強有監督學習的性能。
SSKD 是一種將自監督知識引入有監督知識蒸餾的方法。SSKD 啟發自 SimCLR,它將圖像在經過數據增強之后的對比分布,可以作為有意義的知識形式來進行知識遷移。
上圖就說明了這個知識定義過程,這個過程實際上就是 SimCLR 的過程,針對一個圖像做數據增強之后,盡量去逼近原始圖像和轉換之后圖像的關系,然后進一步遠離其他圖像之間的關系。
由這個關系做一個對比分布,可以由教師網絡遷移給學生網絡,在具體實現層面上的話,是在 Teacher ?Backbone 和 Student Bone 上分別加一個自監督的模塊,文中的自監督模塊,實際上是一個特征向量的轉換模塊,這個特征向量的轉換模塊是由兩個全連接層中間加一個非線性的 ReLU 函數來構成的,目的是將 Backbone 所產生的特征向量做一定程度的轉換,轉換之后的特征向量可以用于進行學習。
實現方法和實驗結果
既然自監督和有監督都能使得網絡學習到更好的特征表示,那么一個很好的研究方向,就是去考慮用自監督任務來增強有監督任務的特征學習。
在這里,考慮聯合兩種方式進行有效的特征表示學習。
上圖中,首先(a)圖是傳統的有監督分類,有監督分類就是針對一張圖像,判斷這個圖像中目標是哪一個類。(b)圖展示的是自監督的分類示意圖,針對一張圖像做了不同的旋轉之后,通過分類器,判斷做的是哪種旋轉。
(c)是聯合有監督和自監督的一種學習方式,在這里它們兩者采用了共享的特征提取器,在分類尾部是采用了不同的分類頭,左側是一個有監督的分類頭,用來分類不同的目標,右側是一個自監督的分類頭,用來分類這個目標的旋轉角度,這兩個分類任務是同時進行、學習共享的特征提取器。
(d)是將這個自監督和有監督的 label 空間進行一定程度的融合,變成一個全連接的空間。舉例子的話是比如說有監督的分類任務是兩個目標,而自監督的角度分類是 4 個 label,通過融合這兩個 label 可以得到一個 2 * 4 的 8 個 label 空間,通過兩個 label 空間做一個笛卡爾積的方式。這里只需要用一個分類頭就可以達到,但是這種方法的缺點就是可能會增大類別空間的類別個數。但是當增大類別個數之后,也會意味著會顯著的編碼更多的有監督和自監督的這種聯合信息。
在研究原始有監督任務的網絡時,不同的 Feature Stage 階段插入輔助分支來學習輔助的有監督和自監督的聯合任務,需要對網絡結構進行了一定的改進。
例如,針對一個 ResNet-56 為例,將主干網絡表達為f(·),并且添加輔助分支C1(·),C2(·),C3(·),主干網絡還是學習原始的有監督的分類任務,輔助分支學習提出的一些有監督結合的一些任務,從而研究無監督任務對有監督任務的一個表現的增益。
不同網絡結構上使用自監督輔助任務來做輔助訓練,對最終有監督分類任務的影響如下表所示:
在不同網絡結構上做了一系列的實驗,發現了如下結論:
(1)使用單一有監督或者自監督任務作為訓練輔助分類器的一些任務以及它們聯合的任務作為輔助任務,都可以使得主干的有監督分類任務的性能有所提升。
(2)采用和主干網絡相同的有監督分類任務,性能提升可能是最小的,可能是因為添加的輔助任務和原有的主干任務是高度同質化的,因此訓練輔助結構并沒有引入額外的一些信息和知識,采用純自監督的任務作為訓練輔助分類器的任務的話,性能增益是要比使用和主干網絡相同的有監督任務要大,表明自監督任務的確是有能有效提升主干有監督任務的性能。
(3)由于自監督和原始的有監督可能是兩種不太相同的任務,自監督也為有監督引入一些額外的表示學習的知識。此外,發現采用一種多任務方式,多任務就是上文提到了采用兩個分類頭的形式,分別學習有監督和自監督的方式,跟原始的自監督之前用自監督的方式性能增益相當。
但是將兩個 label 空間聯合起來,即有監督和次監督這兩個 label 空間采用笛卡爾積的形式變成一個聯合 label 空間,這種 label 空間作為一個輔助的訓練任務,可以使得網絡主干的有監督分類就能得到最大程度的提升,表明學習到了最好的特征表示,因此也啟發采用一種這種聯合的形式作為輔助任務,作為教師學生模型之中的一種 pretext 任務。
在此簡單回顧一下自監督蒸餾知識的定義:
(1)先前最好的方法 SSKD 是基于 SimCLR 框架定義的一種自監督對比分布作為自監督的輔助知識。這里把之前的示意圖做了簡化:假如說這個鳥旋轉 90 度得到一個 Transform 的圖像,然后將這兩個圖像在這個特征空間中逼近,同時遠離這個狗和飛機所產生的特征向量,這個是一個對比分布。
(2)采用了基于一種有監督和無監督的一種聯合分布,產生的軟標簽作為一種知識,該知識包含了有監督和無監督的聯合知識,然后對附加的輔助分類器來學習這些輔助的任務,來遷移有監督增強的一種聯合概率分布。
簡單形式化介紹一下自監督增強的標簽空間分布。針對一個傳統的 N 維的目標分類任務,它的 label 空間為 1~N。針對一個傳統的 M 維的無監督轉換分類任務,label 空間通常是 1~M。然后無監督增強的標簽空間,就可以表達為 label 空間 A 和 label 空間B的一個笛卡爾積的形式,然后新的增強的空間的大小就是 M*N。
上圖展示了自監督增強知識蒸餾框架:HSAKD 的整體的流程。
HSAKD 在教師網絡和學生網絡對應層次上的輔助分類器,進行一個單向的自監督輔助的知識蒸餾。
誤差來源有兩個:
(1)誤差 1:教師網絡產生的自監督增強的分布,q1 到 q3 作為軟標簽來對學生網絡對應層次所產生的自監督增強的分布 q1 到 q3 進行監督;
(2)誤差 2:最終分類器層面的一個概率概率蒸餾。使用原始的 KD,將最終的有監督的類別概率分布 PT 作為軟標簽來監督學生 PS ,這是兩個蒸餾誤差的來源。
具體的實驗方法是由于基于教師學生模式的框架,通常需要預訓練一個教師網絡,將教師結構表達為主干網絡 fT(·)?和 L 個輔助分類器。
(1)給定原始的數據 x,訓練主干網絡 fT(·) 通過傳統的交叉熵來擬合真實的標簽y;
(2)給定旋轉增強的數據 tj(x),其中 tj 為第 j 的旋轉角度,訓練 L 個輔助分類器,通過交叉熵函數來學習自監督增強的 label,這個 K 表示為 KL,然后這兩個 loss 都是基于交叉熵的分類 loss,整體的誤差是將這兩個交叉熵 loss 相加。同時對主干網絡和添加的輔助分支做了一個聯合的訓練。
這里 PT 是指原始的 N 維有監督分類任務上預測概率分布,qT 的話就是為 N*M 維的一個自監督增強的分布,預訓練教師網絡之后,可以用教師網絡進一步指導訓練學生網絡。
將學生結構表達為主干網絡 fS(·),L 個輔助分類器 CS。
(1)對 fS(·) 擬合原始的數據 x,然后通過給定原始的數據 x 來擬合真實的標簽y,這里采用了一種傳統的交叉熵 loss。在這里 ps 的指原始的有監督分類任務上的一個概率分布。除了基礎任務 Loss,還有教師網絡對學生網絡的指導。
(2)給定旋轉增強的數據 tj(x)之后,其中 tj 為第 j 的旋轉角度,訓練L個輔助分類器來學習來自教師網絡產生的自監督增強的軟標簽 qT,啟發自原始的 KD 過程就可以用 KL 散度來實現。通過 KL 散度,可以使得學生網絡產生的監督增強的分布。去擬合教師網絡所產生的分布。在這個公式中,qs 是學生網絡產生的自監督增強的分布,KL 散度中中,取溫度 T=3。
(3)除了自監督增強的一個分布蒸餾,進一步可以考慮最終的有監督的類別概率分布作為前一個知識。由于對輸入端圖像進行了旋轉增強,得到了轉換之后的圖像 tj(x),于是可以遷移這些增強數據的概率分布。這里采用了一種 KL 散度的方式來使得學生網絡去逼近教師網絡產生的概率分布。PT 和 PS 就是分別是有監督分類任務上的概率分布,都是 N 維的。這里同樣取溫度等 3。
整體的誤差就是將上述誤差(1)(2)(3)相加得到最終的誤差。由于誤差來源都是基于概率的熵誤差,所以并沒有引入誤差平衡因子,在實驗中效果表現已經足夠好。
首先在標準的 CIFAR-100 數據集上做了一個實驗對比,主要對比的是先前流行的自蒸餾方法,在不同的教師學生網絡結構上, HSAKD 能一致的超過先前的方法,先前最好的方法是自監督相關的 SSKD,HSAKD 方法在不同網絡結構上超過了 SSKD 2.56%?的平均準確率增益。
進一步的在 ImageNet 和小樣本分類任務上進行了實驗,得到以下結論:
(1)在ImageNet 數據上,在經典的 ResNet-34 和 ResNet-18 結構對上,HSAKD 超過先前最好的 SSKD 0.77% 的準確率增益。
(2)在小樣本分類:驗證各種知識蒸餾方法在數據缺乏場景下的表現。在保留不同比例數據的情況下,HSAKD 獲得了一致的性能提升。表明 HSAKD 在缺乏數據的情況下,依舊能表現出良好的性能。在僅保留訓練集 25% 的情況下,在 ResNet20 依舊能夠達到 68.5% 的準確率,相當于跟原始的 baseline 效果是近似一致。
(2)進一步做了一個遷移學習的實驗,通常在上游的圖像分類數據集之后,可以將訓練好的模型作為一個遷移學習的任務。在 ImageNet 數據集上訓練的 ResNet-18 遷移到下游的目標檢測任務上。在 ImageNet 上訓練的其他的蒸餾方法做一個比較,訓練的 ResNet-18 能在下游的 VOC 上獲得了 78.45 的 mAP,比先前的 KD 超過了 0.85 的 mAP 增益。
(4)將 CIFAR-100 訓練好的模型,固定特征提取器部分,只訓練新添加的分類器,來完成下游的 STL-10 和 TinyImageNet 分類任務。結果顯示,相比于先前最好的 HSAKD 性能增益依舊很大。
消融實驗部分,首先是對一個誤差項做一個消融實驗:
(1)上圖左面基于一個累逐漸累積的方式來做。結果顯示,每一個誤差項都對性能有提升的,其中通過 kl_q 來進行自監督的蒸餾表現增益是最大的。
(2)上圖右邊是一個基于輔助分類器的消融實驗,可以發現就是說從左到右一次使用了單獨分析了 c1、c2、c3,使用每一個分類器來學習并遷移自監督增強的知識都能有性能增益,其中深層的分類器通常比淺層的分類器能夠得到更多的性能增益。原因是深層的分類器所接受的特征通常是比較高階段特征,這些特征通常是具有高階的語義信息,一般圖像分類的話,使用這部分特征所產生的信息一般是更加具有泛化能力的。
最后聚合所有的分類器來進行自蒸餾,能夠最大化增益。
總結與展望
總結包含三個方面:
1. 知識定義:引導了一個分類網絡上的輔助結構來產生自監督增強的概率分布,該分布包含了原始有監督的分類任務和無監督任務的聯合知識;
2. 知識遷移:利用中間的輔助分類器來進行一對一概率化知識的傳遞。使用高度抽象化的概率知識同樣緩解了在特征匹配時語義層次不一致的問題;
3. 實驗結果:HSAKD 顯著地超過了先前基于自監督知識蒸餾方法 SSKD。同時利用自監督表示學習來進行 知識蒸餾通常也能使得網絡學習到更容易泛化的特征表示,有利于下游的語義識別的任務。
隨著自監督表示學習的發展,不同模式的算法被提出,本質上還是能夠使得模型學習到泛化能力更好特征表示。直覺上講,算法越好,其對應的知識表達形式在知識蒸餾領域就越可以 work。
所以,結合自監督的表示學習來進行有監督的知識蒸餾,目前仍然是一個具有前景的研究方向。
特別鳴謝
感謝 TCCI 天橋腦科學研究院對于 PaperWeekly 的支持。TCCI 關注大腦探知、大腦功能和大腦健康。
更多閱讀
#投 稿?通 道#
?讓你的文字被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體,縮短讀者尋找優質內容的成本呢?答案就是:你不認識的人。
總有一些你不認識的人,知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁,促使不同背景、不同方向的學者和學術靈感相互碰撞,迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人,在我們的平臺上分享各類優質內容,可以是最新論文解讀,也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個,讓知識真正流動起來。
📝?稿件基本要求:
? 文章確系個人原創作品,未曾在公開渠道發表,如為其他平臺已發表或待發表的文章,請明確標注?
? 稿件建議以?markdown?格式撰寫,文中配圖以附件形式發送,要求圖片清晰,無版權問題
? PaperWeekly 尊重原作者署名權,并將為每篇被采納的原創首發稿件,提供業內具有競爭力稿酬,具體依據文章閱讀量和文章質量階梯制結算
📬?投稿通道:
? 投稿郵箱:hr@paperweekly.site?
? 來稿請備注即時聯系方式(微信),以便我們在稿件選用的第一時間聯系作者
? 您也可以直接添加小編微信(pwbot02)快速投稿,備注:姓名-投稿
△長按添加PaperWeekly小編
🔍
現在,在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
·
與50位技術專家面對面20年技術見證,附贈技術全景圖總結
以上是生活随笔為你收集整理的IJCAI 2021 | 中科院计算所:自监督增强的知识蒸馏方法的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 想问大家一下目前市面上的门窗十 大品牌有
- 下一篇: 杀虫剂除甲醛资源方面?