Dynamic Pruning for Inference Cost Cuts [Deep Dive]

Inference pruning is one of the few optimization techniques that can cut real transformer cost without retraining the whole model. The trick is timing. If you remove tokens too early, accuracy drops because the model has not built enough context yet. If you prune after a few encoder blocks, later self-attention runs on a much shorter sequence, which reduces compute and memory while keeping the model’s decision boundary nearly unchanged when you calibrate the keep ratio on validation traffic.

Prerequisites

pip install torch transformers

Why this works: Hugging Face documents attention tensors with shape (batch_size, num_heads, sequence_length, sequence_length). That means reducing a later-layer sequence from 512 tokens to 256 cuts the score matrix area by roughly 75%. That is the mechanical reason dynamic pruning can pay off.

Step 1: Choose the pruning policy

For a production-safe first pass, keep the policy boring:

1. Prune only after early contextualization, usually after 25-40% of encoder layers.
2. Never drop [CLS], [SEP], or any token already masked out as padding.
3. Score token importance from the current hidden state, not from raw IDs.
4. Start with a simple score such as hidden-state L2 norm.

This is intentionally not the fanciest scorer. A learned scorer can be better, but a norm-based scorer is easy to validate and keeps the implementation dependency-free. If you want to clean up the script before merging it into your serving repo, run it through TechBytes’ Code Formatter.

Step 2: Implement the pruned forward pass

The implementation below targets a BERT-style classifier. It walks the encoder layers manually, prunes once in the middle, then finishes the remaining layers on the shorter sequence.

import math
import time
import torch
from torch import nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer


class DynamicPrunedBertClassifier(nn.Module):
    def __init__(
        self,
        checkpoint_path,
        prune_after_layer=4,
        keep_ratio=0.6,
        preserve_cls=True,
    ):
        super().__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(checkpoint_path)
        self.bert = self.model.bert
        self.prune_after_layer = prune_after_layer
        self.keep_ratio = keep_ratio
        self.preserve_cls = preserve_cls

    def _select_tokens(self, hidden_states, attention_mask):
        # hidden_states: [batch, seq, hidden]
        # attention_mask: [batch, seq]
        scores = hidden_states.norm(dim=-1)
        scores = scores.masked_fill(attention_mask == 0, float("-inf"))

        batch, seq_len, _ = hidden_states.shape
        keep_k = max(2, math.ceil(seq_len * self.keep_ratio))
        keep_k = min(keep_k, seq_len)

        if self.preserve_cls:
            body_scores = scores[:, 1:]
            topk = max(1, keep_k - 1)
            _, body_idx = torch.topk(body_scores, k=topk, dim=1, sorted=True)
            body_idx = body_idx + 1
            cls_idx = torch.zeros(batch, 1, dtype=torch.long, device=hidden_states.device)
            keep_idx = torch.cat([cls_idx, body_idx], dim=1)
        else:
            _, keep_idx = torch.topk(scores, k=keep_k, dim=1, sorted=True)

        keep_idx, _ = torch.sort(keep_idx, dim=1)

        gather_idx = keep_idx.unsqueeze(-1).expand(-1, -1, hidden_states.size(-1))
        pruned_hidden = torch.gather(hidden_states, dim=1, index=gather_idx)
        pruned_mask = torch.gather(attention_mask, dim=1, index=keep_idx)
        return pruned_hidden, pruned_mask

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        input_shape = input_ids.size()
        hidden_states = self.bert.embeddings(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
        )

        extended_mask = self.bert.get_extended_attention_mask(attention_mask, input_shape)

        for layer_idx, layer_module in enumerate(self.bert.encoder.layer):
            layer_outputs = layer_module(hidden_states, attention_mask=extended_mask)
            hidden_states = layer_outputs[0]

            if layer_idx == self.prune_after_layer:
                hidden_states, attention_mask = self._select_tokens(hidden_states, attention_mask)
                new_shape = attention_mask.shape
                extended_mask = self.bert.get_extended_attention_mask(attention_mask, new_shape)

        pooled_output = self.bert.pooler(hidden_states) if self.bert.pooler is not None else hidden_states[:, 0]
        pooled_output = self.model.dropout(pooled_output)
        logits = self.model.classifier(pooled_output)
        return logits

What the code is doing

• bert.embeddings creates the initial token representations.
• Each encoder block runs normally until pruneafterlayer.
• torch.topk selects the highest-scoring tokens to keep.
• torch.gather applies the same selection to hidden states and masks.
• The remaining encoder blocks run on the shorter sequence only.

This version assumes a BERT-family model with a pooler and classifier head. For other architectures, the pruning pattern is the same, but the submodule names change.

Step 3: Benchmark and verify

Do not ship this based on intuition. You need two checks: latency and task quality.

def benchmark(model, tokenizer, texts, device="cpu", warmup=10, runs=50):
    model.to(device)
    model.eval()

    encoded = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors="pt",
    )
    encoded = {k: v.to(device) for k, v in encoded.items()}

    with torch.no_grad():
        for _ in range(warmup):
            _ = model(**encoded)

        start = time.perf_counter()
        for _ in range(runs):
            _ = model(**encoded)
        elapsed = time.perf_counter() - start

    return elapsed / runs


checkpoint = "<your-finetuned-checkpoint>"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
texts = [
    "The service was fast, but the billing workflow was confusing.",
    "The new release fixed the crash and reduced startup time.",
] * 32

baseline = AutoModelForSequenceClassification.from_pretrained(checkpoint)
pruned = DynamicPrunedBertClassifier(checkpoint, prune_after_layer=4, keep_ratio=0.6)

base_latency = benchmark(baseline, tokenizer, texts)
pruned_latency = benchmark(pruned, tokenizer, texts)

print(f"baseline_avg_s={base_latency:.6f}")
print(f"pruned_avg_s={pruned_latency:.6f}")
print(f"speedup={(base_latency / pruned_latency):.2f}x")

Verification and expected output

On the first pass, you are looking for a result pattern like this:

baseline_avg_s=0.012841
pruned_avg_s=0.009731
speedup=1.32x

Then run your normal validation evaluator and compare the pruned wrapper against the baseline checkpoint:

• If accuracy, F1, or AUC moves less than your release threshold, the configuration is viable.
• If latency improves but quality regresses, increase keep_ratio or prune later.
• If quality is flat but latency barely moves, your sequences may be too short for pruning to matter.

Top 3 troubleshooting

1. Accuracy dropped more than expected

• Move pruneafterlayer deeper into the encoder.
• Raise keep_ratio from 0.6 to 0.7 or 0.8.
• Preserve fixed special tokens explicitly instead of letting them compete in the scorer.

2. You hit mask or shape errors

• Prune hidden_states and attention_mask with the same index tensor.
• Rebuild the extended attention mask after every prune event.
• Keep token order sorted after torch.topk so downstream layers see a stable sequence order.

3. Latency did not improve

• Your batch may be too small, or the sequence length too short.
• Pruning once near the final layer rarely helps because most compute already happened.
• Kernel overhead can erase gains on some hardware, so benchmark on the real serving target.

What's next

Once the simple norm-based version is stable, there are three worthwhile upgrades:

• Replace the heuristic scorer with a tiny learned gate trained to mimic baseline predictions.
• Prune multiple times, for example once at layer 4 and again at layer 8, instead of doing one large cut.
• Pair dynamic pruning with standard inference controls like mixed precision, request bucketing, and cache-aware batching.

The important engineering discipline is unchanged: optimize the later expensive blocks, keep the selection logic deterministic, and prove the quality delta on held-out traffic before rollout.