Title: ENERGY-CONSTRAINED COMPRESSION FOR DEEP NEURAL NETWORKS VIA WEIGHTED SPARSE PROJEC-TION AND LAYER INPUT MASKING

Abstract: Deep Neural Networks (DNNs) are increasingly deployed in highly energyconstrained environments such as autonomous drones and wearable devices while at the same time must operate in real-time. Therefore, reducing the energy consumption has become a major design consideration in DNN training. This paper proposes the first end-to-end DNN training framework that provides quantitative energy consumption guarantees via weighted sparse projection and input masking.
The key idea is to formulate the DNN training as an optimization problem in which the energy budget imposes a previously unconsidered optimization constraint. We integrate the quantitative DNN energy estimation into the DNN training process to assist the constrained optimization. We prove that an approximate algorithm can be used to efficiently solve the optimization problem. Compared to the best prior energy-saving methods, our framework trains DNNs that provide higher accuracies under same or lower energy budgets.

1. INTRODUCTION:

Deep Neural Networks (DNNs) have become the fundamental building blocks of many emerging application domains such as computer vision (Krizhevsky et al., 2012; Simonyan & Zisserman, 2014) , speech recognition (Hinton et al., 2012) , and natural language processing (Goldberg, 2016) . Many of these applications have to operate in highly energy-constrained environments. For instance, autonomous drones have to continuously perform computer vision tasks (e.g., object detection) without a constant power supply. Designing DNNs that can meet severe energy budgets has increasingly become a major design objective.

The state-of-the-art model compression algorithms adopt indirect techniques to restrict the energy consumption, such as pruning (or sparsification) (He et al., 2018; Han et al., 2015a; Liu et al., 2015; Zhou et al., 2016; Li et al., 2016; Wen et al., 2016) and quantization (Gong et al., 2014; Wu et al., 2016; Han et al., 2015a; Courbariaux et al., 2015; Rastegari et al., 2016) . These techniques are agonistic to energy consumption; rather they are designed to reduce the amount of computations and the amount of model parameters in a DNN, which do not truly reflect the energy consumption of a DNN. As a result, these indirect approaches only indirectly reduce the total energy consumption. Recently, Energy-Aware Pruning (EAP) (Yang et al., 2017) proposes a more direct manner to reduce the energy consumption of DNN inferences by guiding weight pruning using DNN energy estimation, which achieves higher energy savings compared to the indirect techniques.

However, a fundamental limitation of all existing methods is that they do not provide quantitative energy guarantees, i.e., ensuring that the energy consumption is below a user-specified energy budget. In this paper, we aspire to answer the following key question: how to design DNN models that satisfy a given energy budget while maximizing the accuracy? This work provides a solution to this question through an end-to-end training framework. By end-to-end, we refer to an approach that directly meets the energy budget without relying heuristics such as selectively restoring pruned weights and layer by layer fine-tuning (Han et al., 2015b; Yang et al., 2017) . These heuristics are effective in practice but also have many hyper-parameters that must be carefully tuned.

Our learning algorithm directly trains a DNN model that meets a given energy budget while maximiz-ing model accuracy without incremental hyper-parameter tuning. The key idea is to formulate the DNN training process as an optimization problem in which the energy budget imposes a previously unconsidered optimization constraint. We integrate the quantitative DNN energy estimation into the DNN training process to assist the constrained optimization. In this way, a DNN model, once is trained, by design meets the energy budget while maximizing the accuracy.

Without losing generality, we model the DNN energy consumption after the popular systolic array hardware architecture (Kung, 1982) that is increasingly adopted in today's DNN hardware chips such as Google's Tensor Processing Unit (TPU) (Jouppi et al., 2017 ), NVidia's Tensor Cores, and ARM's ML Processor. The systolic array architecture embodies key design principles of DNN hardware that is already available in today's consumer devices. We specifically focus on pruning, i.e., controlling the DNN sparsity, as the main energy reduction technique. Overall, the energy model models the DNN inference energy as a function of the sparsity of the layer parameters and the layer input.

Given the DNN energy estimation, we formulate DNN training as an optimization problem that minimizes the accuracy loss under the constraint of a certain energy budget. The key difference between our optimization formulation and the formulation in a conventional DNN training is two-fold. First, our optimization problem considers the energy constraint, which is not present in conventional training. Second, layer inputs are non-trainable parameters in conventional DNN training since they depend on the initial network input. We introduce a new concept, called input mask, that enables the input sparsity to be controlled by a trainable parameter, and thus increases the energy reduction opportunities. This lets us further reduce energy in scenarios with known input data pattern.

We propose an iterative algorithm to solve the above optimization problem. A key step in optimization is the projection operation onto the energy constraint, i.e., finding a model which is closest to the given (dense) model and satisfies the energy constraint. We prove that this projection can be casted into a 0/1 knapsack problem and show that it can be solved very efficiently. Evaluation results show that our proposed training framework can achieve higher accuracy under the same or lower energy compared to the state-of-the-art energy-saving methods.

In summary, we make the following contributions in this paper: 

2. RELATED WORK:

Energy-Agnostic Optimizations Most existing DNN optimizations indirectly optimize DNN energy through reducing the model complexity. They are agonistic to the energy consumption, and therefore cannot provide any quantitative energy guarantees.

Pruning, otherwise known as sparsification, is perhaps the most widely used technique to reduce DNN model complexity by reducing computation as well as hardware memory access. It is based on the intuition that DNN model parameters that have low-magnitude have little impact on the final prediction, and thus can be zeroed-out. The classic magnitude-based pruning (Han et al., 2015b) removes weights whose magnitudes are lower than a threshold. Subsequent work guides pruning using special structures (Liu et al., 2015; Zhou et al., 2016; Li et al., 2016; Wen et al., 2016; He et al., 2017) , such as removing an entire channel, to better retain accuracy after pruning.

Quantization reduces the number of bits used to encode model parameters, and thus reducing computation energy and data access energy (Gong et al., 2014; Wu et al., 2016; Han et al., 2015a) . The extreme case of quantization is using 1-bit to represent model parameters (Courbariaux et al., 2015; Rastegari et al., 2016) . Such binary quantization methods are usually trained from scratch instead of quantizing a pre-trained DNN.

