Introduction

EEG is the most accessible window we have into the living brain. It is non-invasive, relatively cheap, and BCI devices are increasingly wearable. However, real EEG recordings are often messy. They come in various lengths, individual channels sometimes go noisy or drop out mid-session, and montages range from four-electrode consumer headbands to 256-channel research caps. A foundation model that can reconstruct, denoise, and upsample EEG is only useful in practice if it can process real-world data, not just the clean, fixed-format slices it was trained on.

ZUNA1.1 is our step toward making our previous open EEG foundation model ZUNA1 more flexible for practitioners and suitable for real world data. ZUNA1.1 keeps much of the same architecture but we took a different approach to how the model was trained. We expanded and increased the quality of our training corpus to ~3.5M channel hours and trained on a realistic mixture of the ways that real EEG data actually gets corrupted. The result is a more flexible model that can now accept variable-length inputs from 0.5 to 30 seconds and we find that the reconstruction quality is better than, or equal to ZUNA1 while now being able to operate across a far wider range of real-world conditions.

We release ZUNA1.1 under the highly permissive Apache 2.0 license. You can download the model and run it locally or use it instantly in your browser on the Zyphra Cloud in our EEG Playground. Cloud-served inference is routine for large language models but still rare for EEG. Now anyone can try ZUNA1.1 on their data with no GPU, no install, and no coding or machine learning knowledge required.

Making ZUNA Easier to Use

When we released ZUNA1, it did one thing very well: given a subset of EEG channels and their electrode positions, it reconstructed the missing channels, outperforming the ubiquitous spherical-spline interpolation baseline from MNE, with the relative performance of ZUNA increasing as more channels went missing. Since then, feedback from the researchers, clinicians, and BCI developers who put ZUNA to work told us a consistent story.  While raw reconstruction accuracy is obviously important, flexibility and frictionless use is also vital. 

Critical capabilities that make ZUNA1.1 more useful for researchers and practitioners include: The ability to accept recordings that are of various lengths between 0.5 - 30 seconds rather than exactly five seconds long, process channels that are noisy for part of a session but fine for the rest, and accommodate missing-channel patterns that look nothing like the uniform random dropout ZUNA was trained on.

ZUNA Model Architecture

ZUNA is a transformer-based encoder–decoder diffusion autoencoder. It slices each EEG channel into short 0.125 second segments (32 samples at 256 Hz), turns each segment into a continuous-valued token, and serializes them in channel × time order. The key idea is the positional encoding. Each token carries a 4D rotary positional encoding over (x, y, z, t), the electrode's 3D scalp coordinate plus its coarse-time index. Because position, not array index, tells the model where a channel sits, ZUNA is channel-agnostic: it accepts any number of electrodes in any layout and can even generate signals at positions that were never recorded. This capability allows ZUNA1.1 to perform arbitrary channel upsampling by location. The encoder compresses the signal into a latent that conditions the decoder via adaptive-RMS norm; the decoder is trained with a rectified-flow objective. ZUNA1.1 maintains the 380M parameter size of the previous model. For the full architecture details, see the ZUNA1 technical paper. 

Architecture schematic for ZUNA1.1. Like our earlier model, ZUNA1.1 uses a diffusion encoder-decoder trained to reconstruct masked EEG channels. The principal changes we made to the architecture were to improve training stability such as adding additional normalization layers. 

What's New in ZUNA1.1

1. Variable-length inputs (0.5–30 seconds)

While our previous model, ZUNA1, processed only fixed 5 s segments. ZUNA1.1 samples a segment length per training example, snapped to a 0.125 s token grid, drawn across four bins from very short (0.5–1.5 s) to long (10–30 s), with the middle 1.5–10 s range oversampled since it is the most common operating point. Because samples now vary substantially in token count, we pack multiple segments into each batch up to a fixed budget and use flex attention with a sample-aware mask so tokens from one sample never attend to tokens from another. This means ZUNA1.1 can serve both a 0.5 s trial snippet and a 30 s continuous stretch, with no retraining or reconfiguration.

Because each pass takes a different random mask and re-normalizes it locally, the model sees the same underlying recording under many different numerical framings. A nice side-effect of this is that it is essentially free data augmentation. This plus the data augmentation provided by the various dropout schemes lets us train for substantially more epochs over our corpus without overfitting. 

2. A richer mixture of reconstruction tasks

