[论文解读] A flexible, extensible software framework for model compression based on the LC algorithm

Yerlan Idelbayev, Miguel Á. Carreira-Perpiñán|arXiv (Cornell University)|May 15, 2020

Neural Networks and Applications参考文献 45被引用 5

一句话总结

本文提出了一种基于学习-压缩（LC）算法的灵活、可扩展的软件框架，通过在学习（L）和压缩（C）步骤之间进行迭代交替，实现端到端的模型压缩。该框架支持在神经网络的不同层中混合使用剪枝、量化、低秩分解及其组合等压缩技术，在用户干预极少、运行时间与训练原始模型相当的情况下，实现了具有竞争力的准确率和压缩比。

ABSTRACT

We propose a software framework based on the ideas of the Learning-Compression (LC) algorithm, that allows a user to compress a neural network or other machine learning model using different compression schemes with minimal effort. Currently, the supported compressions include pruning, quantization, low-rank methods (including automatically learning the layer ranks), and combinations of those, and the user can choose different compression types for different parts of a neural network. The LC algorithm alternates two types of steps until convergence: a learning (L) step, which trains a model on a dataset (using an algorithm such as SGD); and a compression (C) step, which compresses the model parameters (using a compression scheme such as low-rank or quantization). This decoupling of the "machine learning" aspect from the "signal compression" aspect means that changing the model or the compression type amounts to calling the corresponding subroutine in the L or C step, respectively. The library fully supports this by design, which makes it flexible and extensible. This does not come at the expense of performance: the runtime needed to compress a model is comparable to that of training the model in the first place; and the compressed model is competitive in terms of prediction accuracy and compression ratio with other algorithms (which are often specialized for specific models or compression schemes). The library is written in Python and PyTorch and available in Github.

研究动机与目标

• 为解决在深度神经网络中应用多样化模型压缩技术时缺乏统一、模块化框架的问题。
• 使非专家用户能够轻松地在模型的不同部分应用和比较多种压缩方案，如剪枝、量化和低秩分解。
• 将学习与压缩过程解耦，支持现有训练和压缩算法的即插即用式集成。
• 提供一个单一、可扩展的软件库，支持单独和组合的压缩策略，并具备收敛性保证。
• 通过提供统一、可配置的接口，降低混合压缩策略实验的复杂性。

提出的方法

• 该框架实现了LC算法，该算法在学习（L）步骤（使用SGD或类似优化器训练模型）和压缩（C）步骤（对模型参数应用压缩方案）之间交替进行。
• L步骤最小化复合损失：$ L(\mathbf{w}) + \frac{\mu}{2} \|\mathbf{w} - \boldsymbol{\Delta}(\boldsymbol{\Theta})\|^{2} $，其中 $ \mu $ 控制拟合数据与遵循压缩表示之间的权衡。
• C步骤应用压缩方案（例如，k-means用于量化，SVD用于低秩，或基于大小的剪枝）将当前权重 $ \mathbf{w} $ 投影到低维或离散空间 $ \boldsymbol{\Delta}(\boldsymbol{\Theta}) $。
• 该框架允许对不同层独立应用不同的压缩类型，从而实现细粒度的异构压缩策略。
• 采用指数增长的 $ \mu $ 调度 $ \mu_k = \mu_0 \times a^k $，在迭代过程中逐步强化压缩，其中 $ a \in [1.1, 1.4] $。
• 该库使用Python和PyTorch实现，具有完全模块化设计：更改模型或压缩方法仅需为L或C步骤插入新的子程序。

实验结果

研究问题

• RQ1 一个统一的框架能否在单一算法流程中高效支持多种异构的模型压缩技术（如剪枝、量化和低秩分解）？
• RQ2 学习与压缩步骤的解耦如何实现灵活、模块化和可扩展的压缩，同时不牺牲性能？
• RQ3 当在神经网络的不同层中组合多种压缩方案时，LC算法在压缩比和预测准确率方面能达到多大程度的竞争力？
• RQ4 在混合压缩场景中，哪些关键超参数配置（如 $ \mu $ 调度、学习率衰减）能确保收敛性和性能？
• RQ5 该框架的运行时间和准确率与针对相同模型和压缩类型的专用单体压缩工具相比如何？

主要发现

• 尽管是通用且可扩展的，该框架在压缩比和预测准确率方面仍能与专用算法相媲美。
• 使用LC框架压缩模型的运行时间与从头开始训练原始模型的时间相当。
• 该框架成功支持混合压缩策略，例如在某些层应用量化，在其他层应用剪枝，在又其他层应用低秩分解，如表6所示。
• 对于LeNet300架构，对不同层应用不同压缩类型（如量化、剪枝和低秩）可实现2.18%的测试误差，证明了异构压缩的有效性。
• 采用指数 $ \mu $-调度 $ \mu_k = 9 \times 10^{-5} \times 1.1^k $（低秩时为 $ 1.4^k $）可确保在迭代过程中实现收敛和有效压缩。
• 该框架的模块化设计使得用户只需进行极少代码更改即可插入新的压缩方案，只要压缩例程在C步骤中能正确最小化二次失真损失即可。