Energy-Aware Optimizations Recently, energy-aware pruning (EAP) (Yang et al., 2017) proposes to use a quantitative energy model to guide model pruning. Different from pure magnitude-based pruning methods, EAP selectively prunes the DNN layer that contributes the most to the total energy consumption. It then applies a sequence of fine-tuning techniques to retain model accuracy. The pruning step and fine-tuning step are alternated until the accuracy loss exceeds a given threshold.

Although EAP a promising first-step toward energy-aware optimizations, its key limitation is that it does not provide quantitative energy guarantees because it does not explicitly consider energy budget as a constraint. Our work integrates the energy budget as an optimization constraint in model training.

Latency-Guaranteed Compression Lately, model compression research has started providing guarantees in execution (inference) latency, which theoretically could be extended to providing energy guarantees as well. However, these methods are primarily search-based through either reinforcement learning (He et al., 2018) or greedy-search (Yang et al., 2018) . They search the sparsity setting for every single layer to meet the given budget. Thus, they may require a large number of trials to achieve a good performance, and may not ensure that the resulting model accuracy is maximized.

3. MODELING DNN INFERENCE ENERGY CONSUMPTION:

This section introduces the model of estimating energy consumption of a single DNN inference. We consider the widely-used feed-forward DNNs. Note that our proposed methodology can be easily extended to other network architectures as well. In this section, we first provide an overview of our energy modeling methodology (Section 3.1). We then present the detailed per-layer energy modeling (Section 3.2 and Section 3.3), which allow us to then derive the overall DNN energy consumption (Section 3.4). Our energy modeling results are validated against the industry-strength DNN hardware simulator ScaleSim (Samajdar et al., 2018) .

DNN model sparsity (via pruning) is well recognized to significantly affect the execution efficiency and thus affect the energy consumption of a DNN model (He et al., 2018; Yang et al., 2017; Han et al., 2015a; Liu et al., 2015; Zhou et al., 2016) . We thus use pruning as the mechanism to reduce energy consumption1 . Note, however, that model sparsity is not the end goal of our paper; rather we focus on reducing the energy consumption directly. Many dedicated DNN hardware chips (a.k.a., Neural Processing Units, NPUs) (Jouppi et al., 2017; Chen et al., 2016; Han et al., 2016; Parashar et al., 2017) have been developed to directly benefit from model sparsity, and are already widely available in today's consumer devices such as Apple iPhoneX, Huawei Mate 10, and Microsoft HoloLens. Our paper focuses on this general class of popular, widely-used DNN chips.

3.1. ENERGY MODELING OVERVIEW:

A DNN typically consists of a sequence of convolution (CONV) layers and fully connected (FC) layers interleaved with a few other layer types such as Rectified Linear Unit (ReLU) and batch normalization. We focus mainly on modeling the energy consumption of the CONV and FC layers. This is because CONV and FC layers comprise more than 90% of the total execution time during a DNN inference (Chen et al., 2016) and are the major energy consumers (Han et al., 2015a; Yang et al., 2017) . Energy consumed by other layer types is insignificant and can be taken away from the energy budget as a constant factor. A DNN inference's energy consumption is tied to the underlying hardware that performs the inference.

In particular, we assume a systolic-array-based DNN hardware architecture. Systolic array (Kung, 1982) has long been know as an effective approach for matrix multiplication. Many DNN hardware architectures adopt the systolic array, most notably the Google Tensor Processing Unit (TPU) (Jouppi et al., 2017 ), Nvidia's Tensor Cores in their most recent Volta GPUs, and ARM's ML Processor.

Targeting systolic-array-based DNN hardware ensures that our approach has a wide applicability. However, our modeling and training strategies can generally be applied to other DNN architectures.

Figure 1 shows the overall hardware architecture. The systolic array comprises of several compute units that perform the Multiply-and-Accumulate (MAC) operation, which conducts the following computation: a ← a + (b × c), where b and c are the two scalar inputs and a is the scalar intermediate result called "partial sum." MAC operation is the building block for matrix multiplication. The MAC units are organized in a 2-D fashion. The data is fed from the edges, both horizontally and vertically, which then propagate to the MAC units within the same row and columns.

We decompose the energy cost into two parts: computation energy E comp and data access energy E data . E comp denotes the energy consumed by computation units, and E data denotes the energy consumed when accessing data from the hardware memory. Since we mainly use pruning as the energy reduction technique, we now model how E comp and E data are affected by DNN sparsity.

3.2. ENERGY CONSUMPTION FOR COMPUTATION:

CONV layers perform convolution and FC layer perform matrix-vector multiplication. Both operations can be generalized to matrix-matrix multiplication, which involves only the MAC operation (Chetlur et al., 2014; Jouppi et al., 2017) . Figure 1 illustrates how a matrix-matrix multiplication is carried out on the systolic array hardware. Given X and W , the systolic array computes XW by passing each row of X to each row in the systolic array and passing each column of W to each column in the systolic array. If the width of the systolic array, denoted by s w , is less than the width of W , the hardware will fold W column-wise in strides of s w . Similarly, if the height of X is greater than the height of the systolic array (s h ), X is folded row-size in strides of s h . Figure 1 illustrates a 2 × 2 systolic array multiplying two 4 × 4 matrices. Both matrices are folded twice in strides of 2.

Critically, if either inputs of a MAC operation is zero, we can skip the MAC operation entirely and thus save the computation energy. At a high-level, the total computation energy, E comp , can be modeled as e MAC N MAC , where e MAC denotes the energy consumption of one MAC operation whereas N MAC denotes the total number of MAC operations that are actually performed. The challenge is to identify N MAC for CONV and FC layers, which we discuss below.

Fully connected layer Let X (v) ∈ R 1×c be the input vector and W (v) ∈ R c×d be the weight matrix of the FC layer v. The FC layer performs matrix-vector multiplication X (v) W (v) . The number of MAC operations N MAC is sum(supp(X)supp(W )), where supp(T ) returns a binary tensor indicating the nonzero positions of tensor T . So the computation energy for a fully connected layer v:

EQUATION

where the equality is reached when the input is dense.

Convolution layer:

The CONV layer performs the convolution operation between a 4-D weight (also referred to as kernel or filter) tensor and a 3-D input tensor. Let W (u) ∈ R d×c×r×r be the weight tensor, where d, c, and r are tensor dimension parameters. Let X (u) ∈ R c×h×w be the input tensor, where h and w are the input height and width. The convolution operation in the CONV layer u generates a 3-dimensional tensor:

