Case Study: Optimizing Hugging Face DistilBERT for the ANE

Part of the ANE Knowledge Base. Source: Apple ML Research, "Deploying Transformers on the Apple Neural Engine", code in references/ml-ane-transformers/ane_transformers/huggingface/distilbert.py.

This case study shows how the four principles are retrofitted onto an existing third-party model — distilbert-base-uncased-finetuned-sst-2-english — without retraining and while staying checkpoint-compatible.

1. Strategy: subclass and swap submodules#

Instead of rewriting DistilBERT, Apple subclasses each Hugging Face module and swaps only the offending submodules via setattr, keeping the rest of the upstream logic:

class TransformerBlock(modeling_distilbert.TransformerBlock):
    def __init__(self, config):
        super().__init__(config)
        setattr(self, 'attention', MultiHeadSelfAttention(config))     # ANE version
        setattr(self, 'sa_layer_norm', LayerNormANE(config.dim, eps=EPS))
        setattr(self, 'ffn', FFN(config))                              # 1x1 convs
        setattr(self, 'output_layer_norm', LayerNormANE(config.dim, eps=EPS))

What gets replaced, per the four principles:

HF module	Replacement	Principle
q_lin, k_lin, v_lin, out_lin (nn.Linear)	nn.Conv2d(dim, dim, 1)	P1 (BC1S layout)
FFN.lin1/lin2 (nn.Linear)	nn.Conv2d 1×1	P1
nn.LayerNorm	LayerNormANE	P1 (normalize dim 1)
Fused MHA forward	per-head split + einsum + split softmax	P2, P3
Task heads (pre_classifier, classifier, vocab_projector, qa_outputs)	nn.Conv2d 1×1	P1

All six task variants are covered: DistilBertModel, ForMaskedLM, ForSequenceClassification, ForQuestionAnswering, ForTokenClassification, ForMultipleChoice.

2. Checkpoint compatibility without retraining#

Two load_state_dict pre-hooks make pretrained weights load into the new architecture untouched:

1. linear_to_conv2d_map — unsqueezes every relevant (out, in) Linear weight to (out, in, 1, 1) for Conv2d.
2. correct_for_bias_scale_order_inversion — LayerNormANE applies (x + bias) * weight whereas nn.LayerNorm applies x * weight + bias, so the stored bias is rescaled: bias = bias / weight.

3. FP16 hygiene#

• Epsilon: DistilBERT's original LayerNorm eps of 1e-12 is "not friendly with the float16 precision that ANE uses by default" → EPS = 1e-7.
• Masks: the HF attention mask (bool or int64 (B, S)) is converted to an additive float mask (B, S, 1, 1) with -1e4 for masked positions:

if mask.dtype == torch.bool:
    mask = mask.logical_not().float() * -1e4
elif mask.dtype == torch.int64:
    mask = (1 - mask).float() * -1e4

• Inference-only: the classes raise on self.training or labels is not None — this port is for on-device inference; train with the original HF implementation.
• return_dict must be False: "coremltools does not support dict outputs."

4. Layout details worth noticing#

• Hidden states flow through the whole network as (B, dim, 1, seq_len) (BC1S).
• Pooling the CLS token becomes a slice on the last axis: hidden_state[:, :, :, 0:1] → (B, dim, 1, 1) — no transpose needed.
• Task logits are recovered at the very end with cheap squeeze/final ops, after all heavy compute is done in ANE-friendly form.

5. Results#

Measured on the SST-2 DistilBERT (iPhone 13, iOS 16):

Metric	Value
Latency (seq 128, batch 1)	3.47 ms at 0.454 W (also 9.44 ms at 0.072 W low-power point)
Speedup at seq 128 / batch 1 (Xcode report)	2.84× vs baseline
Speedup at larger workloads (e.g., seq 512, batch 8)	up to 10× latency, 14× peak memory
Reference server-side comparison	AWS c6i/inf1: ~5–6 ms at seq 128
Devices validated	iPhone 12 (iOS 15/16), iPhone 13 (iOS 16), M1 Mac (macOS 13)

Operational notes from the tutorial:

• The optimized model has 606 ops; op count (from chunking) increases load/compile time — a one-time cost, hide it with async loading.
• 4 of 606 ops run on CPU: the embedding-lookup ops, which are simply more efficient on CPU for this configuration. A few CPU ops ≠ failure.
• Latency is ~flat across seq len 32/64/128 at batch 1 → the model is bandwidth-bound there (Principle 4); quantization/pruning headroom remains.

6. The deployment recipe (abridged)#

Full workflow with explanation in doc 05; the essence from the repo README:

baseline_model = transformers.AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english",
    return_dict=False, torchscript=True).eval()

optimized_model = ane_distilbert.DistilBertForSequenceClassification(
    baseline_model.config).eval()
optimized_model.load_state_dict(baseline_model.state_dict())  # hooks do the mapping

tokenized = tokenizer(["..."], return_tensors="pt", max_length=128, padding="max_length")
traced = torch.jit.trace(optimized_model,
                         (tokenized["input_ids"], tokenized["attention_mask"]))

mlpackage = ct.convert(traced, convert_to="mlprogram",
    inputs=[ct.TensorType(f"input_{name}", shape=t.shape, dtype=np.int32)
            for name, t in tokenized.items()],
    compute_units=ct.ComputeUnit.ALL)
mlpackage.save("distilbert_seqLen128_batchSize1.mlpackage")

7. Transferable lessons#

1. You rarely need a new model — an architecture-preserving, math-equivalent re-expression plus state-dict hooks converts existing checkpoints.
2. Subclass + setattr is a clean pattern for porting any HF model family.
3. Fix precision constants (eps, mask values) at the same time you fix layout — FP16 breakage is silent otherwise.
4. Judge success with Xcode performance reports (per-op dispatch), not just "it runs".
5. Expect some ops (embeddings, final squeezes) on CPU; optimize the transformer trunk.

Next: 04 — Vision Transformers on the ANE.