Feature Affinity Assisted Knowledge Distillation and Quantization of
  Deep Neural Networks on Label-Free Data
In this paper, we propose a feature affinity (FA) assisted knowledge
distillation (KD) method to improve quantization-aware training of deep neural
networks (DNN). The FA loss on intermediate feature maps of DNNs plays the role
of teaching middle steps of a solution to a student instead of only giving
final answers in the conventional KD where the loss acts on the network logits
at the output level. Combining logit loss and FA loss, we found that the
quantized student network receives stronger supervision than from the labeled
ground-truth data. The resulting FAQD is capable of compressing model on
label-free data, which brings immediate practical benefits as pre-trained
teacher models are readily available and unlabeled data are abundant. In
contrast, data labeling is often laborious and expensive. Finally, we propose a
fast feature affinity (FFA) loss that accurately approximates FA loss with a
lower order of computational complexity, which helps speed up training for high
resolution image input.
Zhijian Li, Biao Yang, Penghang Yin, Yingyong Qi, Jack Xin
Feature Affinity Assisted Knowledge Distillation and Quantization of Deep Neural Networks on Label-Free Data

###### Abstract

In this paper, we propose a feature affinity (FA) assisted knowledge distillation (KD) method to improve quantization-aware training of deep neural networks (DNN). The FA loss on intermediate feature maps of DNNs plays the role of teaching middle steps of a solution to a student instead of only giving final answers in the conventional KD where the loss acts on the network logits at the output level. Combining logit loss and FA loss, we found that the quantized student network receives stronger supervision than from the labeled ground-truth data. The resulting FAQD is capable of compressing model on label-free data, which brings immediate practical benefits as pre-trained teacher models are readily available and unlabeled data are abundant. In contrast, data labeling is often laborious and expensive. Finally, we propose a fast feature affinity (FFA) loss that accurately approximates FA loss with a lower order of computational complexity, which helps speed up training for high resolution image input.

 Quantization, Convolutional Neural Network, Knowledge Distillation, Model Compression, Image Classification

## I Introduction

Quantization is one of the most popular methods for deep neural network compression, by projecting network weights and activation functions to lower precision thereby accelerate computation and reduce memory consumption. However, there is inevitable loss of accuracy in the low bit regime. One way to mitigate such an issue is through knowledge distillation (KD [10]). In this paper, we study a feature affinity assisted KD so that the student and teacher networks not only try to match their logits at the output level but also match feature maps in the intermediate stages. This is similar to teaching a student through intermediate steps of a solution instead of just showing the final answer (as in conventional KD [10]). Our method does not rely on ground truth labels while enhancing student network learning and closing the gaps between full and low precision models.

### _Weight Quantization of Neural Network_

Quantization-aware training (QAT) searches the optimal model weight in training. Given an objective \(L\), the classical QAT scheme ([6, 21]) is formulated as

\[\begin{cases}w^{t+1}=w^{t}-\nabla_{u}L(u^{t}),\\ u^{t+1}=\text{Quant}(w^{t+1}),\end{cases} \tag{1}\]

where Quant is projection to a low precision quantized space. Yin et al. [28] proposed BinaryRelax, a relaxation form of QAT, which replaces the second update of (1) by

\[\begin{split} u^{t+1}=\frac{w^{t+1}+\lambda^{t+1}\text{Quant}(w^{ t+1})}{1+\lambda^{t+1}},\\ \lambda^{t+1}=\eta\lambda^{t}\quad\text{with }\eta>1.\end{split} \tag{2}\]

Darkhorn et al. [7] further improved (2) by designing a more sophisticated learnable growing scheme for \(\lambda^{t}\) and factoring a learnable parameter into Proj(.). Polino et al. [19] proposed quantized distillation (QD), a QAT framework that leverages knowledge distillation for quantization. Under QD, the quantized model receives supervision from both ground truth (GT) labels and a trained teacher in float precision (FP). The objective function has the generalized form (\(\alpha\in(0,1)\)):

\[\mathcal{L}_{QD}=\alpha\mathcal{L}_{KD}+(1-\alpha)\mathcal{L}_{GT} \tag{3}\]

where \(\mathcal{L}_{KD}\) is Kullback-Leibler divergence (KL) loss, and \(\mathcal{L}_{GT}\) is negative log likelihood (NLL) loss. In order to compare different methods fairly, we introduce two technical terms: end-to-end quantization and fine-tuning quantization. End-to-end quantization is to train a quantized model from scratch, and fine-tuning quantization is to train a quantized model from a pre-trained float precision (FP) model. With the same method, the latter usually lands a better result than the former. Li et al. [15] proposed a mixed quantization (a.k.a. BRECQ) that takes a pre-trained model and partially retrains the model on a small subset of data. We list the performance of some previous works of weight quantization, which will serve as baselines for this work.

\begin{table}
\begin{tabular}{c|c|c|c} Method & 1-bit & 2-bit & 4-bit \\ \hline \hline \multicolumn{4}{c}{Model: ResNet20} \\ \hline QAT ([6, 21]) & 87.07\% & 90.26\% & 91.47\% \\ \hline BinaryRelax [28] & 88.64\% & 90.47\% & 91.75\% \\ \hline QD [19] & 89.06\% & 90.86\% & 91.89\% \\ \hline DSQ [8] & 90.24\% & 91.06\% & 91.92\% \\ \hline BRECQ [15] & N/A & 81.31\% & 83.98\% \\ \end{tabular}
\end{table} TABLE I: Quantization accuracies of some existing quantization-aware training methods on CIFAR-10 dataset. All methods except BRECQ are end-to-end.
### _Activation Quantization_

In addition to weight quantization, the inference of neural networks can be further accelerated through activation quantization. Given a resolution \(\alpha>0\), a quantized ReLU activation function of bit-width \(b\in\mathbb{N}\) is formulated as:

\[\sigma(x,\alpha)=\begin{cases}0&x<0\\ k\alpha&(k-1)\alpha\leq x<k\alpha,\ \ 1\leq k\leq 2^{b}-1\\ (2^{b}-1)\alpha&x\geq(2^{b}-1)\alpha\end{cases} \tag{4}\]

