What Can You Remove From a Transformer?

What Can You Remove From a Transformer?

A decoder-only model is a stack of identical blocks: a causal self-attention sub-layer, an MLP sub-layer, residual connections, and layer normalisation. Structural pruning asks which of those pieces can be removed or narrowed without destroying the model. This is different from zeroing individual weights (unstructured magnitude pruning) and from quantization (reducing numeric precision). Structural removal changes the architecture; the other two keep the architecture and change either which weights are non-zero or how many bits represent them.

This lesson treats structural removal as an architecture decision with three concrete forms:

• Block removal: removing an entire transformer block (attention + MLP + norms + residuals).
• MLP narrowing: reducing the intermediate hidden width of the feed-forward sub-layer inside a block.
• Attention component removal: removing an attention sub-layer entirely, or removing individual heads within multi-head attention, or changing the Q/K/V projection dimensions.

We distinguish all three from unstructured pruning and quantization, explain the shape and weight compatibility rules after each change, describe when recovery training is required, and design one original ablation that separates estimated parameter savings from measured serving latency. No universal ranking exists; any adoption requires measurement on your target model, evaluation set, and hardware.

The one-sentence version

Structural pruning changes the architecture by removing or narrowing components (layers, MLP hidden width, attention sub-layers or heads); it requires shape verification, usually requires recovery training, and does not guarantee lower measured serving latency even when parameter count drops, because latency depends on memory bandwidth, kernel fusion, sequential dependencies, and batch settings, not parameter count alone.

1. Three forms of structural removal — and what they are not

Removal form	What changes in the architecture	What does not change
Block removal (entire decoder block)	One fewer attention + MLP pair; residual stream skips the removed block	Weight matrices of remaining blocks; vocabulary projection; embedding table
MLP narrowing	MLP intermediate dimension d_ff reduced; W_1 and W_2 shapes change	Number of blocks; attention dimensions; embedding and output projection
Attention component removal	Attention sub-layer removed; or individual heads removed (changing W_Q, W_K, W_V, W_O); or projection dimensions narrowed	Residual stream shape if the output projection is kept consistent; remaining MLPs

These are different from:

• Unstructured weight zeroing (magnitude/pruning): keeps all architecture dimensions; sets individual weights to zero. The tensor shapes stay identical; the model can still load the full weight file. Zeroing does not change the number of layers or the MLP width; it only changes which weights are non-zero. It requires sparse kernels or mask storage to yield any speed benefit, and standard dense GPU kernels rarely accelerate sparse weights without specialised implementations.
• Quantization: keeps all dimensions; reduces bit-width (e.g., FP16 to INT8 or 4-bit). The parameter count in terms of elements does not change; the memory footprint and compute format change. Quantization is a precision decision, not a structural decision; it is handled separately in this track’s quantization module.

The distinction matters because each form has different compatibility rules, different quality impacts, and different latency effects.

2. Block removal: the coarsest structural change

Removing a full decoder block removes both the attention sub-layer and the MLP sub-layer, plus their associated layer norms and residual connections. The input and output of the removed block must be compatible with the adjacent residual stream; because the residual connection is x + sublayer(x), removing the sub-layer means the output of the previous block becomes the input of the next block directly.

From He et al. (2024), §3.1, the authors call this Block Drop: removing a full block based on a similarity-based redundancy metric (measured as the cosine similarity between a block’s input and output hidden states). Their finding: block removal degrades performance significantly relative to targeted removal, because it ignores the internal architecture — the attention and MLP sub-layers may have different redundancy patterns (§3.3, Figure 3; §4.1, Table 2).

Key compatibility checks after block removal:

• The hidden dimension d_model does not change; the residual stream is preserved.
• The number of layers L decreases by one; any code that assumes a fixed layer count (e.g., KV-cache pre-allocation by layer index) must be updated.
• Positional encodings (absolute or RoPE) apply to token positions, not layers; removing a layer does not change the positional encoding mechanism, but the representation depth decreases.
• No weight matrix from the removed block is reused; the weights of the remaining blocks stay exactly the same unless recovery training updates them.

