Modulating neural network weights improves error tolerance and reduces storage

LEAN: Loss-Engineered Modulation of the Model Weight-Sensitivity Distribution for Efficient, Error-Tolerant Neural Networks

Machine LearningCryptography and Security

Summary

Not all parts of a neural network's memory are equally important for keeping its accuracy. The authors discovered a way to adjust how sensitive these memory parts are to errors, making the network less likely to break from small bit flips. They do this by fine-tuning the network to spread out sensitivity more evenly across weights, resulting in much better resistance to errors with little loss in performance. This could allow smaller storage formats without hurting how well the network works.

What this means in practice

  • For hardware designers: Design neural network chips with smaller, mixed-bit weight storage to reduce memory costs while maintaining accuracy and tolerance to bit errors.$Commercial implications: Enables commercial low-bit precision neural network accelerators with improved error resilience, benefiting chip manufacturers targeting embedded AI applications.
  • For machine learning engineers: Improve deployed model robustness on error-prone devices by fine-tuning models to reduce sensitivity of critical weights to bit flips.

Authors

Habibur Rahaman, Swarup Bhunia

Abstract

The weights of a trained network do not have uniform sensitivity to model accuracy: the 1% most sensitive weights carry about a fifth of the damage a bit flip can do. This distribution decides the bit-flip tolerance of a model, the fraction of weights a code must cover and the bits each weight needs; existing schemes work with the distribution training leaves. We show that we can systematically modulate or alter the weight sensitivity distribution of a trained model, thereby shifting weights from one sensitivity bin to another. The sensitivity of a stored weight is the product of its gradient and the step a flip takes in the storage format, and LEAN adds a composite regularizer on both factors, a range term that bounds the flip step and a tail term that lowers the gradient of the most sensitive weights, with one coefficient λ setting how far the distribution moves. Ten epochs with the term, as a fine-tune or as the last epochs of training, shrink the sensitive tail 5 to 13 times on nine vision models (ResNet-18, ResNet-50, VGG-16 on CIFAR-10, CIFAR-100, ImageNet) within 0.5 points of accuracy; real flips confirm that the reshaped weights do less damage, a plain gradient-norm penalty does not reproduce the effect, and a preliminary extension to three small language models shows the same tail shaping. With a BCH code over 1% of the weights the model then survives 300 progressive bit-search flips in every training seed and 500 in two of three, where three of four published defenses with the same code fail by 300 and WCAT falls to 75% at 500. The reshaped distribution also cuts the parity for a fixed guarantee 3.7 times and makes a sensitivity-ordered 8/6/4/3-bit format affordable, which stores ImageNet ResNet-50 at a 3.29-bit weight payload, 2.4 times smaller than INT8, above the accuracy of LSQ at 3 bits.