where the resolution parameter \(\alpha\) is learned from data. A plot of \(2\)-bit quantized ReLU is shown in figure 1. However, such quantized activation function leads to vanished gradient during training, which makes the standard backpropagation inapplicable. Indeed, it is clear that \(\frac{\partial\sigma}{\partial x}=0\) almost everywhere. Bengio et al. [2] proposed to use a straight through estimator (STE) in backward pass to handle the zero gradient issue. The idea is to simply replace the vanished \(\frac{\partial\sigma}{\partial x}\) with a non-trivial derivative \(\frac{\partial\sigma}{\partial x}\) of a surrogate function \(\tilde{\sigma}(x,\alpha)\). Theoretical studies on STE and convergence vs. recurrence issues of training algorithms have been conducted in ([17, 27]). Among a variety of STE choices, a widely-used STE is the \(x\)-derivative of the so-called clipped ReLU [3]\(\tilde{\alpha}(x,\alpha)=\min\{\max\{x,0\},(2^{b}-1)\alpha\}\), namely,

\[\frac{\partial\tilde{\sigma}}{\partial x}=\begin{cases}1&0<x<(2^{b}-1)\alpha \\ 0&\text{else}.\end{cases}\]

In addition, a few proxies of \(\frac{\partial\sigma}{\partial\alpha}\) have been proposed ([4, 29]). In this work, we follow [29] and use the three-valued proxy:

\[\frac{\partial\sigma}{\partial\alpha}\approx\begin{cases}0&x\leq 0\\ 2^{b-1}&0<x<(2^{b}-1)\alpha\\ 2^{b}-1&x\geq(2^{b}-1)\alpha.\end{cases} \tag{5}\]

### _Knowledge Distillation_

Several works have proposed to impose closeness of the probabilistic distributions between the teacher and student networks, e.g. similarity between feature maps. A flow of solution procedure (FSP) matrix in [26] measures the information exchange between two layers of a given model. Then \(l_{2}\) loss regularizes the distance between FSP matrices of teacher and student in knowledge distillation. An attention transform (AT) loss [30] directly measures the distance of feature maps outputted by teacher and student, which enhances the learning of student from teacher. Similarly, feature affinity (FA) loss [24] measures the distance of two feature maps. In a dual learning framework for semantic segmentation [24], the FA loss is applied on the output feature maps of a segmentation decoder and a high-resolution decoder. In [25], FA loss on multi-resolution paths also improves light weight semantic segmentation. Given two feature maps with the same height and width (interpolate if different), \(\mathbf{F}^{S}\in\mathbb{R}^{C_{1}\times H\times W}\) and \(\mathbf{F}^{T}\in\mathbb{R}^{C_{2}\times H\times W}\), we first normalize the feature map along the channel dimension. Given a pixel of feature map \(F_{i}\in\mathbb{R}^{C}\), we construct an affinity matrix \(\mathbf{S}\in\mathbb{R}^{WH\times WH}\) as:

\[\mathbf{S}_{ij}=\|\mathbf{F}_{i}-\mathbf{F}_{j}\|_{\theta}:=\cos\theta_{ij}= \frac{\langle\mathbf{F}_{i},\mathbf{F}_{j}\rangle}{||\mathbf{F}_{i}||||\mathbf{ F}_{j}||}.\]

where \(\mathbf{\Theta}_{ij}\) measures the angle between \(F_{i}\) and \(F_{j}\). Hence, the FA loss measures the similarity of pairwise angular distance between pixels of two feature maps, which can be formulated as

\[L_{fa}(\mathbf{F}^{S},\mathbf{F}^{T})=\frac{1}{W^{2}H^{2}}||\mathbf{S}^{T}- \mathbf{S}^{S}||_{2}^{2}. \tag{6}\]

### _Contributions_

In this paper, our main contributions are:

1. We find that using mean squares error (MSE) gives better performance than KL on QAT, which is a significant improvement to QD ([19]).
2. We consistently improve the accuracies of various quantized student networks by imposing the FA loss on feature maps of each convolutional block. We also unveil the theoretical underpinning of feature affinity loss in terms of the celebrated Johnson-Lindenstrass lemma for low-dimensional embeddings.
3. We achieve state-of-art quantization accuracy on CIFAR-10 and CIFAR-100. Our FAQD framework _can train a quantized student network on unlabeled data_.
4. We propose a randomized Fast FA (FFA) loss to accelerate the computation of training loss, and prove its convergence and error bound.

## II Feature Affinity Assisted Distillation and Quantization

### _Feature Affinity Loss_

In quantization setting, it is unreasonable to require that \(F^{S}\) is close to \(F^{T}\), as they are typically in different spaces (\(F^{S}\in\mathcal{Q}\) in full quantization) and of different dimensions. However, \(F^{S}\) can be viewed as a compression of \(F^{T}\) in dimension, and preserving information under such compression has been studied in compressed sensing. Researchers ([20, 22]) have proposed to compress graph embedding to lower dimension so that graph convolution can be computed efficiently. In K-means cluttering problem, several methods ([1, 18]) have been designed to project the data into a low-dimensional space such that

\[||\text{Proj}(\mathbf{x})-\text{Proj}(\mathbf{y})||\approx||\mathbf{x}-\mathbf{ y}||,\ \ \forall\ (\mathbf{x},\mathbf{y}), \tag{7}\]
and so pairwise distances from data points to the centroids can be computed at a lower cost.

In view of the feature maps of student model as a compression of teacher's feature maps, we impose a similar property in terms of pairwise angular distance:

\[||\mathbf{F}_{i}^{S}-\mathbf{F}_{j}^{S}||_{\theta}\approx||\mathbf{F}_{i}^{T}-\mathbf{F}_{j}^{T} ||_{\theta},\ \forall\left(i,j\right)\]

which is realized by minimizing the feature affinity loss. On the other hand, a Johnson-Lindenstrauss (JL [11]) like lemma can guarantee that we have student's feature affinity matrix close to the teacher's, provided that the number of channels of student network is not too small. In contrast, the classical JL lemma states that a set of points in a high-dimensional space can be embedded into a space of much lower dimension in such a way that the _Euclidean_ distances between the points are nearly preserved. To tailor it to our application, we prove the following JL-like lemma in the angular distance case:

**Theorem II.1** (Johnson-Lindenstrauss lemma, Angular Case): _Given any \(\epsilon\in(0,1)\), an embedding matrix \(\mathbf{F}\in\mathbb{R}^{n\times d}\), for \(k\in(16\epsilon^{-2}\ln n,d)\), there exists a linear map \(T(\mathbf{F})\in\mathbb{R}^{n\times k}\) so that_

