首次將「教導主任」引入模型蒸餾，大規(guī)模壓縮優(yōu)于24種SOTA方法

互聯(lián)網(wǎng) 2022-11-30 11:17:25

面對越來越深的深度學習模型和海量的視頻大數(shù)據(jù)，人工智能算法對計算資源的依賴越來越高。為了有效提升深度模型的性能和效率，通過探索模型的可蒸餾性和可稀疏性，本文提出了一種基于 "教導主任 - 教師 - 學生" 模式的統(tǒng)一的模型壓縮技術。

該成果由人民中科和中科院自動化所聯(lián)合研究團隊合作完成，相關論文發(fā)表在人工智能頂級國際期刊 IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI) 上。該成果是首次將 "教導主任" 角色引入模型蒸餾技術，對深度模型的蒸餾與裁剪進行了統(tǒng)一。

論文地址：https://ieeexplore.ieee.org/abstract/document/9804342

目前該項成果已經(jīng)應用于人民中科自主研發(fā)的跨模態(tài)智能搜索引擎 "白澤"。"白澤" 打破圖文音視等不同模態(tài)間信息表達的隔閡，將文字、圖片、語音和視頻等不同模態(tài)信息映射到一個統(tǒng)一特征表示空間，以視頻為核心，學習多個模態(tài)間統(tǒng)一的距離度量，跨越文字、語音、視頻等多模態(tài)內(nèi)容的語義鴻溝，實現(xiàn)大一統(tǒng)的搜索能力。

然而面對海量的互聯(lián)網(wǎng)數(shù)據(jù)尤其是視頻大數(shù)據(jù)，跨模態(tài)的深度模型對計算資源的消耗逐漸提升?；谠擁椦芯砍晒?，"白澤"能夠在保證算法性能的情況下，將模型大小進行大規(guī)模壓縮，從而實現(xiàn)高通量低功耗的跨模態(tài)智能理解和搜索能力。根據(jù)初步的實際應用情況來看，該項技術能夠?qū)⒋竽Ｐ偷膮?shù)規(guī)模壓縮平均四倍以上。一方面能夠極大降低模型對 GPU 服務器等高性能計算資源的消耗，另一方面能夠?qū)o法在邊緣端部署的大模型經(jīng)過蒸餾壓縮后實現(xiàn)邊緣端的低功耗部署。

模型壓縮的聯(lián)合學習框架

深度算法模型的壓縮和加速可通過蒸餾學習或結構化稀疏裁剪實現(xiàn)，但這兩個領域均存在一些局限性。對于蒸餾學習方法，旨在訓練一個輕量化模型（即學生網(wǎng)絡）來模擬復雜龐大的模型（即教師網(wǎng)絡）。在教師網(wǎng)絡的指導下，學生網(wǎng)絡可以獲得比單獨訓練的更優(yōu)性能。

然而，蒸餾學習算法僅僅專注于提升學生網(wǎng)絡的性能，往往忽略了網(wǎng)絡結構的重要性。學生網(wǎng)絡的結構一般是預定義好的，并且在訓練過程中是固定的。

對于結構化稀疏裁剪或濾波器裁剪，這些方法旨在將一個冗余繁雜的網(wǎng)絡裁剪成一個稀疏緊致的網(wǎng)絡。然而，模型裁剪僅僅用于獲得一個緊致的結構。目前已有方法都沒有充分利用原始復雜模型所包含的"知識"。近期研究為了平衡模型性能和大小，將蒸餾學習和結構化稀疏裁剪進行結合。但是這些方法僅限于簡單的損失函數(shù)的結合。

為了深入分析以上問題，該研究首先對模型進行基于壓縮感知訓練，通過分析模型性能和結構發(fā)現(xiàn)，對于深度算法模型，存在兩個重要屬性：可蒸餾性（distillability）和可稀疏性（sparsability）。

具體而言，可蒸餾性指的是能夠從教師網(wǎng)絡中蒸餾出有效知識的密度。它可以通過學生網(wǎng)絡在教師網(wǎng)絡指導下所獲得的性能收益來衡量。例如，擁有更高可蒸餾性的學生網(wǎng)絡可以獲得更高性能?？烧麴s性也可以在網(wǎng)絡層級別上被定量分析。

