Model Compression: Quantization and Pruning Explained

Large models cost memory, download time, latency and money. Model compression trades a little of a model's quality for a lot of those savings, and three techniques do most of the work. Quantization stores numbers with fewer bits. Pruning removes weights. Distillation trains a smaller model to imitate a bigger one.

Each promises smaller and faster models. This guide measures what each actually did to a real sentence encoder, including one result that ran the wrong way.

The three techniques

The classic demonstration of combining them is Deep Compression: pruning, trained quantization and Huffman coding together cut AlexNet from 240 MB to 6.9 MB and VGG-16 from 552 MB to 11.3 MB without losing accuracy, with retraining after each step (Han et al.: Deep Compression).

The tools today

PyTorch is deprecating its built-in torch.ao.quantization module and moving quantization into the separate torchao library (PyTorch: deprecation tracker). torchao's main entry point is quantize_, which swaps a model's layers for quantized versions in place; its own first example runs on a CUDA GPU in bfloat16 and compiles the model with torch.compile (torchao: first example). Pruning lives in torch.nn.utils.prune (PyTorch: pruning tutorial).

The experiment

The model is all-MiniLM-L6-v2, a sentence encoder that maps text to 384-dimensional vectors, released under Apache 2.0 (Hugging Face: all-MiniLM-L6-v2). The task is Banking77: 77 customer intents, with 10,003 training and 3,080 test queries, under CC BY 4.0 (PolyAI: Banking77). A logistic regression classifier is trained once on the float32 encoder's vectors, then each compressed encoder is swapped in without retraining it, the way a compressed model replaces a live one.

import gzip
import io
import time

import pandas as pd
import torch
import torch.nn.utils.prune as prune
from sklearn.linear_model import LogisticRegression
from torchao.quantization import Int8DynamicActivationInt8WeightConfig, Int8WeightOnlyConfig, quantize_
from transformers import AutoModel, AutoTokenizer

NAME = "sentence-transformers/all-MiniLM-L6-v2"                  # Apache 2.0, 384-dimensional sentence embeddings
tokenizer = AutoTokenizer.from_pretrained(NAME)
train, test = pd.read_csv("train.csv"), pd.read_csv("test.csv")  # Banking77 (CC BY 4.0)


@torch.no_grad()
def embed(model, texts, batch=64):                               # mean pooling + normalisation, as the model card shows
    vectors = []
    for start in range(0, len(texts), batch):
        tokens = tokenizer(list(texts[start:start + batch]), padding=True, truncation=True, return_tensors="pt")
        hidden, mask = model(**tokens).last_hidden_state, tokens["attention_mask"].unsqueeze(-1).float()
        vectors.append(torch.nn.functional.normalize((hidden * mask).sum(1) / mask.sum(1), dim=1))
    return torch.cat(vectors)


def megabytes(model, gzipped=False):
    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    data = buffer.getvalue()
    return len(gzip.compress(data) if gzipped else data) / 2**20


def prune_linear(model, share):                                  # zero the smallest weights across all Linear layers
    weights = [(m, "weight") for m in model.modules() if isinstance(m, torch.nn.Linear)]
    prune.global_unstructured(weights, pruning_method=prune.L1Unstructured, amount=share)
    for module, name in weights:
        prune.remove(module, name)                               # make the zeros permanent
    return model


def quantized(model, config):                                  # torchao swaps Linear weights for int8 versions in place
    quantize_(model, config)
    return model


base = AutoModel.from_pretrained(NAME).eval()
table = base.embeddings.word_embeddings.weight
print(f"{sum(p.numel() for p in base.parameters()) / 1e6:.1f}M parameters; the token-embedding table alone is "
      f"{table.numel() * 4 / 2**20:.1f} MB of the {megabytes(base):.1f} MB float32 file")
reference = embed(base, test.text)
classifier = LogisticRegression(max_iter=3000, C=10).fit(embed(base, train.text).numpy(), train.category)

variants = {"float32": lambda m: m,
            "int8 weights": lambda m: quantized(m, Int8WeightOnlyConfig()),
            "int8 weights and activations": lambda m: quantized(m, Int8DynamicActivationInt8WeightConfig()),
            "pruned 30%": lambda m: prune_linear(m, 0.3), "pruned 50%": lambda m: prune_linear(m, 0.5),
            "pruned 70%": lambda m: prune_linear(m, 0.7)}
