[论文解读] FLOPs as a Direct Optimization Objective for Learning Sparse Neural Networks

Raphael Tang, Ashutosh Adhikari+1 more|arXiv (Cornell University)|Nov 7, 2018

Advanced Neural Network Applications参考文献 11被引用 20

一句话总结

本文提出了一种新颖的方法，在神经网络训练过程中通过在L0正则化框架中引入可微分的FLOPs目标，直接优化目标FLOPs数量。通过使用代理目标和得分函数估计器，将FLOPs作为软约束融入优化过程，该方法实现了精确且可调节的模型压缩，在MNIST、CIFAR-10和CIFAR-100数据集上均实现了SOTA的准确率-效率权衡，并具备显式的FLOPs控制能力。

ABSTRACT

There exists a plethora of techniques for inducing structured sparsity in parametric models during the optimization process, with the final goal of resource-efficient inference. However, few methods target a specific number of floating-point operations (FLOPs) as part of the optimization objective, despite many reporting FLOPs as part of the results. Furthermore, a one-size-fits-all approach ignores realistic system constraints, which differ significantly between, say, a GPU and a mobile phone -- FLOPs on the former incur less latency than on the latter; thus, it is important for practitioners to be able to specify a target number of FLOPs during model compression. In this work, we extend a state-of-the-art technique to directly incorporate FLOPs as part of the optimization objective and show that, given a desired FLOPs requirement, different neural networks can be successfully trained for image classification.

研究动机与目标

• 为解决神经网络稀疏化方法中缺乏对FLOPs的直接优化问题，这些方法通常仅将FLOPs作为结果指标报告。
• 使实践者能够在训练过程中指定目标FLOPs值，使模型压缩与移动设备或GPU部署等实际系统约束相匹配。
• 扩展最先进的L0正则化框架，将FLOPs作为优化目标中的软性、可微分约束。
• 证明FLOPs可作为延迟和能耗的合理代理指标，从而在不损失准确率的前提下提升模型效率。

提出的方法

• 该方法引入了一个仅依赖非零参数数量的FLOPs目标 $ L_{\text{flops}}(h, \boldsymbol{\theta}) $，并使用卷积层和全连接层的标准FLOPs公式。
• 通过在二值门上使用硬混凝土分布，将FLOPs约束松弛为可微分形式，从而实现基于梯度的优化。
• 损失函数结合交叉熵与截断的FLOPs惩罚项：$ \mathcal{L} = \mathbb{E}_{p(\tilde{\mathbf{z}}|\boldsymbol{\phi})}[-\log p(\mathcal{D}|\boldsymbol{\theta} \odot \mathbf{z})] + \lambda_f \mathbb{E}_{p(\mathbf{z}|\psi(\boldsymbol{\phi}))} [\max(0, L_{\text{flops}}(h, \boldsymbol{\theta} \odot \mathbf{z}) - T)] $。
• 每层的FLOPs使用标准公式计算：对于卷积层，$ L_{\text{flops}} = (K_wK_hC_{\text{in}} + 1)(I_w - K_w + P_w + 1)(I_h - K_h + P_h + 1) \|\mathbf{z}\|_0 $，对于全连接层，$ L_{\text{flops}} = (I_n + 1) \|\mathbf{z}\|_0 $。
• 使用得分函数估计器反向传播通过FLOPs目标，由于其组合性质，该目标被视为黑箱函数。
• 在推理阶段，根据学习到的概率将门设为0或1，并在剪枝后对最终模型进行微调以恢复准确率。

实验结果

研究问题

• RQ1 是否可以在神经网络训练过程中将FLOPs作为软约束直接优化，以实现精确的模型压缩？
• RQ2 将FLOPs作为目标是否能带来优于标准L0正则化的准确率-效率权衡？
• RQ3 实践者是否能通过在训练期间指定目标FLOPs值，有效控制模型效率？
• RQ4 与基于准确率或参数数量的方法相比，基于FLOPs的剪枝在推理延迟和能效方面表现如何？

主要发现

• 在MNIST上，当目标为100K FLOPs时，该方法仅用153K FLOPs即实现了0.9%的错误率，FLOPs效率优于基线L0方法。
• 在CIFAR-10上，当目标为4B FLOPs时，该方法实现了3.82%的错误率和4.6B FLOPs，相较于原始模型实现了Pareto改进。
• 在CIFAR-100上，当目标为4B FLOPs时，该方法实现了18.93%的错误率和4.6B FLOPs，展示了在准确率与效率之间可接受的权衡，适用于延迟敏感的应用。
• 训练期间的预期FLOPs与推理时的实际FLOPs高度一致，验证了代理目标的有效性。
• 该方法最激进地压缩了卷积层，因为它们在模型中贡献了大部分FLOPs。
• 该方法提供了一个可调节的FLOPs控制旋钮，使实践者能够直接指定并满足目标计算预算。