微軟提出自動化神經(jīng)網(wǎng)絡(luò)訓(xùn)練剪枝框架OTO,一站式獲得高性能輕量化模型
在人工智能時代,如何部署和維護神經(jīng)網(wǎng)絡(luò)是產(chǎn)品化的關(guān)鍵問題考慮到節(jié)省運算成本,同時盡可能小地損失模型性能,壓縮神經(jīng)網(wǎng)絡(luò)成為了 DNN 產(chǎn)品化的關(guān)鍵之一。
DNN 壓縮通常來說有三種方式,剪枝,知識蒸餾和量化。剪枝旨在識別并去除冗余結(jié)構(gòu),給 DNN 瘦身的同時盡可能地保持模型性能,是最為通用且有效的壓縮方法。三種方法通常來講可以相輔相成,共同作用來達到最佳的壓縮效果。
然而現(xiàn)存的剪枝方法大都只針對特定模型,特定任務(wù),且需要很強的專業(yè)領(lǐng)域知識,因此通常需要 AI 開發(fā)人員花費大量的精力才能將這些方法應(yīng)用于自己的場景中,非常消耗人力物力成本。
OTO 概述
為了解決現(xiàn)存剪枝方法的問題并給 AI 開發(fā)者提供便利,微軟團隊提出了 Only-Train-Once OTO 框架。OTO 是業(yè)內(nèi)首個自動化、一站式、用戶友好且通用的神經(jīng)網(wǎng)絡(luò)訓(xùn)練與結(jié)構(gòu)壓縮框架,系列工作陸續(xù)發(fā)表于 ICLR2023 和 NeurIPS2021。
通過使用 OTO,AI 工程師可以方便地訓(xùn)練目標神經(jīng)網(wǎng)絡(luò),一站式地取得高性能且輕量化的模型。OTO 最小化了開發(fā)者的工程時間精力的投入,且全程無需現(xiàn)有方法通常需要的非常耗時的預(yù)訓(xùn)練和額外的模型微調(diào)。
- 論文鏈接:
- OTOv2 ICLR 2023: https://openreview.net/pdf?id=7ynoX1ojPMt
- OTOv1 NeurIPS 2021:https://proceedings.neurips.cc/paper_files/paper/2021/file/a376033f78e144f494bfc743c0be3330-Paper.pdf
- 代碼鏈接:
?https://github.com/tianyic/only_train_once
框架核心算法
理想化的結(jié)構(gòu)剪枝算法應(yīng)該做到:針對通用神經(jīng)網(wǎng)絡(luò),自動化地一站式地從零開始訓(xùn)練,同時達到高性能且輕量化的模型,并無需后續(xù)微調(diào)。但因為神經(jīng)網(wǎng)絡(luò)的復(fù)雜性,實現(xiàn)這一目標是一件極其有挑戰(zhàn)性的事情。為了實現(xiàn)這一最終目的,下面的三個核心問題需要被系統(tǒng)性地解決:
- 如何找出哪些網(wǎng)絡(luò)結(jié)構(gòu)可以被移除?
- 如何在移除網(wǎng)絡(luò)結(jié)構(gòu)的過程中,盡可能不損失模型性能?
- 如何能自動化地完成如上兩點?
微軟團隊設(shè)計并實現(xiàn)了三套核心算法,首次系統(tǒng)性地,全面性地解決這三個核心問題。
自動化 Zero-Invariant Groups (零不變組) 分組
由于網(wǎng)絡(luò)結(jié)構(gòu)的復(fù)雜性和關(guān)聯(lián)性,刪去任意網(wǎng)絡(luò)結(jié)構(gòu)可能會導(dǎo)致剩余的網(wǎng)絡(luò)結(jié)構(gòu)無效。因此自動化網(wǎng)絡(luò)結(jié)構(gòu)壓縮的一個最大的問題之一是如何找到必須要被一起剪枝的模型參數(shù),使得余下的網(wǎng)絡(luò)依然有效。為了解決該問題,微軟團隊在 OTOv1 中提出了零不變組 Zero-Invariant Groups (ZIGs)。零不變組可以理解為一類最小的可移除單元,使得該組對應(yīng)的網(wǎng)絡(luò)結(jié)構(gòu)移除后剩余網(wǎng)絡(luò)依然有效。零不變組的另一大特性是,如果一個零不變組等于零,那么無論輸入值是什么,輸出值永遠是零。在 OTOv2 中,研究者進一步提出并實現(xiàn)了一套自動化算法來解決通用網(wǎng)絡(luò)的零不變組的分組問題。自動化分組算法是由一系列圖算法精心設(shè)計組合而成,整套算法非常高效,具有線性的時間空間復(fù)雜度。
雙半平面投影梯度優(yōu)化算法 (DHSPG)
當劃分目標網(wǎng)絡(luò)所有的零不變組后,接下來的模型訓(xùn)練和剪枝任務(wù)需要找出哪些零不變組是冗余的,哪些又是重要的。冗余的零不變組對應(yīng)的網(wǎng)絡(luò)結(jié)構(gòu)需要被刪除,重要的零不變組需要保留并保證壓縮模型的性能。研究者公式化這個問題為一個結(jié)構(gòu)性稀疏化問題,并提出全新的 Dual Half-Space Projected Gradient (DHSPG) 優(yōu)化算法來解決。
DHSPG 可以非常有效地找出冗余的零不變組并將其投影成零,并持續(xù)訓(xùn)練重要的零不變組來取得跟原始模型相媲美的性能。
與傳統(tǒng)稀疏優(yōu)化算法相比,DHSPG 具有更強更穩(wěn)定地稀疏結(jié)構(gòu)探索能力,且擴展了訓(xùn)練搜索空間并因此通常要獲得更高的實際表現(xiàn)效果。
自動化構(gòu)建輕量壓縮模型
通過使用 DHSPG 對模型進行訓(xùn)練,我們會得到一個服從于零不變組的高結(jié)構(gòu)稀疏性的解,即該解里有很多被投影成零的零不變組,此外該解還會具有很高的模型性能。接下來,研究者把所有對應(yīng)與冗余零不變組的結(jié)構(gòu)刪去來自動化地構(gòu)建壓縮網(wǎng)絡(luò)。由于零不變組的特性,即如果一個零不變組等于零,那么無論輸入值是什么,輸出值永遠是零,因此刪去冗余的零不變組不會對網(wǎng)絡(luò)產(chǎn)生任何影響。所以通過 OTO 得到的壓縮網(wǎng)絡(luò)和完整網(wǎng)絡(luò)會有相同的輸出,無需傳統(tǒng)方法所需要的進一步模型微調(diào)。
數(shù)值實驗
分類任務(wù)
表 1:CIFAR10 中的 VGG16 及 VGG16-BN 模型表現(xiàn)
在 CIFAR10 的 VGG16 實驗中,OTO 將浮點數(shù)減少了 86.6%,將參數(shù)量減少了 97.5%,性能表現(xiàn)令人印象深刻。
表 2:CIFAR10 的 ResNet50 實驗
在 CIFAR10 的 ResNet50 實驗中,OTO 在沒有量化的情況下優(yōu)于 SOTA 神經(jīng)網(wǎng)絡(luò)壓縮框架 AMC 和 ANNC,僅使用了 7.8% 的 FLOPs 和 4.1% 的參數(shù)。
表 3. ImageNet 的 ResNet50 實驗
在 ImageNet 的 ResNet50 實驗中,OTOv2 在不同結(jié)構(gòu)稀疏化目標下,展現(xiàn)出跟現(xiàn)存 SOTA 方法相媲美甚至更優(yōu)的表現(xiàn)。
表 4: 更多結(jié)構(gòu)和數(shù)據(jù)集
OTO 也在更多的數(shù)據(jù)集和模型結(jié)構(gòu)上取得了不錯的表現(xiàn)。
Low-Level Vision 任務(wù)
表 4:CARNx2 的實驗
在 super-resolution 的任務(wù)中,OTO 一站式訓(xùn)練壓縮了 CARNx2 網(wǎng)絡(luò),得到了跟原始模型有競爭力的性能且壓縮了越 75% 的運算量和模型大小。
語言模型任務(wù)
此外,研究者還在 Bert 上針對核心算法之一,DHSPG 優(yōu)化算法,進行了對比試驗,驗證了其相較于其他稀疏優(yōu)化算法的高性能性??梢园l(fā)現(xiàn)在 Squad 上,利用 DHSPG 進行訓(xùn)練所得到的參數(shù)量減小和模型性能要遠遠優(yōu)于其他稀疏優(yōu)化算法。
結(jié)論
微軟團隊提出了一個名為 OTO(Only-Train-Once)的 自動化一站式神經(jīng)網(wǎng)絡(luò)訓(xùn)練結(jié)構(gòu)剪枝框架。它可以將一個完整的神經(jīng)網(wǎng)絡(luò)自動化地壓縮為輕量級網(wǎng)絡(luò),同時保持較高的性能。OTO 大大簡化了現(xiàn)有結(jié)構(gòu)剪枝方法復(fù)雜的多階段流程,適合各類網(wǎng)絡(luò)架構(gòu)和應(yīng)用,且最小化了用戶的額外工程投入,具有通用性,有效性和易用性。