Machine Learning

Mannequin Compression: Make Your Machine Studying Fashions Lighter and Quicker

May 9, 2025
roosho

Whether or not you’re getting ready for interviews or constructing Machine Studying programs at your job, mannequin compression has turn into essential talent. Within the period of LLMs, the place fashions are getting bigger and bigger, the challenges round compressing these fashions to make them extra environment friendly, smaller, and usable on light-weight machines have by no means been extra related.

On this article, I’ll undergo 4 basic compression methods that each ML practitioner ought to perceive and grasp. I discover pruning, quantization, low-rank factorization, and Information Distillation, every providing distinctive benefits. I will even add some minimal PyTorch code samples for every of those strategies.

I hope you benefit from the article!

Mannequin pruning

Pruning might be probably the most intuitive compression approach. The thought may be very easy: take away a few of the weights of the community, both randomly or take away the “much less essential” ones. After all, once we speak about “eradicating” weights within the context of neural networks, it means setting the weights to zero.

Mannequin pruning (Picture by the writer and ChatGPT | Inspiration: [3])

Structured vs unstructured pruning

Let’s begin with a easy heuristic: eradicating weights smaller than a threshold.

[ w’_{ij} = begin{cases} w_{ij} & text{if } |w_{ij}| ge theta_0
0 & text{if } |w_{ij}| < theta_0
end{cases} ]

After all, this isn’t ultimate as a result of we would wish to discover a method to discover the fitting threshold for our downside! A extra sensible method is to take away a specified proportion of weights with the smallest magnitudes (norm) inside one layer. There are 2 widespread methods of implementing pruning in a single layer:

  • Structured pruning: take away whole elements of the community (e.g. a random row from the load tensor, or a random channel in a convulational layer)
  • Unstructured pruning: take away particular person weights no matter their positions and of the construction of the tensor

We are able to additionally use world pruning with both of the 2 above strategies. This may take away the chosen proportion of weights throughout a number of layers, and probably have completely different elimination charges relying on the variety of parameters in every layer.

PyTorch makes this gorgeous easy (by the way in which, you’ll find all code snippets in my GitHub repo).

import torch.nn.utils.prune as prune

# 1. Random unstructured pruning (20% of weights at random)
prune.random_unstructured(mannequin.layer, identify="weight", quantity=0.2)                           

# 2. L1‑norm unstructured pruning (20% of smallest weights)
prune.l1_unstructured(mannequin.layer, identify="weight", quantity=0.2)

# 3. International unstructured pruning (40% of all weights by L1 norm throughout layers)
prune.global_unstructured(
    [(model.layer1, "weight"), (model.layer2, "weight")],
    pruning_method=prune.L1Unstructured,
    quantity=0.4
)                                             

# 4. Structured pruning (take away 30% of rows with lowest L2 norm)
prune.ln_structured(mannequin.layer, identify="weight", quantity=0.3, n=2, dim=0)

Notice: when you’ve got taken statistics courses, you in all probability realized regularization-induced strategies that additionally implicitly prune some weights throughout coaching, through the use of L0 or L1 norm regularization. Pruning differs from that as a result of it’s utilized as a post-Mannequin Compression approach

Why does pruning work? The Lottery Ticket Speculation

Picture generated by ChatGPT

I wish to conclude that part with a fast point out of the Lottery Ticket Speculation, which is each an utility of pruning and an fascinating clarification of how eradicating weights can typically enhance a mannequin. I like to recommend studying the related paper ([7]) for extra particulars.

Authors use the next process:

  1. Practice the complete mannequin to convergence
  2. Prune the smallest-magnitude weights (say 10%)
  3. Reset the remaining weights to their unique initialization values
  4. Retrain this pruned community
  5. Repeat the method a number of instances

After doing this 30 instances, you find yourself with solely 0.930 ~ 4% of the unique parameters. And surprisingly, this community can do in addition to the unique one.

This implies that there’s essential parameter redundancy. In different phrases, there exists a sub-network (“a lottery ticket”) that truly does a lot of the work!

Pruning is one method to unveil this sub-network.

I like to recommend this excellent video that covers the subject!

Quantization

Whereas pruning focuses on eradicating parameters completely, Quantization takes a distinct method: decreasing the precision of every parameter.