Because the architecture is permanently shallower, recovery training is almost always required: the model was trained with L layers of depth; removing one changes the representation trajectory for all tokens that passed through it. Fine-tuning on a representative subset of the original training data (or on a domain-specific evaluation set) is standard practice.

3. MLP narrowing: changing the hidden width

Inside each decoder block, the MLP is usually two linear projections separated by an activation: FFN(x) = W_2 · σ(W_1 · x). The intermediate dimension d_ff is often 4 · d_model (e.g., d_model = 768, d_ff = 3072). Narrowing the MLP reduces d_ff, which changes the shapes of W_1 and W_2:

# Before narrowing: d_ff = 3072
W_1: [d_model, d_ff]  = [768, 3072]
W_2: [d_ff, d_model]  = [3072, 768]

# After narrowing: d_ff' = 1536 (half)
W_1': [768, 1536]
W_2': [1536, 768]

The output of the MLP (W_2 · σ(W_1 · x)) still has shape [batch, seq, d_model]; the residual stream is preserved. The projection weights are smaller; the total parameter count drops linearly with the d_ff reduction. However, the model’s capacity to represent non-linear transformations between attention outputs is reduced; the original MLP width was selected as part of the architecture design, and narrowing it is a capacity tradeoff, not just a size tradeoff.

From He et al. (2024), §3.2 and §4.1: removing MLP layers (or narrowing them) hurts performance more than removing attention layers at similar ratios, because the MLP sub-layer carries a large share of the model’s parameter budget and non-linear capacity. Their Table 2 reports that MLP Drop causes larger perplexity increases than Attention Drop at the same removal ratio.

Compatibility rules after MLP narrowing:

• W_1 and W_2 must be re-initialised or resized; existing weights from a wider MLP cannot be directly mapped to a narrower MLP without truncation or projection.
• The activation function σ does not change; the intermediate tensor shape changes.
• No other layer’s weights need to change, but recovery training is required because the MLP’s learned representations change.

4. Attention component removal: sub-layer or head-level

Removing the attention sub-layer entirely removes W_Q, W_K, W_V, and W_O. The residual stream then skips the attention computation: the MLP sub-layer receives the previous layer’s output directly. This is the most aggressive attention-side structural change and requires careful compatibility checking.

Removing individual heads within multi-head attention is finer-grained. The projection weights W_Q, W_K, W_V, and W_O are reshaped; the head dimension d_head and the number of heads h change together so that d_model = h · d_head holds. If h is reduced and d_head is kept constant, the projection dimensions change; if h is reduced and d_model is preserved by increasing d_head, the projection dimensions stay the same but the attention mechanism’s parallelism changes. Either way, the remaining heads must cover the same output dimension.

From the original multi-head attention definition (Vaswani et al., 2017, §3.2.2): MultiHead(Q, K, V) = Concat(head_1, ..., head_h) W^O. Removing heads changes the Concat dimension and requires W^O to match the new concatenated dimension.

From He et al. (2024), §3.3 and §4.2: Attention Drop (removing attention sub-layers based on a similarity-based redundancy score) achieves the best performance retention relative to Block Drop and MLP Drop at the same removal ratio. Their similarity metric measures cosine similarity between the input and output hidden states of an attention layer; layers with high similarity (low transformation) are considered redundant. Their Figure 3 illustrates the three removal forms: Block Drop removes the full block; Attention Drop removes the attention sub-layer and its preceding LayerNorm; MLP Drop removes the MLP and its preceding LayerNorm.

Important: the similarity metric measures redundancy, not importance. A layer with high input-output similarity may still contribute to the overall representation pathway through its interaction with the residual stream; the metric is a heuristic, not a proof of redundancy.

Compatibility rules after attention component removal:

• If an entire attention sub-layer is removed: no W_Q, W_K, W_V, W_O weights remain for that layer; the MLP sub-layer’s input is the previous block’s output directly. The MLP’s W_1 and W_2 shapes are unchanged.
• If individual heads are removed: W_Q, W_K, W_V output dimensions change; W_O input dimension changes accordingly. The output dimension d_model must stay the same to preserve the residual stream.
• The KV cache (if implemented) no longer stores keys and values for removed attention layers; any code that indexes the cache by layer index must be updated.
• Recovery training is required; removing attention heads changes the representation trajectory through the model.

