E ND - TO -
END
Q UANTIZED T RAINING VIA L OG -
B
ARRIER
E XTENSIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Quantization of neural network parameters and activations has emerged as a successful approach to reducing model size and inference time on hardware that supports native low-precision arithmetic.
Fully quantized training would facilitate further computational speed-ups as well as enable model training on embedded devices, a feature that would alleviate privacy concerns resulting from the transfer of sensitive data and models that is necessitated by off-device training.
Existing approaches to quantization-aware training (QAT) perform “fake” quantization in the forward pass in order to learn model parameters that will perform well when quantized, but rely on higher precision variables to avoid overﬂow in large matrix multiplications, which is unsuitable for training on fully low-precision (e.g. 8-bit) hardware.
To enable fully end-to-end quantized training, we propose Log Barrier
Tail-bounded Quantization (LogBTQ).
LogBTQ introduces a loss term, inspired by the log-barrier for constrained optimization, that enforces soft constraints on the range of values that model parameters can take on.
By constraining and sparsifying model parameters, activations and inputs, our approach eliminates overﬂow in practice, allowing for fully quantized 8-bit training of deep neural network models.
We show that models trained using our approach achieve results competitive with state-of-the-art full-precision networks on the MNIST, CIFAR-10 and
ImageNet classiﬁcation benchmarks.
I NTRODUCTION
As state-of-the-art deep learning models for vision, language understanding and speech grow increasingly large and computationally burdensome (He et al., 2017; Devlin et al., 2018; Karita et al., 2019), there is increasing antithetical demand, motivated by latency, security and privacy concerns, to perform training and inference in these models on smaller devices at the edge rather than in server farms in the cloud.
Model quantization has emerged as a promising approach to enable deployment of deep learning models on edge devices that reduce energy, latency and storage requirements by performing ﬂoating-point computation in low precision (less than 32 bits).
There are two primary strategies for quantization: Post-training approaches quantize the parameters of a model trained in full precision post-hoc, and tend to suffer a heavy penalty on accuracy since their inference graph differs substantially from training (Jacob et al., 2018).
Quantization-aware training (QAT)
(Bhuwalka et al., 2020) combats this discrepancy by simulating quantization during training, so that model parameters are learned that will work well when inference is performed in low precision.
In this work, we focus on the latter setting, suitable for fully quantized training on low-precision (e.g. 8-bit) devices.
Though QAT results in quantized models that perform largely on par with their non-quantized counterparts, current state-of-the-art QAT methods (Wu et al., 2018; Wang et al., 2018; Bhuwalka et al., 2020) are not suitable for training on fully low-precision hardware because they employ fake quantization , meaning each operation is executed using 32- or 16-bit ﬂoating point arithmetic, and its output is quantized to lower precision, e.g. int8.
This results in two key incompatibilities with fully low-precision training, and consequently deployment on real low-precision hardware.
First, existing QAT approaches assume perfect sums in inner product operations, which means that the accumulators used to compute matrix multiplies (the acc row in Table 1) must be higher precision than the values being multiplied (other bit-precision rows in Table 1).
This is to avoid losing res- olution in low-precision additions, also known as swamping (Wang et al., 2018) 1 .
Second, QAT commonly leverages dynamic quantization ranges per-layer, meaning the mapping between highand low-precision values varies by layer, carefully tuned as a function of the network architecture, optimization dynamics and data during training.
While this practice results in higher quantized inference accuracy, it is also a challenge to low-precision training, since it is unclear how to tune those ranges when training on new data in the absence of high-precision arithmetic.
These incompatibilities present a substantial hurdle to quantized training in practice.
For example, an automotive electronics manufacturer may want to deploy a machine learning model on its 8-bit door lock or power window controller to adaptively ﬁt the users’ habits.
In this scenario, existing approaches for quantized training would fail (Sakr et al., 2019).
In response, we propose a new approach for fully quantized training of neural network models, inspired by the barrier method from convex optimization (Boyd & Vandenberghe, 2004).
Log Barrier
Tail-bounded Quantization (LogBTQ) utilizes a log barrier extension loss (Kervadec et al., 2019) to constrain the output of the network, encouraging all model parameters and activations to stay within the same predeﬁned range.
The log barrier function itself is a smooth approximation of the indicator function, which is ideal for selecting the weights that are within the range of quantization (see Figure 1, left).
By ﬁxing a single quantization range throughout the network at the beginning of training, our approach both obviates the need for dynamic ranges, and the limits of the range are set so as to alleviate overﬂow 2 in matrix multiply accumulations.
We combine the log barrier extension loss with an L 1 regularization term (Hoffer et al., 2018) to further reduce the total magnitude of parameters and activations in the model.
To allow for gradients, which tend form a peaky distribution near extremely small values (Zhou et al., 2016; Jain et al., 2020), to be quantized using the same range as the rest of the network, we also adopt the nonlinear µ -law algorithm from audio applications (Deng & Doroslovacki, 2006) to construct a new MU8 codebook that better deals with “swamping” issues compared to the standard IEEE Float Standard.
Experiments show that our approach achieves competitive results compared to state-of-art full-precision models on the MNIST, CIFAR-10 and ImageNet classiﬁcation benchmarks, despite our models being trained end-to-end using only 8 bits of precision. [Table]
B ACKGROUND AND RELATED WORK 2. P OST - TRAINING Q UANTIZATION
There was been a recent surge of interest in quantization research.
In 2020 alone, there were a number of important developments in post-training quantization.
Rusci et al.
Jain et al.
Esser et al.
(2020); Uhlich et al.
(2020) proposed learning-based approaches for determining the quantization ranges of activation and weights at low precision.
Stock et al.
advocates preserving the quality of the reconstruction of the network outputs rather than its weights.
They all show excellent performance compared to full-precision models after quantization.
Sakr &
Shanbhag (2019) presented a detailed analysis of reduced precision training for a feedforward network that accounts for both the forward and backward passes, demonstrating that precision can be greatly reduced throughout the network computations while largely preserving training quality. Our work share the same intuition of preferring small predetermined dynamic range (PDR) and small clipping rate 3 . However, Sakr & Shanbhag (2019)’s approach requires the network ﬁrst be trained to convergence at full 32 bit precision, which is a signiﬁcant limitation.
In this paper, we focus on training rather than inference on low-precision hardware, therefore, we do not assume access to a full-precision high-performing model as a starting point.
2. Q UANTIZATION - AWARE T RAINING
Pioneering works in this domain (Zhou et al., 2016; Courbariaux et al., 2015) looked at quantizing model weights, activations, gradients to lower precision to accelerate neural network training.
The terminology quantization aware training (QAT) was ﬁrst introduced by Jacob et al.
QAT incoporates quantization error as noise during training and as part of the overall loss, which the optimization algorithm tries to minimize.
Hence, the model learns parameters that are more robust to quantization, but QAT is not meant to be performed entirely in low precision, it aims to learn parameters that will work well for low-precision inference.
More recently, several works further pursued the goal of enabling fully low-precision training (Wu et al., 2018; Wang et al., 2018; Das et al., 2018; Sun et al., 2019).
As shown in Table 1, most existing work employs fake quantization, resorting to higher precision values to compensate for the swamping issue, especially during gradient accumulation.
Mixed-precision quantization (Das et al., 2018; Wang et al., 2018; Zhang et al., 2020a), which quantizes a neural network using multiple bit precisions across layers, still relies on higher-precision gradients to preserve model accuracy.
This means it is difﬁcult, if not impossible, to implement these approaches on low-bit (e.g. -bit) hardware
Most similar to our work, Sun et al.
(2019) claim it is possible to do every step in low precision, but the quantization range for the layers in their work is very carefully chosen empirically, which presents great difﬁculty if we were to train models from scratch on low-precision hardware.
Their method also requires a copy of the error ble 1).
In addition to the 9-bit mantissa, the exponent bit in their ﬂoating point format would need to be manually modiﬁed to store the residual due to its small value.
In this paper, we propose a new quantization scheme: log-barrier tail-bounded quantization (LogBTQ) that can perform fully end-to-end low precision training, suitable for deployment on lowprecision hardware.
Our major contributions are the following:
We apply a log barrier extension loss to soft threshold the range of network weights and activations to constrain all the values to be small.
Our quantization scheme also enables global ﬁxed-range quantization which together signiﬁcantly alleviates the overﬂow issue caused by large numbers and dynamic range.
We add an L 1 loss term to encourage sparsity and further reduce overﬂow.
We propose µ -law quantization (MU8) instead of INT8, FP8(1-4-3) or FP8(1-5-2) to construct a more accurate codebook that better compensates for the peaky concentration of network parameters around small values.
L OG B ARRIER T AIL - BOUNDED Q UANTIZATION (L OG BTQ)
The overall diagram of our quantization scheme is shown in Figure 2 (left).
Figure 2 (right) shows the backward pass, where we quantize everything at each layer and all operations including input x , weights w , activations a , errors e , and gradients g (including the gradient accumulation step).
We denote all these values as the set Z = { x, w, a, e, g } .
In this work, different from previous works (Sakr & Shanbhag, 2019; Zhang et al., 2020b) that used adaptive quantization range, we adopt a globally ﬁxed quantization range for every element z ∈ Z , and set z ∈ [ − 2 , 2] .
We do not need to adjust the range and precision during training as in other quantization work that relies on layer-wise dynamic ranges.
This would greatly reduce the overhead for implementation on hardware.
3. C ONSTRAINED F ORMULATION
Let D
=
{ I 1 , ...
I N } denote the labeled set of N training images, and f denote the neural network model, θ here denotes all the parameters of the neural network including weights w .
For task, such as image classiﬁcation, we are usually solving such an optimization problem: min θ L ( f θ ( I ))
where
L is the loss function of our neural network training objective.
In this work, we use the typical cross-entropy loss, and since we are interested in constraining the quantization threshold, we are effectively performing constrained optimization in such a form:
[Equation]
With u our desired barrier (perturbation).
In practice, we set u = 0 .
1 to ensure we can represent as much information as possible within our quantization range (Figure 1, left).
This setting is also explained further in Section 3.3.
3. LOG - BARRIER E XTENSION F UNCTION
Theoretically, problem (1) should be best solved by the log barrier method which is an interior point method that perfectly handles inequality constraints.
(Tibshirani, 2019; Boyd & Vandenberghe, 2004):
In phase I, we would perform Lagrangian-dual optimization to ﬁnd the feasible points:
[Equation]
subject to
[Equation]
where λ ∈ R 1 × N + is the Lagrangian multiplier (dual variable).
After we ﬁnd a feasible set of network parameters, we can use the barrier method to solve equation (1) as an unconstrained problem:
[Equation]
solving problem (3) is Phase II, where ψ t is the standard log-barrier function: ψ t = − 1 t log ( − z ) .
As t approaches inﬁnity, the approximation becomes closer to the indicator function.
Also, for any value of t , if any of the constraints is violated, the value of the barrier approaches inﬁnity.
However, a huge limitation in applicability to practical problems such as ours is that the domain of Eq.
(3) must be the set of feasible points.
The canonical barrier method above is also prohibitively computationally expensive given there are millions of parameters in the network, and we need to alternate the training between primal and dual and do projected gradient ascent for the dual variable.
We are not particularly concerned with the weak duality gap to lower-bound the optimal solution in this work, instead, we are interested in the property of the barrier method to handle inequality constraints.
Therefore, inspired by Kervadec et al.
(2019), we formulate quantization as an unconstrained loss to approximate the constrained optimization problem:
[Equation]
where ˜ ψ t is the log-barrier extension, which is convex, continuous, and twice-differentiable:
[Equation]
in our case, the input z to the log-barrier extension is the same z we deﬁned in the beginning of
Section 3, and t is the scaling parameter.
It basically shares the same property with the standard logbarrier function, {+when t approaches + ∞ , our log-barrier extension would approach a hard indicator: H ( z ) = 
 if z ≤ 0 and + ∞ otherwise. But+} its domain is not restricted to the feasible points only.