Keep in mind that each quantity in a pc is saved as a sequence of bits. A float32 worth makes use of 32 bits (see instance image beneath), whereas an 8-bit integer (int8) makes use of simply 8 bits.

An instance of how float32 numbers are represented with 32 bits (Picture by the writer and ChatGPT | Inspiration: [2])

Most deep studying fashions are skilled utilizing 32-bit floating-point numbers (FP32). Quantization converts these high-precision values to lower-precision codecs like 16-bit floating-point (FP16), 8-bit integers (INT8), and even 4-bit representations.

The financial savings listed below are apparent: INT8 requires 75% much less reminiscence than FP32. However how can we really carry out this conversion with out destroying our mannequin’s efficiency?

The maths behind quantization

To transform from floating-point to integer illustration, we have to map the continual vary of values to a discrete set of integers. For INT8 quantization, we’re mapping to 256 doable values (from -128 to 127).

Suppose our weights are normalized between -1.0 and 1.0 (widespread in deep studying):

[ text{scale} = frac{text{float_max} – text{float_min}}{text{int8_max} – text{int8_min}} = frac{1.0 – (-1.0)}{127 – (-128)} = frac{2.0}{255} ]

Then, the quantized worth is given by

[text{quantized_value} = text{round}(frac{text{original_value}}{text{scale}} ] + textual content{zero_point})

Right here, zero_point=0 as a result of we would like 0 to be mapped to 0. We are able to then spherical this worth to the closest integer to get integers between -127 and 128.

And, you guessed it: to get integers again to drift, we are able to use the inverse operation: [text{float_value} = text{integer_value} times text{scale} – text{zero_point} ]

Notice: in apply, the scaling issue is decided primarily based on the vary values we quantize.

Easy methods to apply quantization?

Quantization may be utilized at completely different levels and with completely different methods. Listed below are just a few methods value realizing about: (beneath, the phrase “activation” refers back to the output values of every layer)

  • Publish-training quantization (PTQ):
    • Static Quantization: quantize each weights and activations offline (after coaching and earlier than inference)
    • Dynamic Quantization: quantize weights offline, however activations on-the-fly throughout inference. That is completely different from offline quantization as a result of the scaling issue is decided primarily based on the values seen to date throughout inference.
  • Quantize-aware coaching (QAT): simulate quantization throughout coaching by rounding values, however calculations are nonetheless accomplished with floating-point numbers. This makes the mannequin study weights which might be extra strong to quantization, which can be utilized after coaching. Below the hood, the concept is to add “faux” operations: x -> dequantize(quantize(x)): this new worth is near x, nevertheless it nonetheless helps the mannequin tolerate the 8-bit rounding and clipping noise.
import torch.quantization as tq

# 1. Publish‑coaching static quantization (weights + activations offline)
mannequin.eval()
mannequin.qconfig = tq.get_default_qconfig('fbgemm') # assign a static quantization config
tq.put together(mannequin, inplace=True)
# we have to use a calibration dataset to find out the ranges of values
with torch.no_grad():
    for knowledge, _ in calibration_data:
        mannequin(knowledge)
tq.convert(mannequin, inplace=True) # convert to a totally int8 mannequin

# 2. Publish‑coaching dynamic quantization (weights offline, activations on‑the‑fly)
dynamic_model = tq.quantize_dynamic(
    mannequin,
    {torch.nn.Linear, torch.nn.LSTM}, # layers to quantize
    dtype=torch.qint8
)

# 3. Quantization‑Conscious Coaching (QAT)
mannequin.prepare()
mannequin.qconfig = tq.get_default_qat_qconfig('fbgemm')  # arrange QAT config
tq.prepare_qat(mannequin, inplace=True) # insert faux‑quant modules
# [here, train or fine‑tune the model as usual]
qat_model = tq.convert(mannequin.eval(), inplace=False) # convert to actual int8 after QAT

Quantization may be very versatile! You’ll be able to apply completely different precision ranges to completely different components of the mannequin. As an illustration, you would possibly quantize most linear layers to 8-bit for max velocity and reminiscence financial savings, whereas leaving important elements (e.g. consideration heads, or batch-norm layers) at 16-bit or full-precision.