5. Why parameter reduction does not mean faster serving

This is the core engineering insight. A structural change that reduces parameter count can still leave measured serving latency unchanged or even worse, for several reasons:

• Memory bandwidth, not parameter count, often dominates: during decode, the bottleneck is often reading the KV cache and loading weights from high-bandwidth memory. Removing a layer reduces the number of weights loaded per token, but the remaining layers’ weights are still loaded; the sequential dependency (each layer depends on the previous) remains; and kernel overhead (launching a GPU kernel per layer) may dominate for small batches.
• Sequential dependencies: a decoder-only model is sequential across layers. Removing a layer reduces the depth of the sequence, which can reduce latency, but the reduction is linear in the number of layers removed, not proportional to the parameter reduction. Removing 25% of layers does not guarantee 25% lower latency, because other overheads (embedding lookup, final projection, sampling, batching) remain.
• Batch and sequence settings: latency depends on batch size, sequence length, prefill length, and whether continuous batching is used. A structural change measured at batch size 1 may behave differently at batch size 8; the measurement must be reported with the exact settings.
• Kernel fusion and framework overhead: modern serving frameworks (vLLM, TensorRT-LLM, etc.) fuse kernels across layers. Removing a layer changes the fusion graph but does not eliminate framework overhead; the actual speedup requires measurement on the target framework and hardware.
• Active vs stored parameters (MoE parallel): from lesson 16, reducing active parameters (e.g., via routing) does not reduce total stored memory or latency unconditionally. The same logic applies to structural removal: fewer parameters does not automatically mean faster inference unless the bottleneck is compute-bound and the removed computation is not replaced by other overhead.

Therefore, any structural pruning claim must distinguish:

• Estimated parameter reduction: calculated from the architecture dimensions.
• Estimated memory reduction: calculated from weight size and KV-cache size.
• Measured latency change: measured on the target hardware, framework, batch size, and sequence length.
• Measured quality change: measured on the target evaluation set.

Only measured values justify adoption.

6. Compatibility check after any structural change

Before running any evaluation, verify the architecture is intact:

Check	How to verify
Residual stream preserved	x + sublayer(x) output shape equals input shape x for all remaining layers
Projection dimensions consistent	W_Q, W_K, W_V, W_O output dimensions match the intended head count and d_model
MLP output dimension preserved	W_2 output dimension equals d_model after narrowing
Layer count updated	Any pre-allocation arrays (KV cache, state buffers) use the new L
Embedding and output projection unchanged	Vocabulary projection W_out shape [d_model, vocab_size] is unchanged unless explicitly changed

A simple shape assertion (pseudocode, inspection only, not executed):

# After structural change, verify one forward pass by inspection
assert block_output.shape == block_input.shape == (B, T, d_model)
assert mlp_output.shape == (B, T, d_model)  # W_2 must restore d_model
assert attention_scores.shape == (B, H_q, T, T)  # if attention kept

These assertions are intended for manual inspection of the architecture definition; they are not a substitute for an actual forward pass with real weights.

7. Recovery training: what it is and why it is required

Recovery training is fine-tuning the structurally changed model on representative data so that the remaining layers learn to compensate for the removed components. It is not optional for structural removal, because:

• The original model was optimised with a specific depth; removing a layer breaks the optimisation trajectory for tokens that depended on it.
• Narrowing an MLP reduces non-linear capacity; the remaining layers must learn a different representation mapping.
• Removing attention heads changes the parallel representation pathways; the remaining heads must learn to cover the representation space that was previously shared.

Recovery training is distinct from initial pretraining: it uses a smaller dataset (often a subset of the original pretraining data or a domain-specific dataset), a lower learning rate, and fewer steps. It is also distinct from instruction fine-tuning: the goal is to restore representation quality, not to teach a new behaviour.

