Neural Network Compression for Noisy Storage Devices

Most memory technologies utilize continuous physical values (e.g., charge) as a means of data storage. The full continuous storage range is often divided into intervals and used to store discrete values. One extreme that maximizes the device’s performance is to allow only two values (high and low) to be written to each memory cell. With this approach, one bit of information can be stored in a memory cell. Storing more values allows more bits to be stored per cell, but also increases the chance of reading an incorrect value from the memory – this represents a natural trade-off between memory density and probability of error. Therefore, storage devices often use error correcting codes (ECCs) in practice to protect data from such memory errors.

Although digital storage (i.e., storing discrete values) is the dominant paradigm in physical memory technology, it is still possible to store continuous values in memory cells, and read them back as continuous values, albeit with noise. This approach, also known as analog storage, has regained attention recently [72] with the emergence of different non-volatile memory (NVM) technologies, such as PCM. NVMs not only retain the stored information even when the power supply is off, but also allow efficient storage of multiple bits per cell. In the extreme case, NVMs can store continuous values [67], making them more efficient than their digital counterparts that require discrete inputs [71, 73]. Among the various NVM technologies, our work focuses on phase-change memory (PCM) technology because it: (i) has faster read/write speed and higher endurance than its competitors; and (ii) can enhance chip performance by reducing the cost of data movement thanks to its on-chip integration.

A major advantage of the PCM memory device over its (digital) competitors is its relatively low cost as compared to its access time (read/write time). Commonly used memory devices such as SRAM (static random access memory) and DRAM (dynamic random access memory) have short access time and high endurance, and are thus used as the cache and main memory in a wide range of technologies. Storage devices such as NOT-AND (NAND) flash, on the other hand, provide low-cost, high density, and non-volatile storage, but their long access time makes them unsuitable to use as memory devices. As shown in Figure 1(a), PCM successfully bridges the gap between memory and storage devices. Such advantages of the PCM device allows for utilizing both memory and storage on the main chip, which, when coupled with its high endurance, makes it a compelling alternative to conventional memories such as NAND flash and DRAM. This is especially true for applications requiring NN inference, where both short access time and non-volatile storage are crucial for real-time deployment.

To formalize the joint compression problem, we consider a supervised learning setup where \(x \in \mathcal {X} \subseteq \mathbb {R}^d\) is the input variable, and \(y \in \mathcal {Y} = \lbrace 1, \ldots , k\rbrace\) is the label. We assume access to samples \(\mathcal {D} = \lbrace (x_i, y_i)\rbrace _{i=1}^n\) drawn from an underlying (unknown) joint distribution \(p_{\rm {data}}(x,y)\) , which are used to train an NN predictor \(f_w: \mathcal {X} \rightarrow \mathcal {Y}\) . This network, parameterized by the weights \(w \in \mathcal {W}\) to be compressed and stored on analog storage devices (where \(\mathcal {W}\) denotes the parameter space of NNs), indexes a conditional distribution \(p_w(y|x)\) over the labels y given the inputs x.

The NN weights w will be exposed to some input dependent device noise \(\epsilon (w)\) when written to the analog storage device, yielding noisy versions of the weights \(\hat{w}=w+ \epsilon (w)\) . In our experiments, we find that PCM noise dramatically hurts classification performance – as a concrete example, the test accuracy of a ResNet-18 model trained on CIFAR-10 drops from 95.10% (using w) to 10% (using \(\hat{w}\) ) after the weights are corrupted (via naive storage on analog PCM). Thus, to preserve NN performance even after it is stored on the analog device, we explore various strategies \(g: \mathcal {W} \rightarrow \mathcal {W}\) for designing reconstructions of the perturbed weights \(g(\hat{w})\) such that the resulting distribution over the output labels \(p_{g(\hat{w})}(y|x)\) is close to that of the original network \(p_w(y|x)\) .