\[\begin{split}(1-\epsilon)||\mathbf{F}_{i}-\mathbf{F}_{j}||_{ \theta}\leq||T(\mathbf{F})_{i}-T(\mathbf{F})_{j}||_{\theta}\\ \leq(1+\epsilon)||\mathbf{F}_{i}-\mathbf{F}_{j}||_{\theta},\ \ \forall\ \ 1\leq i,j\leq n\end{split} \tag{8}\]

_where \(||\mathbf{F}_{i}-\mathbf{F}_{j}||_{\theta}=\frac{\langle\mathbf{F}_{i},\mathbf{F}_{j}\rangle}{ ||\mathbf{F}_{i}||\mathbf{F}_{j}||}\) is the angular distance._

It is thus possible to reduce the embedding dimension down from \(d\) to \(k\), while roughly preserving the pairwise angular distances between the points. In a convolutional neural network, we can view intermediate feature maps as \(\mathbf{F}^{S}\in\mathbb{R}^{HW\times C_{1}}\) and \(\mathbf{F}^{T}\in\mathbb{R}^{HW\times C_{2}}\), and feature affinity loss will help the student learn a compressed feature embedding. The FA loss can be flexibly placed between teacher and student in different positions (encoder/decoder, residual block, etc.) for different models. In standard implementation of ResNet, residual blocks with the same number of output channels are grouped into a sequential layer. We apply FA loss to the features of such layers.

\[\mathcal{L}_{FA}=\sum_{l=1}^{L}L_{fa}(\mathbf{F}_{l}^{T},\mathbf{F}_{l}^{S})\]

where \(\mathbf{F}_{l}^{T}\) and \(\mathbf{F}_{l}^{S}\) are the feature maps of teacher and student respectively. For example, the residual network family of ResNet20, ResNet56, ResNet110, and ResNet164 have \(L=3\), whereas the family of ResNet18, ResNet34, and ResNet50 have \(L=4\).

### _Choice of Loss Functions_

In this work, we propose two sets of loss function choices for the end-to-end quantization and pretrained quantization, where end-to-end quantization refers to having an untrained student model with randomly initialized weights. We investigate both scenarios of quantization and propose two different strategies for each.

The Kullback-Leibler divergence (KL) is a metric of the similarity between two probabilistic distributions. Given a ground-truth distribution \(P\), it computes the relative entropy of a given distribution \(Q\) from \(P\):

\[\mathcal{L}_{KL}(P||Q)=\sum_{x\in\mathcal{X}}P(x)\ln\frac{P(x)}{Q(x)}. \tag{9}\]

While KD is usually coupled with KL loss ([19, 10]), it is not unconventional to choose other loss functions. Kim et al. [14] showed that MSE, in certain cases, can outperform KL in the classic teacher-student knowledge distillation setting. KL loss is also widely used for trade-off between accuracy and robustness under adversarial attacks, which can be considered as self-knowledge distillation. Given a classifier \(f\), an original data point \(\mathbf{x}\) and its adversarial example \(\mathbf{x}^{\prime}\), TRADES [31] is formulated as

\[L_{TRADES}=\mathcal{L}_{CE}(f(\mathbf{x}),\mathbf{y})+\mathcal{L}_{\mathbf{KL}} (\mathbf{f}(\mathbf{x})||\mathbf{f}(\mathbf{x}^{\prime}))\]

Li et al. [16] showed that \(L_{CE}\big{(}f(\mathbf{x}^{\prime}),y\big{)}\) outperforms \(\mathcal{L}_{KL}(f(\mathbf{x})||\mathbf{f}(\mathbf{x}^{\prime}))\) both experimentally and theoretically.

Inspired by the studies above, we conduct experiments on different choices of the loss function. We compare KD on quatization from scratch (end-to-end). As shown in table II, MSE outperforms KL in quantization.

On the other hand, we find that KL loss works better for fine-tuning quantization. One possible explanation is that when training from scratch, the term \(\ln\frac{P(x)}{Q(x)}\) is large. However, the derivative of logarithm is small at large values, which makes it converge slower and potentially worse. On the other hand, when \(\frac{P(x)}{Q(x)}\) is close to 1, the logarithm has sharp slope and converges fast.

### _Feature Affinity Assisted Distillation and Quantization_

Inspired by previous studies ([13, 15, 19]), we propose a feature affinity assisted quantized distillation (FAQD). The end-to-end quantization objective function is formulated as:

\[\begin{split}\mathcal{L}=\alpha\,\mathcal{L}_{KD}+\beta\, \mathcal{L}_{FA}+\gamma\,\mathcal{L}_{GT}\\ =\alpha\,\mathcal{L}_{MSE}\big{(}f^{T}(\mathbf{x}),f^{S}(\mathbf{ x})\big{)}&+\beta\,\sum_{l=1}^{L}\mathcal{L}_{fa}(\mathbf{F}_{l}^{T}, \mathbf{F}_{l}^{S})\\ &+\gamma\,\mathcal{L}_{NLL}(f^{S}(\mathbf{x}),y).\end{split} \tag{10}\]

In fine-tuning quantization, we replace MSE loss in (10) by KL divergence loss. In FAQD, the student model learns not only the final logits of the teacher but also the intermediate extracted feature maps of the teacher using feature affinity norm computed as in [24].

\begin{table}
\begin{tabular}{c|c|c|c|c} Student & Teacher & 1-bit & 2-bit & 4-bit \\ \hline \hline \multicolumn{4}{c}{\(\mathcal{L}_{KD}\)} & \multicolumn{1}{c|}{KL loss in (3)} \\ \hline \hline ResNet20 & ResNet110 & 89.06\% & 90.86\% & 91.89\% \\ \hline \multicolumn{4}{c}{\(\mathcal{L}_{KD}\)} & \multicolumn{1}{c|}{MSE in (3)} \\ \hline ResNet20 & ResNet110 & **90.00\%** & **91.01\%** & **92.05\%** \\ \hline \end{tabular}
\end{table} TABLE II: Comparision of KL loss and MSE loss on CIFAR-10 data set. All teachers are pre-trained FP models, and all students are initial models (end-to-end quantization).
In addition to (10), we also propose a label-free objective which does not require the knowledge of labels:

