FLOPs as a Direct Optimization Objective for Learning Sparse Neural Networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
There exists a plethora of techniques for inducing structured sparsity in parametricmodels during the optimization process, with the ﬁnal goal of resource-efﬁcientinference.
However, to the best of our knowledge, none target a speciﬁc numberof ﬂoating-point operations (FLOPs) as part of a single end-to-end optimizationobjective, despite reporting FLOPs as part of the results.
Furthermore, a one-sizeﬁts-all approach ignores realistic system constraints, which differ signiﬁcantlybetween, say, a GPU and a mobile phone—FLOPs on the former incur less latencythan on the latter; thus, it is important for practitioners to be able to specify a targetnumber of FLOPs during model compression.
In the present work, we extend astate-of-the-art technique to directly incorporate FLOPs as part of the optimizationobjective, and we show that, given a desired FLOPs requirement, different neuralnetworks are successfully trained.
Introduction
Neural networks are a class of parametric models that achieve state of the art across a broad rangeof tasks, but their heavy computational requirements hinder practical deployment on resourceconstrained devices, such as mobile phones, Internet-of-things (IoT) devices and ofﬂine embeddedsystems.
Many recent works focus on alleviating these computational burdens, mainly fallingunder two non-mutually exclusive categories: manually designing resource-efﬁcient models andautomatically compressing popular architectures.
In the latter, increasingly sophisticated techniqueshave emerged [4, 5, 6] which have achieved respectable accuracy–efﬁciency operating points, someeven Pareto-better than that of the original network; for example, network slimming [4] reachesan error rate of 6.20% on CIFAR-10 using VGGNet [10] with a 51% FLOPs reduction—an errordecrease of 0.14% over the original.
However, to the best of our knowledge, none of the methods impose a FLOPs constraint as partof a single end-to-end optimization objective.
MorphNets [1] apply an L 1 norm, shrinkage-basedrelaxation of a FLOPs objective, but for the purpose of searching and training multiple models toﬁnd good network architectures; in this work, we learn a sparse neural network in a single trainingrun.
Other texts directly target device-speciﬁc metrics, such as energy usage [16], but the pruningprocedure does not explicitly include the metrics of interest as part of the optimization objective,instead using them as heuristics.
Falling short of continuously deploying a model candidate andmeasuring actual inference time, as in time-consuming neural architectural search [12], we believethat the number of FLOPs is reasonable to use as a proxy measure for actual latency and energyusage; across variants of the same architecture, it has been suggested to be a stronger predictor ofenergy usage and latency than the number of parameters is [13].
Indeed, there are compelling reasons to optimize for the number of FLOPs as part of the trainingobjective: ﬁrst, it would permit FLOPs-guided compression to occur in a more principled manner.
Second, practitioners can directly specify a wanted target of FLOPs, which is important in deployment.
Thus, our main contribution is to present a novel extension of prior state of the art [7] to incorporatethe number of FLOPs as part of the optimization objective, furthermore allowing practitioners to setand meet a desired compression target.
FLOPs Objective
Formally, we deﬁne the FLOPs objective L flops : f × R m (cid:55)→ N 0 as follows:
[Equation]
where L flops is the FLOPs associated with hypothesis h ( · ; θθθ ) :
= p ( ·| θθθ ) , g ( · ) is a functional with theexplicit dependencies, and I is the indicator function.
We assume L flops to depend only on whetherparameters are non-zero, such as the number of neurons in a neural network.
For a dataset D , ourempirical risk thus becomes
[Equation]
Hyperparameters λ f ∈ R +0 and T ∈ N 0 control the strength of the FLOPs objective and thetarget, respectively.
The second term is a black-box function, whose combinatorial nature preventsgradient-based optimization; thus, using the same procedure in prior art [7], we relax the objectiveto a surrogate of the evidence lower bound with a fully-factorized spike-and-slab posterior as thevariational distribution, where the addition of the clipped FLOPs objective can be interpreted as asparsity-inducing prior p
[Equation]
Let z ∼ p ( z | πππ ) be Bernoullirandom variables parameterized by πππ :
[Equation]
where (cid:12) denotes the Hadamard product.
To allow for efﬁcient reparameterization and exact zeros, ithas been proposed [7] to use a hard concrete distribution as the approximation, which is a stretchedand clipped version of the binary Concrete distribution [8]: if ˆ z ∼ BinaryConcrete ( α, β ) ,
then˜ z := max(0 , min(1 , ( ζ − γ )
ˆ z + γ )) is said to be a hard concrete r.v., given ζ > 1 and γ < 0 .
Deﬁne φφφ :=
( ααα, β ) , and let ψ ( φφφ ) =
Sigmoid (log ααα − β log − γζ ) and z ∼ Bernoulli ( ψ ( φφφ )) .
Then,the approximation becomes
[Equation]
ψ ( · ) is the probability of a gate being non-zero under the hard concrete distribution.
It is more efﬁcientin the second expectation to sample from the equivalent Bernoulli parameterization compared tohard concrete, which is more computationally expensive to sample multiple times.
The ﬁrst termnow allows for efﬁcient optimization via the reparameterization trick [3]; for the second, we applythe score function estimator (REINFORCE; 15), since the FLOPs objective is, in general, nondifferentiable and thus precludes the reparameterization trick.
High variance is a non-issue becausethe number of FLOPs is fast to compute, hence letting many samples to be drawn.
At inference time,the deterministic estimator is ˆ θθθ :=
θθθ (cid:12) max(0 , min(1 , Sigmoid (log ααα )( ζ − γ ) + γ )) for the ﬁnalparameters ˆ θθθ .
2.1 Computing number of FLOPs under group sparsity
In practice, computational savings are achieved only if the model is sparse across “regular” groupsof parameters, e.g., each ﬁlter in a convolutional layer.
Thus, each computational group uses onehard concrete r.v.
[7]—in fully-connected layers, one per input neuron; in 2D convolution layers,one per output ﬁlter.
Under convention in literature where one addition and one multiplication eachcount as a FLOP, the FLOPs for a 2D convolution layer h conv ( · ; θθθ ) given a random draw z is thendeﬁned as L flops ( h conv , z ) =
( K w K h C in + 1)( I w − K w + P w + 1)
( I h − K h
+ P h + 1) (cid:107) z (cid:107) 0for kernel width and height ( K w , K h ) , input width and height ( I w , I h ) , padding width and height( P w , P h ) , and number of input channels C in .
The number of FLOPs for a fully-connected layerh fc ( · ; θ θ θ ) is L flops ( h fc , z ) =
( I n + 1) (cid:107) z (cid:107) 0 , where I n is the number of input neurons.
Note thatthese are conventional deﬁnitions in neural network compression works—the objective can easily useinstead a number of FLOPs incurred by other device-speciﬁc algorithms.
Thus, at each training step,we compute the FLOPs objective by sampling from the Bernoulli r.v.’s and using the aforementioneddeﬁnitions, e.g., L flops ( h conv , · ) for convolution layers.
Then, we apply the score function estimatorto the FLOPs objective as a black-box estimator
Experimental Results
We report results on MNIST, CIFAR-10, and CIFAR-100, training multiple models on each datasetcorresponding to different FLOPs targets.
We follow the same initialization and hyperparameters [7],using Adam [2] with temporal averaging for optimization, a weight decay of 5 × 10 − 4 , and an initialα that corresponds to the original dropout rate of that layer.
We similarly choose β = 2 / 3 , γ = − 0 .
1 ,and ζ = 1 . 1 .
For brevity, we direct the interested reader to their repository 1 for speciﬁcs.
In all of ourexperiments, we replace the original L 0 penalty with our FLOPs objective
[Table]
and we train all models to200 epochs; at epoch 190, we prune the network by weights associated with zeroed gates and replacethe r.v.’s with their deterministic estimators, then ﬁnetune for 10 more epochs.
For the score functionestimator, we draw 1000 samples at each optimization step—this procedure is fast and has no visibleeffect on training time.
[Table]
L
[Table]
We choose λ f = 10 − 6 in all of the experiments for LeNet-5-Caffe, the Caffe variant of LeNet-5.
1 Weobserve that our methods (Table 1, bottom three rows) achieve accuracy comparable to those frompast approaches while using fewer FLOPs, with the added beneﬁt of providing a tunable “knob” foradjusting the FLOPs.
Note that the convolution layers are the most aggressively compressed, sincethey are responsible for most of the FLOPs in this model.
[Table]
Orig.
in Table 2 denotes the original WRN-28-10 model [17]
[Table]
and L 0 -* refers to the L 0 -regularizedmodels [7]; likewise, we augment CIFAR-10 and CIFAR-100 with standard random cropping andhorizontal ﬂipping.
For each of our results (last two rows), we report the median error rate of ﬁvedifferent runs, executing a total of 20 runs across two models for each of the two datasets; we useλ f = 3 × 10 − 9 in all of these experiments.
We also report both the expected FLOPs and actual
FLOPs, the former meaning the number of FLOPs, on average, at training time under stochastic gatesand the latter denoting the number of FLOPs at inference time.
We restrict the FLOPs calculationsto the penalized non-residual convolution layers only.
For CIFAR-10, our approaches result in
Pareto-better models with decreases in both error rate and the actual number of inference-time FLOPs.
For CIFAR-100, we do not achieve a Pareto-better model, since our approach trades accuracy forimproved efﬁciency.
The acceptability of tradeoff depends on the end application.