在人工智能時(shí)代，如何部署和維護(hù)神經(jīng)網(wǎng)絡(luò)是產(chǎn)品化的關(guān)鍵問(wèn)題考慮到節(jié)省運(yùn)算成本，同時(shí)盡可能小地?fù)p失模型性能，壓縮神經(jīng)網(wǎng)絡(luò)成為了 DNN 產(chǎn)品化的關(guān)鍵之一。

DNN 壓縮通常來(lái)說(shuō)有三種方式，剪枝，知識(shí)蒸餾和量化。剪枝旨在識(shí)別并去除冗余結(jié)構(gòu)，給 DNN 瘦身的同時(shí)盡可能地保持模型性能，是最為通用且有效的壓縮方法。三種方法通常來(lái)講可以相輔相成，共同作用來(lái)達(dá)到最佳的壓縮效果。

然而現(xiàn)存的剪枝方法大都只針對(duì)特定模型，特定任務(wù)，且需要很強(qiáng)的專(zhuān)業(yè)領(lǐng)域知識(shí)，因此通常需要 AI 開(kāi)發(fā)人員花費(fèi)大量的精力才能將這些方法應(yīng)用于自己的場(chǎng)景中，非常消耗人力物力成本。

OTO 概述

為了解決現(xiàn)存剪枝方法的問(wèn)題并給 AI 開(kāi)發(fā)者提供便利，微軟團(tuán)隊(duì)提出了 Only-Train-Once OTO 框架。OTO 是業(yè)內(nèi)首個(gè)自動(dòng)化、一站式、用戶(hù)友好且通用的神經(jīng)網(wǎng)絡(luò)訓(xùn)練與結(jié)構(gòu)壓縮框架，系列工作陸續(xù)發(fā)表于 ICLR2023 和 NeurIPS2021。

通過(guò)使用 OTO，AI 工程師可以方便地訓(xùn)練目標(biāo)神經(jīng)網(wǎng)絡(luò)，一站式地取得高性能且輕量化的模型。OTO 最小化了開(kāi)發(fā)者的工程時(shí)間精力的投入，且全程無(wú)需現(xiàn)有方法通常需要的非常耗時(shí)的預(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)該做到：針對(duì)通用神經(jīng)網(wǎng)絡(luò)，自動(dòng)化地一站式地從零開(kāi)始訓(xùn)練，同時(shí)達(dá)到高性能且輕量化的模型，并無(wú)需后續(xù)微調(diào)。但因?yàn)樯窠?jīng)網(wǎng)絡(luò)的復(fù)雜性，實(shí)現(xiàn)這一目標(biāo)是一件極其有挑戰(zhàn)性的事情。為了實(shí)現(xiàn)這一最終目的，下面的三個(gè)核心問(wèn)題需要被系統(tǒng)性地解決：

• 如何找出哪些網(wǎng)絡(luò)結(jié)構(gòu)可以被移除？
• 如何在移除網(wǎng)絡(luò)結(jié)構(gòu)的過(guò)程中，盡可能不損失模型性能？
• 如何能自動(dòng)化地完成如上兩點(diǎn)？

微軟團(tuán)隊(duì)設(shè)計(jì)并實(shí)現(xiàn)了三套核心算法，首次系統(tǒng)性地，全面性地解決這三個(gè)核心問(wèn)題。

自動(dòng)化 Zero-Invariant Groups (零不變組) 分組

由于網(wǎng)絡(luò)結(jié)構(gòu)的復(fù)雜性和關(guān)聯(lián)性，刪去任意網(wǎng)絡(luò)結(jié)構(gòu)可能會(huì)導(dǎo)致剩余的網(wǎng)絡(luò)結(jié)構(gòu)無(wú)效。因此自動(dòng)化網(wǎng)絡(luò)結(jié)構(gòu)壓縮的一個(gè)最大的問(wèn)題之一是如何找到必須要被一起剪枝的模型參數(shù)，使得余下的網(wǎng)絡(luò)依然有效。為了解決該問(wèn)題，微軟團(tuán)隊(duì)在 OTOv1 中提出了零不變組 Zero-Invariant Groups  (ZIGs)。零不變組可以理解為一類(lèi)最小的可移除單元，使得該組對(duì)應(yīng)的網(wǎng)絡(luò)結(jié)構(gòu)移除后剩余網(wǎng)絡(luò)依然有效。零不變組的另一大特性是，如果一個(gè)零不變組等于零，那么無(wú)論輸入值是什么，輸出值永遠(yuǎn)是零。在 OTOv2 中，研究者進(jìn)一步提出并實(shí)現(xiàn)了一套自動(dòng)化算法來(lái)解決通用網(wǎng)絡(luò)的零不變組的分組問(wèn)題。自動(dòng)化分組算法是由一系列圖算法精心設(shè)計(jì)組合而成，整套算法非常高效，具有線(xiàn)性的時(shí)間空間復(fù)雜度。

雙半平面投影梯度優(yōu)化算法 (DHSPG)