From He et al. (2024), §4.3: their experiments evaluate structural removal without full recovery training (they use the original weights with the removed layers simply omitted). The quality degradation reported in their Table 2 reflects this no-recovery scenario; actual adoption would typically include recovery training and would likely show different results. We do not claim any specific recovery-training result; any claim requires measurement on your target setup.

8. Original ablation plan: measuring structural removal defensibly

The following ablation plan is designed to separate estimated savings from measured effects. It is not a recipe that has been executed; it is a proposal that a practitioner could implement and report.

8.1 Fixed baseline

• Model: a decoder-only model with L layers, d_model, H_q query heads, d_head, MLP expansion factor 4, standard pre-norm residual blocks.
• Checkpoint: a pre-trained checkpoint with known evaluation scores on the target task.
• No structural changes applied to the baseline; the baseline provides the reference quality, memory, and latency numbers.

8.2 Controlled evaluation set and hardware

• Evaluation set: a fixed held-out set (not the training set) with a clearly stated metric: perplexity for language modelling, or task-specific accuracy (e.g., MMLU subset, code-generation pass rate) for capability evaluation. The set must be the same for all variants.
• Hardware: a single GPU model (e.g., NVIDIA A100) with fixed driver and CUDA version; no multi-device routing, so communication overhead does not confound latency measurements.
• Serving framework: a standard dense-serving framework (e.g., vLLM) with continuous batching disabled for simplicity; batch settings must be stated explicitly.

8.3 Metrics

Metric category	Specific metric	How measured
Quality	Perplexity on held-out set; or task accuracy	Computed on the fixed evaluation set, same tokenisation
Memory (weights)	Peak GPU memory during inference	Measured by framework profiler or torch.cuda.memory_allocated()
Memory (KV cache)	Per-layer KV-cache bytes at max sequence length	Calculated from formula 2 × L' × H_kv × d_head × S_max × B × b; stated as estimated unless measured
Latency (prefill)	Time from prompt to first token (TTFT)	Measured with a stopwatch or profiler; reported in milliseconds
Latency (decode)	Inter-token latency (ITL) for token 2 through N	Measured with profiler; reported in milliseconds/token
Batch settings	Batch size B, sequence length S, prefill length S_prefill	Stated explicitly; must be identical across variants

8.4 Variants to compare

• Baseline: original L, full MLP width, full attention.
• Variant A (attention removal): remove 25% of attention sub-layers based on similarity metric (He et al. 2024, §3.3); no recovery training for the first measurement; with recovery training for the second measurement.
• Variant B (MLP narrowing): reduce d_ff by 25% (e.g., from 4 · d_model to 3 · d_model); recovery training applied.
• Variant C (block removal): remove 25% of full decoder blocks (He et al. 2024, §3.1); recovery training applied.

8.5 Decision rule

A structural change is adopted only if all of the following conditions are met on the measured results (not on estimates):

1. Quality: the measured quality metric (perplexity or task accuracy) is within the stated acceptable degradation threshold (e.g., perplexity increase ≤ 3%; accuracy decrease ≤ 2% absolute).
2. Memory: the measured peak GPU memory (weights + KV cache at the target sequence length and batch size) is reduced by at least the stated minimum (e.g., 15%).
3. Latency: the measured decode latency (ITL) at the target batch size and sequence length is reduced by at least the stated minimum (e.g., 10%), and prefill latency does not increase beyond the stated maximum (e.g., 5%).
4. Recovery cost: the time and data cost of recovery training is documented and accepted.

If any condition fails, the change is rejected or returned for further tuning (e.g., a different removal ratio, a different recovery dataset, or a different selection metric for attention removal).

8.6 What to label as estimated vs measured

• Estimated: parameter count reduction (calculated from architecture dimensions); memory reduction (calculated from weight size and KV-cache formula); theoretical compute reduction (FLOP estimates).
• Measured: perplexity, accuracy, peak memory, prefill latency, decode latency.
• Any number presented without a measurement source must be labelled “illustrative” or “calculated from stated dimensions, not measured.”

9. Worked calculation: estimated parameter and memory changes

The numbers below are calculated from stated dimensions and are illustrative, not measured from a real run. They show how to separate estimated savings from what must be measured.