Low-Rank Factorization

Now let’s speak about low-rank factorization — a way that has been popularized with the rise of LLMs.

The important thing statement: many weight matrices in neural networks have efficient ranks a lot decrease than their dimensions counsel. In plain English, which means there may be a variety of redundancy within the parameters.

Notice: when you’ve got ever used PCA for dimensionality discount, you’ve gotten already encountered a type of low-rank approximation. PCA decomposes massive matrices into merchandise of smaller, lower-rank elements that retain as a lot info as doable.

The linear algebra behind low-rank factorization

Take a weight matrix W. Each actual matrix may be represented utilizing a Singular Worth Decomposition (SVD):

[ W = USigma V^T ]

the place Σ is a diagonal matrix with singular values in non-increasing order. The variety of constructive coefficients really corresponds to the rank of the matrix W.

SVD visualized for a matrix of rank r (Picture by the writer and ChatGPT | Inspiration: [5])

To approximate W with a matrix of rank ok < r, we are able to choose the ok biggest parts of sigma, and the corresponding first ok columns and first ok rows of U and V respectively:

[ begin{aligned} W_k &= U_k,Sigma_k,V_k^T
[6pt] &= underbrace{U_k,Sigma_k^{1/2}}_{Ainmathbb{R}^{mtimes ok}} underbrace{Sigma_k^{1/2},V_k^T}_{Binmathbb{R}^{ktimes n}}. finish{aligned} ]

See how the brand new matrix may be decomposed because the product of A and B, with the entire variety of parameters now being m * ok + ok * n = ok*(m+n) as an alternative of m*n! This can be a large enchancment, particularly when ok is far smaller than m and n.

In apply, it’s equal to changing a linear layer x → Wx with 2 consecutive ones: x → A(Bx).

In PyTorch

We are able to both apply low-rank factorization earlier than coaching (parameterizing every linear layer as two smaller matrices – not likely a compression methodology, however a design alternative) or after coaching (making use of a truncated SVD on weight matrices). The second method is by far the most typical one and is applied beneath.

import torch

# 1. Extract weight and select rank
W = mannequin.layer.weight.knowledge # (m, n)
ok = 64 # desired rank

# 2. Approximate low-rank SVD
U, S, V = torch.svd_lowrank(W, q=ok) # U: (m, ok), S: (ok, ok), V: (n, ok)

# 3. Kind elements A and B
A = U * S.sqrt() # [m, k]
B = V.t() * S.sqrt().unsqueeze(1) # [k, n]

# 4. Exchange with two linear layers and insert the matrices A and B
orig = mannequin.layer
mannequin.layer = torch.nn.Sequential(
    torch.nn.Linear(orig.in_features, ok, bias=False),
    torch.nn.Linear(ok, orig.out_features, bias=False),
)
mannequin.layer[0].weight.knowledge.copy_(B)
mannequin.layer[1].weight.knowledge.copy_(A)

LoRA: an utility of low-rank approximation

LoRA fine-tuning: W is mounted, A and B are skilled (supply: [1])

I believe it’s essential to say LoRA: you’ve gotten in all probability heard of LoRA (Low-Rank Adaptation) when you’ve got been following LLM fine-tuning developments. Although not strictly a compression approach, LoRA has turn into extraordinarily widespread for effectively adapting massive language fashions and making fine-tuning very environment friendly.

The thought is easy: throughout fine-tuning, moderately than modifying the unique mannequin weights W, LoRA freezes them and study trainable low-rank updates:

$$W’ = W + Delta W = W + AB$$

the place A and B are low-rank matrices. This enables for task-specific adaptation with only a fraction of the parameters.

Even higher: QLoRA takes this additional by combining quantization with low-rank adaptation!

Once more, this can be a very versatile approach and may be utilized at varied levels. Normally, LoRA is utilized solely on particular layers (for instance, Consideration layers’ weights).

Information Distillation

Information distillation course of (Picture by the writer and ChatGPT | Inspiration: [4])

Information distillation takes a basically completely different method from what we have now seen to date. As an alternative of modifying an present mannequin’s parameters, it transfers the “data” from a massive, complicated mannequin (the “trainer”) to a smaller, extra environment friendly mannequin (the “pupil”). The purpose is to coach the coed mannequin to mimic the habits and replicate the efficiency of the trainer, typically a neater process than fixing the unique downside from scratch.

