Zyphra releases ZUNA1.1: Apache 2.0 EEG foundation model with variable length input from 0.5 to 30 seconds

Machine Learning


Zyphra was released this week ZUNA1.1 under Apache 2.0 license. The EEG foundational model reconstructs, denoises, and upsamples data across arbitrary channel layouts. It is built on Zyphra’s initial open EEG foundation model, ZUNA1.

The main change is flexibility, not a quantum leap in raw accuracy. Actual EEG recording is tricky. Session lengths vary, and channels can be noisy or cut out in the middle of the session. The montage ranges from a 4-electrode headband to a 256-channel research cap. ZUNA1 processed only fixed 5-second segments. ZUNA1.1 accepts variable length input from 0.5 seconds to 30 seconds.

What is ZUNA1.1?

To understand its flexibility, start with the model’s features.

ZUNA1.1 is a 380M parameter masked diffusion autoencoder for scalp EEG signals. Specifying a subset of channels denoises existing EEG segments and channels. Rebuild what was lost. It also predicts new channel signals by considering the physical coordinates on the scalp.

The number of parameters is unchanged from ZUNA1. It runs on consumer GPUs and works well on CPUs for many workloads. The weight is placed on the hug face. The inference and preprocessing code can be found on GitHub. Install with pip install zuna. Zyphra also hosts the free browser EEG Playground, all of which are provided for research use only.

How the architecture works

That flexibility relies on tokenization.

ZUNA is a transformer encoder-decoder spreading autoencoder. Slice each channel into 0.125 second segments (32 samples at 256 Hz). Each segment becomes a continuous-valued token. Tokens are serialized in channel x time order.

Positional encoding is the key idea. Each token holds a 4D rotational position encoding over (x, y, z, t). This is the 3D scalp coordinate of the electrode and its coarse temporal index. ZUNA is channel-independent because position, rather than array index, tells the model where the channel should be placed. It accepts any electrode layout and can generate signals at locations never before recorded. This feature allows upsampling of any channel by location.

The encoder compresses the signal into potential. This potentially tunes the decoder via an adaptive RMS norm. The decoder is trained using rectified flow objectives. Architectural changes in ZUNA1.1 target training stability, including the addition of a normalization layer.

Changes from ZUNA1

The architecture remains close, so the difference comes from the training.

1. Variable length input (0.5 to 30 seconds): ZUNA1.1 samples the segment length for each training example and snaps it to a 0.125 second token grid. Length is plotted across four bins, from very short to long. The middle 1.5-10 second range is oversampled because it is the most common operating point. Since the number of tokens varies, Zyphra packs multiple segments per batch up to a fixed budget. Flexible care using sample recognition masks prevents tokens from joining across samples. Therefore, one model provides a 0.5 second snippet and a 30 second stretch without reconfiguration.

2. Richer combination of reconstruction tasks: ZUNA1 was trained with one dropout pattern, i.e. uniformly random across channels. ZUNA1.1 is composed of 4 pieces. The first is a channel-wide dropout, which covers sparse montages and electrode cuts. The second method removes short time ranges across all channels. Third, we remove these stretches from only some channels and cluster the gaps in space and time. The fourth distributes the missing values ​​to individual points.

3. Quality-aware preprocessing and larger corpora: ZUNA1 made channel quality calls at the overall recording level and discarded the usable signal. ZUNA1.1 instead calculates a per-channel, per-second quality score by setting thresholds at load time. This increased the corpus from approximately 2 million channel hours of publicly available EEG data to approximately 3.5 million channel hours. The Zyphra team also precalculates two filter variants for each recording: a 0.1-45 Hz bandpass, and a 0.01 Hz highpass with a notch. Generalizing the overall preprocessing strategy is a stated goal, not a benchmark result.

result

So the question is whether flexibility comes at the cost of accuracy.

For suspended tasks, ZUNA1.1 reaches better or essentially the same reconstructed NMSE than ZUNA1. Both are clearly superior to classical spherical spline interpolation with MNE. To make a fair comparison, these evaluation sets used samples of exactly 5 seconds.

Zyphra also conducted community-based testing. Electrodes from one brain region are removed and reconstructed from the remaining seven regions. This setup is more realistic than random channel drops. There, ZUNA1.1 outperforms both spherical spline and ZUNA1.

interactive explainer

To illustrate how these mechanisms work, the following demo animates the pipeline end-to-end.



Source link