\[\mathcal{L}_{\text{label-free}}=\alpha\mathcal{L}_{MSE}\big{(}f^{T}(\mathbf{x}),f^{S}(\mathbf{x})\big{)}+\beta\sum_{l=1}^{L}\mathcal{L}_{fa}(\mathbf{F}_{l}^{T },\mathbf{F}_{l}^{S}). \tag{11}\]

Despite the pre-trained computer vision models being available from cloud service such as AWS and image/video data abundantly collected, the data labeling is still expensive and time consuming. Therefore, a label-free quantization framework has significant value in the real world. In this work, we verify that the FA loss can significantly improve KD performance. The label-free loss in Eq. (11) can outperform the baseline methods in table I as well as the prior supervised QD in (3).

## III Experimental Results

### _Weight Quantization_

In this section we test FAQD on the dataset CIFAR-10. First, we experiment on fine-tuning quantization. The float precision (FP) ResNet110 teaches ResNet20 and ResNet56. The teacher has 93.91% accuracy, and the two pre-trained models have accuracy 92.11% and 93.31% respectively. While both SGD and Adam optimization work well on the problem, we found KL loss with Adam slightly outperform SGD in this scenario. The objective is

\[\mathcal{L}=\mathcal{L}_{KL}+\mathcal{L}_{FA}\]

for the label-free quantization. When calibrating the ground-truth label, the cross-entropy loss \(\mathcal{L}_{NLL}\) is used as the supervision criterion.

For end-to-end quantization, we found that MSE loss performs better than KL loss. Adam optimization struggles to reach acceptable performance on end-to-end quantization (with either KL or MSE loss). We further test the performance of FAQD on larger dataset CIFAR-100 where an FP ResNet 164 teaches a quantized ResNet110. We report the accuracies for both label-free and label-present supervision. We evaluate FAQD on both fine-tuning quantization and end-to-end quantization.

In the ResNet experiment, the teacher ResNet164 has 74.50% testing accuracy. For the pretrained FAQD, the FP student ResNet110 has 72.96% accuracy. As shown in table VI

Fig. 2: FAQD framework. The intermediate feature maps are supervised by FA loss, and the raw logits by MSE loss.

\begin{table}
\begin{tabular}{c|c|c|c} \hline \hline \multicolumn{4}{c}{Teacher ResNet110: 93.91\%} \\ \hline \hline Method & 1-bit & 2-bit & 4-bit \\ \hline Pre-trained FP student & ResNet20: 92.21\% & \\ \hline Label-free FAQD & 89.97\% & 91.40\% & 92.55\% \\ \hline FAQD with Supervision & 90.92\% & 91.93\% & 92.74\% \\ \hline Pre-trained FP student & ResNet56: 93.31\% & \\ \hline Label-free FAQD & 92.34\% & 92.91\% & 93.52\% \\ \hline FAQD with Supervision & 92.83\% & 93.14\% & 93.77\% \\ \hline \end{tabular}
\end{table} TABLE III: Fine-tuning knowledge distillation for quantization of all convolutional layers.
and table V, FAQD has surprisingly superior performance on CIFAR-100. The binarized student almost reaches the accuracy of FP model, and the 4-bit model surpasses the FP teacher.

### _Full Quantization_

In this section, we extend our results to full quantization where the activation function is also quantized. As shown in table VII, the 4W4A fune-tuning quantization has accuracy similar to float ResNet20. Meanwhile, we fill in the long existing performance gap [9] when reducing precision from 1W2A to 1W1A on CIFAR-10 dataset, as the accuracy drop is linear (with respect to activation precision) and small.

## IV Fast Feature Affinity Loss

### _Proposed Method_

Despite the significant increase of KD performance, we note that introducing FA loss will increase the training time. If we normalize the feature maps by row beforehand, computing FA loss between multiple intermediate feature maps can be expensive.

\[\mathcal{L}_{fa}(F_{1},F_{2})=\|F_{1}F_{1}^{T}-F_{2}F_{2}^{T}\|_{2}^{2}. \tag{12}\]

As we freeze the pre-trained teacher, feature map of the teacher model \(F_{1}=f^{T}(\mathbf{x})\) is a constant, in contrast to student feature map \(F_{1}=f^{S}(\Theta,\mathbf{x})\). Denote \(\mathbf{S}_{1}=F_{1}F_{1}^{T}\in\mathbb{R}^{WH\times WH}\) and \(g(\Theta,\mathbf{x})=f^{S}(\Theta,\mathbf{x})[f^{S}(\Theta,\mathbf{x})]^{T}\). The feature affinity can be formulated as

\[\mathcal{L}_{fa}(\Theta)=\frac{1}{|\mathcal{X}|}\sum_{\mathbf{x}\in\mathcal{X}}\| \mathbf{S}_{1}-g(\Theta,\mathbf{x})\|_{2}^{2}. \tag{13}\]

Computing \(\mathbf{S}_{1}\) and \(g(\Theta,X)\) requires \(\mathcal{O}(W^{2}H^{2}C)\) complexity each (C is the number of channels), which is quite expensive. We introduce a random estimator of \(L_{ffa}(\Theta)\):

\[\mathcal{L}_{fa}(F_{1},F_{2},\mathbf{z})=\frac{1}{|\mathcal{X}|}\sum_{\mathbf{x} \in\mathcal{X}}\|(\mathbf{S}_{1}-g(\Theta,\mathbf{x}))\mathbf{z}\|_{2}^{2}, \tag{14}\]

where \(\mathbf{z}\in\mathbb{R}^{HW}\) is a vector with i.i.d unit normal components \(N(0,1)\). We show below that Eq. (14) is an unbiased estimator of FA loss (13).

**Proposition 1**: \[\mathbb{E}_{\mathbf{x}\sim N(0,1)}[\mathcal{L}_{fa}(F_{1},F_{2},\mathbf{z})]= \mathcal{L}_{fa}(\Theta).\]

This estimator can achieve computing complexity \(\mathcal{O}(HWC)\) by performing two matrix-vector multiplication \(F_{1}\big{(}F_{1}^{T}\mathbf{z}\big{)}\). We define the Fast Feature Affinity (FFA) loss to be the \(k\) ensemble of (14):

