ZUNA1.1 · Core ML

ZUNA1.1 is Zyphra’s masked diffusion autoencoder for scalp EEG. Given a subset of channels and their 3D electrode positions, it denoises what was measured, estimates what was not, and can predict signals at new positions on the scalp.

This release brings the pretrained model to Apple platforms as fixed Core ML profiles. The default profile matches the 14-channel Emotiv EPOC X layout, so a consumer headset’s recording can be cleaned and extended on the device that receives it.

Each channel’s five-second window is cut into 40 coarse time tokens of 32 samples. A token’s index carries its electrode position and time, {x, y, z, t}, which is why coordinates are part of the input contract. The encoder runs once; the decoder then denoises a latent over a fixed number of diffusion steps before the tokens are turned back into signal.

The original ZUNA runs in fp16: its 14-, 16-, 32- and 64-channel fp16 profiles pass their parity gates. ZUNA1.1 does not. The base model was trained in bf16, and its decoder activations exceed the fp16 range, so no fp16 profile passes parity. Every ZUNA1.1 profile ships in fp32; the 14-channel profile is the default.

The port surfaced a bug in coremltools. Its rms_norm translation rescales the input by its maximum before squaring, to protect fp16 on the Neural Engine, but adds epsilon after the rescale. Algebraically the epsilon is inflated by max(|x|)². On spiky activations, which are common in transformer residual streams, that changes outputs by whole percents rather than rounding error, even at fp32 on the CPU.

It was reported as apple/coremltools#2821, with a proposed fix in #2822: rescale epsilon by 1/m², clamp the scale so all-zero rows no longer divide 0 by 0, and add regression tests for both.