🏗️ Structured Pruning: Filters, Channels, and Attention Heads

Table of Contents

• 1. Unstructured vs. Structured Sparsity
• 2. Filter Pruning via ℓ₁ Norms
  • 2.1 Saliency of a Filter
  • 2.2 Dependency Chains and Layer Coupling
  • 2.3 PyTorch: ℓ₁ Filter Pruner
• 3. Channel Pruning via Batch Norm Scaling
  • 3.1 The γ-Sparsity Trick
  • 3.2 PyTorch: BN γ Channel Pruner
• 4. Attention Head Pruning in Transformers
  • 4.1 Head Importance Scoring
  • 4.2 L0 Gate Pruning
  • 4.3 PyTorch: Attention Head Masking
• 5. Rethinking the Value of Network Pruning
• 6. References

1. 💡 Unstructured vs. Structured Sparsity

Classical pruning methods produce unstructured sparsity: individual weights are zeroed at arbitrary positions. While this achieves high compression ratios, the resulting sparse matrices are dense in memory layout — the zero entries still occupy storage and memory bandwidth unless a custom sparse format (CSC, CSR) is used.

On modern hardware (GPU, TPU), dense matrix operations dominate: - cuBLAS SGEMM operates on \(m \times n\) tiles; any zero within a tile still costs a multiply-accumulate. - Conv2d on GPU uses Winograd or FFT-based algorithms that assume dense filter tensors.

Structured pruning removes entire groups of parameters — whole convolutional filters, output channels, or attention heads — so that the resulting network is smaller but still dense. The pruned network requires no custom sparse kernels and runs at full speed on standard hardware.

2. 📐 Filter Pruning via ℓ₁ Norms

Li, Kadav, Durdanovic, Samet, Graf (2016). “Pruning Filters for Efficient ConvNets.” ICLR 2017.

2.1 Saliency of a Filter

For a convolutional layer with weight tensor \(W \in \mathbb{R}^{C_{out} \times C_{in} \times k_H \times k_W}\), each filter \(W_i \in \mathbb{R}^{C_{in} \times k_H \times k_W}\) produces one output channel. Pruning filter \(i\) removes the \(i\)-th output channel entirely and the corresponding input channel in the next layer.

ℓ₁ filter saliency:

\[s_i = \|W_i\|_1 = \sum_{c, h, w} |W_{i,c,h,w}|\]

Filters with small ℓ₁ norm produce activations close to zero on average — they contribute little to the representation. Prune the \(k\) filters with lowest ℓ₁ norm.

Reconstruction error interpretation. Removing filter \(i\) zeroes its output feature map \(\hat{a}_i = 0\). The change in the next layer’s pre-activations is:

