This repository contains Core ML conversions of ZUNA1.1 for Apple-native inference on iPhone, visionOS, and macOS, organized under enumerated profile folders.
ZUNA1.1 is a 380M-parameter masked diffusion autoencoder for scalp EEG reconstruction and superresolution. Given a subset of channels and their 3D electrode coordinates, the model can:
Denoise observed EEG channels
Reconstruct dropped or missing channels
Predict signals at novel scalp positions from physical coordinates
The base model extends ZUNA1 with variable-length masking (windows up to 30 seconds), temporal-segment reconstruction, and a more flexible corruption/dropout training scheme. This Core ML release preserves the pretrained behavior for Apple deployment using the fixed 5-second (seq_len=1280) profile contract described below.
Model Overview
The base model follows the same high-level inference pattern as upstream ZUNA:
Inputs are EEG windows of 30 seconds @ 256 Hz (seq_len=7680).
Signals are tokenized with num_fine_time_pts=32, so each channel produces 7680 / 32 = 240 coarse time tokens.
tok_idx encodes {x, y, z, tc} (electrode position + coarse time index).
Inference performs:
Encoder forward pass (once)
Decoder denoising loop (N diffusion steps)
Token-to-signal reconstruction
As reported in the original ZUNA paper, the base architecture is a ~380M-parameter position-aware diffusion autoencoder trained on a large harmonized public EEG corpus.
This release preserves the base model tensor contract and publishes profile-specific shapes for deterministic Apple deployment.
Preprocessing Contract
For best parity with upstream behavior, keep the same preprocessing assumptions used by ZUNA:
EEG montage must include 3D channel positions
Sampling rate: 256 Hz
Epoch length: 30 seconds (7680 samples)
Token chunk size: 32 (240 coarse tokens per channel)
Normalization aligned with upstream inference (data_norm=10.0)
These assumptions are what the released pretrained weights were optimized for.
Getting Started
Profile artifacts are organized as:
profiles/14ch/fp32/...
Each profile contains:
ZunaEncoder.mlpackage
ZunaDecoderStep.mlpackage
ZunaDecoderStepUpdate.mlpackage
coreml_export_metadata.json
Model split
ZunaEncoder: Encodes tokenized EEG context
ZunaDecoderStep: One denoising step in the diffusion loop
Use DecoderStepUpdate when you want a minimal host-side loop and fewer host tensor ops.
Available profiles
Profile
Channels
Precision
Token Count
Encoder/Decoder Tensor Shape
final-z rel_l2 vs PyTorch
14ch-fp32
14
fp32
3360
[1, 3360, 32]
0.000006
See profiles/index.json for machine-readable profile discovery.
Validation
All published profiles are checked against the original PyTorch weights using a 20-step diffusion parity run.
Profile
MAE
RMSE
max_abs
rel_l2
Threshold
Gate
14ch-fp32
0.000003
0.000005
0.000060
0.000006
0.005000
PASS
Parity Visualization
PyTorch vs Core ML parity summary
Waveform Overlay + Residual
Representative sample from 14ch-fp32 (channel 0), final-z step-loop output.
Waveform overlay and residual
Runtime Notes
14ch/fp32 is included for native Emotiv EPOC X-style 14-channel workflows.
Profiles are FP32-only: the base model is bf16-trained and its decoder residual stream exceeds the fp16 range (max 65504), so fp16 conversions fail parity and are not published.
Throughput/latency depends strongly on diffusion steps.
Lower diffusion steps are useful for rapid iteration; higher steps improve reconstruction quality.