EQUATION

where x, y indicate the position of the output tensor, which has height h = (h + 2p -r)/s + 1 and width w = (w + 2p -r)/s + 1 (p is the convolution padding and s is the convolution stride).

Tensor convolution (2) can be seen as a special matrix-matrix multiplication (Chellapilla et al., 2006; Chetlur et al., 2014) . Specifically, we would unfold the tensor X (u) to a matrix X(u) ∈ R h w ×cr 2 , and unfold the tensor W (u) to a matrix W (u) ∈ R cr 2 ×d . X(u) and W (u) are then multiplied together in the systolic array to compute the equivalent convolution result between X (u) and W (u) .

Nonzero elements in X (u) and W (u) incur actual MAC operations. Thus, N MAC = sum(supp(X (u) ) * supp(W (u) )) ≤ h w W (u) 0 (the equality means the input is dense), resulting in the following computation energy of a CONV layer u:

E (u) comp = e MAC sum(supp(X (u) ) * supp(W (u) )) ≤ e MAC h w W (u) 0 .

(3)

3.3. ENERGY CONSUMPTION FOR DATA ACCESS:

Accessing data happens in every layer. The challenge in modeling the data access energy is that modern hardware is equipped with a multi-level memory hierarchy in order to improve speed and save energy (Hennessy & Patterson, 2011) . Specifically, the data is originally stored in a large memory, which is slow and energy-hungry. When the data is needed to perform certain computation, the hardware will load it from the large memory into a smaller memory that is faster and consume less energy. If the data is reused often, it will mostly live in the small memory. Thus, such a multi-level memory hierarchy saves overall energy and improves overall speed by exploiting data reuse.

Without losing generality, we model a common, three-level memory hierarchy composed of a Dynamic Random Access Memory (DRAM), a Cache, and a Register File (RF). The cache is split into two halves: one for holding X (i.e., the feature map in a CONV layer and the feature vector in a FC layer) and the other for holding W (i.e., the convolution kernel in a CONV layer and the weight matrix in an FC layer). This is by far the most common memory hierarchy in DNN hardware such as Google's TPU (Jouppi et al., 2017; Chen et al., 2016; Zhu et al., 2018; Han et al., 2016) . Data is always loaded from DRAM into cache, and then from cache to RFs.

In many today's DNN hardwares, the activations and weights are compressed in the dense form, and thus only non-zero values will be accessed. This is done in prior works (Chen et al., 2016; Parashar et al., 2017) . Therefore, if the value of the data that is being loaded is zero, the hardware can skip the data access and thereby save energy. There is a negligible amount of overhead to "unpack" and "pack" compressed data, which we simply take away from the energy budget as a constant factor. This is also the same modeling assumption used by Energy-Aware Pruning (Yang et al., 2017) .

To compute E data , we must calculate the number of data accesses at each memory level, i.e., N DRAM , N cache , N RF . Let the unit energy costs of different memory hierarchies be e DRAM , e cache , and e RF , respectively, the total data access energy consumption E data will be e DRAM N DRAM +e cache N cache + e RF N RF . We count the number of data accesses for both the weights and input, then combine them together. The detailed derivation of data access energy is included in the Appendix.

3.4. THE OVERALL ENERGY ESTIMATION FORMULATION:

Let U and V be the sets of convolutional layers and fully connected layers in a DNN respectively. The superscript (u) and (v) indicate the energy consumption of layer u ∈ U and v ∈ V , respectively. Then the overall energy consumption of a DNN inference can be modeled by

EQUATION

where X stacks input vectors/tensors at all layers and W stacks weight matrices/tensors at all layers.

4. ENERGY-CONSTRAINED DNN MODEL:

Given the energy model presented in Section 3, we propose a new energy-constrained DNN model that bounds the energy consumption of a DNN's inference. Different from prior work on model pruning in which energy reduction is a byproduct of model sparsity, our goal is to directly bound the energy consumption of a DNN while sparsity is just used as a means to reduce energy.

This section formulates training an energy-constrained DNN as an optimization problem. We first formulate the optimization constraint by introducing a trainable mask variable into the energy modeling to enforce layer input sparsity. We then define a new loss function by introducing the knowledge distillation regularizer that helps improve training convergence and reduce overfitting.

Controlling Input Sparsity Using Input Mask:

The objective of training an energy-constrained DNN is to minimize the accuracy loss while ensuring that the DNN inference energy is below a given budget, E budget . Since the total energy consumption is a function of X (u) 0 and W (u) 0 , it is natural to think that the trainable parameters are X and W . In reality, however, X depends on the input to the DNN (e.g., an input image to an object recognition DNN), and thus is unknown during training time. Therefore, in conventional DNN training frameworks X is never trainable.

To include the sparsity of X in our training framework, we introduce a trainable binary mask M that is of the same shape of X, and is multiplied with X before X is fed into CONV or FC layers, or equivalently, at the end of the previous layer. For example, if the input to a standard CONV layer is X (u) , the input would now be X (u) M (u) , where denotes the element-wise multiplication. In practice, we do not really do this multiplication but only read X (u) on the nonzero positions of M (u) . Decay the sparsity constraint q = q -∆q; end

W * = W, M * = M .

With the trainable mask M , we can ensure that X (u) M (u) 0 ≤ M (u) 0 , and thereby bound the sparsity of the input at training time. In this way, the optimization constraint during training becomes E(M, W ) ≤ E budget , where E(M, W ) denotes the total DNN inference energy consumption, which is a function of X and W (as shown in Equation ( 4)), and thus a function of M and W .

Knowledge Distillation as a Regularizer Directly optimizing over the constraint would likely lead to a local optimum because the energy model is highly non-convex. Recent works (Mishra & Marr, 2017; Tschannen et al., 2017; Zhuang et al., 2018) notice that knowledge distillation is helpful in training compact DNN models. To improve the training performance, we apply the knowledge distillation loss (Ba & Caruana, 2014) as a regularization to the conventional loss function. Intuitively, the regularization uses a pre-trained dense model to guide the training of a sparse model. Specifically, our regularized loss function is:

EQUATION

