時隔 5 年，network slimming 原作者團隊打造出了 ViT 版的 slimming，入選 CVPR 2022。
卷積網絡版的 network slimming 在 ICCV 2017 上被提出，作為神經網絡剪枝領域代表性的工作之一，目前已被引用超過 1400 次，五年之後的 CVPR 2022，原作者團隊跟 Meta、印度理工學院等機構的研究者聯合打造了 ViT 版的 slimming，據悉該工作得到了四個審稿人一致推薦接收！
論文地址：https://arxiv.org/pdf/2201.00814.pdf
代碼地址：https://github.com/Arnav0400/ViT-Slim
是什麼原因讓 network slimming 的 ViT 版本得到所有審稿人的一致青睞呢，這裡還得說一下 network slimming 這種方法的幾個特點，即：簡單 (simple) + 有效 (effective) + 靈活 (flexible)。通過簡單的引入指示因子，slimming 可以邊訓練邊搜索，類似於隱式的參數共享機制，壓縮後的目標網絡結構通過排序指示因子就可獲得，非常高效方便。ViT-Slim 繼承了這些優點，同時針對 ViT 主幹網絡結構的特性做出了幾個改進。在介紹改進前先回顧一下 network slimming 的方法和原理。用一句話概括就是，對於卷積神經網絡，每層每個卷積核都用一個重要性係數去表徵它對最後預測結果的貢獻程度，如果這個係數接近於 0，就認為它是沒有貢獻並可以被去掉的。作者通過施加一個
正則使得這些係數變得稀疏，同時作者發現由於通常會將卷積核與批量歸一（BN）一起使用，因此可以直接使用 BN 層的縮放係數
作為每層卷積的重要性係數，這樣就不需要引入額外的參數了。對於 ViT Slimming，需要解決的核心問題是如何在沒有 BN 層的情況下設計合適的指示參數來反映不同模塊規模對於整個分類性能的重要程度。文章採用的是通過顯式地定義可導的 soft mask 來確定每個模塊的大小和尺度，整個過程如下圖所示。
ViT Slimming 同時考慮了 ViT 裡面三個模塊：輸入 tokens，MHSA 和 MLP 模塊，它的訓練過程可以看成是一個基於權重共享的子網絡搜索過程，主幹網的權重可以通過加載預訓練好的模型參數來加快搜索。訓練過程中 soft mask 接近 0 對應的特徵相當於動態地被裁剪 / 丟棄了。該方法的優點是只需訓練 / 搜索一次，就可以通過排序 mask 得到無數個子網絡，非常高效靈活。算法優化函數由原始的 cross-entropy 和
稀疏正則組成，如下所示：
其中在 ViT 的 MHSA 模塊內部加 soft mask 的處理如下：
在 MLP 模塊加 soft mask 的處理如下：
下面是 ViT-Slim 跟其他一些 ViT 壓縮或者瘦身方法的區別，主要在於搜索方法以及使用範圍的不同。
相比其他方法，ViT Slimming 可以同時減少模型參數量和計算量，同時由於可以繼承已經正常訓練好的模型參數，搜索過程需要的時間非常少，對資源需要也不是很大。
可以看出在更少的參數及計算量的條件下，ViT-Slim 可以得到更好的性能。
1. 搜索得到的每個 Attention head 結構：
可以看到搜索得到的 attention head 結構數目呈現中間大，兩邊小（高層更加明顯）的現象。
可以看到搜索得到的 MLP 結構同樣呈現中間大，兩邊小的現象，這也是一個非常有意思的地方。
©THE END
轉載請聯繫本公眾號獲得授權
投稿或尋求報道：content@jiqizhixin.com
留言列表