[Paper Review] 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 Applications11 references20 citations

TL;DR

This paper proposes a novel method to directly optimize for a target number of FLOPs during neural network training by extending the L0 regularization framework with a differentiable FLOPs objective. By incorporating FLOPs as a soft constraint via a surrogate objective and score function estimator, the approach enables precise, tunable model compression that achieves state-of-the-art accuracy-efficiency trade-offs on MNIST, CIFAR-10, and CIFAR-100 with explicit FLOPs control.

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.

Motivation & Objective

• To address the lack of direct FLOPs optimization in neural network sparsity methods, which typically report FLOPs only as a result metric.
• To enable practitioners to specify a target FLOPs value during training, aligning model compression with real-world system constraints such as mobile or GPU deployment.
• To extend the state-of-the-art L0 regularization framework to incorporate FLOPs as a soft, differentiable constraint in the optimization objective.
• To demonstrate that FLOPs can serve as a principled proxy for latency and energy usage, improving model efficiency without sacrificing accuracy.

Proposed method

• The method introduces a FLOPs objective $ L_{\text{flops}}(h, \boldsymbol{\theta}) $ that depends only on the non-zero parameter count, using standard FLOPs formulas for convolutional and fully connected layers.
• A surrogate objective is constructed by relaxing the FLOPs constraint into a differentiable form via a hard concrete distribution over binary gates, enabling gradient-based optimization.
• The loss function combines cross-entropy with a clipped FLOPs penalty: $ \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)] $.
• The FLOPs of each layer are computed using standard formulas: for conv layers, $ 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 $, and for FC layers, $ L_{\text{flops}} = (I_n + 1) \|\mathbf{z}\|_0 $.
• A score function estimator is used to backpropagate through the FLOPs objective, which is treated as a black-box function due to its combinatorial nature.
• At inference, gates are set to 0 or 1 based on learned probabilities, and the final model is fine-tuned after pruning to recover accuracy.

Experimental results

Research questions

• RQ1 Can FLOPs be directly optimized as a soft constraint during neural network training to achieve precise model compression?
• RQ2 Does incorporating FLOPs as an objective lead to better accuracy-efficiency trade-offs compared to standard L0 regularization?
• RQ3 Can practitioners effectively control model efficiency by specifying a target FLOPs value during training?
• RQ4 How does FLOPs-based pruning compare to accuracy-focused or parameter-count-focused methods in terms of inference latency and energy efficiency?

Key findings

• On MNIST, the proposed method achieved 0.9% error rate with only 153K FLOPs when targeting 100K FLOPs, outperforming baseline L0 methods in FLOPs efficiency.
• On CIFAR-10, the method achieved 3.82% error rate with 4.6B FLOPs when targeting 4B FLOPs, representing a Pareto improvement over the original model.
• On CIFAR-100, the method achieved 18.93% error rate with 4.6B FLOPs at the 4B target, showing a trade-off between accuracy and efficiency that is acceptable for latency-sensitive applications.
• The expected FLOPs during training closely matched the actual FLOPs at inference time, validating the effectiveness of the surrogate objective.
• The method successfully compressed convolutional layers most aggressively, as they contribute the majority of FLOPs in the network.
• The approach enables a tunable knob for FLOPs control, allowing practitioners to directly specify and meet target computational budgets.