Efficient ML #2 — Neural Net Quantization

This post is a follow-up to Data Types in the Deep Learning Era. If the previous post looked at how INT/FP data types interpret bits as numbers, this one actually takes an already-trained model and shrinks it with those low-bit types, then measures the result. All the code comes from the notebooks in github.com/warpspaceinc/efficient-ml-practice, and the plots below were produced by running that code as-is on Colab (T4).

Do we really need all 32 bits?

There was one thing we noticed at the end of the last post: a neural net’s weights are usually clustered in a bell shape around 0, with long tails. And indeed, if you train a small MNIST MLP (784→256→128→10) for just 2 epochs and plot the histogram of the first layer’s weights, that’s exactly what you get.

If the values are packed into such a narrow range, do we really need to store each one as a separate 32-bit float and multiply in 32 bits? Quantization answers “no.” In this post we run the two canonical approaches directly in code.

• K-Means quantization — group similar values under a shared representative to reduce storage (non-uniform).
• Linear quantization — map onto an integer grid so that both storage and compute become integer (uniform).

Both apply to an already-trained model after the fact — this is PTQ (post-training quantization) — and here we only quantize the weights.

1. K-Means Quantization — Sharing Representatives to Cut Storage

The first idea comes from the 2016 Deep Compression paper¹. The pitch is to lump similar weights together and cover them all with a single representative value. For example, 2.09, 1.92, 1.87 are all “roughly 2.0,” so we replace them with a single 2.00 and, at each position, just record an index for “which representative to use.”

The procedure has three steps.

1. Clustering — K-Means cluster all weights into $k = 2^N$ groups.
2. Codebook — store each group’s centroid as its representative value. This list is the codebook.
3. Indices — at each position, instead of a 32-bit float, store only an $N$-bit integer pointing to “which cluster it is.”

In code we just use scikit-learn’s KMeans. To avoid local optima, we initialize with values that split the min–max range uniformly.

from sklearn.cluster import KMeans
import numpy as np

def kmeans_quantize(w, n_bits, fit_sample=20000):
    """Quantize weights with K-Means. Returns the reconstruction, codebook, and indices."""
    shape = w.shape
    flat = w.reshape(-1, 1)
    k = 2 ** n_bits
    lo, hi = float(flat.min()), float(flat.max())
    init = np.linspace(lo, hi, k).reshape(-1, 1)          # linear initialization
    idx = np.random.choice(len(flat), min(fit_sample, len(flat)), replace=False)
    km = KMeans(n_clusters=k, init=init, n_init=1, max_iter=50).fit(flat[idx])
    codebook = km.cluster_centers_.reshape(-1)            # list of representatives
    indices = km.predict(flat)                            # each weight → cluster index
    recon = codebook[indices].reshape(shape)              # reconstruction
    return recon, codebook, indices

If you run it at 2 bits ($k=4$) and overlay the result on the original distribution, you can see the representatives (vertical lines) land densely around 0, where the values are packed. That’s what “non-uniform” means.

How much do we save?

The compression ratio is bit-accounting. For $M$ parameters and $N$-bit quantization:

num = w.size; k = 2 ** N_BITS
orig_bits = 32 * num
comp_bits = N_BITS * num + 32 * k     # indices + codebook
print(f"compression: {orig_bits/comp_bits:.2f}x")

When $M \gg 2^N$, the codebook term becomes negligible and the compression ratio converges to $32M/NM = \mathbf{32/N}$×. That’s about 16× at 2 bits and 8× at 4 bits. (Measured compression ratio for the first layer fc1: 15.99×)

An important limitation: it only cuts storage

This has to be stated clearly. What K-Means quantization reduces is only the size stored on disk/in memory. At inference time you still have to decode the indices through the codebook to reconstruct the original float weights before multiplying.

• Storage: integer indices (small) ✓
• Compute: still floating-point arithmetic with the reconstructed 32-bit floats ✗

We saved memory bandwidth, but the multiplier (MAC) is still running floats. To actually speed up integer compute too, you need the next approach.

2. Linear Quantization — Inference with Integer Arithmetic Only

Linear quantization ties integers and reals together with a single-line affine mapping. It was proposed by Jacob et al. in 2018², and it’s exactly the method behind TensorFlow Lite’s INT8 quantization.

• $r$ — the original real value
• $q$ — the quantized integer
• $Z$ — the zero point. The integer that maps exactly to real $0$. A device for representing common values like the 0 after ReLU, or padding, with no error.
• $S$ — the scale. How much one integer step is worth in real terms (floating-point).

$S$ and $Z$ come from lining up the two ends of the real range $[r_{\min}, r_{\max}]$ with the two ends of the integer range $[q_{\min}, q_{\max}]$.

