機器之心專欄
作者： 陳天翼-微軟西雅圖-高級研究員
OTO 是業內首個自動化、一站式、用戶友好且通用的神經網絡訓練與結構壓縮框架。
在人工智慧時代，如何部署和維護神經網絡是產品化的關鍵問題考慮到節省運算成本，同時儘可能小地損失模型性能，壓縮神經網絡成為了 DNN 產品化的關鍵之一。
DNN 壓縮通常來說有三種方式，剪枝，知識蒸餾和量化。剪枝旨在識別並去除冗餘結構，給 DNN 瘦身的同時儘可能地保持模型性能，是最為通用且有效的壓縮方法。三種方法通常來講可以相輔相成，共同作用來達到最佳的壓縮效果。
然而現存的剪枝方法大都只針對特定模型，特定任務，且需要很強的專業領域知識，因此通常需要 AI 開發人員花費大量的精力才能將這些方法應用於自己的場景中，非常消耗人力物力成本。
OTO 概述
為了解決現存剪枝方法的問題並給 AI 開發者提供便利，微軟團隊提出了 Only-Train-Once OTO 框架。OTO 是業內首個自動化、一站式、用戶友好且通用的神經網絡訓練與結構壓縮框架，系列工作陸續發表於 ICLR2023 和 NeurIPS2021。
通過使用 OTO，AI 工程師可以方便地訓練目標神經網絡，一站式地取得高性能且輕量化的模型。OTO 最小化了開發者的工程時間精力的投入，且全程無需現有方法通常需要的非常耗時的預訓練和額外的模型微調。
論文連結：
- 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
框架核心算法
理想化的結構剪枝算法應該做到：針對通用神經網絡，自動化地一站式地從零開始訓練，同時達到高性能且輕量化的模型，並無需後續微調。但因為神經網絡的複雜性，實現這一目標是一件極其有挑戰性的事情。為了實現這一最終目的，下面的三個核心問題需要被系統性地解決：
如何找出哪些網絡結構可以被移除？
如何在移除網絡結構的過程中，儘可能不損失模型性能？
如何能自動化地完成如上兩點？
微軟團隊設計並實現了三套核心算法，首次系統性地，全面性地解決這三個核心問題。
自動化 Zero-Invariant Groups (零不變組) 分組
由於網絡結構的複雜性和關聯性，刪去任意網絡結構可能會導致剩餘的網絡結構無效。因此自動化網絡結構壓縮的一個最大的問題之一是如何找到必須要被一起剪枝的模型參數，使得餘下的網絡依然有效。為了解決該問題，微軟團隊在 OTOv1 中提出了零不變組 Zero-Invariant Groups (ZIGs)。零不變組可以理解為一類最小的可移除單元，使得該組對應的網絡結構移除後剩餘網絡依然有效。零不變組的另一大特性是，如果一個零不變組等於零，那麼無論輸入值是什麼，輸出值永遠是零。在 OTOv2 中，研究者進一步提出並實現了一套自動化算法來解決通用網絡的零不變組的分組問題。自動化分組算法是由一系列圖算法精心設計組合而成，整套算法非常高效，具有線性的時間空間複雜度。
雙半平面投影梯度優化算法 (DHSPG)
當劃分目標網絡所有的零不變組後，接下來的模型訓練和剪枝任務需要找出哪些零不變組是冗餘的，哪些又是重要的。冗餘的零不變組對應的網絡結構需要被刪除，重要的零不變組需要保留並保證壓縮模型的性能。研究者公式化這個問題為一個結構性稀疏化問題，並提出全新的 Dual Half-Space Projected Gradient (DHSPG) 優化算法來解決。
DHSPG 可以非常有效地找出冗餘的零不變組並將其投影成零，並持續訓練重要的零不變組來取得跟原始模型相媲美的性能。
與傳統稀疏優化算法相比，DHSPG 具有更強更穩定地稀疏結構探索能力，且擴展了訓練搜索空間並因此通常要獲得更高的實際表現效果。
自動化構建輕量壓縮模型
通過使用 DHSPG 對模型進行訓練，我們會得到一個服從於零不變組的高結構稀疏性的解，即該解里有很多被投影成零的零不變組，此外該解還會具有很高的模型性能。接下來，研究者把所有對應與冗餘零不變組的結構刪去來自動化地構建壓縮網絡。由於零不變組的特性，即如果一個零不變組等於零，那麼無論輸入值是什麼，輸出值永遠是零，因此刪去冗餘的零不變組不會對網絡產生任何影響。所以通過 OTO 得到的壓縮網絡和完整網絡會有相同的輸出，無需傳統方法所需要的進一步模型微調。
數值實驗
分類任務
表 1:CIFAR10 中的 VGG16 及 VGG16-BN 模型表現。
在 CIFAR10 的 VGG16 實驗中，OTO 將浮點數減少了 86.6%，將參數量減少了 97.5%，性能表現令人印象深刻。
表 2:CIFAR10 的 ResNet50 實驗。
在 CIFAR10 的 ResNet50 實驗中，OTO 在沒有量化的情況下優於 SOTA 神經網絡壓縮框架 AMC 和 ANNC，僅使用了 7.8% 的 FLOPs 和 4.1% 的參數。
表 3. ImageNet 的 ResNet50 實驗。
在 ImageNet 的 ResNet50 實驗中，OTOv2 在不同結構稀疏化目標下，展現出跟現存 SOTA 方法相媲美甚至更優的表現。
表 4: 更多結構和數據集。
OTO 也在更多的數據集和模型結構上取得了不錯的表現。
Low-Level Vision 任務
表 4:CARNx2 的實驗。
在 super-resolution 的任務中，OTO 一站式訓練壓縮了 CARNx2 網絡，得到了跟原始模型有競爭力的性能且壓縮了越 75% 的運算量和模型大小。
語言模型任務
此外，研究者還在 Bert 上針對核心算法之一，DHSPG 優化算法，進行了對比試驗，驗證了其相較於其他稀疏優化算法的高性能性。可以發現在 Squad 上，利用 DHSPG 進行訓練所得到的參數量減小和模型性能要遠遠優於其他稀疏優化算法。
結論
微軟團隊提出了一個名為 OTO（Only-Train-Once）的 自動化一站式神經網絡訓練結構剪枝框架。它可以將一個完整的神經網絡自動化地壓縮為輕量級網絡，同時保持較高的性能。OTO 大大簡化了現有結構剪枝方法複雜的多階段流程，適合各類網絡架構和應用，且最小化了用戶的額外工程投入，具有通用性，有效性和易用性。