Not All Bits have Equal Value: Heterogeneous Precisions via Trainable Noise
Abstract
We study the problem of training deep networks while quantizing parameters and activations into low-precision numeric representations, a setting central to reducing energy consumption and inference time of deployed models.
We propose a method that learns different precisions, as measured by bits in numeric representations, for different weights in a neural network, yielding a heterogeneous allocation of bits across parameters.
Learning precisions occurs alongside learning weight values, using a strategy derived from a novel framework wherein the intractability of optimizing discrete precisions is approximated by training per-parameter noise magnitudes.
We broaden this framework to also encompass learning precisions for hidden state activations, simultaneously with weight precisions and values.
Our approach exposes the objective of constructing a low-precision inference-efﬁcient model to the entirety of the training process.
Experiments show that it ﬁnds highly heterogeneous precision assignments for CNNs trained on CIFAR and ImageNet, improving upon previous state-of-the-art quantization methods.
Our improvements extend to the challenging scenario of learning reduced-precision GANs.
Introduction
Scaling up the size of deep neural networks offers returns to accuracy in tasks across a range of application areas, including computer vision and natural language processing.
Expanding network depth from tens of layers [23, 44] to a hundred [13, 52] revolutionized image classiﬁcation and object detection [9, 14].
Following the 340M parameter BERT model [4], the size of state-of-the-art language models has increased by multiple orders of magnitude: GPT-2 with 1.5B [40], GPT-3 with 175B and recently Megatron-Turing with 530B parameters [45] – more than a 1000-fold increase in size over BERT.
Larger models increase the computational cost of inference, posing challenges to their deployment in real-world applications.
To ameliorate these costs, multiple research avenues have explored methods for constructing more compact deep networks while preserving accuracy comparable to larger models.
These include, e.g., architectural building blocks with sparser connectivity [17, 37, 54, 16] or efﬁcient dimensionality reduction stages [20].
Another category of approaches shrinks larger models via pruning [11, 31, 29, 8, 24, 43], reducing precision [12, 19, 51], or binarization [18, 41, 35, 38].
Parallel efforts on neural architecture search [57, 36, 27, 32, 28, 49] optimize coarser aspects of model structure by automating exploration of hyperparameters governing network size and conﬁguration of layers.
[Figure]
Yet, these methods fall short of providing a complete framework for producing efﬁcient networks.
Pruning, quantization, and binarization are often implemented as heuristic post-processing transformations that operate on a trained network [12, 25] or require special training procedures [18, 41, 53, 56].
Section 2 provides a broader survey.
The lack of modular approaches to optimizing the network structure at this ﬁnest level of detail complicates potential combination with neural architecture search techniques for selecting the functional form of layers and building blocks at coarser granularity.
We introduce a principled framework for training compact neural networks, which encompasses pruning, quantization, and binarization.
Each is a special case of learning a more general heterogeneous mixed-precision architecture, wherein a set of auxiliary parameters govern the bit precision of each weight in the original network architecture.
Our approach learns auxiliary parameters simultaneously with learning the original parameters (weights), with training driven entirely by stochastic gradient descent.
In addition to modularity, this strategy provides two crucial beneﬁts:
• Unlike schemes which post-process a trained network, our method makes the entirety of the training process aware of the goal of producing a reduced-precision model.
Learning weights in the same context as quantization avoids divergence between training and inference scenarios.
• The auxiliary parameters can freely participate in the deﬁnition of the loss being optimized during training.
We can craft loss terms that approximate actual
costs ( e.g., in computation or energy usage) of inference as determined by the current conﬁguration of numeric precisions.
Combined with the task-speciﬁc loss term, we can choose how to trade-off accuracy and efﬁciency.
Key to our technical approach is correspondence between a notion of noisy weights and numeric precision.
Speciﬁcally, during training, we instantiate each network weight as a full-precision value plus uniform random noise of a learnable magnitude.
These noise magnitudes are precisely the auxiliary parameters.
The more noise that can be injected into the value of a particular weight without increasing loss, the lower the precision required to represent that weight.
Figure 1 illustrates the intuition behind this correspondence.
During training, noise magnitudes serve as a differentiable proxy for discrete bit precisions.
Upon completion of training, the learned noise magnitudes dictate replacement of noisy weights with ﬁxed, reduced-precision values.
Though related to Bayesian neural networks, our formulation differs substantially from prior work on Bayesian approaches to network pruning and quantization [30].
Section 3 provides details, with complete derivations in Appendix A.
Our method is able to allocate precisions under any level of granularity ( e.g., per-parameter, per-layer) by sharing the trainable noise magnitude variables across different parameters.
Although a per-layer precision scheme is a more natural ﬁt for current GPUs, new hardware designs can leverage more ﬁne-grained and heterogeneous schemes [34, 46, 21]; hence, we focus on the limit of granularity where each parameter has its own precision.
This setting has the potential to achieve signiﬁcant energy efﬁciency and run-time advantages over the per-layer setting without sacriﬁcing accuracy, as suggested by our empirical evidence.
Future commercially available accelerators, based on the latest hardware research, could make these advantages real.
Figure where different precisions should be assigned to each parameter due to their distinct perturbation limits.
Note the connection between each parameter’s perturbation limit, the width of the level curve in the parameter’s axis, and the allocated precision for each parameter.
Results in Section 4 show that, on standard benchmarks such as creating weight-quantized CNNs for image classiﬁcation on CIFAR-10 and ImageNet, our approach outperforms competitors by achieving higher accuracy using fewer bits per parameter.
Further experiments show similar improvements when training GANs.
We also show empirically that, when extended to quantize hidden activations of a network, our method is able to ﬁnd heterogeneous precision schemes and outperforms well-established activation quantization methods.
The source of our advantages may span both the ability to simultaneously account for precision and task objectives within the training dynamics, as well as the ability to learn heterogeneous weight precisions at the ﬁnest level of granularity.
Related Work
Most approaches in the quantization literature assume that a predeﬁned precision is given prior to training [55, 53, 10, 51, 19, 33, 3, 26, 12, 6], designating the number of bits used to represent each parameter of the network.
The focus of these works is typically to circumvent the obstacle posed by the non-differentiability of the quantization function, as this makes ﬁrst-order methods inapt to train quantized models due to the lack of meaningful gradients.
DoReFa [55] updates real-valued weights using gradients w.r.t.
their quantized forms ( i.e., a straightthrough estimator [1]), enabling training over quantized parameter values with gradient descent.
LQ-Nets [53], PACT [3], and LSQ [6] introduce further ﬂexibility to quantization by also optimizing the quantization step size, which we refer by quantization scale throughout our work, i.e., the real value that each bit corresponds to is also learned.
DSQ [10] proposes a smooth approximation to the quantization mapping, which allows weights to be directly trained with SGD without straight-through estimators.
SLB [51] introduces per-parameter auxiliary variables that induce distributions over values that quantized weights can take – these variables are trained with gradient descent and a ﬁnal quantized model is created by approximating each weight’s distribution by a point-mass.
More recently, some works on quantization have explored ways to assign different precisions to parameter groups in a neural network [48, 47, 5, 50] – such methods, however, remain a minority in the quantization literature.
DNAS [48] uses neural architecture search to map candidate precisions to different parameter groups.
HAQ [47] uses reinforcement learning to learn a quantization policy.
HAWQ [5] uses second-order information to allocate precisions to different parameter groups.
Closest to our work is BSQ [50], which introduces new variables that are trained with gradient-based methods, which are then used to allocate precisions throughout the network.
It operates by ﬁrst mapping each real-valued weight to a bitstring variable whose length is chosen a-priori.
These variables are then iteratively trained with gradient descent and pruned based on their magnitudes, resulting in a decrease in precision whenever a component (a bit) of the bitstring is removed.
Finally, the bitstrings are mapped back to weights, and the quantized model undergoes ﬁne-tuning by adopting straight-through estimators for the quantized weights.
Lastly, Fan et al.
[7] improve the performance of quantized networks by randomly selecting a subset of weights to be quantized during training while keeping the remaining parameters in their original
1: procedure SMOL ( w ∈ R d , L : R d → R , λ ∈ R , p init ∈ N ) 2:
Instantiate parameters s ∈ R d and initialize s init =
− ln (cid:0) 2 p init − 1 − 1 (cid:1) 3: for t = 1 → T 1 do 4:
Sample (cid:15) ( t ) ∼ U d ( ± 1) 5:
Compute L ( w, s ) =
L ( w + σ ( s ) (cid:12) (cid:15) ( t ) ) + λ (cid:107) log 2 (1 + e − s ) (cid:107) 1 and ∇ w,s L ( w, s ) 6:
Update w and s using ∇ w L ( w, s ) and ∇ s L ( w, s ) 7:
Clip w to ± (2 − σ ( s )) 8: end for 9:
Set p = 1 + round (log 2 (1 + e − s )) 10:
ZPA (Optional): If | w | < | w − q ( w, p ) | then set p =
(element-wise) 11: for t =
T 1 → T 2 do 12:
Compute L ( w q ) and ∇ w q L ( w q ) , where w q = q ( w, p ) 13:
Update w using ∇ w q L ( w q ) instead of ∇ w L ( w q ) 14: end for 15: end procedure format.
While our approach also uses randomness to improve quantization in some fashion, the form of noise samples, along with how and why they are used, are vastly different.
Method 3.1 Allocating Precisions by Optimizing Noise Magnitudes
We consider the problem of allocating a precision p (a positive integer which represents the number of bits) to each weight w of a network under a total precision budget (the total number of bits that can be used to represent the model).
If we let q ( w, p ) denote the quantization of w with p many bits, then our goal is to minimize L ( q ( w, p )) subject to p ≤ C , where C is our total budget measured in bits and L captures the loss incurred by the weights – typically the cross-entropy loss over a training set.
Although our method is the result of a rigorous sequence of approximations to the objective described above, it admits an intuitive description and motivation that relies on a connection between quantization and perturbations.
More speciﬁcally, if we assume that a weight w can be perturbed in any direction by at least (cid:15) without degrading the performance of the network, then we can safely represent w with p bits as long as the quantization error does not exceed (cid:15) – that is, | w − q ( w, p ) | ≤ (cid:15) .
In other words, if we know the largest perturbation that each weight admits (its perturbation limit ) without yielding performance degradation, then we can easily assign a precision to each weight: it sufﬁces to choose the smallest number of bits such that the quantization error falls below the perturbation limit of the corresponding weight.
Our method works by estimating the perturbation limit of each parameter to be quantized, which is achieved by directly optimizing the magnitude of the weight perturbations through a novel loss function.
We call our method SMOL : S earching for M ixed-Precisions by O ptimizing L imits for
Perturbations.
Once the perturbation limit of each weight has been estimated through optimization by our method, we assign per-weight precisions by mapping each perturbation limit to a number of bits.
More speciﬁcally, for a model with weight parameters w and training loss L , our method ﬁrst introduces new parameters s of the same shape as w , and at each iteration t aims to update w, s to decrease an augmented loss function:
[Equation]
where (cid:15) ( t ) has the same size as s and whose components are independently sampled from an uniform distribution over the interval [ − 1 , +1] , and λ controls the trade-off between performance and representational cost – a larger value for λ will place more importance on using less bits to represent weights compared to achieving small loss.
Note that the noise samples (cid:15) ( t ) are independent across different iterations, i.e., we draw a new sample (cid:15) ( t ) at each iteration t .
Here, the term σ ( s ) represents the perturbation limit estimates of the weights w : the distance between w and w + σ ( s ) (cid:12) (cid:15) ( k ) is at most σ ( s ) since (cid:15) k is between − 1 and +1 .
A key aspect of our proposed loss function lies in the fact that it is fully differentiable w.r.t.
to s ; hence, gradient-based methods like SGD and Adam can be applied off-the-shelf to optimize both the original weights w and the auxiliary variables s .
This enables the parameters s to be seen as being part of the model itself, allowing for our method to be easily combined with higher-order optimizers, neural architecture search, and other algorithms that operate on networks.
After training, we map the ﬁnal values of s to precisions via some function (cid:100)·(cid:99) that outputs integers e.g., rounding or truncation:
[Equation]
The resulting tensor p will have the same shape as w , and its components represent the number of bits assigned to each parameter in w .
Having allocated a precision to each weight parameter, we can discard the auxiliary variables s and use quantized weights w q = q ( w, p ) to compute the model’s predictions instead of applying perturbations.
This will typically lead to non-negligible changes in the model’s activations, which we circumvent by ﬁne-tuning the model by further optimizing w to minimize the training loss.
Following prior work, we use straight-through estimators to optimize w since q is non-differentiable, i.e., we set ∇ w L = ∇ w q L and perform gradient descent on w .
Lastly, note that inverting the s → p mapping offers a way to initialize s given a desired initial precision p init , more speciﬁcally we set s init =
− ln (cid:0) 2 p
[Equation]
−
[Equation]
(cid:1) .
For example, adopting an initial precision p init = 8 results in s init =
− ln(2 8 − 1 − 1) = − ln(127) , and the perturbation limit estimate will start as σ ( s init )
= 2 1 − p init =
2 − 7 for all weights w of the network.
Our method is also applicable if groups of parameters must share the same precision value e.g., the layer-wise setting where all parameters of each layer are represented with the same number of bits.
In this case, we assign each parameter group i to a single component
s i of s , and when applying perturbations to the weights in the group we scale all the noise samples by the same scalar σ ( s i ) .
At the end of training, we map s i to a single integer p i via (2) which is shared across the group.
3.2 Zero Precision Allocation
A fundamental limitation of having each i ’th bit map to either +2 1 − i or − 2 1 − i is that zero weights cannot be represented: the possible values that a 1-bit weight can assume are {− 1 , +1 } , for a 2bit weight {− 1 . 5 , − 0 .
5 , +0 . 5 , +1 . 5 } , and moreover a quantized weight cannot assume the valueregardless of its precision.
However, in our setting – where we can assign a different precision to each weight – we can directly re-deﬁne the quantization function q to map any w to 0 whenever p = 0 , hence introducing the notion of zero-precision weights .
This is analogous to the procedure of assigning zero precision to a ﬁlter in BSQ, which leads to the corresponding convolution to be completely skipped when computing a model’s outputs.
We propose to assign zero precision to a weight w whenever | w | ≤ | w − q ( w, p ) | .
Since the quantization function q maps w to the closest value that is representable with p bits, whenever a weight’s precision is set to zero due to our proposed strategy, the induced quantization error is guaranteed not to increase.
This follows since prior to changing p the quantization error is | w − q ( w, p ) | ; but once we set p = 0 , it becomes | w − q ( w, 0) | = | w − 0 | = | w | , which cannot lead to an increase since the condition is precisely | w | ≤ | w − q ( w, p ) | .
For example, q ( w = 0 . 2 , p = 2) = 0 .
5 since it is the value in {− 1 . 5 , − 0 .
5 , +0 . 5 , +1 .
5 } that is closest to w = 0 .
2 , while q ( w = 0 .
2 , p = 0)
=
following our re-deﬁnition of q yields a quantization error of | 0 .
2 − 0 | = 0 .
On the other hand, for p =
we have | 0 .
5 − 0 .
2 | = 0 .
Hence, assigning zero precision in this case not only frees up 2 bits but also decreases the quantization error by 0 . 1 .
One advantage of adopting this procedure to assign zero precisions to weights is that it only changes q and hence the quantized network, therefore not affecting the procedure described in the previous section.
In other words, one can allocate precisions to weights by optimizing perturbations as described previously, and then quantize the model separately with and without zero precision allocation.
This results in two quantized models with different precision assignments, hence we can obtain two networks by training the auxiliary variables s only once.
The notion of zero precisions also uniﬁes quantization and pruning, since in practice assigning zero precision is equivalent to pruning a weight.
Results
We evaluate our method in the tasks of quantizing CNNs trained for image classiﬁcation and generation.
For all experiments, we adopt an initial precision p init = 8 , which corresponds to initializing each component of the new parameter s as − ln(2 7 − 1)
≈ − 4 .
[Table] [Table] [Table]
[Table] [Table]
SMOL-
[Table]
[Table]
[Table]
We use the ﬂoor operation to map real-valued precisions to integers: based on our preliminary experiments this more aggressive rounding operation yields more compact models and rarely results in performance degradation.
For all experiments we train the auxiliary parameters s with Adam [22], using the default learning rate of 10 − 3 and no weight decay – all its other hyperparameters are set to their default values.
We run SMOL with and without Zero Precision
Allocation: allowing for zero precisions results in more compact models, and we use SMOL* to denote results with zero precision allocation.
When allocating per-layer precisions, we refer to our method as SMOL-L.
To compare different methods we measure the performance and the size of the resulting network.
We report the average number of bits assigned to the model’s weights (the sum of all precisions divided by the number of weights), which we refer to as ‘average bpp’, along with the compression ratio relative to a full-precision model, which equals to 32 divided by the average bpp.
Appendix D.4 provides preeliminary experiments with Transformers.
4.1 Image Classiﬁcation on CIFAR-
We ﬁrst compare SMOL against different quantization methods on the small-scale CIFAR-10 dataset, adopting different networks to evaluate how aggressively each method can quantize weights without degrading the model’s generalization performance.
Since prior works typically rely on other methods to quantize activations ( e.g., using PACT [3] for activations while applying a new method to the weights), we only consider results with full-precision activations for CIFAR-10 as this isolates each method’s performance to its ability to quantize weights, leading to a more direct comparison.
[Table]
We adopt the standard data augmentation procedure of applying random translations and horizontal ﬂips to training images, and train each network for a total of 650 epochs: the precisions are trained with
SMOL for the ﬁrst 350 while the remainingare used to ﬁne-tune the weights while the precisions remain ﬁxed.
Note that this training budget assigned to our method is considerably smaller than
BSQ’s 1000 total epochs which are split between pre-training, precision allocation, and ﬁne-tuning.
Following prior work, we keep the batch normalization parameters in full-precision as they represent a small fraction of the network’s total parameters.
Like in BSQ, weights of parameterized shortcut connections are also kept in full-precision – namely, shortcuts that consist of 1 × 1 convolutions followed by normalization in ResNet-20 and MobileNetV2.
To train the weights we use SGD with a momentum of 0 .
9 and an initial learning rate of 0 . 1 , which is decayed at epochs 250, 500, and 600.
We use a batch size of 128 and a weight decay of 10 − 4 for
ResNet-20, 4 · 10 − 5 for MobileNetV2, and 5 · 10 − 4 for ShufﬂeNet.
ResNet-20 results are given in Table 1: SMOL comfortably outperforms BSQ and other competing methods by offering higher performance at lower precision.
Comparing against BSQ’s 91 .
9% accuracy at 2 .
7 bpp, SMOL provides 0 .
6% higher accuracy at 0 .
6 lower bpp ( 92 . 5% at 2 .
1 bpp).
With zero-precision allocation (SMOL*), our method outperforms BSQ by 0 .
7% at 1 . 0 lower bpp ( 92 .
6% at 1 . 7 bpp), and matches the performance of the full-precision model with a 18 . 8 × compression ratio.
When allocating layer-wise precisions (SMOL-L), we observe a 0 .
5% higher accuracy at 0 .
3 lower bpp compared to BSQ, showing that although per-parameter precisions improve efﬁciency, our method outperforms the state-of-the-art even when constrained to the less ﬂexible, layer-wise setting (more details in Appendix D.1).
Figure 3 shows efﬁciency curves for BSQ and SMOL-L.
Table 2 presents results for MobileNetV2 and ShufﬂeNet: SMOL also comfortably outperforms
BSQ, offering 0 .
7% higher accuracy at 1 . 1 lower bpp on MobileNetV2 ( 94 . 8% at 1 .
7 bpp, compared to 94 .
1% at 2 . 8 bpp), and 0 .
2% higher performance at 0 .
5 lower bpp on ShufﬂeNet ( 92 . 0% at 2 .bpp, compared to 91 .
8% at 3 . 4 bpp), while outperforming the full-precision model in both cases.
Additional empirical studies are provided in Appendix D.
4.2 ImageNet
For the large-scale ImageNet classiﬁcation task, we quantize the ResNet-18 and ResNet-50 models which allow for comparisons against LQ-Nets, DSQ, SLB, and BSQ.
For both networks we train the weight parameters with SGD, a momentum of 0.9, a weight decay of 10 − 4 and a batch size ofwhich is distributed across 4 GPUs.
We follow LQ-Nets in terms of data augmentation.
We train ResNet-18 for a total of 180 epochs: the ﬁrst 120 are used for precision training and the last 60 for ﬁne-tuning, and SGD has an initial learning rate of 0.1 which is decayed by 10 at epochs 45, 90, 150, and 165.
For ResNet-50, we start from a pre-trained, full-precision model and train for 100 epochs: allocating the ﬁrst 60 for precision training and the remaining 40 for ﬁne-tuning.
An initial learning rate of 0 .
01 is decayed by 10 at epoch 30 for the precision training phase, while ﬁne-tuning starts with the same learning rate of 0 .
01 which is decayed by 10 at epochs 15 and 30.
Note that, as in the
CIFAR-10 experiments, our training budget is considerably smaller than BSQ’s, which also starts from a pre-trained model but has a budget of 180 additional epochs.
Results in Table 3 show that SMOL outperforms competing methods while achieving higher performance than the full-precision baselines.
On ResNet-18, our method offers 0 .
6% higher accuracy than
LQ-Nets at 1 .
7 lower bpp ( 69 . 9% at 2 .
3 bpp compared to 69 .
3% at 4 . 0 bpp), while also outperforming the full-precision model by 0 .
[Table]
On ResNet-50, SMOL outperforms LQ-Nets at 2 .
2 lower bpp, with an average number of bits of 0 . 8 – lower than a binary network.
With 4 . 1 bpp, our method provides a 1 . 0% improvement over the full-precision baseline
[Table]
suggesting that the noise injection adopted by our method has additional regularizing effects that can further improve generalization.
Performance by average precision plots for BSQ, SMOL
[Table]
and LQ are given in Figure 4.
4.3 Image Generation with DCGAN
We train a DCGAN [39] model on the CIFAR-10 dataset to perform unconditional image generation at a resolution of 32 × 32 .
The generator is implemented as four transposed convolutional layers with batch norm and ReLU activations, while the discriminator consists of four strided convolutional layers with batch norm and LeakyReLU activations.
Models are trained with the binary cross-entropy loss using Adam with a learning rate = 2 · 10 − 4 and ( β 1 , β 2 ) = (0.5, 0.999).
We only quantize the weights of the generator since it is the network used for deployment.
For BSQ, we ﬁrst train a full-precision DCGAN for 30 epochs, quantize its weights to 8-bits, and conduct 100 epochs of precision learning with a pruning interval of 10 epochs and λ = 2 · 10 − 3 .
The models are then ﬁne-tuned for 30 epochs, with a 10 × decayed learning rate.
Similarly, for SMOL we train for a total of 160 epochs: 130 for precision training followed by 30 epochs for ﬁne-tuning.
We randomly generate 10,000 samples for all methods, each assessed with Inception Score(IS)
[42] and Fréchet Inception Distance (FID) [15].
As shown in Table 4, our method consistently outperforms
BSQ, with higher generation quality at lower bpp, demonstrating generalization capability to the challenging task of quantizing GANs.
We drastically improve BSQ’s FID from 37 .
1 at 3 .
9 bpp to 29 .
9 at only 3 . 5 bpp.
For BSQ at 2 . 8 bpp,
SMOL* achieves 7 . 2 better FID with only 1 . 0 bpp.
Even when evaluated at an extremely low bpp of 0 . 2 , our method still generates images with good quality.
4.4 Heterogeneous Precisions for Quantized Activations
In order to extend SMOL to train precisions for hidden activations, we introduce new trainable parameters to estimate the perturbation limit of activation outputs instead of weights.
For an activation tensor u , we instantiate a new variable s of the same shape which will be be trained jointly with the network’s original parameters.
[Table] [Table]
Similarly to the weight quantization case described in Section
[Table]
at each training iteration t we sample a tensor (cid:15) ( t ) with the same shape as u , where each component is drawn uniformly from [ − 1 , +1] , and generate perturbed activations u + M 2 · σ ( s ) (cid:12) (cid:15) ( t ) , which are used in place of u throughout the next layers of the network.
The scalar M denotes the range of activation function used to compute u , i.e., M = 1 for sigmoid and M = 2 for tanh activations, and is used to match the magnitudes of the activations and its perturbations.
Since ReLU activations are unbounded, we use PACT [3] which clips values to [0 , α ] where α is a new trainable parameter – this yields M = α .
Once s has been trained, we map its components to integer precisions which are used to quantize the activations u .
Table 5 shows the performance of a ResNet-20 trained on CIFAR-10 whose activations are quantized by PACT and SMOL (referred as SMOL-A): our method provides off-the-shelf improvements over
PACT, offering higher accuracy at lower average bits per activation (bpa).
All networks are trained for 200 epochs, following the training settings from Section 4.1.
4.5 Computational Efﬁciency
A key question is whether per-parameter precisions can result in inference energy cost reductions.
The power required to multiply a 2-bit and 3-bit weight with a 4-bit activation is 2 .
41 × and 3 .
83 × higher than what is required for a 1-bit weight; and the latency is 1 .
91 × and 2 .
10 × higher, respectively.
These numbers are estimated using ripple-carry adder based multiplier designs (which are suitable for low-precision operations) for the corresponding precision settings [2].
For the accumulation operations, we assume that the power and latency values are the same for different precision settings.
We take as an example the last convolutional layer of ResNet-20, for which BSQ assigns 2 bits for all parameters while our method assigns 1, 2, and 3 bits to 74%, 24.5%, and 1.5% of the parameters, respectively.
In this case, the layer with precisions assigned by SMOL requires only 57.5% of computation power while improving latency by 36% compared to BSQ, which amounts to an energy cost reduction (power x latency) of 62.7%.
Note that in real hardware designs, additional control overheads, e.g., around 25% [34], are required to perform ﬁne-grained mixed-precision operations.
Although estimates, these suggest that a ﬁne-grained precision allocation scheme can result in signiﬁcant energy savings on hardware designed to support the corresponding arithmetic operations.
Discussion
With advances in hardware design and signiﬁcant overparameterization of recent deep networks, learning compact model representations that can be leveraged by specialized hardware has become a key problem – one that is not restricted to the machine learning and systems communities.
Our work goes a step further in terms of network quantization by proposing SMOL, a modular and principled method that is able to learn heterogeneous precisions at the ﬁnest granularity, where different weights of the same model can be represented in distinct precisions regardless of how the weights are connected in the network’s topology.
[Figure]
Moreover, SMOL operates by introducing auxiliary variables to the target model, thus being agnostic to the network’s topology, the choice of optimizer, or the underlying training setting, and hence it can be easily combined with other algorithms such as neural architecture search.
Our procedure to allocate precisions lies in the optimization of a trade-off between the model’s performance and its cost – a cost function that can be chosen according to the underlying application or the user’s concern, e.g., computational or environmental cost.
Finally, our method offers an uniﬁcation of pruning, binarization, and quantization, while at the same time achieving state-of-the-art results in image classiﬁcation and image generation with quantized models.
Our results show that adopting heterogeneous precisions yields highly compact models that not only perform well, but can also result in signiﬁcant computational advantages when deployed on specialized hardware.
Acknowledgments and Disclosure of Funding
This work was supported in part by the University of Chicago CERES Center.
The authors have no competing interests.
Appendix
We present additional experiments and discussions in the following sections.
Appendix A provides the derivation of our precision-training framework, including steps and approximations taken to arrive at our proposed regularized loss.
Appendix B presents a discussion on how to accurately compare methods that use different bit values, which is crucial when comparing our method against BSQ [50] as it adopts unsigned bits and a mutual notion of precision must be established.
Appendix C has additional details on our main experiments, including how λ was chosen and how activations were quantized.
Finally, Appendix D offers additional experimental results with the goal of studying key questions regarding our method, including whether it is able to assign precisions in a layer-wise fashion instead of per-parameter.
A Framework Derivation A.1 Preliminaries
We start by deﬁning a function family F parameterized by Θ ⊆ R d : for any θ ∈ Θ , we have a function f θ ∈ F , f θ : R d in → R d out that is induced by θ .
In our practical setting, the variables θ denote a conﬁguration for the parameters of neural network while f θ represents the function that the network implements under such parameter values.
For the sake of convenience, we directly deﬁne a loss functional L : Θ → R that measures the (average) instant loss of a function f θ ∈ F induced by some θ ∈ Θ .
For example, L can be explicitly given by L ( θ ) = 1 N (cid:80)
Ni =1 (cid:96) ( f θ ( x i ) , y i ) , where (( x i , y i ))
Ni =1 is a ﬁxed dataset and (cid:96) is an instant loss ( i.e., squared loss or cross-entropy).
Now, consider the problem of ﬁnding a parameter setting θ ∈ Θ where each component θ i is represented by exactly p
i many (signed) bits given a bit-value map ( v 1 , v 2 , . . . )
that assigns a scalar v j ∈ R to each position j ∈ N in a bit string.
In other words, we have
[Equation]
where b i,j ∈ {± 1 } is the j ’th bit used to represent θ i and ( v j ) Pj =1 is the bit-value map.
W.l.o.g.
we will assume that the bit-value map is (countably) inﬁnite and
that V d ⊆ Θ , where V = { (cid:80) ∞
j =1 b j · v j | b j ∈ {± 1 }} is the set of all possible values that each parameter can assume.
We can then formalize the problem of ﬁnding a set of bit precisions p ∈ N d and bit strings b ∈ {± 1 } d ×∞ such that the total precision (cid:80) di =1 p i is minimized and the parameter setting θ ∈ Θ induced by b achieves ‘small loss’:
Deﬁnition A.1.
(Mixed Precision Problem)
[Equation]
where L ∗ = min θ ∈ V d L ( θ ) is the smallest achievable loss by F over V d and δ ∈ R + is a scalar that captures tolerance in terms of suboptimality.
Here, V is the function that maps bit strings b ∈ {± 1 } d ×∞ and precision values p ∈ N d to parameter values θ ∈ V d element-wise:
[Equation]
The above problem poses a few obstacles: ﬁrst, it involves a search over multiple binary values for each of the d parameters of the network; second, the derivative w.r.t.
each integer precision p i is undeﬁned, making it ill-suited for gradient-based methods.
Alternatively, we can directly optimize θ ∈ V d under the constraint that its ﬁnite-precision representation is used to evaluate the loss L :
Deﬁnition A.2.
(Mixed Quantization Problem)
[Equation]
where Q maps θ ∈ V to a ﬁnite-precision representation of itself with exactly p many bits.
To deﬁne
Q precisely we will ﬁrst introduce a notion of inverse for V (a ‘notion’ since V is not generally bijective).
For any θ i ∈ V , we let V − 1 ( θ i ) denote a setting ( b i , p i ) , with p i ∈ N and b i ∈ {± 1 } ∞ such that
V ( b i , p i )
=
θ i and, moreover, for any ( b (cid:48)
i , p (cid:48) i ) such that V ( b (cid:48)
i , p (cid:48) i )
=
θ i , it follows that p (cid:48)
i ≥ p i .
That is, V − 1 maps each element θ i ∈ V to a precision-bitstring tuple that perfectly represents θ i using the least number of bits.
We also deﬁne V − 1 p ( θ i ) =
p i and V − 1 b ( θ i )
=
b i where, as before,
V − 1 ( θ i ) =
( b i , p i ) .
Then, we can deﬁne Q by Q ( θ i , p i )
=
(cid:80) p i j =1 V − 1 b ( θ i )
j · v j , i.e., it is the value induced by the shortest bit string that represents θ , when considering only its ﬁrst p i bits.
Note that this yields
Q ( θ i , p i ) =
V ( V − 1 b ( θ i ) , p i ) , which will be a key property in the proofs below.
While it might be tempting to see Q as a quantization mapping, it not always acts as such.
For example, consider the value map given by v j = 2 1 − j , where θ
i = 1 can be represented by p (1) = 1 , b (1) i = (1 , 1 , 1 , 1 , . . . )
, by p (2)
= ∞ , b (2)
i =
(1 , − 1 , 1 , 1 , . . . ) , and so on.
In this case, we could let
V − 1 (1) =
( b (1) , p (1) ) , but not V − 1 (1) =
( b (2) , p (2) ) .
For the former, we would have Q (1 , 2) = 32 , which is not a quantization – in the common sense – of θ i = 1 under the considered value map.
With this deﬁnition of Q we can show that the optimization problems in Equations (3) and (4) are equivalent (under V and V − 1 ).
Lemma A.3.
For any setting ( b, p ) that satisﬁes L ( V ( b, p ))
≤ L ∗ + δ in Equation (3) , we have that ( θ = V ( b, p ) , p ) satisﬁes L ( Q ( θ, p ))
≤ L ∗ + δ in Equation (4) .
Proof.
By deﬁnition, we have:
[Equation]
Therefore, the parameter setting induced by V ( p, b ) in Equation (3) will be the same as the one induced by Q ( θ, p ) in Equation (4), and hence the losses will be the same and satisﬁability follows.
Lemma A.4.
For any setting ( θ, p ) that is a minimizer of the problem in Equation (4) , we have that ( b = V − 1 b ( θ ) , p )
satisﬁes
L ( V ( b, p )) ≤ L ∗ + δ in Equation (3) .
Proof.
Since ( θ, p ) minimizes (cid:80) di =1 p i by assumption, we have that no θ i can be represented with less than p i many bits and hence p =
V − 1 b ( θ ) .
Therefore
[Equation]
Therefore, the parameter setting induced by Q ( θ, p )
= θ in Equation (4) will be the same as the one induced by V ( b, p ) in Equation (3), and hence the losses will be the same and satisﬁability follows.
With the two Lemmas we can prove that the two problems are indeed equivalent:
Corollary A.5.
The optimization problems in Equations (3) and (4) are equivalent under V and
V − 1 .
That is, for any ( b, p ) that minimizes (3) , we have that ( θ = V ( b, p ) , p ) minimizes (4) , and for every ( θ, p ) that minimizes (4) we have that ( b = V − 1 b ( θ ) , p ) minimizes (3) .
Moreover, the induced function f θ will be the same in all cases.
Proof.
• (4) → (3): assume for the sake of contradiction that ( b = V − 1 b ( θ ) , p ) is suboptimal in (3), then there exists a satisﬁable ( b (cid:48) , p (cid:48) ) with (cid:107) p (cid:48)
(cid:107) <
(cid:107) p (cid:107) , but from Lemma A.3 it follows that ( θ = V ( b (cid:48) ,
p (cid:48) ) , p (cid:48) ) is satisﬁable in (4), which contradicts the optimality of ( θ, p ) .
• (3) → (4): assume for the sake of contradiction that ( θ = V ( b, p ) , p ) is suboptimal in (4), then there exists a satisﬁable ( θ (cid:48) , p (cid:48) ) with (cid:107) p (cid:48)
(cid:107) <
(cid:107) p (cid:107) , but from Lemma A.4 it follows that ( b = V − 1 b ( θ (cid:48) , p (cid:48) ) , p (cid:48) ) is satisﬁable in (3), which contradicts the optimality of ( b, p ) .
A.2 Stochastic Approximation
We now present a sequence of steps, involving approximations, that yield our ﬁnal method.
We can ﬁrst write
[Equation] [Equation]
which allows us to show that
[Equation]
by relying on the deﬁnition of Q and V − 1 alone.
Next, we consider the following assumption: if ( θ i , p i ) is satisﬁable, then ( θ i , p (cid:48)
i ) is satisﬁable for any p (cid:48)
i > p
i (i.e. assigning more bits to a parameter preserves its satisﬁability).
This assumption can be more simply framed as
[Equation]
and is the core assumption for our method.
Using the deﬁnition of R and w.l.o.g. denoting Q ( θ, p ) simply by θ (cid:48) , we get from the above equation that
[Equation]
i.e., using the trinary system, this means that if a bitstring (1 , 1 , 0 , 0 , 0 , . . . )
is satisﬁable, then any bitstring (1 , 1 , ± 1 , ± 1 , ± 1 , . . . )
will also be satisﬁable.
Now, we can frame the following optimization problem which has inﬁnitely many constraints:
[Equation]
Note that R ( θ i , p (cid:48) i ) includes V − 1 b ( θ i ) which admits any conﬁguration for all bits after V − 1 b ( θ i ) -pth one due to the fact that Q operates via performs truncation.
We further assume that { R ( θ i , p (cid:48) i ) | p (cid:48)
i > V − 1 p ( θ i ) } is dense in some interval
I ( θ i , p i ) =
[ l, u ] ⊂ V .
We then frame a stochastic approximation for the problem above
[Equation]
which is designed by sampling from the inﬁnitely many constraints that were introduced:
[Equation]
Next, we pose the constraints using
Lagrange
multipliers to yield a regularized
objective
[Equation]
Note that the derivation so far is agnostic to the underlying value map.
We now assume the value map is given by v j = 2 1 − j , where we get P ( I ( θ i , p i ))
=
U ( ± 2 1 − p i ) under a bit distribution of U ( ± 1) for all bits after the V − 1 p ( θ i ) -th one:
[Equation]
To allow for gradients to ﬂow through the sampling procedure
we apply the reparameterization trick and a continuous relaxation on p :
[Equation]
Finally, we reparameterize both p and θ to yield variables that can be optimized over the reals – although not necessary, this results in a procedure that does not require variables to be projected to a domain after each update
We deﬁne s = σ − 1 (2 1 − p ) , θ = τ ( w ) , and frame the problem as optimization over ( s, w ) instead, with the advantage that, unlike ( p, θ ) , they both belong to R d :
[Equation]
where ideally we have τ : R → V be such that τ ( R ) =
V .
This concludes the derivation of our regularized loss.
Note that instead of adopting the reparameterization θ = τ ( w ) , we can directly optimize θ while guaranteeing that θ ∈ V : this can be easily achieved by projecting the parameters θ to an interval, for example to [ − 1 , +1] when using signed bit representations.
B Bit Counting: Implicit vs Explicit Sign Bits
The process of quantizing data is heavily dictated by whether bits are chosen to be signed (taking values { +1 , − 1 } ) or unsigned (values { 0 , 1 } ), and comparing methods that use different settings might require adopting notions of ‘precision’ that diverge from the ones used in their original papers.
Most methods on quantizing neural networks [53, 51, 10] – ours included – adopt signed bits, which is often done implicitly by splitting the [ − 1 , +1] interval uniformly in 2 p bins to extract the values that quantized data can assume.
In this case, assigning a value of 2 1 − i to each i -th bit results in a bitstring ( b 1 , b 2 , . . .
, b p ) being mapped to (cid:80) pi =
b i · 2 1 − i , a variable that can assume 2 p different values uniformly distributed over an interval centered at the origin i.e., [ − M, + M ] for positive M .
On the other hand, adopting unsigned bits results in the possible values for (cid:80) pi =1 b i · 2 1 − i to be uniformly distributed over a [0 , M (cid:48) ] interval instead (for positive M (cid:48) ), hence bitstrings can only be mapped to non-negative values.
An additional ‘sign bit’ can be adopted to allow the representation of negative numbers: in particular, points are encoded to unsigned bitstrings of the form ( b s , b 1 , . . .
, b p (cid:48) ) , which are then mapped to s ( b s ) · (cid:80) p (cid:48)
i =
1 b i · 2 1 − i where s ( b s ) =
+1 if b s = 1 and s ( b s )
= −
if b s = 0 : note that such mapping allows for a total of 2 p (cid:48)
+1 different outputs.
In this case the term ‘precision’ p can be used to either denote p (cid:48) + 1 (the total number of bits used to encode each point) or p (cid:48)
(number of bits except for the ‘sign bit’ b s ).
However, when comparing a method that relies on such representation to one that adopts signed bits instead, we must count the additional ‘sign bit’ i.e., p = p (cid:48)
+ 1 , as in this case both methods can represent exactly 2 p values given a precision of p bits.
When considering results reported by BSQ [50], we account for an additional ‘sign bit’ that is implicitly used when quantizing full-precision weights with bitstrings.
More speciﬁcally, once a model has been ﬁne-tuned, BSQ converts its full-precision weights W s ∈ R d to binary tensors
W (0) p , . .
, W ( n − 1) p , W (0) n , . . .
, W ( n − 1) n , each ∈ { 0 , 1 } d , where ( W ( b ) p ) n − 1 b =0 encodes the positive components of W s while ( W ( b ) n ) n − 1 b =0 encodes the absolute value of its negative components.
These binary tensors are then mapped to quantized weights W final = 1 2 n − 1 (cid:80) n − 1 b =0 ( W ( b ) p −
W ( b ) n )2 n − b − 1 , whose components can take values 1 2 n − 1 { 0 , 1 , . .
, 2 n − 1 } by setting W ( b ) n =and enumerating all conﬁgurations for ( W ( b ) p ) n − 1 b =0 , while setting W ( b ) p =
and considering all conﬁgurations for ( W ( b ) n ) n − 1 b =0 yields the values − 1 2 n − 1 { 0 , 1 , . .
, 2 n − 1 } : in total, we have that elements of W final can assume 2 n +1 − 1 distinct values ( 2 n − 1 positive, 2 n − 1 negative, and 0 ).
An implicit ‘sign bit’ is also used in the forward computation of quantized models, which maps weights W to2 n − 1 · sign( W ) (cid:12) (cid:100)| W | · (2 n − 1) (cid:99) : the implicit representation of n bits does not encode the sign of the weights, which is recovered after quantizing its absolute value over 2 n bins.
To summarize, BSQ uses unsigned bit representations to encode weights but without an explicit ’sign bit‘: the weight signs are either bypassed and hence not encoded (as in the forward pass), or are implicitly encoded by allocating two bitstrings for each weight (as when converting full-precision weights to bitstrings) 1 – where which one is used is dictated the the weight’s sign.
To account for the mismatch in how many values can be represented with p many bits by BSQ compared to our method, we account for one more bit when reporting BSQ results, which designates the implicit ‘sign bit’ in addition to the unsigned n -bit representations.
C Experimental Details
Values for λ : When training ResNet-20 networks on CIFAR-10, we used λ ∈ { 10 − 6 , 7 · 10 − 7 , 5 · 10 − 7 } when running SMOL, where larger values for λ result in less bits per parameter: λ = 10 −resulted in a network with 2 . 1 bpp while λ
= 5 · 10 − 7 yielded a model with 2 . 8 bpp.
For MobileNetV we used λ ∈ { 10 − 7 , 2 · 10 − 7 } , while for ShufﬂeNet we adopted λ ∈ { 5 · 10 − 8 , 7 · 10 − 10 }
For our ImageNet experiments, we adopted { 10 − 6 , 10 − 8 } and { 10 − 7 , 10 − 8 } as values for λ when training ResNet-18 and ResNet-50 networks, respectively.
Different values for λ were required (compared to the ones we used for CIFAR-10) due to the ImageNet ResNets having considerably more parameters than the ResNet-20 model we adopted for CIFAR-10: note that our regularizer is a sum over all precisions instead of, for example, an average.
Finally, for image generation we used λ ∈ { 2 · 10 − 5 , 10 − 5 , 5 · 10 − 6 , 2 · 10 − 6 , 10 − 6 } , where the most compact model (0.2 bpp) was achieved with λ = 2 · 10 − 5 along with zero precision allocation.
In most cases we saw a decrease of between 1 and 2 bpp when enabling zero precision allocation: training a DCGAN with λ = 10 − 5 results in 0.5 bpp with zero precision allocation and 1.7 without it.
Quantized activations for ImageNet Experiments: We follow BSQ [50] and replace all ReLU modules by ReLU6 throughout the network, resulting in outputs constrained to the [0 , 6] interval.
Activations are not quantized prior to ﬁne-tuning i.e., precision training with SMOL uses full-precision activations, which are only quantized once precisions have been allocated and ﬁne-tuning has started.
All our experiments that quantize activations consider 4-bit precisions i.e., the activation tensors are quantized using 2 4 = 16 different values.
Also following BSQ, we use straight-through estimators to allow for gradient ﬂow through the activation quantization procedure.
The quantization values are uniform over half of the activation range – typically [0 , 3] , which results in the full-precision activations being quantized to one of thevalues in the set { 0 , 15 , 25 , . .
D Additional Experimental Questions D.1 What Precision does SMOL Allocate to Each Layer?
In Section 4.1 we present results for SMOL when training layer-wise precisions on a ResNettrained on CIFAR-10 (SMOL-L in Table 1), where higher performance under lower bpp is achieved compared to BSQ.
Figure 6 shows the precisions allocated to each layer of the ResNet-20 under three different values for λ : in all cases, later convolutional layers are assigned low precisions compared to earlier stages of the network.
The ﬁrst convolution, along with the fully-connected layer at the end of the network, are assigned signiﬁcantly higher precision than other layers.
D.2 More Precision Training or More Fine-tuning?
In our main experiments we allocate roughly half of the training budget to precision training and the other half to ﬁne-tuning i.e., further optimization of the weight parameters while applying quantization and adopting straight-through estimators.
Here we show how the performance and size (measured in average number of bits per parameter) of the ﬁnal model behaves under different ratios for the number of epochs allocated to precision training over the total training budget.
We use the settings described in Section 4.1, training models with λ = 7 · 10 − 7 but changing what how many of the total 650 epochs are used for precision training – the remaining ones are always allocated to ﬁne-tuning the model.
Results are shown in Figure 7: there is signiﬁcant change in the model’s bpp when the fraction of epochs allocated to precision training is lower than 250 650 , indicating that SMOL requires betweenand 250 epochs to fully optimize the auxiliary variables s .
For ratios larger than 250650 our method yields networks with similar bits-per-parameter (around 2 .
5 bpp)
, but the test accuracy starts to decrease as the ratio increases due less epochs being used to ﬁne-tune the model.
The best performance is achieved by allocating 350 epochs for precision training, which results in the remaining 300 epochs being assigned for ﬁne-tuning.
[Table]
D.3 Does SMOL Learn Structured Precision Maps?
To evaluate whether the per-parameter precision assignments learned by SMOL have an underlying structure or not, we set up an additional experiment which aims to re-train and evaluate a network after destroying any possible structure in its precision map.
In particular, we start from a ResNetmodel that has been fully trained and ﬁne-tuned by SMOL, and proceed to re-train it under two different settings.
For the ﬁrst, we re-initialize the weight parameters (using a different random seed) but fully preserve the precision assignments generated by SMOL i.e., each i -th weight w i with precision p i is randomly re-initialized ( w i ∼ D init ), but p i is kept the same.
As for the second setting, we also apply a random permutation on the precision tensor of each layer i.e., weights w i and w j of the same layer have their precisions p
i and p j swapped – with the goal of destroying any possible structure that the tensor p might have.
We then re-train the two networks using the same training settings adopted in Section 4.1, but skipping precision training and hence allocating a budget of 300 epochs to weight training under quantization and straight-through estimators, akin to ﬁne-tuning.
6 and Figure 8 present
) the precisions yields in signiﬁcantly lower performance ( 89 . 3% compared to 91 .
5% ), suggesting that the precision assignments generated by SMOL do have an underlying structure, which can be some form of per-layer organization i.e., allocating high or low precision to all elements in the same convolutional ﬁlter, or cross-layer structure i.e., feature maps (layer outputs) have low or high precision connections to both the previous and the next layer.
We leave further the characterization of such structure to future work.
D.4 Can SMOL Quantize Transformer Models?
[Table]
Since our results in Section 4 focus on CNN-like networks like
ResNets and MobileNetV2, an immediate question is whether SMOL is able to efﬁciently quantize architectures of different families, for example Transformers which do not use convolutions and whose main component is the attention mechanism.
To answer this question, we run preeliminary experiments where we train a standard encoder-decoder Transformer with 6 attention blocks on the neural machine translation
IWSLT’14 German to English task, where we quantize all weights except for the layer normalization parameters.
A total of 50 epochs were used for training: for SMOL, the ﬁrst 30 were used to train the precisions and the remaining 20 for ﬁne-tuning.
Results in 7 show that SMOL can match the performance of the full-precision model at 5 .
8 bpp, which amounts to a 5 . 5 × compression ratio.
At a lower bpp of 3 . 9 ( 8 . 2 × ratio), the BLEU score is lower by only 0 . 2 .
Note that we use the same training budget that is commonly employed when training a full-precision model for this task, which is likely suboptimal for quantization.
Therefore, although promising, these results are very preeliminary and can likely be improved by simply training for more epochs.