In code we only need to compute $S$ and $Z$; quantization (round·clamp) and reconstruction (dequant) are handled by PyTorch’s built-in fake_quantize_per_tensor_affine. Arbitrary bit widths are set via quant_min/quant_max.

import torch

def linear_quantize(w, n_bits, symmetric=False):
    """Linear (affine) quantization via PyTorch's built-in fake_quantize. Returns the reconstruction and S, Z."""
    t = torch.as_tensor(w, dtype=torch.float32)
    if symmetric:                                   # Z=0 fixed (simpler compute)
        qmax = 2 ** (n_bits - 1) - 1; qmin = -qmax
        S = float(t.abs().max()) / qmax; Z = 0
    else:                                           # asymmetric: match the range exactly
        qmin = -2 ** (n_bits - 1); qmax = 2 ** (n_bits - 1) - 1
        S = (float(t.max()) - float(t.min())) / (qmax - qmin)
        Z = int(round(qmin - float(t.min()) / S))
    recon = torch.fake_quantize_per_tensor_affine(t, S, Z, qmin, qmax)
    return recon.numpy(), S, Z

Run it at 2 bits and the representatives (vertical lines) land at uniform spacing $S$. This is exactly where it contrasts with K-Means — here the grid is uniform no matter where the data is packed.

Why does “integer arithmetic all the way” become possible?

The real value of Linear quantization isn’t storage but compute. If you rewrite each value of the matrix product $Y = WX$ through the affine mapping:

Everything inside the parentheses, including $q_W q_X$, is integer multiply-and-add, and terms independent of the input like $Z_W Z_X$ can be precomputed. The only real quantity is the leading scale ratio $\frac{S_W S_X}{S_Y}$, and empirically this value is always in $(0,1)$, so it’s handled as a fixed-point multiply + bit shift of the form $2^{-n} M_0$. In the end no floating-point unit is needed at all — the “integer arithmetic is tens of times cheaper than float” from the energy table in Part 1 becomes reality right here.

Since weights are usually symmetric about 0, in practice we often use symmetric quantization with $Z_W = 0$. Then all the $Z_W$-related terms vanish outright and the equation gets even cleaner.

3. So, How Much Do We Lose? — Measure It Directly

Words alone don’t give a feel for it, so we compared the three methods (Linear asymmetric / Linear symmetric / K-Means) across bit widths using MNIST test accuracy. We quantized the weights of all three Linear layers of the model (replacing them with the reconstructions) and then measured accuracy.

def eval_linear(model, n_bits, symmetric):
    m = copy.deepcopy(model)
    with torch.no_grad():
        for layer in [m.fc1, m.fc2, m.fc3]:
            wl = layer.weight.detach().cpu().numpy()
            recon, *_ = linear_quantize(wl, n_bits, symmetric=symmetric)
            layer.weight.copy_(torch.tensor(recon, dtype=torch.float32, device=device))
    return accuracy(m)

bits_list = [2, 3, 4, 8]
acc_asym = [eval_linear(model, b, False) for b in bits_list]
acc_sym  = [eval_linear(model, b, True)  for b in bits_list]
# K-Means compared the same way via eval_kmeans(model, b)

The big picture to take away is this.

• At 8 bits it’s essentially free. All three methods hold accuracy so well you can’t tell them apart from float32.
• The lower the bit width, the more K-Means wins. Because it places representatives to match the data distribution (non-uniform), at the same 2 bits its error is smaller than Linear’s uniform grid. This is even clearer in the MSE curve below.

But accuracy isn’t everything. Even if K-Means has smaller error at 2 bits, its inference compute is still float. Linear accepts a bit more error in exchange for making the compute itself integer, which is what actually speeds things up on real hardware. The two methods aren’t competitors — they’re tools for different purposes.

Try It Yourself

All the plots and tables above came from running the two notebooks below on Colab. Runtime → Run all is all it takes, and you can apply it to your own model just by uploading its .safetensors and changing the path.

• 📓 linear-quantization.ipynb — affine mapping, symmetric/asymmetric
• 📓 kmeans-quantization.ipynb — codebook/indices, compression ratio

Wrap-up

• Quantization = reducing the number of representable values. Since weights are packed around 0, using all 32 bits is wasteful.
• K-Means (non-uniform) places representatives to match the distribution, cutting storage but keeping compute in float. Its error is small at low bit widths.
• Linear (uniform) turns both storage and compute into integers via the affine mapping $r=S(q-Z)$. It’s the foundation of hardware acceleration.
• 8 bits is nearly free; push all the way to 2 bits and the choice of method decides your accuracy.

In the next experiment note, we’ll take on activation quantization and QAT (quantization-aware training), which we couldn’t cover here, and build the integer inference pipeline all the way to the end.