CBraMod-CoreML-Apple

CBraMod-CoreML-Apple

CBraMod as a native Core ML embedding extractor for Apple platforms (macOS / iOS / visionOS), fp32, validated to rel-L2 1.3e-5 against the PyTorch reference.

This repository packages the pretrained CBraMod EEG foundation model (Wang et al., ICLR 2025) — via the braindecode re-hosted checkpoint braindecode/cbramod-pretrained — as a single .mlpackage that maps a 14-channel EEG window to a 14,000-dimensional embedding, entirely on-device via CoreML.framework.

It is published as a companion to the ZUNA Core ML profiles: in our 109-subject motor-imagery evaluation, CBraMod embeddings were the strongest foundation-model feature for sustained motor-imagery decoding, so this artifact is the MI engine of an on-device BCI stack.

What is in the package

The channel count (14) matches the Emotiv EPOC X montage (AF3, F7, F3, FC5, T7, P7, O1, O2, P8, T8, FC6, F4, F8, AF4), but nothing in the backbone is montage-specific beyond the fixed input shape: CBraMod treats channels symmetrically at the patch level, so any 14-channel montage at 200 Hz can use this package. For other channel counts, re-export from source (script linked below).

Expected preprocessing

The parity and task-level validation below used this exact chain (matching the upstream evaluation convention of z-scored inputs):

1. Window the raw EEG to 5 s.
2. Average-reference across the available channels.
3. Global z-score the window (single mean/std over all channels and samples).
4. Resample to 200 Hz (polyphase), crop/zero-pad to exactly 1000 samples.

Validation

All gates were run against the original PyTorch checkpoint on real EEG (PhysioNet EEGBCI motor-imagery recordings mapped to the 14-channel montage), not random tensors.

Downstream context (not a property of this artifact, but of the underlying checkpoint): on 109 PhysioNet EEGBCI subjects with held-out-run evaluation, CBraMod embeddings + logistic regression reached 63.7% ± 13.8% left/right-hand sustained motor-imagery accuracy on cue-offset windows (a control that excludes visual-cue-evoked confounds), significantly above every classical and FM baseline we tested (paired test vs best prior, p ≈ 2e-6).

Conversion notes (for reproducers)

Two standard approaches fail on this architecture; both failures are worth knowing:

• torch.jit.trace is non-deterministic here unless torch.backends.mha.set_fastpath_enabled(False) is set (nn.MultiheadAttention's fastpath produces divergent traces), and even then Core ML const-folding fails on ~349 symbolic-int (aten::Int) nodes arising from CBraMod's criss-cross reshape arithmetic.
• torch.export.export(...).run_decompositions({}) converts cleanly: the modern frontend specializes static shapes to constants, eliminating the symbolic-int nodes. This is the recommended path for reshape-heavy transformers.

The export script (including the parity and task-equivalence gates) is open source: scripts/port_cbramod_coreml.py.

Usage

Python (coremltools)

import numpy as np
import coremltools as ct
from huggingface_hub import snapshot_download

# NOTE: use local_dir — the default symlinked HF cache breaks the Core ML
# compiler, which cannot resolve a symlinked weight.bin inside an .mlpackage.
repo = snapshot_download("oraculumai/CBraMod-CoreML-Apple", local_dir="CBraMod-CoreML-Apple")
model = ct.models.MLModel(f"{repo}/CBraModEmbedder.mlpackage")

eeg = np.random.randn(1, 14, 1000).astype(np.float32)  # preprocessed as above
embedding = model.predict({"eeg": eeg})["embedding"]   # (1, 14000)

Swift (CoreML.framework)

import CoreML

// Compile once: xcrun coremlcompiler compile CBraModEmbedder.mlpackage <outdir>
let model = try MLModel(contentsOf: compiledURL)
let input = try MLMultiArray(shape: [1, 14, 1000], dataType: .float32)
// ... fill input with the preprocessed window ...
let out = try model.prediction(from: MLDictionaryFeatureProvider(dictionary: ["eeg": input]))
let embedding = out.featureValue(for: "embedding")!.multiArrayValue!  // [1, 14000]

A typical decoder is a small linear head (e.g. logistic regression) trained on these embeddings; the 63.7% MI result above is exactly that.

Limitations

• Fixed input shape [1, 14, 1000]. Other montages/window lengths require re-export.
• fp32 only. We have not published a compressed variant; palettization/int8 were not validated for this model.
• Research artifact. Not validated for medical diagnosis, treatment, or clinical decision-making. Use at your own risk and follow the base model's license (BSD-3-Clause).

Provenance & credit

• Original model: CBraMod — "CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding", Wang et al., ICLR 2025. Repository: https://github.com/wjq-learning/CBraMod
• Pretrained weights: braindecode/cbramod-pretrained (safetensors, BSD-3-Clause), loaded through the braindecode CBraMod implementation. The classification head (absent from the pretrained checkpoint) is replaced with Identity; all backbone tensors load cleanly.
• This conversion: oraculumai — Core ML export, parity/task-level validation, and packaging. Evaluation harness: https://github.com/nschlaepfer/oraculum-gpt-mk1

If you use this model, please cite the original CBraMod paper and credit braindecode for the checkpoint distribution.