While ZUNA1 was trained entirely on a single dropout pattern where we simply removed a uniformly-random subset of whole channels, we aimed to cover a substantially larger range of different dropout schemes with ZUNA1.1. This is vital to further enable ZUNA1.1’s  extremely flexible reconstruction abilities for more realistic settings. For instance, motion artifacts corrupt clusters of nearby electrodes at once; electrodes drop out in bursts correlated over time; transmission glitches can knock out a stretch of the whole recording but without taking an entire channel offline completely.

ZUNA1.1 was trained on a mixture of four dropout schemes, each capturing a different way EEG gets corrupted or goes missing. The first is whole-channel dropout where entire channels are removed, covering sparse montages and dead electrodes. The second removes short stretches of time across every channel at once, mimicking moments when the whole signal drops out or a head movement corrupts everything briefly. The third removes those same time stretches from only some channels rather than all of them, producing gaps clustered in both space and time. And the fourth scatters missing values randomly across individual points in the recording, reflecting the temporary, localised noise of a muscle twitch or a momentary channel failure. By training across this mixture of dropout schemes, ZUNA1.1 can handle almost arbitrary reconstructions.

ZUNA1.1’s new set of dropout schemes is significantly more diverse than ZUNA1, which dropped entire channels over all time per sample. Here, we show the combinations of time and channel masking used in the training of ZUNA1.1.

3. Quality-aware preprocessing and a bigger corpus

Our original ZUNA1 data pipeline made channel-quality decisions at the whole-recording level. If a channel was noisy in any long stretch, it was zeroed for the entire recording, and epochs with too many bad channels were thrown out. This was wasteful since a channel that is clean for 90% of a session still carries a partial, usable signal.

For the training of ZUNA1.1, we instead compute a per-channel, per-second quality score that allows us to threshold the data used for training at load time. This lets us recover signal from partially-noisy channels and ingest datasets that our old pipeline could not, growing the training corpus from ~2M to roughly 3.5M channel-hours of public EEG data.

We also precompute two filter variants for every recording. One variant was a bandpass filtered from 0.1 to 45HZ, resembling the frequency band in which EEG data is most commonly analyzed. We then added a second variant with less stringent preprocessing, applying only a highpass filter at 0.01Hz as well as a notch filter to remove power line noise. Training the model on data with different preprocessing allows it to be more versatile which is important given that real-world EEG recordings often have heterogeneous preprocessing applied. One of the goals for our ZUNA1.1 model is the ability for it to generalize across different preprocessing strategies. 

Results

We found that incorporating this flexibility does not have obvious costs in reconstruction quality. On our held-out evaluation tasks, ZUNA1.1 reaches better or essentially the same reconstruction NMSE as ZUNA1, and both models clearly outperform classical spline interpolation. In other words, ZUNA1.1 keeps or surpasses the accuracy of the original model while extending it to a much wider range of real-world conditions including variable-length inputs and realistic, correlated missingness, at no cost to fidelity. 

Performance of ZUNA1.1 vs ZUNA1 and MNE Spherical Spline interpolation on a range of reconstruction evaluation datasets. To compare with ZUNA1, we restricted the evaluation sets to those with exactly 5 second samples. 

We also see that ZUNA1.1 outperforms spherical-spline and ZUNA1 on reconstructions for given brain regions. 

We delete electrodes from a specific brain region, and then reconstruct that region given the remaining seven regions, a more experimentally realistic setup than typical masked autoencoders which drop channels at random. ZUNA1.1’s new training regime yields tangible benefits on more realistic scenarios for missing data.

Try ZUNA1.1: Download it or Run it in Your Browser

As with ZUNA1, everything is open-source. The ZUNA1.1 model weights are available on HuggingFace. The inference and MNE-compatible preprocessing code are available on GitHub. ZUNA1.1 can be installed using pip install zuna. The model runs fast on a consumer GPU and acceptably on CPU. Because it's Apache 2.0 you can build on it freely, including self-hosting for sensitive or clinical data that can't leave your premises.

In addition, ZUNA1.1 is available on the Zyphra Cloud. This means you can now run ZUNA1.1 without installing anything, simply upload your EEG recordings (in .fif file format) to Zyphra Cloud and run the model on our servers, directly in the browser. 

Demo of ZUNA1.1 EEG Playground in Zyphra Cloud.

Disclaimer: This website and related services (“Services”) are provided for research use only and is not intended for use in the diagnosis, cure, mitigation, treatment, or prevention of any disease or health condition. The Services have not been validated for any medical or clinical use. The information provided through the Services is for general informational purposes only and is not a substitute for any professional medical or healthcare advice. We do not warrant that any information provided through the Services is accurate, complete, or useful to you. Any reliance you place on such information is strictly at your own risk.