print(f"\n{'variant':<29} {'MB':>5} {'gzip MB':>8} {'cosine to float32':>18} {'accuracy':>9}")
for name, compress in variants.items():
    model = compress(AutoModel.from_pretrained(NAME).eval())
    start = time.perf_counter()
    vectors = embed(model, test.text)
    seconds = time.perf_counter() - start
    accuracy = (classifier.predict(vectors.numpy()) == test.category).mean()
    print(f"{name:<29} {megabytes(model):>5.1f} {megabytes(model, gzipped=True):>8.1f} "
          f"{(vectors * reference).sum(1).mean().item():>18.4f} {accuracy:>9.2%}")
    print(f"{'':<29} embedded {len(test):,} queries in {seconds:.1f} s (varies by machine)")

22.7M parameters; the token-embedding table alone is 44.7 MB of the 86.7 MB float32 file

variant                          MB  gzip MB  cosine to float32  accuracy
float32                        86.7     79.3             1.0000    92.95%
                              embedded 3,080 queries in 11.4 s (varies by machine)
int8 weights                   56.0     50.7             0.9996    92.99%
                              embedded 3,080 queries in 12.2 s (varies by machine)
int8 weights and activations   56.0     50.7             0.9976    92.92%
                              embedded 3,080 queries in 1428.3 s (varies by machine)
pruned 30%                     86.7     70.9             0.9387    92.47%
                              embedded 3,080 queries in 24.5 s (varies by machine)
pruned 50%                     86.7     63.8             0.7506    86.04%
                              embedded 3,080 queries in 24.6 s (varies by machine)
pruned 70%                     86.7     56.6             0.4065    47.40%
                              embedded 3,080 queries in 20.1 s (varies by machine)

What the output shows

• Half the file was one table. The token-embedding table took 44.7 MB of the 86.7 MB float32 file. The int8 methods here quantize Linear layers only, so they could never take the file below about 56 MB; embedding tables need their own treatment.
• Int8 weights were free. The file fell from 86.7 MB to 56.0 MB, the new embeddings matched the old ones with an average cosine similarity of 0.9996, and accuracy was 92.99% against 92.95%, within noise. Speed did not change on this CPU.
• Quantizing activations as well was a trap here. Accuracy held at 92.92%, but embedding the 3,080 queries took 1,428 seconds instead of 11, about 125 times slower. torchao's own example compiles the model on a CUDA GPU; on this CPU, in eager mode and uncompiled, the same configuration ran on slow fallback code. Quantization speeds a model up only where fast kernels exist.
• Pruning without retraining degraded quickly. Removing the smallest 30% of Linear weights cost half a point, 92.47%. Removing 50% cost almost seven points, 86.04%. At 70%, accuracy collapsed to 47.40% and the embeddings barely resembled the originals, with a cosine similarity of 0.41.
• Pruned files were no smaller. Every pruned file stayed at 86.7 MB, because the zeros are still stored. Only compression exposed them: 70% pruning brought the gzip size to 56.6 MB, still more than the int8 file's 50.7 MB, at far worse accuracy. Deep Compression retrained after pruning for exactly this reason.

Distillation: the compression you choose up front

The encoder itself is the product of distillation. It is built on MiniLM, whose authors trained small students to mimic a large teacher's self-attention, and kept more than 99% of the teacher's accuracy on SQuAD 2.0 and several GLUE tasks with half its Transformer parameters and computation (Wang et al.: MiniLM). Choosing a distilled model before training anything is often the largest saving available, and quantization can still be applied on top.

Choosing

• Limited by size or memory: quantize the weights first. It is cheap and, here, cost nothing measurable.
• Limited by latency: measure on the target hardware with its compiled or vendor kernels before promising a speed-up.
• Willing to retrain: distil, or prune and then fine-tune, as Deep Compression did.
• Always: check accuracy on the task and on important slices, not only the size.

For running compressed models on phones and devices, see edge ML deployment. For saving and shipping model files, see model packaging, and for whether a GPU would help more, see GPU vs CPU for ML.

Common mistakes

• Assuming quantized means faster.
• Pruning without fine-tuning and accepting the accuracy drop.
• Reporting dense file sizes for pruned models as if zeros were free.
• Checking only overall accuracy after compression.
• Ignoring layers the method skips, such as embedding tables.

Model serialisation and deployment are covered in the Model Deployment module of the MLOps course.

Sources

• Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding (arXiv) — Han, Mao and Dally
• MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers (arXiv) — Wang et al.
• torch.ao.quantization deprecation tracker (pytorch/ao, GitHub) — PyTorch
• First Quantization Example (torchao documentation) — PyTorch
• Pruning Tutorial (PyTorch) — PyTorch
• all-MiniLM-L6-v2 (Hugging Face) — sentence-transformers
• Task-specific datasets, including Banking77 (GitHub, PolyAI) — PolyAI