apple-aiml-research/ml-ane-transformers

Apple Neural Engine (ANE) Transformers

What it is – ane_transformers is a small open‑source library that provides a reference PyTorch implementation of Transformer models optimized for Apple’s Neural Engine (ANE). It ships two packages:

• ane_transformers.reference – a stand‑alone, easy‑to‑read implementation that shows the low‑level changes needed to run efficiently on ANE.
• ane_transformers.huggingface – drop‑in replacements for selected Hugging Face model classes (e.g., DistilBert) that incorporate those optimizations while keeping the same Python API.

Why it matters – Apple’s A14, A15, M1, M2, etc., include a dedicated neural‑processing unit. By restructuring the computation graph (fusing ops, moving embedding look‑ups to the CPU, and tailoring the model for the ANE’s instruction set) the library claims up to 10× lower latency and 14× lower peak memory compared with a vanilla PyTorch/torch‑script model when run on‑device via Core ML.

How to use it – The README walks through a complete tutorial for the Hugging Face distilbert-base-uncased-finetuned-sst-2-english model:

1. Load the standard model with transformers.AutoModelForSequenceClassification.
2. Replace it with ane_transformers.huggingface.distilbert.DistilBertForSequenceClassification, copying the original state dict.
3. Tokenize some input, trace the model with torch.jit.trace, and convert the traced model to a Core ML mlprogram package using coremltools.convert.
4. Drop the resulting .mlpackage into an Xcode project and use Xcode’s Performance tab to verify the speed‑up on an iPhone or Mac.

Installation – The package is published on PyPI:

pip install ane_transformers          # pre‑built wheels (fastest)
# or, for development:
pip install -e .

If the build fails on the tokenizers dependency, the README points to a known fix involving a Rust compiler.

Supported hardware – The optimizations are tuned for Apple silicon that includes an ANE (iPhone 12/A14 or newer, iPad with A14+, Macs with M1 or newer). Unit tests will warn when run on older hardware, but the generated Core ML model will still run (just without the advertised speed‑up).

Limitations / caveats

• Only a handful of Hugging Face models are currently wrapped (the example shows DistilBERT). Extending to other architectures requires additional work.
• The optimized model adds a small one‑time compilation cost; the first load is slower.
• Four of the 606 operations remain on the CPU (embedding look‑ups), which is intentional for this model size.
• Performance gains depend on sequence length and batch size; the README cites up to 10× latency reduction for longer sequences (e.g., 512 tokens, batch 8).

Who should look at this – iOS/macOS developers who want to ship NLP models locally and need the best possible on‑device latency and memory footprint, as well as researchers interested in how to adapt transformer graphs for Apple’s ANE.