where W dense is the original dense model, and L(M, W ) is the original loss, e.g., cross-entropy loss for classification task. φ(X; W ) is the network's output (we use the output before the last activation layer as in Ba & Caruana (2014) ), |φ(•; W )| is the network output dimensionality and 0 ≤ λ ≤ 1 is a hyper parameter similar to other standard regularizations.

Thus, training an energy-constrained DNN model is formulated as an optimization problem:

EQUATION

5. OPTIMIZATION:

This section introduces an algorithm to solve the optimization problem formulated in (6). The overall algorithm is shown in Algorithm 1. Specifically, the algorithm includes three key parts:

• Initialization by training a dense model. That is,

EQUATION

• Fix M and optimize W via approximately solving (using W dense initialization):

EQUATION

• Fix W and optimize M by approximately solving :

EQUATION

After the initialization step (Line 1 in Algorithm 1), the training algorithm iteratively alternates between the second (Line 3-6 in Algorithm 1) and the third step (Line 8-13 in Algorithm 1) while gradually reducing the sparsity constraint q (Line 14 in Algorithm 1) until the training accuracy converges. Note that Equation ( 7) is the classic DNN training process, and solving Equation ( 9) involves only the well-known L 0 norm projection P M 0≤q (Q) := arg min M 0≤q M -Q2 . We thus focus on how Equation ( 8) is solved.

Optimizing Weight Matrix W To solve (8), one can use either projected gradient descent or projected stochastic gradient descent. The key difficulty in optimization lies on the projection step

P Ω(Ebudget) (Z) := arg min W ∈Ω(Ebudget) W -Z 2 (10)

where Z could be W -η∇ W L(W, M ) or replacing ∇ W L(W, M ) by a stochastic gradient ∇W L(W, M ). To solve the projection step, let us take a closer look at the constraint Equation ( 4). We rearrange the energy constraint Ω(E budget ) into the following form with respect to W :

EQUATION

where W stacks all the variable {W (u) } u∈U ∪V , and α

EQUATION

4 and k are properly defined nonnegative constants. Note that α

(u) 1 ≤ α (u)

2 and k is a positive integer. Theorem 1 casts the energyconstrained projection problem to a 0/1 knapsack problem. The proof is included in the Appendix. Theorem 1. The projection problem in (10) is equivalent to the following 0/1 knapsack problem:

EQUATION

where Z stacks all the variables {Z (u) } u∈U ∪V , A and ξ are of the same shape as Z, and the j-th element of A (u) for any u ∈ U ∪ V is defined by

A (u) j = α (u) 1 + α (u) 3 , if Z (u) j

is among the top k elements of Z (u) in term of magnitude; α

(u) 2 + α (u) 3 , otherwise. (13)

The optimal solution of (10) is Z ξ * , where ξ * is the optimal solution to the knapsack problem (12).

The knapsack problem is NP hard. But it is possible to find approximate solution efficiently. There exists an approximate algorithm (Chan, 2018) that can find an -accurate solution in O(n log(1/ ) + -2.4 ) computational complexity. However, due to some special structure in our problem, there exists an algorithm that can find an an -accurate solution much faster.

In the Appendix, we show that an (1 + )-approximate solution of problem (10) can be obtained in Õ(n + 1

2 ) time complexity ( Õ omits logarithm), though the implementation of the algorithm is complicated. Here we propose an efficient approximate algorithm based on the "profit density." The profit density of item j is defined as Z 2 j /A j . We sort all items based on the "profit density" and iteratively select a group of largest items until the constraint boundary is reached. The detailed algorithm description is shown in the Appendix (Algorithm 2). This greedy approximation algorithm also admits nice property as shown in Theorem 2. Theorem 2. For the projection problem (10), the approximate solution W ∈ Ω(E budget ) to the greedy approximation algorithm admits

W -Z 2 ≤ P Ω(Ebudget) (Z)-Z 2 +Top W 0+1 ((Z Z) A)•min((max(A)-gcd(A)), R(W ))

(14) where max(A) is the maximal element of A, which is a nonnegative matrix defined in (13); Top k (•) returns the k-th largest element of •; denotes the element-wise division. gcd(•) is the largest positive rational number that divides every argument, e.g., gcd(0, 1/3, 2/3) = 1/3. In (14), gcd(A) denotes the greatest common divisor of all elements in A 2 , and R(W ) denotes the remaining budget

R(W ) = E budget - u∈U ∪V α (u) 4 -A, supp(W ) .

The formal proof is in the Appendix. W is the optimal projection solution to (10) if either of the following conditions holds:

1. (The remaining budget is 0.) It means that the greedy Algorithm 2 runs out of budget; 2. (The matrix A satisfies max(A) = gcd(A).) It implies that all elements in A have the same value. In other words, the weights for all items are identical.

6. EVALUATION:

The evaluations are performed on ImageNet (Deng et al., 2009) , MNIST, and MS-Celeb-1M (Guo et al., 2016) datasets. For the MS-Celeb-1M, we follow the baseline setting reported in the original paper (Guo et al., 2016) , which selects 500 people who have the most face images. We randomly sample 20% images as the validation set. We use both classic DNNs, including AlexNet (Krizhevsky et al., 2012) and LeNet-5 (LeCun et al., 1998) , as well as recently proposed SqueezeNet (Iandola et al., 2016) and MobileNetV2 (Sandler et al., 2018) .

We compare our method mainly with five state-of-art pruning methods: magnitude-based pruning (MP) (Han et al., 2015b; a) , structured sparsity learning (SSL) (Wen et al., 2016) , structured bayesian pruning (SBP) (Neklyudov et al., 2017) , bayesian compression (BC) (Louizos et al., 2017) and energy-aware pruning (EAP) (Yang et al., 2017) . Filter pruning methods (Li et al., 2016; He et al., 2017) require a sparsity ratio to be set for each layer, and these sparsity hyper-parameters will determine the energy cost of the DNN. Considering manually setting all these hyper-parameters in energy-constrained compression is not trivial, we directly compare against NetAdapt (Yang et al., 2018) which automatically searches such sparsity ratios and use filter pruning to compress DNN models. We implement an energy-constrained version of NetAdapt, which is originally designed to restrict the inference latency. Note that MobileNetv2 and SqueezeNet have special structures (e.g. residual block) that are not fully supported by NetAdapt. Thus, we show the results of NetAdapt only for AlexNet and LeNet-5. Hyper-parameters In the experiment, we observe that knowledge distillation can improve the performance of MP and SSL, so we apply knowledge distillation to all methods including the baseline for a fair comparison. The results of removing knowledge distillation on MP and SSL are included in the Appendix. We choose the distillation weight λ = 0.5. EAP proposes an alternative way to solve the overfitting issue, so we directly use their results. For all the DNNs, we turn off the dropout layers since we find the knowledge distillation regularization will perform better. In all the experiments, we choose ∆q = 0.1|M | where |M | is the number of all mask elements. For optimizing W , we use a pre-trained dense initialization and update W by SGD with the learning rate η 1 = 0.001 and weight decay 10 -4 . For optimizing input mask parameters M , we use the Adam optimizer (Kingma & Ba, 2014) with η 2 = 0.0001 and weight decay 10 -5 (MNIST)/10 -6 (MS-Celeb-1M). To stabilize the training process, we exponentially decay the energy budget to the target budget, and also use this trick in MP training (i.e. decaying the sparsity budget) for fair comparisons.