We note that this notion of “closeness” between the original weights w and the reconstructed weights \(g(\hat{w})\) has several interpretations. In Section 3.1.1, we explore methods to minimize the distance between w and \(g(\hat{w})\) in Euclidean space:In Sections 3.1.2–3.3, we study ways to minimize the Kullback-Leibler (KL) divergence between these two output distributions:where we learn the appropriate transformation \(g(\cdot)\) .

In this section, we devise several novel coding strategies for \(g(\cdot)\) that can be applied to the model post-training to mitigate the effect of perturbations in the weights. For each strategy, we require a pre-mapping process to remove the input dependence in the mean of the PCM response in Figure 2(c).

Pre-mapping:. Let C represent the PCM channel. The mean function of C, denoted as \(\mu : \mathcal {X} \rightarrow \mathbb {R}\) , in Figure 2(c) can be learned via a k-nearest neighbor classifier on the channel response measurements. We can also learn an inverse function \(h = \mu ^{-1}\) using the measurement data and use it to remove the input dependence in the mean. More precisely, we pre-map the data with h prior to channel usage, which yields an identity function with zero-mean noise \(\phi = C \circ h\) (Figure 3), i.e., \(\phi (x) = x + \epsilon _0(x)\) where \(\epsilon _0(x)\) is a zero-mean noise with input dependent std \(\sigma (x)\) due to the PCM channel. Thus, the relationship between input weights \(w_{in}\) to be stored and output weights \(w_{out}\) to be read is:where \(\alpha\) and \(\beta\) are scale and shift factors, respectively. Since the noise is zero-mean, we have:If we use the noisy channel \(\phi\) r times (i.e., store the same weight on r “independent” cells and average over the outputs – much like repetition codes in coding theory), the relationship between \(w_{in}\) and \(w_{out}\) becomes:

Let us define a new random variable \(\tilde{\epsilon }(w_{in}, r) =\frac{1}{r}\sum _{i=1}^r \epsilon _{0,i}(w_{in})\) . Notice that the standard deviation of \(\tilde{\epsilon }(w_{in}, r)\) is \(\frac{\sigma (w_{in})}{\sqrt {r}}\) . Then the standard deviation of \(\hat{\epsilon }(w_{in}, r, \alpha)= \frac{\tilde{\epsilon }((w_{in}, r))}{\alpha }\) in Equation (2) is given by: \(\frac{\sigma (w_{in})}{\alpha \sqrt {r}}.\) More precisely, we have:where the standard deviation of \(\hat{\epsilon }(w_{in},r,\alpha)\) is \(\frac{\sigma (w_{in})}{\alpha \sqrt {r}}\) . This provides us two tools to protect the NN weights against zero-mean noise.

For the first method, we observe in Figure 3 that the response becomes less noisy as we increase the number of PCM cells used (r). However, we desire to keep r at a minimum for an efficient storage. The second method allows us to make use of the full analog range by scaling the weights. This is particularly useful when storing weights with small magnitude. But we cannot increase \(\alpha\) without bound because of the device constraint, since we must satisfy \(-1 \le \alpha \cdot w_{in} -\beta \le 1\) (in practice, this constraint is \(-0.65 \le \alpha \cdot w_{in} -\beta \le 0.75\) since the remaining portion of the response is non-invertible, and therefore unusable for our purposes). Such limitations leave more to be desired for a general-purpose robust coding strategy for NN weights.

The following set of techniques for constructing \(g(\cdot)\) are designed to correct the errors that the robust coding strategies in the previous section fail to address.