\[\delta z_{i'}^{(l+1)} \approx W^{(l+1)}_{i', i} * 0 = 0 \quad \forall i'\]

The reconstruction error in the next layer’s feature maps depends on how much subsequent filters relied on the now-deleted channel. Filters with small activations (small \(\|W_i\|_1\)) contribute less to this error.

2.2 Dependency Chains and Layer Coupling

Pruning filter \(i\) from layer \(l\) forces us to also remove the \(i\)-th input channel from layer \(l+1\). This propagates through skip connections and concatenations:

• Sequential layers: Remove filter \(i\) from layer \(l\) and input channel \(i\) from layer \(l+1\). Straightforward.
• ResNet skip connections: Layer \(l\) and the skip branch must produce the same number of channels. If we prune from the residual branch, we must also prune from the identity branch (or use projection shortcuts).
• Concatenation (DenseNet): Pruning channels in a dense block requires coordinating which channels to remove from all concatenated inputs.

In practice, structured pruning tools handle these dependencies by building a layer dependency graph and propagating pruning decisions through it.

2.3 💻 PyTorch: ℓ₁ Filter Pruner

import torch
import torch.nn as nn
from typing import Optional


def l1_filter_saliency(conv: nn.Conv2d) -> torch.Tensor:
    """
    Compute ℓ₁ saliency for each output filter of a Conv2d layer.
    Returns a (C_out,) tensor of filter saliencies.
    """
    # Sum |W| over C_in, k_H, k_W dimensions
    return conv.weight.data.abs().sum(dim=[1, 2, 3])


def prune_conv_filters(
    conv: nn.Conv2d,
    next_conv: Optional[nn.Conv2d],
    next_bn: Optional[nn.BatchNorm2d],
    n_keep: int,
) -> tuple[nn.Conv2d, Optional[nn.Conv2d], Optional[nn.BatchNorm2d]]:
    """
    Prune a Conv2d layer to keep `n_keep` filters (output channels),
    and update the next Conv2d layer's input channels accordingly.

    Returns new (conv, next_conv, next_bn) modules with pruned weights.
    """
    saliency = l1_filter_saliency(conv)
    keep_idx = saliency.topk(n_keep).indices.sort().values

    # Prune current conv: keep only selected output filters
    new_weight = conv.weight.data[keep_idx]
    new_bias = conv.bias.data[keep_idx] if conv.bias is not None else None
    new_conv = nn.Conv2d(
        in_channels=conv.in_channels,
        out_channels=n_keep,
        kernel_size=conv.kernel_size,
        stride=conv.stride,
        padding=conv.padding,
        bias=conv.bias is not None,
    )
    new_conv.weight.data = new_weight
    if new_bias is not None:
        new_conv.bias.data = new_bias

    # Update BN layer if present (prune running stats and γ/β)
    new_bn = None
    if next_bn is not None:
        new_bn = nn.BatchNorm2d(n_keep)
        new_bn.weight.data = next_bn.weight.data[keep_idx]
        new_bn.bias.data = next_bn.bias.data[keep_idx]
        new_bn.running_mean = next_bn.running_mean[keep_idx]
        new_bn.running_var = next_bn.running_var[keep_idx]

    # Update next conv's input channels
    new_next_conv = None
    if next_conv is not None:
        new_next_conv = nn.Conv2d(
            in_channels=n_keep,
            out_channels=next_conv.out_channels,
            kernel_size=next_conv.kernel_size,
            stride=next_conv.stride,
            padding=next_conv.padding,
            bias=next_conv.bias is not None,
        )
        new_next_conv.weight.data = next_conv.weight.data[:, keep_idx]
        if next_conv.bias is not None:
            new_next_conv.bias.data = next_conv.bias.data

    return new_conv, new_next_conv, new_bn

3. 🎚️ Channel Pruning via Batch Norm Scaling

Liu, Li, Shen, Huang, Yan, Zhang (2017). “Learning Efficient Convolutional Networks through Network Slimming.” ICCV 2017.

3.1 The γ-Sparsity Trick

Batch normalization scales each channel’s activations by a learned parameter \(\gamma_i\):

\[\hat{a}_i = \gamma_i \cdot \frac{a_i - \mu_i}{\sqrt{\sigma_i^2 + \epsilon}} + \beta_i\]

If \(\gamma_i \approx 0\), channel \(i\)’s contribution to subsequent layers is negligible — it can be removed with little accuracy impact. Network Slimming adds an ℓ₁ sparsity regularizer on all BN \(\gamma\) parameters during training:

\[L_{total} = L_{task} + \lambda \sum_{i \in BN} |\gamma_i|\]

This is the lasso penalty applied to the BN scaling factors. After training, channels with \(|\gamma_i| < \tau\) (threshold) are pruned globally.

Why BN γ instead of weights directly? BN γ is a single scalar per channel, making it an ideal “gate.” The ℓ₁ regularizer drives small-γ channels to exactly zero (as in LASSO), creating a natural threshold. Pruning a channel requires no architectural assumptions about how weights were organized — any channel with \(\gamma \approx 0\) can be removed, regardless of its position.

The training objective. Let \(\Theta\) be all network parameters and \(\Gamma = \{\gamma_i\}\) the BN scaling parameters:

\[\min_{\Theta, \Gamma} L_{task}(\Theta, \Gamma) + \lambda \|\Gamma\|_1\]

The ℓ₁ penalty on \(\Gamma\) induces a sparse \(\hat{\Gamma}\) at convergence. After training, threshold at a global \(\tau\) determined by the desired sparsity budget.

3.2 💻 PyTorch: BN γ Channel Pruner

import torch
import torch.nn as nn


def sparsity_regularized_loss(
    model: nn.Module,
    task_loss: torch.Tensor,
    lam: float,
) -> torch.Tensor:
    """
    Add ℓ₁ penalty on all BN γ (weight) parameters to the task loss.
    Used as the training loss to drive BN scales toward zero.
    """
    bn_l1 = sum(
        module.weight.abs().sum()
        for module in model.modules()
        if isinstance(module, (nn.BatchNorm2d, nn.BatchNorm1d))
    )
    return task_loss + lam * bn_l1


def global_bn_threshold(model: nn.Module, sparsity: float) -> float:
    """
    Find the global γ threshold such that `sparsity` fraction of BN channels
    fall below it. Uses the global distribution of |γ| values.
    """
    all_gamma = torch.cat([
        module.weight.data.abs()
        for module in model.modules()
        if isinstance(module, (nn.BatchNorm2d, nn.BatchNorm1d))
    ])
    return float(torch.quantile(all_gamma, sparsity))


def prune_bn_channels(
    model: nn.Module,
    threshold: float,
) -> dict[str, torch.Tensor]:
    """
    Return a dict of channel masks for each BN layer.
    mask[name][i] = False means channel i is pruned (|γ_i| < threshold).
    Does NOT physically remove channels — call rebuild_pruned_model for that.
    """
    masks = {}
    for name, module in model.named_modules():
        if isinstance(module, (nn.BatchNorm2d, nn.BatchNorm1d)):
            masks[name] = module.weight.data.abs().ge(threshold)
    return masks


def zero_out_pruned_channels(model: nn.Module, masks: dict[str, torch.Tensor]) -> None:
    """
    Zero out γ and β for pruned channels in-place (soft mask application).
    Zeroed channels produce zero output — effectively removed at inference.
    """
    for name, module in model.named_modules():
        if name in masks:
            mask = masks[name].float()
            module.weight.data.mul_(mask)
            module.bias.data.mul_(mask)

4. 🤖 Attention Head Pruning in Transformers

Michel, Levy, Neubig (2019). “Are Sixteen Heads Really Better than One?” NeurIPS 2019. Voita, Talbot, Moiseev, Sennrich, Titov (2019). “Analyzing Multi-Head Self-Attention: Specialized Heads do the Heavy Lifting, the Rest can be Pruned.” ACL 2019.

4.1 Head Importance Scoring

A multi-head attention layer computes:

\[\text{MHA}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W^O\]

\[\text{head}_h = \text{softmax}\!\left(\frac{Q W^Q_h (K W^K_h)^\top}{\sqrt{d_k}}\right) V W^V_h\]

Can we remove heads? Michel et al. introduce a head importance score based on the sensitivity of the loss to masking head \(h\):

\[I_h = \left|\mathbb{E}_{x \sim \mathcal{D}}\left[\frac{\partial L}{\partial \xi_h} \cdot \xi_h\right]\right|\]

where \(\xi_h\) is a scalar gate multiplying head \(h\)’s output. This is the magnitude of the expected gradient times gate value — equivalent to the first-order Taylor term for removing the head.

Key empirical finding: On WMT En-De, most of BERT’s 144 heads can be removed with minimal BLEU drop. The distribution of importance scores is heavy-tailed: a few “specialized” heads carry most of the load.

4.2 L0 Gate Pruning

Voita et al. train a differentiable binary gate \(z_h \in \{0, 1\}\) for each head using the Hard Concrete distribution (a continuous relaxation of discrete gates):

\[z_h \sim \text{HardConcrete}(\alpha_h) \quad \in [0, 1]\]

During training, the attention output is:

\[\text{MHA}(Q, K, V) = \text{Concat}(\text{head}_1 \cdot z_1, \ldots, \text{head}_H \cdot z_H) W^O\]

with an ℓ₀ regularizer on the expected number of active heads:

\[L_{total} = L_{task} + \lambda \sum_h \mathbb{E}[z_h]\]

At inference, gates are rounded to \(\{0, 1\}\) — any head with \(z_h = 0\) is permanently removed.

Results on WMT En-Ru (6-layer Transformer): 38 out of 48 encoder heads pruned with only \(-0.15\) BLEU. The surviving specialized heads encode positional, syntactic, and rare token information.

4.3 💻 PyTorch: Attention Head Masking

import torch
import torch.nn as nn
import math
from typing import Optional


class MaskedMultiheadAttention(nn.Module):
    """
    Multi-head attention with per-head binary masks.
    Masked heads (mask[h] = 0) produce zero output.
    """

    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

        # Learnable head gates (used for importance-based pruning training)
        self.head_mask = nn.Parameter(torch.ones(n_heads), requires_grad=False)

    def forward(
        self,
        x: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        B, T, _ = x.shape
        H, d_k = self.n_heads, self.d_k

        Q = self.W_q(x).view(B, T, H, d_k).transpose(1, 2)  # (B, H, T, d_k)
        K = self.W_k(x).view(B, T, H, d_k).transpose(1, 2)
        V = self.W_v(x).view(B, T, H, d_k).transpose(1, 2)

        # Scaled dot-product attention per head
        scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)  # (B, H, T, T)
        if attn_mask is not None:
            scores = scores + attn_mask
        attn = scores.softmax(dim=-1)
        out = attn @ V  # (B, H, T, d_k)

        # Apply head mask: zero out pruned heads
        # head_mask: (H,) → broadcast to (1, H, 1, 1)
        out = out * self.head_mask.view(1, H, 1, 1)

        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_o(out)

    def compute_head_importance(
        self,
        loader: torch.utils.data.DataLoader,
        criterion: nn.Module,
        device: str = "cuda",
    ) -> torch.Tensor:
        """
        Estimate head importance I_h = |E[dL/d(gate_h) * gate_h]|
        by temporarily making head_mask a differentiable gate.
        """
        self.head_mask.requires_grad_(True)
        importance = torch.zeros(self.n_heads, device=device)
        n_batches = 0

        for inputs, targets in loader:
            inputs, targets = inputs.to(device), targets.to(device)
            out = self(inputs)  # forward with current mask
            loss = criterion(out, targets)
            loss.backward()

            if self.head_mask.grad is not None:
                # I_h ≈ |grad_h * mask_h|
                importance += (self.head_mask.grad * self.head_mask).abs().detach()
                self.head_mask.grad.zero_()
            n_batches += 1

        self.head_mask.requires_grad_(False)
        return importance / n_batches

    def prune_heads(self, n_prune: int) -> list[int]:
        """Prune the `n_prune` least important heads (set mask to 0)."""
        importance = self.compute_head_importance(...)
        _, prune_idx = importance.topk(n_prune, largest=False)
        with torch.no_grad():
            self.head_mask[prune_idx] = 0.0
        return prune_idx.tolist()

5. 🔄 Rethinking the Value of Network Pruning

Liu, Sun, Zhou, Huang, Darrell (2019). “Rethinking the Value of Network Pruning.” ICLR 2019.

This paper challenges the conventional wisdom that the inherited weights of a pruned network are the source of its value.

Standard view: Train a large network → prune → fine-tune the surviving weights. The value of pruning is that the surviving weights are “better” (they were selected by the training process and survived the saliency filter).

Liu et al.’s finding: For structured pruning methods (filter pruning, channel pruning, Network Slimming), fine-tuning the pruned model with its inherited weights offers no significant advantage over: - Training the pruned architecture (same structure, same parameter count) from random initialization. - The pruned architecture’s performance is determined by its structure, not the inherited weights.

Implication: Structured pruning is equivalent to Neural Architecture Search — it discovers a compact architecture. The training of the unpruned model is just a expensive architecture search procedure, not a source of good weights.

6. 📚 References