Comparing Fisher Information Regularization with Distillation for DNN Quantization
Abstract
A large body of work addresses deep neural network (DNN) quantization and pruning to mitigate the high computational burden of deploying DNNs.
We analyze two prominent classes of methods; the ﬁrst class uses regularization based on the Fisher Information Matrix (FIM) of parameters, whereas the other uses a student-teacher paradigm, referred to as Knowledge Distillation (KD).
The Fisher criterion can be interpreted as regularizing the network by penalizing the approximate KL -divergence (KLD) between the output of the original model and that of the quantized model.
The KD approach bypasses the need to estimate the FIM and directly minimizes the KLD between the two models.
We place these two approaches in a uniﬁed setting, and study their generalization characteristics using their loss landscapes.
Using CIFAR-10 and CIFAR-100 datasets, we show that for higher temperatures, distillation produces wider minima in loss landscapes and yields higher accuracy than the Fisher criterion.
Introduction
Deep Neural Network (DNN) networks are compressed or quantized to reduce hardware complexity, memory usage and energy usage.
A DNN model, F ( x ; θ ) , has output distribution p ( y | x ; θ ) .
Upon quantization, the output of the network is perturbed to p ( y | x ; θ − δ θ ) , where θ − δ θ ≡ Q ( θ ) and
Q ( . ) is the quantization operation.
Quantization has a deleterious effect on the performance of the model, particularly in the low-bit regime, where the number of bits N ≤ 4 [1, 2, 3, 4, 5].
The performance degradation can be mitigated using a regularizer that minimizes the KL -divergence (KLD) between a high performance FP32 (32 bit ﬂoating point) network, and a quantized network.
We consider the FP32 DNN, F ( x ; φ ) , with output distribution p ( y | x ; φ ) , and a quantized network, F ( x ; Q ( θ )) with output distribution p ( y | x ; Q ( θ )) .
We use p φ and p Q ( θ ) to denote the output distributions.
The regularized loss for training the quantized network is ˜
[Equation]
where L CE is the standard cross entropy loss.
Note that we consider two different sets of parameters φ and θ here for generality.
Broadly speaking, there are two paradigms for operationalizing the regularizer in (1): (a) Setting φ
= θ and using a second order expansion around θ of the KLdivergence in terms of the Fisher information matrix (FIM), then empirically estimating the FIM and solving for θ ; (b) Using distillation by pre-training and ﬁxing an FP32 teacher model with parameters φ ; then optimizing the student quantized model with parameters θ .
In this paper, we place these two approaches in a uniﬁed setting and study their generalization characteristics using their loss landscapes.
We show an advantage for implementing (1) via distillation as increasing the temperature hyperparameter in distillation ﬂattens the loss function at the minimum.
We empirically validate that the ﬂatter minima result in improved model accuracy on benchmark data for different bit widths.
Methods
In this section we provide an overview of the two families of methods: Fisher-based regularization and KD.
The details of regularizing using Fisher Information for DNN quantization is outlined in
Section 2.1; Quantization based on KD is covered in Section 2.2.
In Section 2.3 we derive bound on the trace of the Hessian, inversely relating KD temperature and loss surface curvature.
2.1 Fisher Information regularization for quantization
We return to the KL divergence in (1) and set φ = θ (where θ are the FP32 parameters).
This loss function identiﬁes a broad optimum for the network such that small perturbations δ θ = θ − Q ( θ ) around the optimum do not change the output distribution of the network.
For sufﬁciently small δ θ , a second order Taylor expansion of the KL divergence can be written in terms of the FIM:
D KL ( p θ || p θ − δ θ )
≈ δ
Tθ F δ θ .
Therefore, we implement (1) using the FIM,
[Equation]
The FIM, F , deﬁnes a metric on a Riemannian manifold of probability densities on the space of network parameters.
The FIM is the expected value of the negative Hessian of the log-loss (standard cross entropy for DNNs), or equivalently, the covariance of the score 1 .
Many approaches for reducing the complexity of DNNs can be interpreted as means of implementing the regularizer in (2).
For network pruning, unnecessary weights can be eliminated using salience criteria based on the FIM (estimated using Hessian) [6, 7, 8, 9, 10, 11].
These methods are heuristics for discarding weights that induce a large δ θ for parameters that have a small entry in F .
Several methods quantize weights based on the FIM.
In [12] weights are clustered by FIM diagonal value, and high FIM clusters are assigned larger bit depth.
In [13] a non-uniform quantizer is learned using a FIM-weighted k -means scheme.
Matrix-free power iteration is used for mixed-precision layer-wise quantization based on the Hessian eigenspectrum [14, 15, 16].
These methods shape the quantization error, δ θ , such that it is inversely proportional to the entries in F .
Some methods directly regularize δ Tθ δ θ [17] (equivalent to using an identity matrix for FIM), or constrain (cid:96) ∞ -norm of δ θ [5].
Implementing FIM-based methods: Approaches that optimize a variant of (2) require estimation of the FIM.
We estimate F using the diagonal approximation in [4, 13, 18, 12].
The second moment, ˆ v ( θ ) , of the ADAM optimizer [19] is an exponential moving average of squared gradients, and can be used to estimate the FIM diagonal.
Then, the gradient update from the regularizer for weight θ j is:
[Equation]
Before regularization, F is ﬁxed after ﬁrst ﬁnding a good minimum.
The regularizer gradient is not used when computing momentum, but only added when calculating the ﬁnal parameter update [20].
2.2 Distillation as KL -Divergence minimization
We return to the original loss function based on the KL divergence in (1).
In contrast to FIM regularization for quantization, KD uses a pre-trained FP32 teacher, F ( x ; φ ) , to guide the training of the quantized, reduced capacity student model, F ( x ; Q ( θ )) .
We expand (1), D KL ( p φ || p Q ( θ ) ) =
H ( p φ , p Q ( θ ) )
− H ( p φ ) ; since the teacher model is pre-trained and ﬁxed, the second term H ( p φ ) is a constant.
Thus, (1) reduces to the traditional distillation loss,
[Equation]
KD has been previously used for DNN quantization; because KD is parameter agnostic, it can be used with other compression techniques [10, 21, 22, 23, 24, 25].
Quantization-aware KD [26] distills a compressed student model onto an FP32 teacher trained from scratch, before ﬁne tuning the student using the teacher.
FIM-based approaches have even been used in conjunction with KD [10] for model pruning.
KD has been used extensively to reduce the size of transformer models [22, 23, 24, 25, 27].
These variants of KD compress BERT [28] to 7 MB applying layer-wise knowledge transfer [22], or apply KD every L -layers [29].
2.3 Flatness of minima
The ﬂatness of the loss function at a minimum is encoded in the eigenvalues of the Hessian and is used by [30, 31] to connect improvements in generalization performance to ﬂatter minima.
The study in [30] characterizes loss curvature at a minimum using a diagonal FIM to estimate the Hessian eigenspectrum.
The FIM-based methods directly make use of the Hessian during optimization; however KD bypasses the need to compute the Hessian altogether.
As we detail in the ensuing section, KD consistently outperforms FIM-based methods on multiple benchmark problems and for all bit widths considered.
We hypothesize that the increased test accuracy of KD corresponds to reduced curvature of the loss at the minimum.
We show that curvature is inversely proportional to
KD temperature by establishing an upper bound on the the trace of the Hessian.
The Hessian’s trace has been applied as a useful measure of curvature in previous studies [15].
To establish the connection between curvature and the temperature parameter, consider the last layer (softmax) of a DNN.
KD uses a temperature softmax: q k = [Equation] e (
[Equation] ) to smooth predicted class probabilities as a function of T (output probability q k for the k th class).
For a K class problem, the softmax layer has weights θ (cid:96)
∈ R K × M , layer inputs x (cid:96) ∈ R M , and the pre-softmax logit for the k -th class is z k
= θ (cid:96) T k, : x (cid:96) , where θ (cid:96)k, : is a row of the weight matrix.
KD temperature plays an important role for quantization as it directly impacts the width at the optimum and, as a result, the allowable quantization error.
The Hessian H ∈ R MK × MK has block structure with each M × M block,
[Equation]
where I k,j is an element of the identity matrix, k, j ∈ { 1 , . .
. , K } , and additional the subscript n denotes the sample index out of N data points.
The trace of the Hessian can be expressed as
[Equation]
where the bound the arises from q k (1 − q k ) ≤ K − 1 K 2 .
This bound shows an inverse relationship between the curvature of the loss at the minimum and the temperature parameter, T 2 .
That is, as the temperature is increased, the minimum becomes ﬂatter.
This is validated experimentally in the next section where we visualize the loss at the minimum.
Furthermore, we also show that the ﬂatter minimum results in improved accuracy.
Results and Discussion
In this section we empirically compare methods based on FIM regularization with distillation on
CIFAR-10 and CIFAR-100.
We ﬁrst visualize the loss landscapes for a 4-bit quantized ResNet-18 on
CIFAR-10, then compare accuracy results across different bit widths on CIFAR-10 and CIFAR-100.
For all bit width levels on CIFAR-10 and CIFAR-100, networks were trained for 300 epochs with a straight through estimator (STE) to ﬁnd a good minimum θ ∗ , and then subsequently trained with the regularizers (and STE) for 100 epochs.
We also include mean squared quantization error (MSQE) regularization as a useful comparison as it is identical to Fisher regularization, as it assumes that
the
FIM is the identity matrix (i.e. assuming all parameters are equally important).
3.1 Loss landscape visualization
The ﬁlter normalization method of [32] is used to compare the loss surfaces at the minimum for each regularization method.
The loss is plotted along two random vectors β and γ , sampled from a Gaussian distribution and rescaled per ﬁlter.
Qualitatively, Figure 1 shows that the loss surface exhibits local convexity along β and γ , and the improvement in test accuracy correlates with minimum width.
This is also seen in [32] where increases in test accuracy (due to modiﬁcations such as L regularization, skip connections, selection of hyperparameters), correspond to ﬂatter minima.
To quantitatively characterize minima ﬂatness, we ﬁt a paraboloid to the computed loss surface as it is locally convex.
The trace of the Hessian of the ﬁtted paraboloid is used to measure the curvature of the loss.
For consistency, β and γ are ﬁxed across all regularization methods, rescaled with ﬁlter normalization for every network.
Figure 1 (h), demonstrates that as temperature is increased, test accuracy increases, and the trace of the Hessian of the ﬁtted curve decreases.
This empirically consistent with the bound in (6).
This is repeated for 10 randomly chosen γ and β .
3.2 Accuracy results
Tables 1 (a) and (b) show the effect of regularization on CIFAR-10 and CIFAR-100 across all bit widths.
MSQE and Fisher regularization not only result in similar test accuracy, but have very similar
[Figure]
[Figure]
STE
MSQE Fisher
[Table]
[Table]
[Figure]
[Table]
[Figure]
minima ﬂatness.
Since networks trained with STE already use Q ( θ ∗ ) in the forward pass, reducing the perturbation δ θ ∗ does not change the operating point of the network (for both Fisher and MSQE).
3.2.1 Why does distillation outperform Fisher-based methods?
We posit that the difference in performance between distillation, Fisher, and MSQE regularization is a consequence of the difference in minima ﬂatness observed in Figure 1 (h).
The temperature, T , in distillation widens the minimum, as measured by the trace of the Hessian, which translates to improved accuracy at all bit widths.
Another possible reason for the improved performance is that distillation bypasses FIM estimation altogether, while still minimizing (1).
Estimating the FIM is challenging.
Approaches often use diagonal approximations [14, 15, 18, 12, 13] (the Hessians of
DNN loss surfaces are non-diagonal [33, 34]), or computationally intensive methods [11, 14, 15].
These approaches have tenuous theoretical foundation, since the “empirical Fisher approximation,” the widely used FIM estimate, computes expectation over the training labels y n , and is distinct from estimating the true FIM, an expectation over model outputs, y .
Weight updates preconditioned with empirical Fisher can be orthogonal or opposite to updates preconditioned with the true Fisher [35].
Conclusion
In this work we studied two regularization methods for DNN quantization that constrain the KLdivergence between the output distribution of the full precision model and the output of the model under a quantization perturbation.
The ﬁrst method uses an approximation to the KL -divergence in terms of the FIM to constrain the solution space.
The second method uses distillation as an alternative to directly regularizing the KL -divergence between an FP32 and quantized model without assumptions of closeness in parameter space or having to estimate the FIM.
Experiments on CIFAR and CIFAR100 demonstrated distillation outperformed Fisher and MSQE regularization.
We provided evidence that this improvement is due to ﬂatter minima obtained from increasing the temperature parameter in KD.
Broader Impact
An extremely broad array of applications demand the use of low power DNNs across computer vision, natural language processing, and speech recognition.
Computer vision systems such as autonomous driving, human activity detection, robotic systems require efﬁcient DNN inference.
Voice driven smart home systems also require low latency, low power deep learning inference for understanding spoken commands.
Beyond mobile devices, efﬁcient and accurate DNN inference is likely to be crucial for datacenters as well.
Datacenter workloads can involve DNN inference for recommender systems, computer vision, speech processing, and language.
An increasing and large fraction of datacenter workload is expected to come from DNN inference [36] while the model size of state of the art models, in key application domains such as natural language processing, are simultaneously also increasing [28, 37, 38, 39].
DNN compression could become a necessity for low-latency inference of these complex models, and for managing energy density and ecological impact.
Several ethical concerns arise in the application areas of low power DNN.
Pervasive sensor networks could utilize efﬁcient, high performance DNN models for to track individual behaviors, movements, and intents without the knowledge or consent of the surveilled.
For example, The authors in [10] use the Fisher information and Knowledge distillation to prune a network for faster gaze prediction.