如圖 1-(a)所示，柱形圖表示蒸餾學習損失梯度和真值分類損失梯度之間的余弦相似度（Cosine Similarity）。更大的余弦相似度說明當前蒸餾的知識對于模型性能更有幫助。這樣，余弦相似度也可以成為可蒸餾性的一種度量。由圖 1-(a)可得，可蒸餾性隨著模型層數(shù)變深逐漸增大。這也解釋了為什么常規(guī)使用蒸餾學習的監(jiān)督均施加在模型最后幾層中。并且，在不同的訓練輪次，學生模型也有不同的可蒸餾性，因為隨著訓練時間變化余弦相似度也在改變。因此，在訓練過程中對不同層進行可蒸餾性的動態(tài)分析十分必要。

另一方面，可稀疏性指的是模型在有限精度損失下能夠獲得的裁剪率（或稱壓縮率）。更高的可稀疏性對應更高裁剪率的潛力。如圖 1-(b)所示，網(wǎng)絡的不同層或模塊展現(xiàn)了不同的可稀疏性。類似于可蒸餾性，可稀疏性也可以在網(wǎng)絡層級別和時間維度進行分析。然而，目前沒有方法去探索和分析可蒸餾性和可稀疏性。現(xiàn)有方法常常使用一種固定的訓練機制，這樣很難達到一個最優(yōu)結果。

圖 1 深度神經(jīng)網(wǎng)絡的可蒸餾性和可稀疏性示意圖

為了解決以上問題，該研究分析了模型壓縮的訓練過程，從而獲得有關可蒸餾性和可稀疏性的相關發(fā)現(xiàn)。受這些發(fā)現(xiàn)啟發(fā)，該研究提出了一種基于動態(tài)可蒸餾性與可稀疏性聯(lián)合學習的模型壓縮方法。它能動態(tài)結合蒸餾學習和結構化稀疏裁剪，通過學習可蒸餾性和可稀疏性，自適應地調(diào)節(jié)聯(lián)合訓練機制。

與常規(guī)的 "教師 - 學生（Teacher-Student）" 框架不同，本文提出的方法能夠被描述成 "在學校學習（Learning-in-School）" 框架，因為它包含三大模塊：教師網(wǎng)絡，學生網(wǎng)絡和教導主任網(wǎng)絡。

具體而言，與之前相同，教師網(wǎng)絡教導學生網(wǎng)絡。而教導主任網(wǎng)絡負責控制學生網(wǎng)絡學習的強度以及學習的方式。通過獲取當前教師網(wǎng)絡和學生網(wǎng)絡的狀態(tài)，教導主任網(wǎng)絡可以評估當前學生網(wǎng)絡的可蒸餾性和可稀疏性，然后動態(tài)地平衡和控制蒸餾學習監(jiān)督和結構化稀疏裁剪監(jiān)督的強度。

為了優(yōu)化本文方法，該研究還提出一種基于交替方向乘子法的蒸餾學習 & 裁剪的聯(lián)合優(yōu)化算法，來更新學生網(wǎng)絡。為了優(yōu)化和更新教導主任網(wǎng)絡，本文提出一種基于元學習的教導主任優(yōu)化算法。借助動態(tài)調(diào)節(jié)監(jiān)督信號，反過來可蒸餾性也能被影響。如圖 1-(a)所示，本文方法證明能夠延緩可蒸餾性的下降趨勢，并且通過合理利用蒸餾的知識，提升了整體的可蒸餾性。

本文方法的整體算法框架和流程圖如下圖所示。該框架包含三大模塊，教師網(wǎng)絡，學生網(wǎng)絡和教導主任網(wǎng)絡。其中，初始的待壓縮裁剪的復雜冗余網(wǎng)絡被看作教師網(wǎng)絡，而在后面的訓練過程中，逐漸被稀疏的原始網(wǎng)絡被看作是學生網(wǎng)絡。教導主任網(wǎng)絡是一個元網(wǎng)絡，它輸入教師網(wǎng)絡和學生網(wǎng)絡的信息來衡量當前可蒸餾性和可稀疏性，從而控制蒸餾學習和稀疏的監(jiān)督強度。

這樣，在每一時刻，學生網(wǎng)絡都能被動態(tài)地蒸餾知識指導和被稀疏。例如，當學生網(wǎng)絡有更高的可蒸餾性，則教導主任會讓更強的蒸餾監(jiān)督信號指導學生網(wǎng)絡（見圖 2 中粉色箭頭信號）；與此相反，當學生網(wǎng)絡有更高的可稀疏性，教導主任會讓更強的稀疏監(jiān)督信號施加于學生網(wǎng)絡中（見圖 2 中橙色箭頭信號）。