Hessian-Aware Training for Enhancing Model Resilience for In-Memory Computing

We thank reviewer 2ZTo for the time and effort in reading and evaluating our manuscript. Below, we answer the reviewer’s concerns and questions.

[Concern 1: Representation of noise]

We appreciate the reviewer’s concerns regarding the representation of noise specific to In-Memory Computing (IMC) hardware. We respectfully disagree with the reviewer and want to clarify that our study is more concerned about improving the parameter-level robustness of a DNN against bitwise error in model parameters. IMC noise models often apply random noise across parameters and report average performance losses, offering a general view of model resilience under typical noise conditions [A, B, C]. However, such average-case models may obscure cases where individual bitwise errors cause catastrophic accuracy drops, which is critical for robustness in fault-prone hardware environments like IMC. If we randomly apply noise to parameters, we barely find cases of accuracy drop. Our approach aims to identify and quantify worst-case scenarios, showing the impact of specific parameter manipulations on model accuracy. This allows us to understand and address parameters especially sensitive to bitwise faults, regardless of hardware type. On top of that, by systematically evaluating the sensitivity of each parameter to single-bit flips, we provide a more granular resilience analysis.

[Concern 2 : Clarification on figure 1 and 2]

A significant percentage of bit flips result in a negligible improvement in accuracy for Hessian-aware trained models. For example, a model with hessian-aware training with 92.67% baseline validation accuracy might exhibit a 92.68% validation accuracy following particular bit flips, resulting in a slight negative RAD score. Figures 1 and 2 do not include these instances, as the RAD metric only reflects more than 10% accuracy drops.

[Concern 3: Generalizability of our approach in pruning applications]

We appreciate the reviewer’s recognition of the advantages of Hessian-based regularization in enhancing pruning resilience. We would like to clarify that, contrary to the reviewer's impression, we conducted evaluations across multiple model and dataset combinations. These additional results, which support the generalizability of our approach in pruning applications, are included in the appendix section (B.6) of the manuscript. Our analysis demonstrates consistent pruning resilience benefits across different models and tasks, underscoring the broad applicability of our Hessian-aware training technique.

We thank the reviewer xzTf for the time and effort in evaluating our manuscript. Below, we provide answers to the reviewer’s concerns and questions. We also updated our manuscript to address the feedback. These changes are highlighted in blue.

[Concern 1 and question 1: overhead]

We appreciate the reviewer’s concern and agree that the reliance on the Hessian trace present scalability challenges for very large models due to the computational costs of Hessian approximation. To address this concern, we conducted additional experiments on reducing overhead. Instead of computing hessian trace across the entire parameter space, we can focus on the most sensitive layers to reduce the overhead. We applied Hessian regularization incrementally, starting with only the last layer and extending it to the last 2, 3, and finally 4 layers of the model and compared the runtime with baseline training. We run our training 5 times and report per epoch training time in Pytorch. Our results are summarized in the table below.

Our results demonstrate that training overhead increases as we increase the “layers involved in hessian-regularization.” However, using only the last 1 layer of the model, we can reduce the overhead to almost the same as baseline training, making our method efficient for very large models. We added these results in the appendix section (B.7) for enhancing the clarity on overhead of our method.

[Concern 2 and Question 2: Adopting our approach to large models with many skip connections]

We thank the reviewer for raising his concerns about the applicability of our method to models with a large percentage of skip connections. We respectfully want to point out that our method is compatible with models that have a large percentage of skip connections, such as visual transformers. Our experiment with the Diet-tiny model (Appendix B.4) shows that, in alignment with our results in Table 3, our method reduces erratic bits by ~7%, hence providing evidence that our approach is compatible with models with large percentage of skip connections.

[Concern 3 : Providing hyperparameter selection guidelines]

We agree with the reviewer that including an additional hessian regularization term can introduce tuning complexity, particularly for larger models. We clarify that, we use a computationally inexpensive approach (section 4.3, methodology) to determine suitable values for the regularization coefficient $\lambda$ across different models. Using a small mini-batch of data, we train models 5 epochs and observe the model's sensitivity. We use the regularization coefficient $\lambda$ that minimizes sensitivity. Our experiments show that for larger models, $\lambda$ value of $10^{-2}$ works quite well, and for small MNSIT models, $\lambda$ value 1 is suitable. To address the reviewer's concern, we updated our manuscript with an additional explanation of our hyperparameter selection. These changes can be found in the section 4.3, lines 262 and 272.