The distillation loss

Let’s clarify some ideas within the case of a classification downside:

  • The trainer mannequin is normally a big, complicated mannequin that achieves excessive efficiency on the duty at hand
  • The pupil mannequin is a second, smaller mannequin with a distinct structure, however tailor-made to the identical process
  • Mushy targets: these are the trainer’s mannequin predictions (possibilities, and never labels!). They are going to be utilized by the coed mannequin to imitate the trainer’s behaviors. Notice that we use uncooked predictions and never labels as a result of in addition they comprise details about the boldness of the predictions
  • Temperature: along with the trainer’s prediction, we additionally use a coefficient T (referred to as temperature) within the softmax operate to extract extra info from the smooth targets. Growing T softens the distribution and helps the coed mannequin give extra significance to improper predictions.

In apply, it’s fairly easy to coach the coed mannequin. We mix the standard loss (normal cross-entropy loss primarily based on laborious labels) with the “distillation” loss (primarily based on the trainer’s smooth targets):

$$ L_{textual content{complete}} = alpha L_{textual content{laborious}} + (1 – alpha) L_{textual content{distill}} $$

The distillation loss is nothing however the KL divergence between the trainer and pupil distribution (you may see it as a measure of the space between the two distributions).

$$ L_{textual content{distill}} = D{KL}(q_{textual content{trainer}} | | q_{textual content{pupil}}) = sum_i q_{textual content{trainer}, i} log left( frac{q_{textual content{trainer}, i}}{q_{textual content{pupil}, i}} proper) $$

As for the opposite strategies, it’s doable and inspired to adapt this framework relying on the use case: for instance, one can even examine logits and activations from intermediate layers within the community between the coed and trainer mannequin, as an alternative of solely evaluating the ultimate outputs.

Information distillation in apply

Much like the earlier methods, there are two choices:

  • Offline distillation: the pre-trained trainer mannequin is mounted, and a separate pupil mannequin is skilled to imitate it. Each fashions are utterly separate, and the trainer’s weights stay frozen through the distillation course of.
  • On-line distillation: each fashions are skilled concurrently, with data switch taking place through the joint coaching course of.

And beneath, a straightforward method to apply offline distillation (the final code block of this text 🙂):

import torch.nn.useful as F

def distillation_loss_fn(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # Commonplace Cross-Entropy loss with laborious labels
    student_loss = F.cross_entropy(student_logits, labels)

    # Distillation loss with smooth targets (KL Divergence)
    soft_teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    soft_student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)

		# kl_div expects log possibilities as enter for the primary argument!
    distill_loss = F.kl_div(
        soft_student_log_probs,
        soft_teacher_probs.detach(), # do not calculate gradients for trainer
        discount='batchmean'
    ) * (temperature ** 2) # elective, a scaling issue

    # Mix losses based on components
    total_loss = alpha * student_loss + (1 - alpha) * distill_loss
    return total_loss

teacher_model.eval()
student_model.prepare()
with torch.no_grad():
     teacher_logits = teacher_model(inputs)
	 student_logits = student_model(inputs)
	 loss = distillation_loss_fn(student_logits, teacher_logits, labels, temperature=T, alpha=alpha)
	 loss.backward()
	 optimizer.step()

Conclusion

Thanks for studying this text! Within the period of LLMs, with billions and even trillions of parameters, mannequin compression has turn into a basic idea, important in nearly each state of affairs to make fashions extra environment friendly and simply deployable.

However as we have now seen, mannequin compression isn’t nearly decreasing the mannequin measurement – it’s about making considerate design choices. Whether or not selecting between on-line and offline strategies, compressing the whole community, or focusing on particular layers or channels, every alternative considerably impacts efficiency and value. Most fashions now mix a number of of those methods (try this mannequin, for example).

Past introducing you to the principle strategies, I hope this text additionally conjures up you to experiment and develop your personal inventive options!

Don’t neglect to take a look at the GitHub repository, the place you’ll discover all of the code snippets and a side-by-side comparability of the 4 compression strategies mentioned on this article.

Take a look at my earlier articles:

References