Structured Pruning › Worked Example: Pruning a Transformer Encoder

To make these concepts concrete, consider a small BERT-like encoder with 6 layers, 8 attention heads per layer, a model dimension of 512, and a feedforward dimension of 2048. We will walk through each pruning strategy step by step with specific numbers.

Starting parameters:

The full model has approximately:

• Attention weights: $6 \times 4 \times 512^2 = 6{,}291{,}456$ parameters (Q, K, V, O projections per layer)
• Feedforward weights: $6 \times 2 \times (512 \times 2048) = 12{,}582{,}912$ parameters
• Total (excluding embeddings): approximately 18.9 million parameters

With $H = 8$ heads and $d_k = 64$ per head, each attention layer contains 4 projection matrices, each $512 \times 512 = 262{,}144$ parameters, for a total of about 1,048,576 per layer. The feedforward sublayer with $d_{\text{ffn}} = 2048$ has two projection matrices totalling $2 \times 512 \times 2048 = 2{,}097{,}152$ parameters per layer. Across 6 layers, the feedforward sublayers collectively dominate the parameter count.

Step 1: Compute layer similarities

Run 1000 representative inputs through the model and compute input-output cosine similarity for each of the 6 layers. Suppose you find:

Layers 3 and 4 have near-identity behavior and are candidates for removal. Notice that the high-similarity layers are not at the ends but in the middle, which is the typical pattern for BERT-style encoders. Removing layers 3 and 4 eliminates 2 of the 6 transformer blocks, reducing the transformer parameter count by 33% and cutting latency by roughly the same fraction (assuming attention and feedforward layers dominate computation, which is true for longer sequences).

Step 2: Head importance scoring

For each head in each of the remaining 4 layers (after considering layer removal), compute gradient-based importance scores. With 4 layers and 8 heads per layer, you have 32 heads in total. Suppose that roughly 6 heads across the 4 layers have importance scores more than one standard deviation below the mean. These 6 heads are candidates for removal, reducing the head count from 32 to 26 and shrinking the attention weight matrices proportionally.

The heads with the lowest importance scores are often concentrated in one or two specific layers. In the typical pattern, layers just after the removed layers (layer 5 in our example, which was previously layer 3 but is now adjacent to where layers 3 and 4 were removed) sometimes show reduced head importance because they had been receiving partially pre-processed representations from the removed layers and now receive less-processed input.

Step 3: Feedforward width assessment

For each feedforward sublayer in the remaining 4 layers, compute mean absolute activation for each of the 2048 intermediate neurons over your calibration set. If you find that roughly 20% of neurons have mean absolute activation below a threshold of 0.01, you can remove those 400 neurons per layer across 4 layers, totalling 1,600 neuron removals. Each removal eliminates one column from $\mathbf{W}_1$ and one row from $\mathbf{W}_2$ in that layer, proportionally reducing feedforward computation.

Step 4: Prune and recover

Remove layers 3 and 4 (leaving 4 layers), then remove the 6 low-importance heads from the 4 remaining layers, then remove the low-activation feedforward neurons. Fine-tune the resulting model for 5% of the original training steps with a learning rate 10 times smaller than the original. The resulting model has roughly:

• 4 transformer layers instead of 6: 33% reduction in layer count
• 26 heads across 4 layers instead of 32: 19% reduction in attention heads
• 80% of feedforward neurons retained: 20% reduction in feedforward width

The combined effect on parameters is multiplicative across the architectural dimensions, typically producing a model around 40-50% smaller than the original with quality loss of 1-3 points on standard benchmarks after recovery fine-tuning.