Robust Training:. For robust training, we regularize the standard cross-entropy loss with the KL divergence in Equation (1). Specifically, the loss function is as follows:In our experiments, we add a noise (with a carefully adjusted standard deviation) as a perturbation \(\delta w\) (i.e., \(g(\hat{w})=w+\delta w\) is a noisy weight) during robust training, and we observe that the trained network is more robust to PCM noise and also to pruning effects. This is because the final perturbation \(\delta w\) on the weights ( \(g(\hat{w})=w+\delta w\) ) can be thought of as a noise or the effect of pruning on the w. Although our framework does not involve a pruning step (pruning can be performed independently as we show in Section 4.1), we believe that making NN weights more robust to pruning effects is an important additional feature of our strategy. In particular, when we apply pruning, we have: \(\delta w_j = 0\) for a non-pruned weight \(w_j\) , and \(\delta w_j = -w_j\) for a pruned weight \(w_j\) . Recall that in magnitude pruning, only the small weights are set to zero, that is \(\delta w_j\) is equal to 0 (for large weights) or \(-w_j\) (for small weights). In other words, \(\delta w_j\) is always a small value, just like a noise, therefore this strategy could provide robustness against pruning effects as well. Our experiments on CIFAR-10 and ImageNet verify that robust training has a protective effect against PCM noise where robust coding strategies are not enough.

Robust Distillation:. Distillation is a well-explored NN compression technique [33]. The idea is to first train a teacher network to capture a smooth probability distribution on labels, and then train a smaller student network by distilling the output probability information from the teacher. We use distillation to optimize a compressed model to be robust to PCM noise via the noisy student loss:where \(\lambda \in [0,1]\) is a scalar, \(w_t\) and \(w_s\) are teacher and student weights, \(g(\hat{w}_s)= g(w_s+\epsilon (w_s))=w_s+\delta w_s\) with \(\epsilon (w_s)\) being the PCM noise and \(\delta w_s\) being the noise injected onto the student network’s weights. Although [74] provides an initial exploration into distillation for noisy storage, we leverage experimental data collected from real storage devices to build a realistic model of the PCM noise.

Finally, we explore a probabilistic, end-to-end learning approach for \(g(\cdot)\) that jointly optimizes over the model compression and the known characteristics of the noisy storage device. We assume access to a set of weights \(\lbrace W_j\rbrace _{j=1}^K\) from K models that have been trained on the same dataset \(\mathcal {D}\) – this serves as an empirical distribution over NN weights, as different initializations of NNs typically converge to different local minima. Additionally, we assume that each weight is a sample from a normal distribution \(W_{ij} \sim \mathcal {N}(0,1)\) [20, 76].

To learn the representation of the NN weight, we maximize the mutual information (MI) between the original network weight W and the compressed weight Z that has been corrupted by the noisy channel. Following [13], our coding scheme can be represented by the following Markov chain²:where \(\hat{Z}\) denotes the compressed weight and \(Z = \hat{Z} \,+\, \epsilon (Z)\) where \(\epsilon (Z)\) denotes the input-dependent PCM noise. Then, we obtain the following lower bound to the true MI:where \(q_\phi (Z|W)\) and \(p_\theta (W|Z)\) denote approximations to the true posteriors \(p(Z|W)\) and \(p(W|Z)\) , respectively [4]. We train an autoencoder \(g_{\phi ,\theta }(\cdot)\) with a stochastic encoder and decoder to learn these variational posteriors. The decoder is trained to output a set of reconstructed weights such that its predictions are close to those of the original network. Notably, our approach differs from KD in that we learn the weight compression scheme rather than using a fixed student network architecture.

We show the effect of sparsity- and sensitivity-driven protection in Figure 5 on ResNet-50 for ImageNet (Figure 8 for CIFAR-10 in Appendix). The exact numerical results are given in Table 1.

In CIFAR-10 experiments, the number of small weights was \(N_{small}=11\) M and the number of large weights was \(N_{large}=5.6\) K and number of cells per weight on average for adaptive redundancy is not more than 0.02 above the listed numbers in the table. Similarly, in ImageNet experiments, \(N_{small}=25.4\) M, \(N_{large}=15\) K and the difference between the number of cells per weight on average and the listed number is always smaller than 0.06.