Assume a toy decoder-only model with:

• d_model = 768
• L = 12 layers
• d_ff = 3072 (4 · d_model)
• H_q = 12, d_head = 64, so d_model = H_q · d_head holds
• Vocabulary size V = 50257
• Batch size B = 1, max sequence length S_max = 1024, precision b = 2 bytes (FP16)

9.1 Baseline parameter count (calculated, illustrative)

For one layer:

• Attention projections (W_Q, W_K, W_V, W_O): 4 · d_model · d_model (simplified; exact shapes depend on H_q and d_head, but the total attention parameter budget per layer is 4 · d_model² for standard MHA).
• MLP (W_1, W_2): d_model · d_ff + d_ff · d_model = 2 · d_model · d_ff.
• Embedding table: V · d_model.
• Output projection (W_out): V · d_model (often tied to embedding; if tied, count once).

Simplified total (untied output projection, no bias terms):

# Attention params per layer ≈ 4 · 768² = 2,359,296
# MLP params per layer ≈ 2 · 768 · 3072 = 4,718,592
# Embedding + output projection = 2 · 50257 · 768 ≈ 77,194,752

# Total ≈ 12 · (2,359,296 + 4,718,592) + 77,194,752 ≈ 162,129,408

This is an illustrative parameter estimate; the actual model’s parameter count depends on exact projection shapes, bias terms, layer norm parameters, and weight tying.

9.2 Variant A: remove 25% of attention sub-layers

Remove 3 of 12 attention sub-layers. Attention parameters removed: 3 · 2,359,296 ≈ 7,077,888. Relative reduction: 7,077,888 / 162,129,408 ≈ 4.4%. The MLP parameters, embedding, and output projection are unchanged. The KV cache for the removed layers is no longer stored; at S_max = 1024, B = 1, b = 2, the estimated KV-cache reduction per removed layer is 2 · 12 · 64 · 1024 · 2 = 3,145,728 bytes = 3.00 MiB (calculated, not measured). The actual latency reduction must be measured on the target framework and hardware.

9.3 Variant B: narrow MLP by 25% (d_ff from 3072 to 2304)

MLP parameter reduction per layer: 2 · 768 · (3072 - 2304) = 1,179,648. For 12 layers: 14,155,776. Relative reduction: 14,155,776 / 162,129,408 ≈ 8.7%. This is larger than the attention removal reduction because the MLP carries a larger share of the parameter budget per layer.

9.4 Variant C: remove 25% of full blocks

Remove 3 full blocks (attention + MLP). Per-block parameter reduction: 2,359,296 + 4,718,592 = 7,077,888. For 3 blocks: 21,233,664. Relative reduction: 21,233,664 / 162,129,408 ≈ 13.1%. This is the largest parameter reduction but also the most aggressive architecture change.

9.5 The gap between estimated and measured

These calculations show the estimated parameter reduction only. They say nothing about:

• Actual perplexity or accuracy change (requires measurement on the target evaluation set).
• Actual memory reduction (requires measurement on the target framework, because framework overhead, activation storage, and caching policies change with architecture).
• Actual latency reduction (requires measurement with the target batch size, sequence length, and serving framework).

Any claim that relies only on these estimates is incomplete. The ablation plan in section 8 requires measured values for adoption.

10. Exit task: a defensible ablation report

Your task is not to execute the ablation, but to write the report that would justify or reject a structural change. Use the framework from section 8.

Write a short report (maximum 500 words) with the following sections:

1. Baseline and variants: state the model (L, d_model, H_q, d_ff), the structural change applied (block, MLP, or attention), the removal ratio, and the selection method (e.g., similarity-based for attention; fixed ratio for MLP or block).
2. Measured results (or stated as not measured): for each variant, report the measured perplexity (or accuracy), measured peak memory, measured prefill latency (TTFT), and measured decode latency (ITL) at the stated batch size and sequence length. If you have not executed the measurement, state clearly: “Not measured; the numbers below are estimated only.”
3. Estimated savings: calculate the parameter and memory reduction using the formulas in section 9, with the exact dimensions of your model. Label these as estimated.
4. Decision: state whether the change is adopted, rejected, or conditional, using the decision rule from section 8.5. Justify from the measured deltas (or explain why measurement is still required before any adoption).
5. Verification gaps: state what has not been measured (e.g., multi-batch latency, long-context behaviour, quality on out-of-distribution inputs, recovery-training cost).

