Structured Pruning › Code Implementation
Language AI HandbookAll chapters Structured Pruning › Code ImplementationOne section · Sections in this chapter7 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We implement structured pruning for a transformer model, focusing on head pruning and layer pruning with gradient-based importance scores.
Setting Up the Model
We start by building a simple multi-layer transformer encoder that we can prune. This implementation makes head pruning explicit by treating each head independently.
from copy import deepcopy
import numpy as np
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads, d_k=None):
super().__init__()
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_k if d_k is not None else (d_model // num_heads)
d_k_total = self.num_heads * self.d_k
self.W_q = nn.Linear(d_model, d_k_total, bias=False)
self.W_k = nn.Linear(d_model, d_k_total, bias=False)
self.W_v = nn.Linear(d_model, d_k_total, bias=False)
self.W_o = nn.Linear(d_k_total, d_model, bias=False)
def forward(self, x):
B, T, _ = x.shape
H, dk = self.num_heads, self.d_k
Q = self.W_q(x).view(B, T, H, dk).transpose(1, 2)
K = self.W_k(x).view(B, T, H, dk).transpose(1, 2)
V = self.W_v(x).view(B, T, H, dk).transpose(1, 2)
scores = (Q @ K.transpose(-2, -1)) / (dk**0.5)
attn = torch.softmax(scores, dim=-1)
out = (attn @ V).transpose(1, 2).contiguous().view(B, T, H * dk)
return self.W_o(out)
class TransformerBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
self.attn = MultiHeadAttention(d_model, num_heads)
self.ff1 = nn.Linear(d_model, d_ff)
self.ff2 = nn.Linear(d_ff, d_model)
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.ff2(torch.relu(self.ff1(self.ln2(x))))
return x
class SimpleTransformerEncoder(nn.Module):
def __init__(
self,
d_model=128,
num_heads=8,
d_ff=256,
num_layers=6,
num_classes=2,
vocab_size=1000,
max_len=32,
):
super().__init__()
self.embed = nn.Embedding(vocab_size, d_model)
self.blocks = nn.ModuleList(
[
TransformerBlock(d_model, num_heads, d_ff)
for _ in range(num_layers)
]
)
self.classifier = nn.Linear(d_model, num_classes)
def forward(self, x):
h = self.embed(x)
for block in self.blocks:
h = block(h)
return self.classifier(h.mean(dim=1))
torch.manual_seed(42)
model = SimpleTransformerEncoder(
d_model=128, num_heads=8, d_ff=256, num_layers=6
)Original model parameters: 920,066 Number of layers: 6 Heads per layer: 8
The model starts with a fixed number of parameters spread across 6 transformer blocks, each with 8 attention heads. Even at this small scale, the structural relationships between layers and heads are the same as in production-scale models: each block contains an attention sublayer and a feedforward sublayer connected by residual connections.
Computing Layer Importance Scores
We compute the cosine similarity between input and output for each layer. Layers where the hidden state barely changes are candidates for removal. The function below passes a batch through each block sequentially, saving the input before each block and comparing it to the output.
def compute_layer_similarities(mdl, data_loader, device="cpu"):
"""Compute mean cosine similarity between input and output of each block."""
mdl.eval()
sims = [[] for _ in range(len(mdl.blocks))]
with torch.no_grad():
for x, _ in data_loader:
x = x.to(device)
h = mdl.embed(x)
for i, block in enumerate(mdl.blocks):
h_in = h.detach().clone()
h = block(h)
cos_sim = nn.functional.cosine_similarity(
h_in.reshape(-1, h_in.shape[-1]),
h.reshape(-1, h.shape[-1]),
dim=-1,
)
sims[i].append(cos_sim.mean().item())
return [float(np.mean(s)) for s in sims]
def make_synthetic_data(
n=512, seq_len=16, vocab_size=1000, num_classes=2, seed=42
):
rng = np.random.default_rng(seed)
X = rng.integers(0, vocab_size, (n, seq_len))
y = rng.integers(0, num_classes, n)
dataset = torch.utils.data.TensorDataset(
torch.tensor(X, dtype=torch.long), torch.tensor(y, dtype=torch.long)
)
return torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=False)
data_loader = make_synthetic_data()
layer_sims = compute_layer_similarities(model, data_loader)Layer-wise cosine similarity (input vs output): Layer 1: 0.9670 ############################# Layer 2: 0.9714 ############################# Layer 3: 0.9699 ############################# Layer 4: 0.9688 ############################# Layer 5: 0.9722 ############################# Layer 6: 0.9722 ############################# Pruning candidates (similarity > 0.85): layers [1, 2, 3, 4, 5, 6]
Layers with similarity above the threshold transform the representation very little and are the best candidates for removal. The text bar chart in the output gives an immediate visual sense of which layers are active and which are near-identity.
The bar chart below shows the cosine similarity for each layer, making the near-identity layers visually obvious.
Computing Head Importance Scores
For head pruning, we use gradient-based importance scores. We run a forward pass to compute the loss, backpropagate to get gradients, and compute the importance of each head using the gradient-weight product norm. This is the Taylor expansion criterion: it approximates how much the loss would change if that head's output were zeroed out.
import torch.nn.functional as F
def compute_head_importance(mdl, data_loader, device="cpu"):
"""Compute importance of each attention head using gradient-weight product norm."""
mdl.eval()
num_layers = len(mdl.blocks)
num_heads = mdl.blocks[0].attn.num_heads
importance = torch.zeros(num_layers, num_heads)
for x, y in data_loader:
x, y = x.to(device), y.to(device)
mdl.zero_grad()
logits = mdl(x)
loss = F.cross_entropy(logits, y)
loss.backward()
for l, block in enumerate(mdl.blocks):
W_o = block.attn.W_o.weight
grad = block.attn.W_o.weight.grad
if grad is None:
continue
dk = block.attn.d_k
for h in range(num_heads):
start = h * dk
end = start + dk
score = (grad[:, start:end] * W_o[:, start:end]).norm()
importance[l, h] += score.item()
importance = importance / max(1, len(data_loader))
return importance
head_importance = compute_head_importance(model, data_loader)Head importance scores (gradient-weight norm): Layer 1: min=0.0018, max=0.0020 (least important: heads [1, 4]) Layer 2: min=0.0019, max=0.0025 (least important: heads [5, 8]) Layer 3: min=0.0021, max=0.0025 (least important: heads [2, 6]) Layer 4: min=0.0023, max=0.0032 (least important: heads [4, 5]) Layer 5: min=0.0024, max=0.0029 (least important: heads [7, 6]) Layer 6: min=0.0022, max=0.0027 (least important: heads [5, 3])
Each row shows the range of head importance scores in a layer. The least important heads have scores well below the layer average and are removed first. Notice that the ratio between the minimum and maximum importance scores within a single layer can be substantial, confirming that heads within a layer are not equally important.
The heatmap below visualizes importance scores across all layers and heads, making it easy to see which heads cluster at the bottom of the importance distribution and should be removed first.
Performing Layer Pruning
With layer similarities computed, we remove the least important layers by creating a new model with a reduced number of transformer blocks. The pruning function deep-copies the model to preserve the original, then filters out the blocks at the specified indices.
def prune_layers(mdl, layers_to_remove):
"""Create a new model with the specified layers removed."""
pruned = deepcopy(mdl)
remaining = [
block
for i, block in enumerate(pruned.blocks)
if i not in layers_to_remove
]
pruned.blocks = nn.ModuleList(remaining)
return pruned
high_sim_threshold = 0.85
layers_to_remove = sorted(
[i for i, s in enumerate(layer_sims) if s > high_sim_threshold],
reverse=True,
)[:2]
model_layer_pruned = prune_layers(model, layers_to_remove)Layers removed: [6, 5] Remaining layers: 4 Original parameters: 920,066 Layer-pruned parameters: 656,130 Parameter reduction: 28.7%
Removing layers that apply near-identity transformations eliminates their parameters while preserving the model's functional behavior as closely as possible. The parameter reduction is proportional to the number of layers removed divided by the total number of layers, which for two layers out of six is roughly 33%.
Performing Head Pruning
We implement head pruning by physically removing the head-corresponding slices from the attention weight matrices. The function below handles the index arithmetic carefully: each head occupies a contiguous slice of columns in the Q, K, V matrices and a contiguous slice of rows in the output projection.
def prune_heads_in_layer(attn_module, heads_to_remove):
"""
Physically remove heads from a MultiHeadAttention module by slicing weight matrices.
Ensures at least 1 head is retained to avoid zero-dimension matrices.
"""
H = attn_module.num_heads
dk = attn_module.d_k
d_model = attn_module.d_model
keep_heads = [h for h in range(H) if h not in heads_to_remove]
# Guard: always keep at least 1 head
if len(keep_heads) == 0:
keep_heads = [0]
new_H = len(keep_heads)
keep_cols = []
for h in keep_heads:
keep_cols.extend(range(h * dk, (h + 1) * dk))
# Pass d_k explicitly so __init__ does not recompute it from d_model // num_heads
new_attn = MultiHeadAttention(d_model, new_H, d_k=dk)
with torch.no_grad():
new_attn.W_q.weight.copy_(attn_module.W_q.weight[keep_cols, :])
new_attn.W_k.weight.copy_(attn_module.W_k.weight[keep_cols, :])
new_attn.W_v.weight.copy_(attn_module.W_v.weight[keep_cols, :])
new_attn.W_o.weight.copy_(attn_module.W_o.weight[:, keep_cols])
return new_attn
def prune_heads(mdl, importance, prune_fraction=0.25):
"""Prune the least important heads globally across all layers."""
pruned = deepcopy(mdl)
num_layers = len(pruned.blocks)
num_heads = pruned.blocks[0].attn.num_heads
total_heads = num_layers * num_heads
n_prune = max(1, int(total_heads * prune_fraction))
flat = [
(importance[l, h].item(), l, h)
for l in range(num_layers)
for h in range(num_heads)
]
flat_sorted = sorted(flat, key=lambda x: x[0])
to_remove_flat = flat_sorted[:n_prune]
heads_by_layer = {}
for _, l, h in to_remove_flat:
heads_by_layer.setdefault(l, []).append(h)
for l, heads in heads_by_layer.items():
pruned.blocks[l].attn = prune_heads_in_layer(
pruned.blocks[l].attn, heads
)
return pruned, heads_by_layer
model_head_pruned, removed_heads = prune_heads(
model, head_importance, prune_fraction=0.25
)Removed heads per layer: Layer 1: heads [1, 2, 3, 4, 5, 6, 7, 8] Layer 2: heads [2, 5, 8] Layer 3: heads [2] Original parameters: 920,066 Head-pruned parameters: 829,954 Parameter reduction: 9.8%
The pruned model has fewer parameters in its attention layers, with the size reduction proportional to the number of heads removed. The global selection strategy means the removed heads are concentrated in layers where importance scores are most unevenly distributed, which is exactly the desired behavior.
Combining Layer and Head Pruning
In practice, you apply layer pruning and head pruning together to maximize compression while minimizing quality loss. The combined function first removes low-similarity layers, then recomputes head importance on the smaller model and removes low-importance heads.
def prune_combined(
mdl,
data_loader,
layer_sim_threshold=0.85,
head_prune_fraction=0.20,
max_layers_to_remove=2,
):
"""Apply both layer pruning and head pruning in sequence."""
sims = compute_layer_similarities(mdl, data_loader)
layers_to_drop = sorted(
[i for i, s in enumerate(sims) if s > layer_sim_threshold], reverse=True
)[:max_layers_to_remove]
model_lp = prune_layers(mdl, layers_to_drop)
importance_lp = compute_head_importance(model_lp, data_loader)
model_comb, removed = prune_heads(
model_lp, importance_lp, prune_fraction=head_prune_fraction
)
return model_comb, layers_to_drop, removed
model_combined, removed_layers, removed_heads_info = prune_combined(
model, data_loader, layer_sim_threshold=0.85, head_prune_fraction=0.20
)Combined structured pruning results: Layers removed: [6, 5] (2 layer(s)) Remaining layers: 4 Original parameters: 920,066 Combined pruning: 606,978 Total reduction: 34.0%
Combining both strategies achieves substantially more compression than either alone. The sequential order matters: by removing layers first and then recomputing head importance on the smaller model, we ensure the head importance scores reflect the actual model state after layer removal, rather than importance scores computed on a different (larger) architecture.
The three pruning strategies target different parts of the model architecture, and their parameter savings compound when applied together. The visualization below shows how each strategy contributes to the total reduction.
Key Parameters
The key parameters for structured pruning are:
- layer_sim_threshold: Cosine similarity threshold above which a layer is considered near-identity. Typical values range from 0.80 to 0.95. Setting this too low removes important layers; setting it too high removes nothing. Start at 0.90 and adjust based on observed quality on a validation set.
- head_prune_fraction: Fraction of total attention heads to remove globally. Start conservatively at 10-20% and increase only after verifying quality recovery at lower fractions.
- max_layers_to_remove: Upper bound on how many layers to prune in a single pass. Removing too many at once makes recovery harder. For 12-layer models, start with at most 2 layers per pruning round.
- recovery_steps: Number of fine-tuning steps after pruning. Typically 5-20% of original training steps, with a learning rate 5-10 times smaller than the original.