In Table 1, we provide detailed experimental results on the effect of each sparsity- and sensitivity-driven strategy with ResNet-10 on CIFAR-10; and ResNet-50, MobileNet-v2, and EfficientNet-B2 on ImageNet. We provide the required number of cells per weight to store the continues value of the weight (3rd-12th columns) and the additional bits to store for each strategy (last column) separately. In Table 2, we report some of the results from Table 1 again, this time by providing the number of cells per parameter required to store (1) the continuous weight value, and (2) the additional bits. Recall that we can store 2 bits or 1.8 bits digitally in one PCM cell with the ideal or realistic baselines, respectively. It is seen from the two tables that sparsity driven protection strategies are enough for ResNet-18 on CIFAR-10 to preserve the classification accuracy ( \(95.10 \%\) ) with 4 cells per weight. Moreover the accuracy with 2.5 cells per weight. is only \(0.1 \%\) less than the baseline accuracy ( \(95.10 \%\) ). Similar observation can be made for the ImageNet results as well. Table 1 also shows that sign protection is a critical step: for instance, the AM+AR and SP+AM+AR rows indicate that sign protection can increase the accuracy from \(58 \%\) AM+AR to \(94.4 \%\) SP+AM+AR on CIFAR-10.³

We also combine sparsity driven protection with pruning. Results on \(90 \%\) pruned ResNet-18 (on CIFAR-10) are given in Table 15 in Appendix B.2. We follow the standard pruning procedure: first train the model as usual, then prune \(90\%\) of the weights, and then retrain the remaining non-pruned weights by keeping the pruned weights frozen at value zero. In our experiment, we retrain the pruned model for 20 epochs. Then for the storage of this pruned model, we assume that the pruning mask could be stored using the compressed sparse row (CSR) or compressed sparse column (CSC) format and Huffman coding, as detailed in [28]. This would require at most 0.5 bits per weight. Then we need to store the continuous weight values of the non-pruned weights. Note that this is the standard technique for storing a pruned network: (1) first store the pruning mask, (2) then store the values of the non-pruned weights. In Table 15, the number of cells we report are computed as follows:translating the number of cells per non-pruned weight to number of cells per (pruned or non-pruned) weight and also considering additional 0.25 cells per parameter to store the pruning mask. Note that the accuracy of the pruned model without noise is \(94.8 \%\) and we get \(94.2 \%\) test accuracy with PCM noise using only 1.35 cells per weight. Compared to the realistic digital baseline (no compression and each 32-bit weight is digitally stored on 18 PCM cells), we provide \(\frac{18}{1.35} = 13.3\times\) more efficient storage with analog storage combined with our strategies and \(90\%\) pruning.

We compare robust training with naive training (training with no noise) in Table 3 on ResNet-18 (on CIFAR-10) and on ResNet-50 (on ImageNet). In our experiments, we use \(\lambda = 0.5\) as the coefficient of the KL term in the loss. We experimented with different combinations of the sparsity and sensitivity driven strategies, and in all cases, robust training provides better robustness against PCM noise than the naive training. Robust training provides robustness against pruning as well. When ResNet-18 trained without noise is pruned with \(90\%\) sparsity, the accuracy drops from \(95.1 \%\) to \(\mathbf {90.2 \%}\) . However, the same model trained with \(N(0,0.006)\) gives \(\mathbf {95.0 \%}\) accuracy after \(90\%\) pruning.

Table 4 shows the distillation results on ResNet-20 distilled from ResNet-18 (on CIFAR-10) with PCM noise applied at test time. We compare three networks: (1) a student ResNet-20 distilled without noise injection, (2) a student ResNet-20 distilled with Gaussian noise injection, and (3) a teacher ResNet-20 (a baseline) trained without noise. As shown in Table 4, student ResNet-20 distilled with noise injection outperforms both teacher ResNet-20 and student ResNet-20 distilled without noise when weights are perturbed by PCM at test time. (see Appendix B.4 for additional results.)