The report must not claim benchmark improvements that have not been reproduced. Any number without a measurement source must be labelled estimated or illustrative.

Exercise: Design and defend a structural-pruning decision

1. 01Using the ablation framework in section 8, propose one structural change (block removal, MLP narrowing, or attention removal) for a decoder-only model with L = 16, d_model = 1024, H_q = 16, d_head = 64, d_ff = 4096. State the removal ratio, the selection method, the evaluation metric, the batch/sequence settings, and the decision thresholds.

   Show a hint

   Keep the baseline fixed. State the exact dimensions for the changed architecture. Label any estimated parameter reduction clearly.

2. 02Calculate the estimated parameter reduction for your proposed change using the formulas from section 9. State the exact arithmetic and label the result as estimated, not measured.

   Show a hint

   Use the per-layer formulas: attention ≈ 4 · d_model²; MLP ≈ 2 · d_model · d_ff. Scale by L.

3. 03State the decision rule from section 8.5 and apply it to your proposed change. If you have not measured quality, memory, or latency, state clearly what measurement is required before adoption.

   Show a hint

   Do not claim benchmark gains you have not reproduced. A defensible decision requires measured deltas, not only ratios.

11. Related lessons

• Lesson 15 (mha-mqa-gqa): attention architecture tradeoffs; the head-grouping mechanics that any attention-side structural change must preserve.
• Lesson 16 (mixture-of-experts): active vs stored parameters, and why fewer parameters does not guarantee faster serving — the same logic applies to structural removal.
• Lesson 20 (architecture-comparison, planned): the framework for comparing configurations defensibly; the ablation plan in this lesson feeds directly into that planned comparison framework.

Related lessons

• LLM Internals & Architecture · Modern LLM architecturesMHA vs. MQA vs. GQA: Attention Architecture Tradeoffs14 min
• LLM Internals & Architecture · Modern LLM architecturesMixture of Experts: Routing, Dispatch, and Load Balancing15 min

12. Sources and claim-to-source mapping

Source	Exact section	What it supports	Scope and limits
Vaswani et al. (2017), “Attention Is All You Need” (arXiv:1706.03762)	§3.2.2 (multi-head attention definition)	Original Concat + W^O mechanics; removing heads requires matching W^O input dimension	Universal architecture definition; does not address pruning
He et al. (2024), “What Matters in Transformers? Not All Attention is Needed” (arXiv:2406.15786)	§3.1 (Block Drop), §3.2 (MLP Drop), §3.3 (Attention Drop), Figure 3 (removal forms), §4.1 (quality results, Table 2), §4.3 (recovery training not measured)	Three structural removal forms; similarity-based redundancy metric; attention removal retains quality better than block or MLP removal at same ratio; quality results reported without full recovery training	Measured on the authors’ model family and evaluation setup; does not claim universal quality retention; results without recovery training
PyTorch nn.Linear docs (torch.nn.Linear)	Weight shape documentation (in_features, out_features)	Projection dimensions must match; structural removal changes these dimensions	Implementation reference; does not address architecture design

Verification gaps stated:

• The similarity-based redundancy metric from He et al. (2024) measures input-output similarity, not functional importance; it is a heuristic.
• The quality results in Table 2 of He et al. (2024) are measured without full recovery training; actual adoption would typically include recovery training and would show different results.
• The parameter estimates and memory calculations in section 9 are calculated from stated dimensions; they have not been executed against a real model checkpoint.
• The latency effects (section 5) are discussed from first principles; any specific latency claim requires measurement on the target framework, GPU, batch size, and sequence length. No universal latency improvement is claimed.
• The ablation plan (section 8) is a design proposal, not a completed experiment. Any adoption decision requires the measured values described in the decision rule.