PyTorch

Starting from LeCun’s 1989 Optimal Brain Damage, this is the story of shrinking a model by cutting weights away (pruning). What to cut, how, and how much; the surprise of compression that survives even a 95% cut thanks to retraining; and the point that you only get a real payoff when the hardware (NVIDIA 2:4) backs you up — all with plots we measured ourselves on an MNIST MLP.

We take the trained 32-bit weights of a neural net and shrink them down to 2–8 bits using two approaches: K-Means (non-uniform) and Linear (uniform, integer arithmetic). From the affine mapping r=S(q−Z) all the way to the compression-ratio vs. accuracy trade-off, with plots measured by actually running the code on an MNIST MLP.