This removes the demanding requirement for explicit Lagrangian optimization.
We are doing approximated Lagrangian optimization with implicit dual variables.
Our strict positive gradient of ˜ ψ t will get higher when z approaches violation and effectively push back into the feasible value set.
Because our penalty does not serve as a strict barrier of the feasible set, there is possibility of overﬂow.
However, our goal is to achieve the practical goal of fully-quantized training on low-bit hardware, as long as the majority of values stay within a high conﬁdence interval, the approach will work in practice, as we show in experimental results ( § 4).
Recall the scenario in Figure 1(left).
3. T AIL B OUND OF D ISTRIBUTION
In this section, we demonstrate that the probability of overﬂowing the quantization range can be controlled for a standard ResNet model.
The ResNet model has L layers, each layer contains a CNN operation and ReLU as an activation.
Consider at layer l , for each output element O ∈ R in the layer, a CNN operation with the kernel size k and channel size c can be viewed as a rectiﬁed dot product between the input feature I ∈ R ck 2 with w ∈ R ck 2 .
[Equation]
Suppose w i ∼ N (0 , σ 2( l ) ) and σ 2( l ) is dependent on the layer l and the parameter u we choose.
The second moment of any element z ( l +1) ∈ R from layer l + 1 can be connected with the second moment of element z ( l ) from layer l as follows:
[Equation]
In the He initialization, we could set σ w = (cid:113) 2 ck 2 to cancel the ﬁrst two terms.
In this work, we want to control σ 2( l ) in order to reduce the chance of overﬂow.
We choose σ 2( l ) to depend on u which was deﬁned in optimization problem (1): σ ( l )
=
(cid:113) 2 u ck 2 .
Then, we can simplify Eq.
[Equation]
Suppose the input feature to the ﬁrst layer has second moments E (1) =
E [ z 21 ] , then the second moments at layer l can be estimated as: E
[Equation]
l ) =
E [ z 2( l ) ] =
u l − 1 E (1) .
Next, for each element z ( l ) at layer l , its tail distribution outside the quantization region can be bounded by
[Equation]
where the second inequality is established by Chebyshev’s inequality, here, z l is guaranteed to be nonnegative since it’s the output of the ReLU layer.
We can further estimate the worst upper bound for the right hand side by considering the following optimization problem:
[Equation]
As z ( l ) is the rectiﬁed value, E [ z ( l ) ] > 0 .
One upper bound can be established as follows:
[Equation]
l [Equation] because E [ z ( l ) ] < (cid:112) E ( l ) and Var( z ( l ) ) < E ( l ) .
Notice that the last term can be approximated to E ( l ) when E ( l ) is small enough, and we can establish our probability bound of overﬂowing the quantization range:
[Equation]
Finally, consider the average overﬂow probability of z ∈ R from any layer:
[Equation]
Formula (12) is derived because in deep neural networks, L is usually a large number allowing us to effectively ignore u L .
Therefore, the overﬂow probability is determined by the input’s second moments E , the layer size L and the barrier parameter u .
In practice, we can adjust both u and L to control the overﬂowing tails.
In this work, E (1) is set around 1.0 because the input features are normalized, and L is around 50 in the ResNet-50.
By choosing u = 0 .
1 , the number of overﬂowing parameters can be controlled under 2 .
3. S PARSITY
In order to achieve the goal of practical implementation on 8-bit hardware, we aggressively ﬁx the range of quantization to z ∈ [ − 2 , 2] .
This leaves us with the task of mapping millions of parameters to this range.
Therefore, we desire a sparse solution, and naturally, as is pointed out by Hoffer et al.
(2018), we use L 1 penalty | w i | to encourage sparsity.
Our unconstrained loss then becomes:
[Equation]
where γ > 0 and is tuning variable, and ˜ ψ t is the log-barrier extension as proposed.
3. µ - LAW QUANTIZATION CODEBOOK (MU8)
Even after sparsifying the weight, we still are left will many parameters which are non-linear distributed.
Luckily, they should all be small enough by now thanks to our log-barrier constraint.
Uniform quantization (e.g. INT8) would cause huge information loss in this case.
Inspired by µ -law’s
(Deng & Doroslovacki, 2006) application in audio encoding, we construct a non-linear codebook accordingly, which could be implemented on 8-bit hardware directly.
As is shown in
Figure 1(right), we computed all possible values of FP8(1-4-3) or FP8(1-5-2) 4 used by Sun et al.
(2019), and we can see our MU8 encoding are better at handling very small values e.g. FP8 (1-43)
’s smallest possible number is 1 .
9 × 10 − 3 , FP8 (1-5-2) is 1 .
5 × 10 − 5 , whereas our MU8 could handle 5 × 10 − 6 with the cost of getting sparser in larger numbers, which we have almost eliminated using log-barrier.
Let’s denote our the quantization function as Q µ 8 ( X ) which take x as input and output quantized value x (cid:48) .
First, following the range of µ -law encoding, we set the input to Equation 15 x = z/ 2 , since z ∈ [ − 2 , 2] in this paper.
[Equation]
where µ is a tuning variable, the larger µ is, the more non-linear it becomes, and sgn ( x ) is the sign function.
Then, we get the y ∈ [ − 1 , 1] as the output of F ( x ) , then we perform Stochastic Rounding step introduced by Gupta et al.
e.g. y × (2 7 − 1) ∈ [15 , 16] , if y = 15 . 5 / (2 7 − 1) , then P (ˆ y = 16 / (2 7 − 1)) = P (ˆ y = 15 / (2 7 − 1))
= 0 .
5 ; if y = 15 .
1 / (2 7 − 1) , then P ((ˆ y = 15 / (2 7 − 1))
= 0 .
{+9 , P (ˆ y = 
+}
(2 7 − 1))
= 0 .
1 , as shown in the following equation.
[Equation]
where UL indicates upper limit and LL indicates lower limit.
At last, we get the quantized value x (cid:48) via the inverse function of encoding function (15):
[Equation]
3. O THER U SEFUL T ECHNIQUES
Straight-through Estimator (STE) : We adopt the same STE as Zhou et al.
(2016) which is critical to the convergence of our models.
The gradient of the quantization steps in the quantization model is based on the straight through estimator.
Chunked Updates : We also use the same gradient accumulation strategy as Sakr et al.
Wang et al.
During updates, since we have quantized the gradient and weight parameters using Q µ 8 ( X ) , the gradient will be lower bounded by the smallest possible value of the MU encoding( 10 − 6 ), which is small enough to preserve model accuracy within the 8-bit constraint.
However, directly adding this small gradient to a large number (weight) in our MU8 quantization would cause this small gradient to be rounded-off and thus lose information due to the growing intervals of MU8 encoding at larger values.
Performing chunked gradient updates, updating the weights only after k steps of gradient descent (in this work, we set k = 20 ), helps to accumulate the gradients to be large enough to be rounded up, aiding convergence.
E XPERIMENTAL R ESULTS
Table.
2 shows the performance of models trained with LogBTQ from scratch on the MNIST,
CIFAR-10, and ImageNet datasets.
Our performance on CIFAR-10 is even higher than the FP baseline, and our results using ResNet-50 trained from scratch on ImageNet are competitive against both of the works in comparison.
This is particularly impressive given that LogBTQ uses strictly 8-bits to store gradients whereas previous work resorted to FP16 or hybrid 8/16.
Our accuracy loss of ResNet-50 on ImageNet is only 2.48% which for many use cases is worth the beneﬁt of keeping everything in only 8 bits of precision.
For practical purposes, 2-5% accuracy loss should be tolerable if calculations can be kept strictly in 8 bits. Training Scheme MNIST CIFARImageNet
ResNetResNet-
ResNetMobileNet
FP 8-bit FP 8-bit FP 8-bit FP 8-bit
Wang et al.
HFP8 (Sun et al., 2019) — — — — 76.
LogBTQ FP8(1-4-3) 99.
LogBTQ FP8(1-5-2) 99.
[Table]
D ISCUSSION
Hou et al.
(2019) also provide a bound for the quantized gradient of QAT models, giving the intuition that quantized gradients would slow down convergence, which we also observe in our training.
For each arithmetic operation, we perform our Q µ 8 ( X ) quantization right after each operation to ensure the number still falls into 8 bits.
For out of range numbers, we clamp.
We verify our assump- tions in our experiments: Figure 3 (left) shows a layer’s weight distribution during training, and we can see that they are nicely bounded by our log-barrier to within [-1.5, 1.5].
Figure 3 (right) shows the weight distribution is concentrating to smaller values during training.
Since we only have 10 − 6 precision to handle the gradients, we are inevitably losing some accuracy compared with higher precision training schemes.
FP16 can handle 10 − 8 precision and FP32 can handle 10 − 32 precision.
As we can see in Table 2, LogBTQ can perform almost perfectly on easier tasks such as MNIST and CIFAR-10, but degrades when we are training on the more challenging
ImageNet dataset.
Our Mu8 encoding performs better than 1-5-2 and 1-4-3 with our LogBTQ quantization scheme, showing µ -law encoding can handle small ﬁxed quantization range better than 1-5-2 or 1-4-3.
As is pointed out by Sheng et al. , MobileNet architecture has some layers that are unfriendly for quantization, e.g. An outlier in one channel could cause a huge quantization loss for the entire model due to an enlarged data range. We employed the same techniques proposed by Sheng et al. (2018), though MobileNet benchmark is still not as good as ResNet-50 due to the agressive reduction of parameters.
Overall, it still depends on the use case of the deep learning models to consider the trade-off between accuracy and precision to choose which quantization scheme to adopt.
As far as we are aware,
LogBTQ is the ﬁrst full 8-bit quantization scheme with competitive performance on a large-scale dataset such as ImageNet.
This opens up the prospect of training fully-quantized models on lowprecision hardware.
C ONCLUSION
Motivated by the limitations of fake quantization, we propose Log Barrier Tail-bounded Quantization (LogBTQ) which introduces a log-barrier extension loss term that enforces soft constraints on the range of values that model parameters can take on at every operation.
Our approach eliminates overﬂow in practice, sparsifying the weights and using µ -law non-uniform quantization , allowing for fully quantized 8-bit training of deep neural network models
By constraining the neural net - work parameters driven by theoretical motivations, this work enables the possibility for the ﬁrst time fully-quantized training on low-precision hardware.