\[\mathcal{L}_{ffa,k}(\Theta)=\frac{1}{|\mathcal{X}|}\sum_{\mathbf{x}\in\mathcal{X} }\frac{1}{k}\|(\mathbf{S}_{1}-g(\Theta,\mathbf{x}))Z_{k}\|_{2}^{2} \tag{15}\]

where \(Z_{k}\in\mathbb{R}^{HW\times k}\) with i.i.d \(\mathcal{N}(0,1)\) components, and we have \(k\ll WH\). The computational complexity of \(\mathcal{L}_{ffa,k}(\Theta)\) is \(\mathcal{O}(kWHC)\).

Finally, we remark that FFA loss can accelerate computation of pairwise Euclidean distance in dimensional reduction such as in (7). The popular way to compute the pairwise distance of rows for a matrix \(A\in\mathbb{R}^{n\times c}\) is to broadcast the vector of row norms and compute \(AA^{T}\). Given the row norm vector \(v=(\|A_{1}\|^{2},\cdots,\|A_{n}\|^{2})\), the similarity matrix \((\mathbf{S}_{ij})\), \(\mathbf{S}_{ij}=\|A_{i}-A_{j}\|^{2}\), is computed as

\[\mathbf{S}=\mathbf{1}\otimes v-2AA^{T}+v\otimes\mathbf{1}.\]

The term \(2AA^{T}\) can be efficiently approximated by FFA loss.

