ICLR 2023 | 微软提出自动化模型训练剪枝框架OTO，一站式获得轻量级架构

2023-03-22 05:03

©PaperWeekly 原创 · 作者 | 陈天翼

单位 | 微软

研究方向 | 模型压缩

大型神经网络 DNN 对资源的巨大需求限制了其在现实世界的部署，如何在尽可能小地损失性能的前提下，压缩神经网络是 DNN 产品化的关键。

剪枝是最常见的 DNN 压缩方法之一，旨在减少冗余结构，给 DNN 模型瘦身的同时保持模型性能。然而，现有的剪枝方法通常指针对特定模型，特定任务，并需要 AI 工程师投入大量的工程和时间精力来应用这些方法到自己的任务上。

为了解决这些问题，微软团队提出了 OTOv2 框架，并发表于 ICLR 2023。OTOv2 是业内首个自动化，一站式，用户友好，且通用的神经网络训练与结构压缩框架。通过使用 OTOv2，AI 工程师可以非常方便地训练目标神经网络，一站式地取得高性能且轻量化的模型。OTOv2 最小化了开发者额外的工程时间精力的投入，且全程无需现有方法通常需要的非常耗时的预训练和额外的模型微调。

论文链接：

https://openreview.net/pdf?id=7ynoX1ojPMt

代码链接：

https://github.com/tianyic/only_train_once

该研究的主要贡献概括如下：

1. 一站式自动化神经网络结构剪枝架构。研究者提出了一个名为 OTOv2（Only-Train-Once）的 one-shot 训练和剪枝框架。它可以将一个完整的神经网络压缩为轻量级网络，同时保持较高的性能。OTO 大大简化了现有剪枝方法复杂的多阶段流程，适合各种架构和应用，且最小化了用户的额外工程投入，因具有通用性，有效性和易用性。

2. 自动化 Zero-Invariant Groups（零不变组）分组。由于网络结构的复杂性和关联性，删去任意网络结构可能会导致剩余的网络结构无效。因此自动化网络结构压缩的一个最大的问题之一是如何找到必须要被一起剪枝的模型参数，使得余下的网络依然有效。

为了解决该问题，研究者提出神经网络的可移除单元和零不变组 Zero-Invariant Groups（ZIGs）。零不变组可以理解为一类最小的可移除单元，使得该组对应的网络结构移除后剩余网络依然有效。研究者进一步提出并实现了一套自动化算法来解决通用网络的 ZIGs 分组问题。

3. 双半平面投影梯度优化算法（DHSPG）。接下来的模型训练和剪枝需要找出冗余的和重要的零不变组。这个问题可以归纳为一个结构稀疏化问题，研究者提出 Dual Half-Space Projected Gradient（DHSPG）优化算法来解决该问题。DHSPG 可以非常有效地找出冗余的零不变组并将其投影成零，并持续训练重要的零不变组来取得跟原始模型相媲美的性能。