當(dāng)劃分目標(biāo)網(wǎng)絡(luò)所有的零不變組后，接下來(lái)的模型訓(xùn)練和剪枝任務(wù)需要找出哪些零不變組是冗余的，哪些又是重要的。冗余的零不變組對(duì)應(yīng)的網(wǎng)絡(luò)結(jié)構(gòu)需要被刪除，重要的零不變組需要保留并保證壓縮模型的性能。研究者公式化這個(gè)問(wèn)題為一個(gè)結(jié)構(gòu)性稀疏化問(wèn)題，并提出全新的 Dual Half-Space Projected Gradient  (DHSPG) 優(yōu)化算法來(lái)解決。

DHSPG 可以非常有效地找出冗余的零不變組并將其投影成零，并持續(xù)訓(xùn)練重要的零不變組來(lái)取得跟原始模型相媲美的性能。

與傳統(tǒng)稀疏優(yōu)化算法相比，DHSPG 具有更強(qiáng)更穩(wěn)定地稀疏結(jié)構(gòu)探索能力，且擴(kuò)展了訓(xùn)練搜索空間并因此通常要獲得更高的實(shí)際表現(xiàn)效果。

自動(dòng)化構(gòu)建輕量壓縮模型

通過(guò)使用 DHSPG 對(duì)模型進(jìn)行訓(xùn)練，我們會(huì)得到一個(gè)服從于零不變組的高結(jié)構(gòu)稀疏性的解，即該解里有很多被投影成零的零不變組，此外該解還會(huì)具有很高的模型性能。接下來(lái)，研究者把所有對(duì)應(yīng)與冗余零不變組的結(jié)構(gòu)刪去來(lái)自動(dòng)化地構(gòu)建壓縮網(wǎng)絡(luò)。由于零不變組的特性，即如果一個(gè)零不變組等于零，那么無(wú)論輸入值是什么，輸出值永遠(yuǎn)是零，因此刪去冗余的零不變組不會(huì)對(duì)網(wǎng)絡(luò)產(chǎn)生任何影響。所以通過(guò) OTO 得到的壓縮網(wǎng)絡(luò)和完整網(wǎng)絡(luò)會(huì)有相同的輸出，無(wú)需傳統(tǒng)方法所需要的進(jìn)一步模型微調(diào)。

數(shù)值實(shí)驗(yàn)

分類(lèi)任務(wù)

表 1:CIFAR10 中的 VGG16 及 VGG16-BN 模型表現(xiàn)

在 CIFAR10 的 VGG16 實(shí)驗(yàn)中，OTO 將浮點(diǎn)數(shù)減少了 86.6%，將參數(shù)量減少了 97.5%，性能表現(xiàn)令人印象深刻。

表 2:CIFAR10 的 ResNet50 實(shí)驗(yàn)

在 CIFAR10 的 ResNet50 實(shí)驗(yàn)中，OTO 在沒(méi)有量化的情況下優(yōu)于 SOTA 神經(jīng)網(wǎng)絡(luò)壓縮框架 AMC 和 ANNC，僅使用了 7.8% 的 FLOPs 和 4.1% 的參數(shù)。

表 3. ImageNet 的 ResNet50 實(shí)驗(yàn)

在 ImageNet 的 ResNet50 實(shí)驗(yàn)中，OTOv2 在不同結(jié)構(gòu)稀疏化目標(biāo)下，展現(xiàn)出跟現(xiàn)存 SOTA 方法相媲美甚至更優(yōu)的表現(xiàn)。

表 4: 更多結(jié)構(gòu)和數(shù)據(jù)集

OTO 也在更多的數(shù)據(jù)集和模型結(jié)構(gòu)上取得了不錯(cuò)的表現(xiàn)。

Low-Level Vision 任務(wù)

表 4:CARNx2 的實(shí)驗(yàn)

在 super-resolution 的任務(wù)中，OTO 一站式訓(xùn)練壓縮了 CARNx2 網(wǎng)絡(luò)，得到了跟原始模型有競(jìng)爭(zhēng)力的性能且壓縮了越 75% 的運(yùn)算量和模型大小。 

語(yǔ)言模型任務(wù)

此外，研究者還在 Bert 上針對(duì)核心算法之一，DHSPG 優(yōu)化算法，進(jìn)行了對(duì)比試驗(yàn)，驗(yàn)證了其相較于其他稀疏優(yōu)化算法的高性能性。可以發(fā)現(xiàn)在 Squad 上，利用 DHSPG 進(jìn)行訓(xùn)練所得到的參數(shù)量減小和模型性能要遠(yuǎn)遠(yuǎn)優(yōu)于其他稀疏優(yōu)化算法。

結(jié)論

微軟團(tuán)隊(duì)提出了一個(gè)名為 OTO（Only-Train-Once）的 自動(dòng)化一站式神經(jīng)網(wǎng)絡(luò)訓(xùn)練結(jié)構(gòu)剪枝框架。它可以將一個(gè)完整的神經(jīng)網(wǎng)絡(luò)自動(dòng)化地壓縮為輕量級(jí)網(wǎng)絡(luò)，同時(shí)保持較高的性能。OTO 大大簡(jiǎn)化了現(xiàn)有結(jié)構(gòu)剪枝方法復(fù)雜的多階段流程，適合各類(lèi)網(wǎng)絡(luò)架構(gòu)和應(yīng)用，且最小化了用戶(hù)的額外工程投入，具有通用性，有效性和易用性。