Structured Pruning › Layer Pruning

Layer pruning removes entire transformer blocks from the model. A transformer block includes the full self-attention sublayer, the feedforward sublayer, and their associated layer normalization layers. Removing a block reduces the number of sequential operations in every forward pass and eliminates all associated parameters. For a 12-layer BERT model, each layer represents roughly 8% of the total transformer parameters (excluding embeddings), so removing one layer provides approximately that fraction of total parameter reduction.

The computational benefit of layer removal is particularly significant for autoregressive generation. In generation, each new token requires a full forward pass through all transformer layers. Removing two layers from a 12-layer model reduces per-token generation time by roughly 17%, a meaningful speedup when you are generating hundreds of tokens per request.

The conceptually cleanest importance score for a layer measures its contribution to the residual stream. Modern transformers use residual connections: each block adds its output to the input before passing to the next block. You can measure how much a layer changes the representation by computing the cosine similarity between the input and output of that layer:

$$\text{sim}^{(l)} = \frac{1}{N} \sum_{x \sim \mathcal{D}} \cos\left( \mathbf{h}^{(l)}(x), \mathbf{h}^{(l+1)}(x) \right)$$

where:

• $\mathbf{h}^{(l)}(x)$ is the hidden state entering layer $l$ for input $x$
• $\mathbf{h}^{(l+1)}(x)$ is the hidden state leaving layer $l$
• $\cos(\cdot, \cdot)$ is cosine similarity, ranging from -1 to 1
• $N$ is the number of tokens evaluated

A layer with similarity close to 1.0 barely modifies the representation; it is essentially a near-identity transformation. Such layers are strong candidates for removal. A layer with similarity much less than 1.0 applies a substantial transformation and is important to retain. This metric is intuitive because it measures what matters to the residual stream: whether the layer changes the information flowing through the model.

This block influence (BI) criterion was investigated by several research groups studying large language models and consistently found that some middle layers have very high input-output similarity, suggesting they contribute little. The distribution of layer similarities is not uniform: early layers typically apply large transformations (encoding syntactic and positional structure into the embeddings), late layers also apply significant transformations (projecting representations toward output token probabilities), and some middle layers are nearly identity. This U-shaped distribution of layer importance has been replicated across model families from BERT to GPT to Llama.

Another approach uses gradient-based sensitivity. Compute the gradient of the task loss with respect to the output of each transformer block. Layers where the gradient magnitude is small contribute little to the final prediction and can be removed without large quality drops:

$$S^{(l)} = \left\| \frac{\partial \mathcal{L}}{\partial \mathbf{h}^{(l+1)}} \right\|_F$$

where:

• $S^{(l)}$ is the sensitivity score for layer $l$
• $\frac{\partial \mathcal{L}}{\partial \mathbf{h}^{(l+1)}}$ is the gradient of the task loss with respect to the layer's output hidden state
• $\|\cdot\|_F$ is the Frobenius norm, measuring the overall magnitude of the gradient signal

Layers with large sensitivity scores are critical to the model's predictions because backpropagation flows strongly through them. Layers with small scores carry little gradient and can be removed with minimal quality impact. The gradient-based score is computationally more expensive than the cosine similarity score because it requires a backward pass, but it is more directly connected to the training objective.

A more refined version of the cosine similarity criterion uses angular distance rather than cosine similarity directly. For each layer, you compute how much the hidden state rotates on average as it passes through:

$$\theta^{(l)} = \arccos\left(\text{sim}^{(l)}\right)$$

where:

• $\theta^{(l)}$ is the angular distance (in radians) for layer $l$
• $\text{sim}^{(l)}$ is the mean cosine similarity between the input and output hidden states of layer $l$
• $\arccos$ converts cosine similarity to an angle, with 0 meaning no rotation and $\pi$ meaning a complete reversal

Layers with small $\theta^{(l)}$ (small rotation, near-identity) are pruned first. This angular perspective is useful because it is invariant to the scale of the hidden state norms, focusing purely on the geometric transformation. A layer that scales all hidden states by a constant factor but does not rotate them would have $\theta^{(l)} = 0$ and would be correctly identified as removable, whereas a norm-based criterion might incorrectly flag such a layer as active because its output magnitude differs from its input.

The angular distance criterion is the foundation of the ShortGPT method, introduced by Ma and colleagues in 2024. They applied it to Llama 2 and found that removing the 25% of layers with the smallest angular distances produced models that retained over 95% of the original performance on standard benchmarks while running substantially faster. The layers identified as most removable were consistently in the middle of the network, confirming the U-shaped importance profile observed in earlier work.

When you drop a transformer block, the layers before and after it are now adjacent. The hidden state that was computed by layer $l-1$ flows directly into layer $l+1$. If the removed layer was performing a useful but subtle transformation, the model can often compensate through fine-tuning because the surviving layers can adjust their parameters to approximately recover the removed layer's contribution.

Recovery from layer pruning benefits significantly from targeted fine-tuning where only the layers adjacent to the removed block are updated initially, then the whole model is fine-tuned at a lower learning rate. This staged approach stabilizes training and typically achieves better final quality than fine-tuning the entire pruned model from the start. The intuition is that the adjacent layers need to make the largest adjustments and giving them a focused fine-tuning phase before introducing global optimization pressure prevents conflicting gradient signals.

When removing multiple layers, the order of removal matters. Removing the single least-important layer first, then recomputing similarities on the resulting model, and repeating this process gives better results than selecting the $k$ least-important layers simultaneously and removing them all at once. The iterative approach allows the model to adapt after each removal, which changes the importance scores for the remaining layers. This iterative single-layer removal is more compute-intensive but consistently produces higher-quality pruned models.

A practical consideration when deploying layer-pruned models is that layer count affects KV-cache size in autoregressive models. Each transformer layer contributes entries to the KV cache for each token in the context. Removing two layers from a 32-layer model reduces KV-cache memory by 6.25% for long contexts, a meaningful saving in memory-constrained deployments.