We thank reviewer SkKx for the time and effort in reading and evaluating our manuscript. Below, we answer the reviewer’s concerns and questions about understanding our result, fairness in the experiment setup and novelty. To address the concerns, we revised our manuscript and highlighted the changes in blue.

Concern 1: Understanding results in Table 3

We thank the reviewer for raising concerns over the results in Table 3. We respectfully disagree with the reviewer on the point where he/she mentions, “baseline even has a better accuracy.” We show in Table 3 that the validation accuracies of models trained with our method are comparable to baseline training, and in some instances, even better (see results of ResNet18 on Cifar10).

We do not mention the model accuracy under bitwise error to DNN model parameters. The systematic resilience measurement framework we use [A] alters single-bits, hence, a model with a million parameters will have 1x32(bits per parameter) = 32 Million bits flips. Instead, we show how many parameters are in the model who’s flipping leads to a more than 10% RAD, to measure the resilience of a DNN under bitwise error. We want to further clarify that there is no error ratio in our study. Also, figures 1,2 and 6 show the distribution plot of accuracy degradation after bit-flips and how these ranges are getting reduced in models trained using hessian-aware training.

[Concern 2 : Fairness in experiment setup]

We clarify that our experimental setup balances computational feasibility with meaningful bitwise error to the model parameters across models of varying sizes. Due to the prohibitive computational time required for complete bitwise error analysis on large models, we leverage findings from MNIST and LeNet experiments—where 99.99% of erratic bits are exponents—and follow prior studies [A] to selectively analyze exponent bits flips for larger models like ResNet18 and MSB for ResNet50. This approach ensures practical yet rigorous evaluation, as a full analysis would take ~503 days for ResNet18 (11M parameters, CIFAR-10) and ~1172 days for ResNet50 (25.6M parameters, ImageNet) with our current resources (4-8 GPUs and 60-120 CPUs). It is important to note that prior works have not performed a complete analysis of bitwise errors on large-scale models simply due to the prohibitive computational demands [A, B, C]. To address the reviewer's concern, we have incorporated these additional explanations into Section 5.1 (lines 297–302) to improve clarity on the fairness of our experiment setup.

