最近，快手 Y-Tech 西雅圖 AI lab 聯合羅切斯特大學等研究者提出了一種基於能耗建模的壓縮方法，他們一脈相承的兩篇論文分別被 ICLR 2019 和 CVPR 2019 接收。在這篇文章中，我們將介紹這種新型模型壓縮的核心思想及主要做法，神經網絡壓縮也許該走向有目標的前進之路了。
模型壓縮的應用
目前快手的手機端應用，基本上都會進行模型壓縮，因為不管是能耗、效率，還是推斷速度，都需要滿足很多條件才能投入使用。不管是將深度學習模型部署到雲端還是移動端，模型壓縮都必不可少。
快手 Y-Tech 西雅圖 AI 團隊的負責人劉霽教授表示部署到手機端的模型，例如美顏相機、視頻的人臉識別、姿態識別等應用都會經過模型壓縮，其重點在於減小模型體積、降低模型能耗，以及保證推斷效率。
部署到伺服器也面臨類似問題，因為每一秒都會有成千上萬條調用請求，而這些請求需要在規定時間內得到響應。例如更新一次推薦視頻，快手的伺服器在接到請求以後需要立即計算，把最合適的 Top-N 視頻反饋到客戶端，這種響應甚至要控制在幾毫秒才能滿足用戶良好的體驗。
正因為有很多實際需求，快手嘗試構建更高效和有目的性的模型壓縮算法。在本文介紹的這兩篇論文中，他們以能耗為約束構建更高效和合理的壓縮方法。雖然論文以能耗為目標，但實際上可以把它換成推斷時間等不同的度量，因此這類神經網絡壓縮方法有更廣泛的應用。
除了這種適用於各種任務的模型壓縮，快手還會提出新方法以構建更緊湊的神經網絡，這本質上也是一種模型壓縮。例如在 Interspeech 2018 的 oral 論文中，語音組提出了一種能使用下文信息的門控循環單元，該模型能在極低延遲的情況下為快手提供語音識別、語音特效和語音評論等應用。
神經網絡壓縮新觀點
在模型壓縮方法中，最常見的就是剪枝、量化和知識蒸餾等，我們認為當參數和運算量減少時，模型的能耗和推斷延遲等也都會下降。但實際上，它們與能耗等度量並非簡單的正比關係，因此這種「想當然」的做法並非是最優方案。
例如模型剪枝，我們希望刪除單個不重要的神經元連接，或者直接刪除一組不重要的連接（channel 級的剪枝）。通常情況下，這種剪枝只關注減少參數量（或者模型大小）和運算量，並不考慮硬體條件或具體能耗等約束。劉霽表示經典的模型壓縮有一個隱含的假設：「在不同層刪除一條邊所節省的能量，或者說所提升的效率是等價的。」然而，由於物理硬體的複雜性，這個基本假設實際上並不絕對正確。有時候大模型可能比小模型的能量消耗更少或者推斷時間更短。
物理硬體實際上很複雜，如果我們用能耗作為約束，它會由幾部分組成：計算和數據加載等產生的能耗。算法工程師一般只關注計算複雜度，且通常不太關注數據加載的能耗。因此劉霽表示：「以前的研究工作對硬體本身缺乏足夠的關注，不同的硬體都是用同一個壓縮的模型，現在我們需要在充分理解硬體的基礎上，設計符合各種硬體的模型。」
劉霽等研究者認為目前的模型壓縮缺少一種導向，即為某個明確的目的進行壓縮，這種導向對當前模型壓縮研究很重要。他表示：「與標準剪枝或量化等壓縮方法相比，我們的研究重點在於找到指導模型進行剪枝或量化的目標。標準方法只需要模型變小變瘦就行了，但我們需要一種約束來引導優化模型，例如降低能耗或減少模型推斷時間等。」
也就是說，我們應該在給定能耗或延遲的約束下，找到滿足約束的最「精準」的模型。為了引入這種約束，我們需要解決兩大問題：如何對模型能耗或延遲建模；如何解這種帶約束的最優化問題。
針對這兩個問題，快手 Y-Tech 聯合羅切斯特大學近來提出了兩套解決方案，它們一脈相承且分別被 ICLR 2019 和 CVPR 2019 兩大頂會接收。以能耗為約束，表明這種模型壓縮框架能更有效地剪枝，從而滿足各種不同的應用。
基於能耗的模型壓縮
總的而言，這兩篇論文都在思考如何設計一種深度神經網絡，它能在滿足給定能耗的情況下最大化準確率。這兩篇論文都將這種能耗約束加入到最優化過程，並通過端到端的方法學習如何進行更高效的剪枝。
論文：Energy-Constrained Compression for Deep Neural Networks via Weighted Sparse Projection and Layer Input Masking（ICLR 2019）論文地址：https://openreview.net/pdf?id=BylBr3C9K7
論文：ECC: Energy-Constrained Deep Neural Network Compression via a Bilinear Regression Model（CVPR 2019）論文地址：https://arxiv.org/pdf/1812.01803.pdf從方法上來說，劉霽教授表示我們既可以通過數學分析的形式對模型能耗建模，也可以通過數據驅動的方式建模能耗，為了對硬體的行為了解更加深刻和清晰，此該工作與羅切斯特大學計算機體系結構教授緊密配合。ICLR 那篇就通過研究核心的底層硬體確定近似能耗，這種對 GPU 等晶片的分析方法可以快速計算不同剪枝網絡的能耗，從而指導剪枝過程。CVPR 那篇通過我們更熟的數據驅動方法，希望模型能學習不同網絡與具體能耗的關係。
對於帶約束的最優化方法，其解法並非被廣泛熟知。目前深度學習大多數都是無約束優化，直接調用各種優化器就能解決。反觀帶約束的最優化方法，我們沒有成熟的優化器，很多都需要根據針對具體問題實現新的最優化方法。正如劉霽教授所言：「並不是大家沒從約束的角度考慮模型壓縮，而是帶約束的最優化確實需要更多技巧與方法，我們需要更深入的工作來解決這些問題。」
總體而言，基於能耗的模型壓縮過程如 ECC 那篇論文的圖 1 所示：ECC 通過數據驅動離線建模能耗，再將其應用到帶約束的在線最優化過程中。
能耗的直接建模
前面已經介紹能耗的建模可以通過數學分析，但這需要我們對硬體的理解比較充分。在 ICLR 的那篇論文中，劉霽等研究者重點關注如何對卷積層與全連接層的能耗建模，因為它們在推斷過程中佔了 90% 以上的執行時間與能耗。其它模塊的能耗比較少，可以近似地看作一個很小的常量。
神經網絡的推斷能耗與硬體底層架構有關，該論文研究了 TPU 和 Volta GPU 等硬體都採用的矩陣運算方法脈動陣列（systolic array）。對該硬體進行能耗建模最終就能計算出不同模型的能耗，如下所示為計算矩陣乘法時的主要能耗，其中脈動陣列包含一些能執行乘加運算（MAC）的計算單元。
總能耗主要可以分為 MAC 的計算能耗與數據讀寫的能耗，這篇論文的總體建模即弄清楚模型稀疏性如何影響這兩部分能耗。
由於硬體的多樣性和複雜性，對能耗進行分析建模通常並不容易，機器學習研究者對底層硬體也不是那麼了解。為此在最近 CVPR 2019 提出的 ECC 中，Y-Tech 團隊提出另一種「更優美」的數據驅動方法。劉霽教授表示這種方法的核心思想即：「給定一個網絡，我們隨機選擇每一層的節點數，然後把這樣的模型放到硬體測試它的能耗或延遲，這樣就能構造一個很全面的數據集。如果再用某個模型根據這個數據集建模就能學會神經網絡與能耗之間的關係。」
值得注意的是，這兩種方法雖然都在建立神經網絡稀疏性與能耗之間的關係，但第一篇 ICLR 2019 論文主要通過權重級的細粒度剪枝獲得稀疏性，第二篇 CVPR 2019 論文主要通過 Channel 級的粗粒度剪枝獲得稀疏性。除了剪枝，量化等方法在這裡也可以同時使用。不過 Channel 級的剪枝更適用於各種硬體，我們不需要為它準備特殊的硬體架構。
在第二篇 ECC 中，研究者通過雙線性模型和構造的數據集對稀疏性和能耗之間的關係進行建模。這種 DNN 能耗模型能移植到不同的硬體架構中，它只需要將硬體平臺視為黑盒就行了。此外，因為雙線性模型是可微的，所以學習過程也極其簡單，使用傳統基於梯度的方法就夠了。
具體而言，對於能耗模型 ε(s)，可微的目標函數可以表示為以下：
其中ε(s) 表示在給定稀疏性 s 情況下的真實能耗，我們希望用雙線性函數 f(s) 逼近真實能耗。一般而言，DNN 層級的能耗是受輸入通道數與輸出通道數影響的，它們又相當於當前層和後一層的稀疏性。因此可以簡單地用雙線性模型為 s 建模：
其中 a_0 到 a_j 為可學習參數，|U| 表示層級數量。注意 a_0 這個偏置項常數，劉霽教授表示：「我們採集的硬體平臺數據會包含一些基礎能耗，這一部分是和模型推斷無關的，因此雙線性模型中的常數項很好地對這些能耗進行了建模。」
使用雙線性建模 Channel 級剪枝的能耗也是很有道理的，因為 DNN 在推斷過程中總的算術運算數大致上是一種雙線性形式。劉霽教授表示其它壓縮方法可能並不適合雙線性模型，但是我們可以用神經網絡等非線性模型對能耗建模。只不過在這個案例中因為雙線性模型的簡單和緊湊，它的效果非常好。
最後，給定硬體和神經網絡的組合，ECC 中能耗模型的建立完全是離線完成的，它沒有運行時開銷。
帶約束的最優化
如果我們將能耗模型作為剪枝操作的指導，那麼就可以構建一個帶約束的最優化過程。這兩篇論文的解法並不相同，劉霽教授表明：「一般帶約束的最優化方法有幾種常見解法，首先如同 ICLR 那篇採用的是 Projection 方法，即先執行 SGD，然後執行 Projection 將解映射到約束中以令它強行滿足條件。如果約束比較複雜，那麼 Projection 可能很難計算，這個時候就需要 CVPR 那篇採用的 ADMM 算法了。」
這一部分主要介紹第二篇 CVPR 所採用的 ADMM+ Adam 解法，當然這裡只簡要介紹核心思想，詳細的更新過程可參考原論文。對於 Channel 級的權重剪枝，首先我們需要形式化地表達整個帶約束的最優化問題：
因為 ADMM 方法並不能直接使用，我們需要把複雜約束換為多個簡單約束。其中我們希望在給定權重 W 的情況下最小化損失函數，且每一層的權重 w 都需要滿足某種程度的稀疏性，而這些稀疏性又需要滿足總體的能耗約束。
劉霽教授表示，ADMM 本質上會根據約束的滿足程度給出一個懲罰：如果解滿足約束，那麼就不會給解加上懲罰，如果解不滿足約束，那麼懲罰會根據不滿足程度變大而強行令解滿足約束。所以一定意義上，ADMM 可以看作是一種特殊的增廣拉格朗日乘子法。拉格朗日乘子法算我們比較熟悉的了，那麼上面方程 2 可以等價地構造為增廣拉格朗日函數的極小極大問題：
其中新引入的 y 和 z 分別是方程 2b 和 2c 的對偶變量，L 為增廣拉格朗日函數：即 L(W, s, y, z) := l(W) + L_1(W, s, y) + L_2(s, z)。
對於 ADMM 而言，基本步驟還是使用 SGD 類的迭代方法如 Adam 更新權重，不過這一更新會加上一個「罰項」L_1，即在更新權重時希望它滿足某個權重稀疏性的約束。其次在更新稀疏性 s 時，因為 L_1 與 L_2 兩項（2b 和 2c）都與稀疏性 s 相關，所以 SGD 可直接極小化這兩項而更新 s。最後，按照拉格朗日乘子法的常規解法，再更新對偶變量 y 和 z 就行了。我們可以直觀理解為，如果不滿足程度越大，那麼對偶變量會增大，它們給出的懲罰也會變大。
如下所示為增廣拉格朗日函數兩個子項 L_1 與 L_2 的表達式：
劉霽教授表示，L_1 和 L_2 的第一項可以看成二階的懲罰項，其中「+」號表示滿足條件下是沒任何懲罰的，只懲罰不滿足條件的解。而第二項可以看作一階的罰項，它在滿足條件下甚至還有「獎勵」（罰項為負）。另外，我們也可以從拉格朗日乘子法的角度理解，前一項表示增廣拉格朗日函數項，後一項一般表示拉格朗日函數項。
解決這個帶約束的最優化問題後，有目的地剪枝就不成問題了。最後，兩篇論文的實驗結果表明，在不同的硬體上對不同的神經網絡進行壓縮，能耗都要顯著低於其它方法，感興趣的讀者可以查閱原論文。