微软自动化神经网络训练剪枝框架OTO，一站式高性能轻量化模型

机器之心专栏

作者： 陈天翼-微软西雅图-高级研究员

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 大大简化了现有结构剪枝方法复杂的多阶段流程，适合各类网络架构和应用，且最小化了用户的额外工程投入，具有通用性，有效性和易用性。