[...cont'd from the rebuttal]

[Question 3: Major Machine Learning Contributions]

First, we clarify that there has been no prior work on improving the “model’s inherent resilience” to bitwise error in their parameters. Second-order objectives have been utilized in prior works as a measure of the flatness of the loss landscape [B]. Few other studies have used hessian to improve accuracy or generalization in quantization [A, C, D]. Our study proposes a novel hessian-aware DNN model training algorithm that utilizes the trace of the hessian and top-p% eigenvalues to train DNN models resilient to bitwise error in their parameters.

Second, our study identifies that naturally flat layers (layers with skip connections in residual blocks) are more resilient to bitwise error in parameters due to their naturally flat loss landscape. This is a significant ML contribution because it identifies the most sensitive layers in a DNN model against bitwise error in parameters (e.g., Conv and Fully Connected), laying the groundwork for future studies, which could be particularly beneficial for designing resilient architectures for error-prone computing environments and emerging devices.

Lastly, our method benefits extreme model compression, enabling more aggressive quantization (2-bit) and pruning with higher performance retention than traditional training. This contribution is important for resource-constrained applications and positions our work as a versatile solution that enhances parameter resilience against bitwise error and efficiency against extreme parameter compression techniques.

[Summary]

We thank the reviewer for raising concern about the novelty and significant contribution of our paper to the ML field. In our rebuttal, we clarified our novelty and established the major ML contribution of our work. We also clarify that our work does not propose a defense against targeted bit-flip attacks. We therefore kindly request the reviewer to adjust the rating if our answer addresses the concerns and answers the questions. If there are any remaining concerns/questions, we are happy to address them during the discussion phase.

[References]

A. Yang, H., Yang, X., Gong, N. Z., & Chen, Y. (2022, July). Hero: Hessian-enhanced robust optimization for unifying and improving generalization and quantization performance. In Proceedings of the 59th ACM/IEEE Design Automation Conference (pp. 25-30).

B. Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020.

C. Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021.

D. Zhen Dong, Zhewei Yao, Daiyaan Arfeen, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Hawq-v2: Hessian aware trace-weighted quantization of neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18518–18529. Curran Associates, Inc., 2020.

We thank reviewer bQFx for the time and effort in reading and evaluating our manuscript. Below, we answer the reviewer’s concerns and questions one by one. The manuscript has been updated to address the reviewers' concerns, with the changes highlighted in blue. The line numbers are based on the revised manuscript.

[Concern 1 and Question 1: Novelty and differentiation of our work from prior work]

Thank you for the careful review and enlightening questions about the novelty of our work. After carefully reading the HERO [A] paper, we found that our work differs from this prior work in terms of perturbation bounds, hessian approximation and DNN model training methodology using hessian. These differences are listed below.

• Hero and Hessian aware training uses different weight perturbations to model parameters.

  • Hero uses a theoretically bounded perturbation bounded by $l_{\infty}$ norm, defined as $\parallel \delta \parallel = {\parallel W_q - W \parallel}_{\infty} \leq \bigtriangledown/2$, where $W$ and $W_q$ denote the original and quantized weights, and \bigtriangledown/2 denotes the maximum quantization of each weight for a quantization bit width of $\bigtriangledown$.

  • Our proposed Hessian-aware training uses single bitwise error to model parameters, which are not theoretically bounded and can change the model parameter dramatically.

• Hessian-approximation and calculation of regularization term is different in our work and prior work [A].

  • For hessian matrix $H$’s eigenvalues $\lambda$, HERO’s regularization term is calculated as $L_r = \sum_{i}^{} \mathrm{\lambda}_{i}^{2} = E_z \parallel H_z \parallel^2, z\sim N(0, I).$ They calculate the hessian term $H_z$ using finite difference method, given by $H_z \approx \frac{\nabla L (W + hz) - \nabla L(W)}{h}$, h is a small positive number and h is a random vector along the direction of the gradient.

  • Our study utilizes Hutchinson's method to approximate the trace of the Hessian, which is a faster method [D]. For a set of random vectors $v$ drawn from the Rademacher distribution, the trace is approximated as $Tr (H) = \mathbb{E} [v^T Hv]$ = $\frac{1}{N_v} = \sum_{i = 1}^{N_v} v^{T}_{i} Hv_i$. We obtain $v^T Hv$ by computing the gradient of the loss function twice, given by $v^t Hv = v^T \cdot \frac{\partial }{\partial \theta} \left( \frac{\partial L}{\partial x} \right) \cdot v$.

• Our study and prior work differ in model training methods.

  • Hero applies the hessian regularization term in each training step.

  • Our algorithm computes the top-p eigenvalues for each minibatch and keeps a running mean of the eigenvalues for all minibatches. We only regularize the model if the current minibatch mean of eigenvalues is greater than the mean eigenvalue observed previously across all minibatches. Lines 234-239 of our paper discuss these unique approaches (lines 15-20 in algorithm 1).

We believe these are sufficient novelty to differentiate our work from prior work. To address reviewer's concern, we have included the prior work in section 4.1 (lines 160-163) of the revised manuscript.

[Concern 2 and Question 2: Defending against Bit-flip Attacks]

We thank the reviewer for raising this concern, and we would like to clarify that we do not propose a defense against bit-flip attack, nor do we claim anywhere in our paper that ours is a defense against bit-flip attack. In lines 108-109 of our paper, we mention that the focus of our study is not defending against malicious bit-flip attacks. Malicious bit-flip attacks are targeted attacks on the DNN model’s bitwise representation, aiming towards high accuracy degradation. Instead, we are more concerned about the naturally occurring bit-flips in emerging computing platforms. Thus we believe that the reviewer’s concern falls outside the scope of our current study.