6.1. IMAGENET:

We set an energy budget to be less than the minimal energy consumption among the three baseline methods. We use the same performance metric (i.e. top-5 test accuracy) and hardware parameters, i.e., e MAC , e DRAM , e cache , e RF , s h , s w , k W , k X , as described in the EAP paper (Yang et al., 2017) . We initialize all the DNNs by a pre-trained dense model, which is also used to set up the knowledge distillation regularization. The top-5 test accuracies on the dense models are 79.1% (AlexNet), 80.5% (SqueezeNet), and 90.5% (MobileNetV2). We use batch size 128 and train all the methods with 30 epochs. For SSL and NetAdapt, we apply 20 additional epochs to achieve comparable results. We implement the projection operation P Ω(Ebudget) on GPU, and it takes < 0.2s to perform it in our experiments. The detailed wall-clock result is included in the Appendix.

Table 1 shows the top-5 test accuracy drop and energy consumption of various methods compared to the dense model. Our training framework consistently achieves a higher accuracy with a lower energy consumption under the same energy budget. For instance on AlexNet, under a smaller energy budget (26% < 27%), our method achieves lower accuracy drop over EAP (0.5% vs. 0.8%). The advantage is also evident in SqueezeNet and MobileNetV2 that are already light-weight by design.

EAP does not report data on MobileNetV2. We observe that weight sparsity is not a good proxy for energy consumption. Our method achieves lower energy consumption despite having higher density.

Figure 2 comprehensively compares our method with prior work. Solid markers represent DNNs trained from our framework under different energy budgets (x-axis). Empty markers represent DNNs produced from previous techniques. DNNs trained by our method have lower energies with higher accuracies (i.e., solid markers are closer to the bottom-left corner than empty markers). For instance on SqueezeNet, our most energy-consuming DNN still reduces energy by 23% while improves accuracy by 0.2% compared to EAP. MNIST and MS-Celeb-1M (Guo et al., 2016) represent datasets where inputs have regular patterns that are amenable to input masking. For instance, MS-Celeb-1M is a face image dataset and we use its aligned face images where most of the facial features are located in the center of an image. In such scenarios, training input masks lets us control the sparsity of the layer inputs and thus further reduce energy than merely pruning model parameters as in conventional methods. We do not claim that applying input mask is a general technique; rather, we demonstrate its effectiveness when applicable.

We compare our method with MP and SSL using LeNet-5 and MobileNetV2 for these two datasets, respectively. The pre-trained dense LeNet-5 has 99.3% top-1 test accuracy on MNIST, and the dense MobileNetV2 has 65.6% top-5 test accuracy on MS-Celeb-1M. EAP does not report data on these two datasets. Similar to the evaluation on ImageNet, we set the energy budget to be lower than the energy consumptions of MP and SSL. We use batch size 32 on MNIST and 128 on MS-Celeb-1M, and number of epochs is set the same as the ImageNet experiments. Table 2 compares the energy consumption and accuracy drop. Our method consistently achieves higher accuracy with lower energy under the same or even smaller energy budget. We visualize the sparsity of the learned input masks in Figure 3 . 

7. CONCLUSION:

This paper demonstrates that it is possible to train DNNs with quantitative energy guarantees in an end-to-end fashion. The enabler is an energy model that relates the DNN inference energy to the DNN parameters. Leveraging the energy model, we augment the conventional DNN training with an energy-constrained optimization process, which minimizes the accuracy loss under the constraint of a given energy budget. Using an efficient algorithm, our training framework generates DNNs with higher accuracies under the same or lower energy budgets compared to prior art.

APPENDICES DETAIL OF ENERGY CONSUMPTION FOR DATA ACCESS FULLY CONNECTED LAYER:

To multiply v) is used once but loaded three times, once each from DRAM, cache and RF, respectively. Thus, the number of DRAM, cache, and RF accesses for weight matrix W (v) is:

X (v) ∈ R c and W (v) ∈ R c×d , each nonzero element of W (

EQUATION

Input X (v) is fed into the systolic array d/s w times, where s w denotes the the systolic array width. Thus, the number of cache accesses for X (v) is:

EQUATION

Let k X be the cache size for input X (v) . If k X is less than X (v) 0 , there are X (v) 0 -k X elements that must be reloaded from DRAM every time. The rest k X elements need to load from only DRAM once as they will always reside in low-level memories. Thus, there are d/s w ( X (v) 0 -k X ) + k X DRAM accesses for X (v) . In addition, the output vector of the FC layer (result of X (v) W (v) ) needs to be written back to DRAM, which further incurs d DRAM accesses. Thus, The total number of DRAM accesses to retrieve X (v) is:

EQUATION

Each input element is loaded from RF once for each MAC operation, and there are two RF accesses incurred by accumulation for each MAC operation (one read and one write). Thus, the total number of RF accesses related to X (v) is:

EQUATION

In summary, the data access energy of a fully connected layer v is expressed as follows, in which each component follows the derivations in Equation (15) through Equation ( 18): ). (20)

E (v) data =

The notations are the same as in FC layer. We now show how the different components are modeled.

To convolve W (u) ∈ R d×c×r×r with X (u) ∈ R c×h×w , each nonzero element in the weight tensor W (u) is fed into the systolic array h w /s h times, where s h denotes the height of the systolic array and h and w are dimension parameters of X (u) . Thus,

EQUATION