We also consider quantization as a way to improve the efficiency of digital storage. Table 5 shows a comparison between analog storage improved with our strategies and digital storage improved with quantization techniques. We consider both the ideal (2 bits per cell) and realistic baselines (1.8 bits per cell). For digital storage, we apply quantization using different techniques from the literature [3, 41]. We find the number of cells to store one quantized weight in both ideal and realistic cases. Then we adjust the number of cells used to store one weight in analog PCMs to be the same as the digital PCMs for both baselines, by adapting the redundancy for large and less sensitive weights. For instance, when [3] performs 4-bit quantization, each parameter is represented by 4 bits. In the ideal baseline, this would require 2 cells per parameter. We adjust the parameters in our robust strategies so that the number of cells required to store one parameter is 2 and report the result (76.02) in the “PCM (analog) + our robust strategies” row under “Ideal” column. We repeat the same procedure for the realistic baseline which requires 2.22 cells per parameter to store 4 bits. We report the result (76.08) in the “PCM (analog) + our robust strategies” row under “Realistic” column. As shown in Table 5, noisy analog storage improved by our strategies outperforms digital storage of quantized weights. We do not compare against more aggressive quantization techniques [2, 14, 23, 43, 54, 66] that can achieve higher efficiency in digital storage since they incur a huge complexity with multiple retraining stages.