\begin{table}
\begin{tabular}{
### _Experimental Results_

We test Fast FA loss on CIFAR-15. As mentioned in the previous section, ResNet-20 has 3 residual blocks. The corresponding width and height for feature maps are 32, 16, and 8, \(H=W\) for all groups, so the dimension (\(HW\)) of similarity matrices are 1024, 256, and 64. We test the fast FA loss with dimensions 1, 6, and 10. The results are shown in table VIII. Meanwhile, the FFA has much training time for each step. When \(k=1\), the accuracies are inconsistent due to large variance. With too few samples in the estimator, the fast FA norm is too noisy and jeopardize the distillation. When \(k=6\), the fast FA loss stabilizes and shows a significant improvement from the baseline, \(\mathcal{L}=\mathcal{L}_{MSE}+\mathcal{L}_{CE}\) as in table II. When \(k\) increases to \(10\), the performance fast FA loss is comparable with the exact FA loss (table IV). Moreover, we experiment with the time consumption for computing FA loss and FFA loss. We plot the time in log scale vs. \(H\), (\(H=W\)) for feature maps. The theoretical time complexity for computing exact FA loss is \(\mathcal{O}(H^{4})\) and that for computing FFA loss is \(\mathcal{O}(H^{2})\). We see that figure 4(a) agrees with the theoretical estimate.

The larger the \(H\), the more advantage the FFA loss. For (medical) images with resolutions in the thousands, the FFA loss will have significant computational savings.

### _Theoretical Analysis of FFA Loss_

As shown in proposition 4.1, the FFA loss is a \(k\) ensemble unbiased estimator of FA loss. By the strong law of large numbers, the FFA loss converges to the exact FA loss with probability 1.

**Theorem IV.1**: _For given \(\Theta\), suppose that \(|\mathcal{L}_{fa}(\Theta)|\leq\infty\), then_

\[\forall\epsilon>0,\exists N\ \ s.t.\ \ \forall k>N,\ \ |\mathcal{L}_{ffa,k}( \Theta)-\mathcal{L}_{fa}(\Theta)|<\epsilon.\]

_Namely, the FFA loss converges to FA loss pointwise:_

\[\forall\Theta,\ \lim_{k\rightarrow\infty}\mathcal{L}_{ffa,k}(\Theta)=\mathcal{ L}_{fa}(\Theta).\]

_We also establish the following error bound for finite \(k\)._

**Proposition 2**: \[\mathbb{P}\big{(}|\mathcal{L}_{ffa,k}(\Theta)-\mathcal{L}_{fa}(\Theta)|> \epsilon\big{)}\leq\frac{C}{\epsilon^{2}k},\]

_where \(C\leq 3\,\|\mathcal{L}_{fa}(\Theta)\|_{2}^{4}\)._

Proposition IV.2 says that the probability that the FFA estimation has an error beyond a target value decays like \(\mathcal{O}(\frac{1}{k})\). The analysis guarantees the accuracy of FFA loss as an efficient estimator of FA loss. Another question one might ask is whether minimizing the FFA loss is equivalent to minimizing the FA loss. Denote \(\Theta^{*}=\arg\min L_{fa}(\Theta)\) and \(\Theta^{*}_{k}=\arg\min L_{ffa,k}(\Theta)\), and assume the minimum is unique for each function. In order to substitute FA loss by FFA loss, one would hope that \(\Theta^{*}_{k}\) converges to \(\Theta^{*}\). Unfortunately, the point-wise convergence in Theorem IV.1 is not sufficient to guarantee the convergence of the optimal points, as a counter-example can be easily constructed. In the rest of this section, we show that such convergence can be established under an additional assumption.

**Theorem IV.2** (Convergence in the general case): _Suppose that \(\mathcal{L}_{ffa,k}(\Theta)\) converges to \(\mathcal{L}_{fa}(\Theta)\) uniformly, that is_

\[\forall\epsilon>0,\exists N\ \ s.t.\ \ \forall k>N,\ \ |\mathcal{L}_{ffa,k}( \Theta)-\mathcal{L}_{fa}(\Theta)|<\epsilon\]

\(\forall\Theta\) _and_

\[|\mathcal{L}_{fa}(\Theta)|\leq\infty,\ \forall\Theta.\]

_Then_

\[\lim_{k\rightarrow\infty}||\Theta^{*}_{k}-\Theta^{*}||^{2}=0. \tag{16}\]

_The uniform convergence can be relaxed if \(\mathcal{L}_{fa}\) is convex in \(\Theta\). We would like to present a consequence of Theorem IV.2._

**Corollary IV.2.1** (Convergence in the convex case): _Suppose that \(L_{fa}:\mathbb{R}^{n}\rightarrow\mathbb{R}\) is convex and \(L\)-smooth, and that there \(\exists M\) such that \(||\Theta^{*}_{k}||\leq M,\,\forall k\). Then, \(\forall k,\ \ \mathcal{L}_{ffa,k}\) is also convex, and \(\lim_{k\rightarrow\infty}||\Theta^{*}_{k}-\Theta^{*}||^{2}=0\)._

\begin{table}
\begin{tabular}{c|c|c|c|c} Model & k & 2W32A & 4W32A & 4W4A \\ \hline \hline \multicolumn{5}{c}{Fast FA Loss Accuracy} \\ \hline ResNet20 & 1 & 90.83\(\pm\)2.75\% & 91.83\(\pm\)3.01\% & 90.32\(\pm\) 2.35\% \\ ResNet20 & 5 & 91.37\% & 92.14\% & 91.12\% \\ ResNet20 & 15 & 91.45\% & 92.39\% & 91.53\% \\ \hline \multicolumn{5}{c}{Fast FA Loss Computing Time Per Step} \\ \hline ResNet20 & 1 & 44 ms & 65 ms & 187 ms \\ ResNet20 & 5 & 46 ms & 66 ms & 188 ms \\ ResNet20 & 15 & 48 ms & 68 ms & 190 ms \\ \hline \multicolumn{5}{c}{Exact FA Loss Computing Time Per Step} \\ \hline ResNet20 & N/A & 60 ms & 80 ms & 205 ms \\ \hline \end{tabular}
\end{table} TABLE VIII: Test FFA loss for different \(k\) on CIFAR-10 (end-to-end quantization). As \(k\) increases, the performance is approaching the exact FA loss. The \(4\)-bit projection is more time consuming than \(2\)-bit. STE and \(\alpha\) update in full-quantization add extra time in milliseconds (ms).

Fig. 4: Plots for inference time of FA loss and FFA loss with \(k=1\).
## V Conclusion

We presented FAQD, a feature assisted (FA) knowledge distillation method for training-aware quantization. It couples MSE loss with FA loss and significantly improves the accuracy of the quantized student. FAQD works for both weight only and full quantization, and outperforms baseline Resnets on CIFAR-10 and CIFAR-100. We also analyzed an efficient randomized approximation (FFA) to the FA loss for feature maps with large dimensions. This theoretically founded FFA loss benefits training models on high resolution images.

## VI Appendix

**Proof of Theorem 2.1:** It suffices to prove that for any set of \(n\) unit vectors in \(\mathbb{R}^{d}\), there is a linear map nearly preserving pairwise angular distances, because the angular distance is scale-invariant.

Let \(T\) be a linear transformation induced by a random Gaussian matrix \(\frac{1}{\sqrt{k}}A\in\mathbb{R}^{k\times d}\) such that \(T(\mathbf{F})=\mathbf{F}A^{T}\). Define the events \(\mathcal{A}^{-}_{ij}=\{T:(1-\epsilon)\|\mathbf{F}_{i}-\mathbf{F}_{j}\|^{2}\leq\|T(\mathbf{F })_{i}-T(\mathbf{F})_{j}\|^{2}\leq(1+\epsilon)\|\mathbf{F}_{i}-\mathbf{F}_{j}\|^{2}\) fails\(\}\) and \(\mathcal{A}^{-}_{ij}=\{T:(1-\epsilon)\|\mathbf{F}_{i}+\mathbf{F}_{j}\|^{2}\leq\|T(\mathbf{F}) _{i}+T(\mathbf{F})_{j}\|^{2}\leq(1+\epsilon)\|\mathbf{F}_{i}+\mathbf{F}_{j}\|^{2}\) fails\(\}\).

Following the proof of the classical JL lemma in the Euclidean case [23], we have:

\[P(\mathcal{A}^{-}_{ij})\leq 2e^{-\frac{(\epsilon^{2}-\epsilon^{3})k}{4}},\ \ P( \mathcal{A}^{+}_{ij})\leq 2e^{-\frac{(\epsilon^{2}-\epsilon^{3})k}{4}}. \tag{17}\]

Let \(\mathcal{B}_{ij}=\{T:|\mathbf{F}_{i}\cdot\mathbf{F}_{j}-T(\mathbf{F})_{i}\cdot T(\mathbf{F})_{ j}|>\epsilon\}\), where \(\cdot\) is the shorthand for inner product. We show that \(\mathcal{B}_{ij}\subset A^{-}_{ij}\cup\mathcal{A}^{+}_{ij}\) for \(\|F_{i}\|=\|F_{j}\|=1\) by showing \(\mathcal{A}^{-C}_{ij}\cap\mathcal{A}^{+C}_{ij}\subset\mathcal{B}^{C}_{ij}\).

If \(\mathcal{A}^{-C}_{ij}\cap\mathcal{A}^{+C}_{ij}\) holds, we have

\[4T(\mathbf{F})_{i}\cdot T(\mathbf{F})_{j}\] \[= \|T(\mathbf{F})_{i}+T(\mathbf{F})_{j}\|^{2}-\|T(\mathbf{F})_{i}-T(\mathbf{F})_{j} \|^{2}\] \[\leq (1+\epsilon)\|\mathbf{F}_{i}+\mathbf{F}_{j}\|^{2}-(1-\epsilon)\|\mathbf{F}_ {i}-\mathbf{F}_{j}\|^{2}\] \[= 4\mathbf{F}_{i}\cdot\mathbf{F}_{j}+2\epsilon(\|F_{i}\|_{2}^{2}+\|F_{j} \|^{2})\] \[= 4\mathbf{F}_{i}\cdot\mathbf{F}_{j}+4\epsilon.\]

Therefore, \(\mathbf{F}_{i}\cdot\mathbf{F}_{j}-T(\mathbf{F})_{i}\cdot T(\mathbf{F})_{j}\geq-\epsilon\). By a similar argument, we have \(\mathbf{F}_{i}\cdot\mathbf{F}_{j}-T(\mathbf{F})_{i}\cdot T(\mathbf{F})_{j}\leq\epsilon\). Then we have \(\mathcal{A}^{-C}_{ij}\cap\mathcal{A}^{+C}_{ij}\subset\mathcal{B}^{C}_{ij}\), and thus

\[\mathbb{P}(\mathcal{B}_{ij})\leq\mathbb{P}(\mathcal{A}^{-}_{ij}\cup A^{+}_{ij} )\leq 4\exp\{-\frac{(\epsilon^{2}-\epsilon^{3})k}{4}\}\]

and

\[\mathbb{P}(\cup_{i<j}\mathcal{B}_{ij})\leq\sum_{i<j}\mathbb{P}(\mathcal{B}_{ ij})\leq 4n^{2}\exp\{-\frac{(\epsilon^{4}-\epsilon^{3})k}{4}\}.\]

This probability is less than \(1\) if we take \(k>\frac{16\ln n}{\epsilon^{2}}\). Therefore, there must exist a \(T\) such that \(\cap_{i<j}\mathcal{B}^{C}_{ij}\) holds, which completes the proof.

**Proof of Proposition 4.1:** Letting \(N=WH\), \(a_{ij}=(F_{1}F_{1}^{T})_{ij}\), and \(b_{ij}=(F_{2}F_{2}^{T})_{ij}\) in equation (14), we have:

\[\mathbb{E}_{z}\mathcal{L}_{ffa}(F_{1},F_{2};2)=\mathbb{E}_{z}\sum _{i=1}^{N}(\sum_{j=1}^{N}|a_{ij}-b_{ij}|z_{j})^{2}\] \[=\mathbb{E}_{z}\sum_{i=1}^{N}(\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{2}z_ {j}^{2}+2\sum_{j\neq k}|a_{ij}-b_{ij}||a_{ik}-b_{ik}|z_{j}z_{k})\] \[=\mathbb{E}_{z}\sum_{i=1}^{N}\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{2}z_ {j}^{2}+2\sum_{i=1}^{N}\sum_{j\neq k}|a_{ij}-b_{ij}||a_{ik}-b_{ik}|z_{j}z_{k}\] \[=\sum_{i=1}^{N}\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{2}\mathbb{E}_{z}z_ {j}^{2}\] \[+2\sum_{i=1}^{N}\sum_{j\neq k}|a_{ij}-b_{ij}|a_{ik}-b_{ik}| \mathbb{E}_{z}z_{j}z_{k}\] \[=\sum_{i=1}^{N}\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{2}=\mathcal{L}_{fa} (F_{1},F_{2};2).\]

**Proof of Theorem 4.1:** Given a Gaussian matrix \(Z_{k}=[\mathbf{z}_{1},\cdots,\mathbf{z}_{k}]\in\mathbb{R}^{n\times k}\),

\[\mathcal{L}_{ffa,k}(\Theta)=\frac{1}{k}\sum_{l=1}^{k}\mathcal{L}_{ffa}(F_{1}, F_{2},\mathbf{z}_{l}).\]

For any fixed \(\Theta\), \(\mathcal{L}_{ffa}(F_{1},F_{2},\mathbf{z}_{l})\), \(l=1,\cdots,k\), are i.i.d random variables. Suppose the first moment of each random variable is finite, by the strong law of large numbers, \(\mathcal{L}_{ffa,k}(\Theta)\) converges to \(\mathbb{E}[\mathcal{L}_{ffa}(F_{1},F_{2},\mathbf{z}_{1})]\) almost surely. In other words, \(\lim_{k\rightarrow\infty}\mathcal{L}_{ffa,k}(\Theta)=\mathcal{L}_{fa}(\Theta)\) with probability 1.

**Proof of Proposition 4.2:** By Chebyshev's inequality, we have

\[\mathbb{P}\big{(}\big{|}\mathcal{L}_{ffa,k}(\Theta)-\mathbb{E}[ \mathcal{L}_{ffa,k}(\Theta)]\big{|}>\epsilon\big{)}\leq\] \[\frac{\text{Var}(\mathcal{L}_{ffa,k}(\Theta))}{\epsilon^{2}}=\frac{ \text{Var}(\mathcal{L}_{ffa}(F_{1},F_{2},\mathbf{z}_{1}))}{\epsilon^{2}k}. \tag{18}\]

In order to estimate

\[\text{Var}(\mathcal{L}_{ffa}(F_{1},F_{2},\mathbf{z}_{1})=\] \[\mathbb{E}[\mathcal{L}_{ffa}^{2}(F_{1},F_{2},\mathbf{z}_{1})]- \big{(}\mathbb{E}[\mathcal{L}_{ffa}(F_{1},F_{2},\mathbf{z}_{1})]\big{)}^{2}, \tag{19}\]

it suffices to estimate

\[\mathbb{E}[\mathcal{L}_{ffa}^{2}(F_{1},F_{2},\mathbf{z}_{1})]=\] \[\mathbb{E}_{z}\big{(}\sum_{i=1}^{N}\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{ 2}z_{j}^{2}+\sum_{i=1}^{N}\sum_{j\neq k}|a_{ij}-b_{ij}||a_{ik}-b_{ik}|z_{j}z_ {k}\big{)}^{2}\]

which equals (as cross terms are zero):

\[=\mathbb{E}_{z}\big{(}\sum_{i=1}^{N}\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{2}z_{j}^{2} \big{)}^{2}\]
Direct computation yields:

\[\sum_{i=1}^{N}\sum_{j=1}^{N}|a_{ij}-b_{ij}|^{4}z_{j}^{4}+\sum_{i=1}^{ N}\sum_{j=1}^{N}\sum_{l\neq i}^{N}|a_{ij}-b_{ij}|^{2}|a_{lj}-b_{lj}|^{2}z_{j}^{4}\] \[+2\sum_{i=1}^{N}\sum_{j=1}^{N}\sum_{l\neq j}^{N}|a_{ij}-b_{ij}|^{ 2}|a_{il}-b_{il}|^{2}z_{j}^{2}z_{l}^{2}\] \[+\sum_{i=1}^{N}\sum_{j=1}^{N}\sum_{k=1}^{N}\sum_{l\neq j}^{N}|a_{ ij}-b_{ij}|^{2}|a_{kl}-b_{kl}|^{2}z_{j}^{2}z_{l}^{2}\]

Notice that \(\mathbb{E}[z_{i}^{4}]=3\). Taking \(\mathbb{E}[\cdot]\), we derive the upper bound \(3\|\mathcal{L}_{fa}\|_{2}^{4}\).

**Proof of Theorem 4.2:** Since \(\lim\limits_{k\to\infty}\mathcal{L}_{ffa,k}(\Theta^{*})=\mathcal{L}_{fa}( \Theta^{*})\), it suffices to show that

\[\lim\limits_{k\to\infty}\ \inf\limits_{\Theta}\mathcal{L}_{ffa,k}(\Theta)= \mathcal{L}_{fa}(\Theta^{*}).\]

Note that

\[\forall\Theta,\lim\limits_{k\to\infty}\mathcal{L}_{ffa,k}(\Theta)=\mathcal{L}_ {fa}(\Theta)\leq\mathcal{L}_{fa}(\Theta^{*}).\]

Then,

\[\mathcal{L}_{fa}(\Theta^{*})\geq\lim\limits_{k\to\infty}\inf\limits_{\Theta}L _{ffa,k}(\Theta).\]

On the other hand, for arbitrary \(\epsilon>0\), we have:

\[\exists N\ \ s.t.\ \ \forall k>N\ \ |\mathcal{L}_{ffa,k}(\Theta)-\mathcal{L}_{ fa}(\Theta)|<\frac{\epsilon}{2},\,\forall\Theta\]

and there exists a sequence \(\{\Theta_{k}\}\) s.t.

\[\mathcal{L}_{ffa,k}(\Theta_{k})<\inf\limits_{\Theta}\mathcal{L}_{ffa,k}( \Theta)+\frac{\epsilon}{2}.\]

Note that \(|\mathcal{L}_{ffa,k}(\Theta_{k})-\mathcal{L}_{fa}(\Theta_{k})|<\frac{\epsilon }{2}\) for \(k>N\), so:

\[\mathcal{L}_{fa}(\Theta^{*})-\epsilon\leq\mathcal{L}_{fa}(\Theta_{k})- \epsilon<\inf\limits_{\Theta}\mathcal{L}_{ffa,k}(\Theta),\ \ \forall k>N.\]

Since \(\epsilon\) is arbitrary, taking \(k\to\infty\), we have

\[\mathcal{L}_{fa}(\Theta^{*})\leq\lim\limits_{k\to\infty}\inf\limits_{\Theta} \mathcal{L}_{ffa,k}(\Theta).\]

**Proof of Corollary 4.2.1:** For readability, we shorthand: \(\mathcal{L}_{ffa,k}=f_{k}\) and \(\mathcal{L}_{fa}=f\). Let

\[\mathbf{H}=\frac{\nabla^{2}f}{\nabla\Theta\nabla\Theta^{T}}\succcurlyeq\mathbf{ 0}\in\mathcal{R}^{n\times n}\]

be the Hessian matrix of FA loss, which is positive semi-definite by convexity of \(L_{fa}\). Then,

\[\frac{\nabla^{2}f_{k}}{\nabla\Theta\nabla\Theta^{T}}=Z_{k}^{T}\mathbf{H}Z_{k} \succcurlyeq\mathbf{0}\in\mathbb{R}^{k\times k}\]

which implies the convexity of \(f_{k}\) for all \(k\). Moreover, it is clear that \(f_{k}\) is smooth for all \(k\) since

\[\|\nabla f_{k}(\mathbf{x})-\nabla f_{k}(\mathbf{y})\|=\|Z_{k}(\nabla f (\mathbf{x})-\nabla f(\mathbf{y}))\|\\ \leq L\cdot\|Z_{k}\|\cdot\|\mathbf{x}-\mathbf{y}\|. \tag{20}\]

We note that \(f_{k}\) is also smooth. Although we cannot claim equi-smoothness since we cannot bound \(\|Z_{k}\|\) uniformly in \(k\), the above is sufficient for us to prove the desired result.

For \(\forall k\), given any initial parameters \(\Theta^{0}\), by smoothness and convexity of \(f_{k}\), it is well-known that

\[\|\Theta_{k}^{t}-\Theta_{k}^{*}\|\leq\|\Theta^{0}-\Theta_{k}^{*}\|\]

where \(\Theta_{k}^{t}\) is the parameter we arrive after \(t\) steps of gradient descent. Hence, we can pick a compact set \(\mathbf{K}=\overline{B_{R}(\Theta^{*})}\) for \(R\) large enough such that \(\{\Theta_{k}\}_{k=1}^{\infty}\subset\mathbf{K}\) (denote \(\Theta_{\infty}^{*}=\Theta^{*}\)). Now, it's suffices to prove \(f_{k}\) converges to \(f\) uniformly on \(K\). In fact, \(f_{k}\) converges to \(f\) on any compact set. To begin with, we state a known result from functional analysis ([5, 12]):

**Lemma 6.1**: _(Uniform boundness and equi-Lipschitz) Let \(\mathcal{F}\) be a family of convex function on \(\mathbb{R}^{n}\) and \(K\subset\mathbb{R}^{n}\) be a compact subset. Then, \(\mathcal{F}\) is equi-bounded and equi-Lipschitz on \(K\)._

This result is established in any Banach space in [12], so it automatically holds in finite dimensional Euclidean space. By Lemma 6.1, we have that the sequence \(\{f_{k}\}_{k=1}^{\infty}\), where \(f_{\infty}=f\), is equi-Lipschitz. \(\forall>0\), \(\exists\,\delta>0\) s.t. \(|f_{k}(x)-f_{k}(y)|<\epsilon\) for all \(k\) and \(x,y\in K\) when \(|x-y|<\delta\). Since \(\{B(x,\delta)\}_{x\in K}\) forms an open cover for \(K\), we have a finite sub-cover \(\{B(x_{j},\delta)\}_{j=1}^{m}\) of \(K\). Since there are finitely many points \(x_{j}\), there exists \(N_{\epsilon}\) such that

\[\forall k>N_{\epsilon},\ \ |f_{k}(x_{j})-f(x_{j})|<\epsilon,\text{ for }j=1, \cdots,m.\]

For any \(x\in K\), \(x\in B(x_{j^{*}},\delta)\) for some \(j^{*}\). For all \(k>N_{\epsilon}\), we have

\[|f_{k}(x)-f(x)|\leq\\ |f_{k}(x)-f_{k}(x_{j^{*}})|+|f_{k}(x_{j^{*}})-f(x_{j^{*}})|+|f(x _{j^{*}})-f(x)|\\ \leq(2\tilde{L}+1)\epsilon \tag{21}\]

where \(\tilde{L}\) is the Lipschitz constant for equi-Lipschitz family. Therefore, \(f_{k}\) converges to \(f\) uniformly on \(K\).