Similar to the FC layer, the number of RF accesses for W (u) during all the MAC operations is:

EQUATION

)

Let k W be the cache size for the weight matrix W (u) . If W (u) 0 > k W , there are k W nonzero elements of W (u) that would be accessed from DRAM only once as they would reside in the cache, and the rest W (u) 0 -k W elements would be accessed from DRAM by h w /s h times. Thus,

EQUATION

Let k X be the cache size for input X (u) . If every nonzero element in X (u) is loaded from DRAM to cache only once, N input DRAM would simply be X (u) 0 . In practice, however, the cache size k X is much and generates a solution W ∈ Ω(E budget ) that admits

EQUATION

)

where min(A + ) is the minimum of the positive elements in A.

|U | and |V | denote the number of CONV and FC layers, respectively. They are very small numbers that can be treated as constants here. Thus, the computational complexity for our problem is reduced to Õ(n + 1 2 ), where Õ omits the logarithm term. In the following, we will prove this theorem by construction.

PROBLEM FORMULATION:

Definition 1. Inverted knapsack problem. Given n objects I := {(v i , w i )} n i=1 each with weight w i > 0, and value v i ≥ 0, define h I (x) to be the smallest weight budget to have the total value x:

h I (x) := min ξ∈{0,1} n n i=1 w i ξ i (29) s. t. n i=1 v i ξ i ≥ x

We are more interested in the case that the weights of n objects are in m clusters, i.e. there are only m distinct weights,

|{w i } n i=1 | = m.

In our case, m is proportional to the number of layers in DNN, and n is the number of all the learnable weights in W , so m n. Definition 2. Inverse of step function. The inverse of the step function f is defined as the maximal x having the function value y:

f -1 (y) := max

EQUATION

Observation The inverse of the step function h -1 I (y) is just the maximal value we can get given the weight budget, i.e. the original knapsack problem:

EQUATION

Observation Given a step function with l breakpoints, its inverse can be generated with O(l) time complexity, and vice versa.

Thus, given the step function of h I in (29) which has l breakpoints, we can get h -1 I (i.e. the original knapsack problem) within O(l) time complexity. Definition 3. w-uniform.

Step function f is w-uniform if the ranges of f is from -∞, 0, w, 2w, ..., lw.

Observation If all the objects in I have the same weight w, i.e. m = 1, then the function h I (x) is nondecreasing and w-uniform. Moreover, its breakpoints are:

(0, 0), (v 1 , w), (v 1 + v 2 , 2w), ..., n i=1 v i , nw ,

if the objects' indices follows the decreasing order in terms of the values, i.e. v 1 ≥ v 2 ≥ ... ≥ v n . Thus we can get all possible function values of h I (x):

h I (x) = kw, ∀x ∈ k-1 i=1 v i , k i=1 v i .

Definition 4. (min, +)-convolution. For functions f, g, the (min, +)-convolution is:

(f ⊕ g)(x) = min x (f (x ) + g(x -x )). Observation If object sets I 1 ∩ I 2 = ∅, then f I1∪I2 = f I1 ⊕ f I2 .

Observation The inverse of (min, +)-convolution between w-uniform function f and w-uniform function g is the (max, +)-convolution between f -1 and g -1 :

EQUATION

Lemma 4. For any f and g nonnegative step functions, given an arbitrary number b, we always have min{f ⊕ g, b} = min{min{f, b} ⊕ min{g, b}, b}

Proof. Given any x, let z ∈ Arg min

x f (x ) + g(x -x ) and z ∈ Arg min x min(f (x ), b) + min(g(x -x ), b), so we have (f ⊕ g)(x) = f (z) + g(x -z) and (min{f, b} ⊕ min{g, b})(x) = min(f (z), b) + min(g(x -z), b).

Consider the following cases: 

(x -z) < b, so min(f (z), b) + min(g(x -z), b) ≤ min(f (z), b) + min(g(x -z), b) = f (z) + g(x -z) = (f ⊕ g)(x) < b.

Since both f and g are nonnegative, we have f (z) < b and g(x -z) < b which imply min

(f (z), b) + min(g(x -z), b) = f (z) + g(x -z) ≥ (f ⊕ g)(x). Therefore, we have min(f (z), b) + min(g(x -z), b) = f (z) + g(x -z) ⇔ (min{f, b} ⊕ min{g, b})(x) = (f ⊕ g)(x).

EFFICIENCY OF (MIN, +)-CONVOLUTION Lemma 5. Let f and g be nondecreasing w-uniform functions with O(l) breakpoints, the (min, +)-convolution f ⊕ g (having O(l) breakpoints) can be generated with O(l 2 ) time complexity.