We test our strategies on a regression setting as well. For the regression task, we consider the Neural Radiance Fields (NeRF) framework [52] that contains a neural network to model the relation between (1) the 3D coordinates of a point in a given 3D scene and the view direction; and (2) the view-dependent continuous color⁴ and volume density at that 3D point. More precisely, a NeRF model takes 3D positions \((x,y,z)\) of a point in 3D space and a particular direction to view the 3D scene \((\theta , \phi)\) ; and outputs a view-dependent RGB color and a volume density at that point. We note that this is a nontrivial task that has been attracting significant interest from the computer graphics and vision communities since the first paper in 2020 [52]. We report the PSNR of the NeRF model on fern dataset (https://github.com/bmild/nerf) that was generated by the authors of [52] under the effect of PCM noise in Table 6. Though the PSNR is affected severely by the PCM noise when there is no protection, our robust strategies help to preserve the PSNR. In particular, the PSNR is 24.75 dB with an average number of 4.5 cells, which is even higher than the baseline (without noise) PSNR 24.73 dB. Additionally, we note that we would normally expect to see a more severe effect on PSNR by the PCM noise. This is because the exact values of the model outputs matter more in a regression setting than in a discriminative one, since the classifier’s accuracy can remain unchanged even with small perturbations to the logits. However, as can be seen from Table 6, PSNR is still within an acceptable range once we apply the robust strategies. We believe this is due to the discretization at the evaluation phase, which may cancel out some noise on the output.

As a final note, the storage density for the NeRF model or other 3D regression models could be improved further through model compression techniques such as [7, 35, 36].

Finally, we explore the effectiveness of the end-to-end learning scheme in which we train an autoencoder on a set of model weights. We generate two sets of 2-D Gaussian mixtures: an “easy” task in which we first sample a two-dimensional mean vector \(\mu _1 \in \mathbb {R}^2\) where \(\mu _{1i} \sim \mathcal {U}[-1,0)\) and \(\mu _2 \in \mathbb {R}^2\) where \(\mu _{2i} \sim \mathcal {U}[0,1)\) for \(i=\lbrace 1,2\rbrace\) . After sampling these means, we draw a set of 50K points for the two mixture components: \(x_1 \sim \mathcal {N}(\mu _1, I)\) and \(x_2 \sim \mathcal {N}(\mu _2,I)\) . This ensures that the two mixture components are well-separated. For the “hard” task, we sample overlapping means: both components of \(\mu _1\) and \(\mu _2\) are drawn from \(\mathcal {U}[-1,+1]\) before sampling from their respective mixture components. We generate 50 datasets per task, where each dataset has 50K data points.

We train 50 different logistic regression models on each of the datasets for both the “easy” and “hard” tasks (each model has three parameters). After training the logistic regression models, we use an autoencoder with MLP encoder and MLP decoder architectures and ReLU nonlinearities, as shown in Table 7. The autoencoder for both the “easy” and “hard” tasks is trained for 10 epochs with a batch size of 100 using the Adam optimizer with learning rate = 0.001, \(\beta _1 = 0.9\) , \(\beta _2 = 0.999\) , and early stopping on a held-out validation set.

We test the autoencoder on both Gaussian and PCM noise: that is, we corrupt the 1-dimensional encoder output (z-representation of the classifier weights) with the appropriate perturbation before passing in the encoded representation into the decoder. Figure 6 provides an illustration of the autoencoding process for the PCM array, which is analogous to the Gaussian noise setup. For the easy task with Gaussian noise, the autoencoder achieves an accuracy of 95.2%; for PCM noise, it achieves 91.8% accuracy. For the hard task with Gaussian noise, the autoencoder achieves an average accuracy of 78.8% across all classifiers; for PCM channel noise, it achieves 78.6% on average across all 50 classifiers.

Interestingly, we find that the 1-D classifier representations in the autoencoder’s latent space is semantically meaningful. In Figure 7, we qualitatively analyze the learned representations of the logistic regression classifier weights. In Figure 7(a), we plot all 50 datasets of the Gaussian mixtures (“easy task”) as well as the true decision boundaries for each of the 50 logistic regression models, each boundary colored by the magnitude of its z-representation as learned by the autoencoder. We plot the same for the “hard task” in Figure 7(b). For the easy task, we note that the classifiers are encoded by their relative location in input-space: those that are in the lower left corner of the scatter plot have smaller magnitudes than those on the upper right. For the hard task, however, the z-representations appear to be fairly random – at a first glance, there does not appear to be a particular correlation between the magnitudes of the z-encodings and the original classifier weights.

We further explore this phenomenon for the hard task in Figure 7(c-d). For two particular datasets (though the trend holds across all 50 datasets), we color the original data points by their mixture component as well as the true decision boundary by the magnitude of its z-encoding. As shown in Figure 7(c) and (d), we find that the autoencoder has learned to map all classifiers with the positive class to the left of the decision boundary to z-representations with large magnitude; conversely, those with the positive class to the right of the decision boundary are encoded to z-representations with smaller magnitudes. Through this preliminary investigation, we demonstrate that the end-to-end approach for learning both the compression scheme while taking the physical constraints of the storage device into account shows promise.

The architectures for LeNet and the MLP are shown in Tables 8 and 9, respectively:

For both the LeNet and MLP classifiers, we use a batch size of 100 and train for 100 epochs, early stopping at the best accuracy on the validation set. We use the Adam optimizer with learning rate \(=0.001\) , and \(\beta _1=0.9, \beta _2=0.999\) with weight decay \(=5e^{-4}\) . For knowledge distillation, we use a temperature parameter of \(T = 1.5\) and equally weight the contributions of the student network’s cross entropy loss and the distillation loss ( \(\lambda =0.5\) ).

We provide the architectural details for the ResNet-18 and the slim ResNet-20 used in our experiments in Tables 10 and 11 below:

For both ResNet-18 and slim ResNet-20, we use a batch size of 128 and train for 350 epochs, early stopping at the best accuracy on the validation set. We use SGD with learning rate \(=0.1\) , and momentum \(=0.9\) and weight decay \(=5e^{-4}\) .

We provide the architectural details for the ResNet-50 used in our experiments in Table 12.

We use the pretrained ResNet-50 from PyTorch (https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py), with a batch size of 64. For the robust training experiments, we retrain the model for 20 epochs with early stopping at best accuracy. We use SGD with learning rate \(=0.001\) , and momentum \(=0.9\) and weight decay \(=5e^{-4}\) .

We give the full results of sparsity-driven and sensitivity-driven protection experiments on CIFAR-10 and ImageNet with confidence intervals in Tables 13 and 14. In Figure 8, we present experimental results with ResNet-18 on CIFAR-10. As can be seen from Figure 8(a), our strategies, namely sign protection, adaptive mapping, and adaptive redundancy, reduce the number of PCM cells per weight required to preserve accuracy from 1,024 to 1. We also test our strategies against Gaussian noise. In the Gaussian experiments, we consider hypothetical storage devices with white Gaussian noise where the standard deviation of the overall noise on the output ( \(\hat{\epsilon }(r, \alpha)\) “not input-dependent”) can be reduced by using the channel multiple times (Method #1) and by using the channel in the its power limit (Method #2). Figure 8(b) shows that our strategies increase the standard deviation threshold, where the NN performance sharply drops, from 0.005 to 0.2, i.e., robustness increases by 40 times. We note that we present results of Gaussian noise experiments to be comparable to future work although this scenario is not realistic.

We give the results of robust pruning experiment in Table 15. We apply one-shot pruning followed by 20 epochs of retraining. When sparsity- and sensitivity-driven strategies are applied, 2.05 cells per weight are enough to preserve the original accuracy of the pruned model ( \(94.8 \%\) ). Moreover, if we use only 1.35 cells per weight on average, the accuracy drop is only \(0.1 \%\) , which is insignificant considering the fact that pruning reduces the accuracy by \(0.3\%\) (from \(95.1\%\) to \(94.8\%\) ). We note that this performance degradation due to pruning can be eliminated with robust training, explained in Section 3.2. With this strategy, we can reduce the number of cells required to store each weight from 16 (in digital storage) to 1.35 with analog storage combined with our robust strategies and pruning. This corresponds to an 11.85 times more efficient storage.

We give the full results of robust training experiments on CIFAR-10 and ImageNet with confidence intervals in Tables 16 and 17. We use \(\lambda =0.5\) as the coefficient of the KL regularization term in the loss function in Section 3.2. The level of the noise (standard deviation) injected during the training is adjusted according to r and \(\alpha\) values to be used at storage time since the noise at storage time has a standard deviation of \(\frac{\sigma (w_{in})}{\alpha \sqrt {r}}\) . It is seen from Tables 16 and 17 that robust training improves the robustness of the network against noise. We have also observed that the pruned network reaches higher accuracy after robust training compared to naive training without noise injection (an increase from \(90.2\%\) to \(95.0\%\) without retraining).

Knowledge distillation (KD) is a well-established NN compression method where a large teacher network is trained with \(\mathcal {L}_t\) loss given bywhere \(w_t\) is the weights of the teacher network, and probability of each class i is the output of the high temperature softmax activation applied to logits:where \(z_i\) is the logit for class i. Temperature \(T\gt 1\) helps the output probabilities of the teacher network be softer. Using the same temperature, a smaller student network can be trained with the following student loss function \(\mathcal {L}_s\) :where \(w_s\) is the weights of the student network. It has been shown that distilled student network achieves a test accuracy that a teacher network with the same architecture cannot achieve. In other words, a student network distilled from a teacher network performs comparable to a larger network. This suggests that knowledge distillation can be regarded as a promising compression method. In addition to compression, distillation has been shown to be an effective method for other desired NN attributes such as generalizability and adversarial robustness [55]. In this work, we define “robustness” as preserving a network’s downstream classification accuracy when noise is added to the weights. This is achieved in part by robust training in the previous section where a trained network is robust to pruning and noise on the weights. Here, we present a student loss function that would make a (compressed) student network more robust to noise with no change in the teacher network training:

In the experiments, we use a temperature parameter of \(T = 1.5\) and equally weight the contributions of the student network’s cross entropy loss and the KL term ( \(\lambda =0.5\) ). Similar to robust training, the noise level (standard deviation) during training is adjusted according to r and \(\alpha\) values to be used at storage time. We give the full results on CIFAR-10 with confidence intervals in Tables 18 and 19 with weights perturbed by PCM cells and Gaussian noise, respectively. We also present the same set of experiments on MLP distilled from LeNet (on MNIST) in Tables 20 and 21. As in CIFAR-10 experiments, noise injection during distillation makes the student network more robust to both PCM and Gaussian noise in MNIST experiments.