Proof. Firstly, we compute the inverse representation of f and g, i.e. compute f -1 and g -1 from Equation (30). The inverse representation can be computed in O(l) time (proportional to the number of breakpoints). From Equation (32), we can compute the inverse of f ⊕ g. For each y ∈ {0, 1w, ..., 2lw}, function (f ⊕ g) -1 (y) can be computed in O(l) time by brute force. Thus a total O(l 2 ) is enough to get (f ⊕ g) -1 which has O(l) breakpoints. We can get f ⊕ g via (f ⊕ g) -1 by the inverse definition (30) in O(l) time. Theorem 9. For the knapsack problem defined in Equation ( 12), if all the n objects have m distinct weights, there exists an approximate algorithm with computational complexity O((n + Z Z, ξ .

From Theorem 9, we can directly get Theorem 3.

All the elements of the solution are either 0 or 1 except the last picked one which corresponds to Top ξ 0+1 ((Z Z) A). Let the ( ξ 0 + 1)-th largest element in (Z Z) A be indexed by t.

We have 0 ≤ ξt ≤ 1 -gcd(A)/A t . Therefore, comparing with the original 0/1 knapsack problem, we have 

SUPPLEMENTARY EXPERIMENT RESULTS:



RESULTS OF BASELINE WITHOUT KNOWLEDGE DISTILLATION:

Table 3 shows the energy and accuracy drop results of the baseline methods MP and SSL when the knowledge distillation is removed from their loss function. By using knowledge distillation, the results in Table 1 are much better. Therefore, we use knowledge distillation in all the experiments when it is applicable. The projection operation P Ω(Ebudget) in Algorithm 1 can be implemented on GPU. We measured its wall-clock time on a GPU server (CPU: Xeon E3 1231-v3, GPU: GTX 1080 Ti), and the result is shown in Table 4 (the time is averaged over 100 iterations). 

:

Quantization is another useful mechanism to reduce energy consumption. It is orthogonal to the pruning mechanism and they could be combined. This paper specifically focuses on the pruning mechanism.

Here we assume A only contains rational numbers since gcd is used.

Figures and Tables:

Figure 1: Illustration of the energy cost of computing matrix multiplication XW .

Energy-Constrained DNN Training. Input: Energy budget E budget , learning rates η 1 , η 2 , mask sparsity decay step ∆q. Result: DNN weights W * , input mask M * . Initialize W = W dense , M = 1, q = M 0 -∆q; while True do // Update DNN weights while W has not converged do W = W -η 1 ∇W L(M, W ) ; // SGD step W = P Ω(Ebudget) (W ) ; // Energy constraint projection for weights W end If previous_accuracy > current_accuracy, exit loop with previous W and M ; // Update input mask while M has not converged do M = M -η 2 ∇M L(M, W ) ; // SGD step Clamp values of M into [0, 1]: assign 1 (or 0) to the values if they exceeds 1 (or negative); M = P M 0≤q (M ) ; // L 0 constraint projection for input mask M end Round values of M into {0, 1};

Figure 2: Accuracy drops under different energy consumptions on ImageNet.

Figure 3: Input mask learned on MNIST and MS-Celeb-1M. For MNIST, the input mask for the first layer (with one channel) is shown. For MS-Celeb-1M, we show the input mask averaged across the 96 channels in the 7 th layer. 0 indicates a pixel is masked off, and 1 indicates otherwise.

1. (f ⊕ g)(x) ≥ b. In this case, we claim that (min{f, b} ⊕ min{g, b})(x) ≥ b. We prove it by contradiction. Suppose (min{f, b} ⊕ min{g, b})(x) < b which implies min(f (z), b) + min(g(x -z), b) < b. Because both f and g are nonnegative, we have f (z) < b and g(x -z) < b which imply min(f (z), b) + min(g(x -z), b) = f (z) + g(x -z) < b, However, this contradicts (f ⊕ g)(x) ≥ b. Therefore, we have min((f ⊕ g)(x), b) = min((min{f, b} ⊕ min{g, b})(x), b) = b. 2. (f ⊕ g)(x) < b. In this case, we have f (z) < b and g

Let f and g be nondecreasing step functions with l breakpoints in total, min{f ⊕ g, b} can be approximated by a step function φ b with O(l + 1 2 ) complexity and 2 b additive error, i.e. min{f ⊕ g, b} ≤ φ b ≤ min{f ⊕ g, b} + 2 b. The resultant function φ b has O(1/ ) breakpoints. 1. Construct a set B := {2 i n max(w) ∈ [min(w), n max(w)]; i ∈ Z ≤0 }, where min(w) and max(w) are the minimum and maximum weight of items respectively, and Z ≤0 is the nonpositive integer set. We have |B| = O(log n max(w) min(w) ). 2. For every b ∈ B, construct ψ b to approximate min{h I1 ⊕ ... ⊕ h Im , b} based on Lemma 7if b/2 < y ≤ b and y > min(B); ψ -1 min(B) (y), if y ≤ min(B). where min(B) is the minimum element in B. The resultant function h-1 I (or hI ) has at most O( 1 log n max(w) min(w) ) breakpoints. 4. Compute the original function hI from h-1 I . According to the above procedure, for any h I (x) ∈ (b/2, b], hI (x) approximate h I (x) with additive error O(m b), so we have h I (x) ≤ hI (x) ≤ (1 + O(m ))h I (x). The algorithm takes O((n + m/ 2 ) log n max(w) min(w) ), if we require the approximation factor to be 1 + O( ), i.e., h I (x) ≤ hI (x) ≤ (1 + O( ))h I (x), ∀x, we need O (n + m 3 / 2 ) log n max(w) min(w) time complexity.

≤ h -1 I (y), ∀y.Proof. From Lemma 8, we have hI (x) ≤ (1 + O( ))h I (x) which implies{x | (1 + O( ))h I (x) ≤ y} ⊆ {x | hI (x) ≤ y}. So max h I (x)≤y/(1+O( )) can get {x | hI (x) ≤ y} ⊆ {x | h I (x) ≤ y} from Lemma 8, so we have max h I (x)≤y x ≥ max hI (x)≤y x ⇔ h -1 I (y) ≥ h-1 I (y).Let I be the set of objects whose weights are nonzero elements in A and values are the corresponding elements in Z Z, i.e.I + = {(Z 2 i , A i ) | ∀i ∈ {1, 2, ..., |A|} and A i > 0}, ξ+ be the solution corresponding to h-1 I (E budget -u∈U ∪V α (u)4 ). Let ξI c + = 1 and ξI+ = ξ+ , whereI c + = {(Z 2 i , A i ) | ∀i ∈ {1, 2, ..., |A|} and A i = 0}is the complement of I + . Here we have m ≤ 2|U | + |V | distinct values in A. According to Theorem 9, we have Z Z, ξ ≥ max ξ Z Z, ξ , s.t. A, ξ ≤ Ebudget-u∈U ∪V α (u) 4 1+O( ) , which implies Z Z, ξ ≥ max ξ∈Ω(Ebudget/(1+O( )))

ξ ≤ Z Z, ξ + (Z Z) t • (1 -gcd(A)/A t ) = Z Z, ξ + Top ξ 0 +1 ((Z Z) A) • A t • (1 -gcd(A)/A t ) = Z Z, ξ + Top ξ 0 +1 ((Z Z) A) • (A t -gcd(A)) ≤ Z Z, ξ + Top ξ 0 +1 ((Z Z) A) • (max(A) -gcd(A))Since {ξ | ξ is binary} ⊆ Γ, we have Z Z, ξ * ≤ max ξ∈Γ Z Z, ξ . So we have the inequality (35).

To the best of our knowledge, this is the first end-to-end DNN training framework that provides quantitative energy guarantees; • We propose a quantitative model to estimate the energy consumption of DNN inference on TPU-like hardware. The model can be extended to model other forms of DNN hardware; • We formulate a new optimization problem for energy-constrained DNN training and present a general optimization algorithm that solves the problem.

Table content: <table/>

Energy consumption and accuracy drops compared to dense models on ImageNet. We set the energy budget according to the lowest energy consumption obtained from prior art.

Table content: <table><tr><td>DNNs</td><td>AlexNet</td><td/><td>SqueezeNet</td><td>MobileNetV2</td></tr><tr><td>Energy Budget</td><td>26%</td><td/><td>38%</td><td>68%</td></tr><tr><td>Methods</td><td colspan="4">MP SSL EAP NetAdapt Ours MP SSL EAP Ours MP SSL Ours</td></tr><tr><td>Accuracy Drop</td><td>0.7% 2.4% 0.8%</td><td>4.4%</td><td colspan="2">0.5% 1.7% 2.7% 0.1% 0.4% 1.7% 2.0% 1.0%</td></tr><tr><td>Energy</td><td>34% 32% 27%</td><td>26%</td><td colspan="2">26% 44% 50% 76% 38% 70% 72% 68%</td></tr><tr><td colspan="2">Nonzero Weights Ratio 8% 35% 9%</td><td>10%</td><td colspan="2">31% 34% 61% 28% 48% 52% 63% 63%</td></tr></table>

Energy consumptions and accuracy drops on MNIST and MS-Celeb-1M.

Table content: <table><tr><td>DNNs@Dataset</td><td/><td/><td colspan="2">LeNet-5@MNIST</td><td/><td/><td colspan="2">MobileNetV2@MS-Celeb-1M</td></tr><tr><td>Energy Budget</td><td/><td/><td>17%</td><td/><td/><td/><td/><td>60%</td></tr><tr><td>Methods</td><td>MP</td><td colspan="3">SSL NetAdapt SBP</td><td>BC</td><td>Ours</td><td>MP</td><td>SSL</td><td>Ours</td></tr><tr><td colspan="3">Accuracy Drop 1.5% 1.5%</td><td>0.6%</td><td colspan="5">1.5% 2.2% 0.5% 1.1% 0.7%</td><td>0.2%</td></tr><tr><td>Energy</td><td colspan="2">18% 20%</td><td>18%</td><td colspan="5">22% 26% 17% 66% 72%</td><td>60%</td></tr></table>

AN (1 + )-APPROXIMATE SOLUTION FOR PROBLEM(10)    Theorem 3. For the projection problem (10), there exists an efficient approximation algorithm that has a computational complexity of O n + (|U |+|V |) 3

Table content: <table><tr><td>2</td><td>log n max(A) min(A+)</td></tr></table>

Energy consumption and accuracy drops compared to dense models on ImageNet. Knowledge distillation is removed from the loss function.

Table content: <table><tr><td>DNNs</td><td colspan="2">AlexNet</td><td colspan="2">SqueezeNet</td><td colspan="2">MobileNetV2</td></tr><tr><td>Methods</td><td>MP</td><td>SSL</td><td>MP</td><td>SSL</td><td>MP</td><td>SSL</td></tr><tr><td>Accuracy Drop</td><td>2.6%</td><td>17.6%</td><td>1.9%</td><td>16.0%</td><td>2.0%</td><td>1.9%</td></tr><tr><td>Energy</td><td>34%</td><td>35%</td><td>44%</td><td>52%</td><td>71%</td><td>77%</td></tr><tr><td colspan="4">ENERGY-CONSTRAINED PROJECTION EFFICIENCY</td><td/><td/><td/></tr></table>

Wall-clock time of the projection operation P Ω(Ebudget) .

Table content: <table><tr><td>DNNs</td><td>AlexNet</td><td>SqueezeNet</td><td>MobileNetV2</td></tr><tr><td>Time (seconds)</td><td>0.170</td><td>0.023</td><td>0.032</td></tr></table>

Peer review (iclr2019):
 The paper proposes a method for neural network training under a hard energy constraint (i.e. the method guarantees the energy consumption to be upper bounded). Based on a systolic array hardware architecture the authors model the energy consumption of transferring the weights and activations into different levels of memory (DRAM, Cache, register file) during inference. The energy consumption is therefore determined by the number of nonzero elements in the weight and activation tensors. To minimize the network loss under an energy constraint, the authors develop a training framework including a novel greedy algorithm to compute the projection of the weight tensors to the energy constraint. Pros: The proposed method allows to accurately impose an energy constraint (in terms of the proposed model), in contrast to previous methods, and also yields a higher accuracy than these on some data sets. The proposed solution seems sound (although I did not check the proofs in detail, and I am not very familiar with hardware energy consumption subtleties). Questions: The experiments in Sec. 6.2 suggest that the activation mask is mainly beneficial when the data is highly structured. How are the benefits (in terms of weight and activation sparsity) composed in the experiments on Imagenet? How does the weight sparsity of the the proposed method compare to the related methods in these experiments? Is weight sparsity in these cases a good proxy for energy consumption? How does the activation sparsity (decay) parameter (.delta) q affect the accuracy-energy consumption tradeoff for the two data sets? The authors show that the weight projection problem can be solved efficiently. How does the guarantee translate into wall-clock time? Filter pruning methods [1,2] reduce both the size of the weight and activation tensors, while not requiring to solve a complicated projection problem or introducing activation masks. It would be good to compare to these methods, or at least comment on the gains to be expected under the proposed energy consumption model. Knowledge distillation has previously been observed to be quite helpful when constraining neural network weights to be quantized and/or sparse, see [3,4,5]. It might be worth mentioning this. Minor comments: - Sec. 3.4. 1st paragraph: subscript -> superscript - Sec. 6.2 first paragraph: pattens -> patterns, aliened -> aligned [1] He, Y., Zhang, X., & Sun, J. (2017). Channel pruning for accelerating very deep neural networks. ICCV 2017. [2] Li, H., Kadav, A., Durdanovic, I., Samet, H., & Graf, H. P. Pruning filters for efficient convnets. ICLR 2017. [3] Mishra, A., & Marr, D. Apprentice: Using knowledge distillation techniques to improve low-precision network accuracy. ICLR 2018. [4] Tschannen, M., Khanna, A., & Anandkumar, A. StrassenNets: Deep learning with a multiplication budget. ICML 2018. [5] Zhuang, B., Shen, C., Tan, M., Liu, L., & Reid, I. Towards effective low-bitwidth convolutional neural networks. CVPR 2018.