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 can become corrupted by noise 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, with varied electrode positions, sample lengths and corruption profiles.

ZUNA1.1 is a significant step toward making our previous open EEG foundation model ZUNA1 more flexible for practitioners and suitable for real world data. ZUNA1.1 uses a similar  architecture that improves training stability and allows us to train for 4X as many steps. We increased both quality and quantity of our training corpus to ~3.5M channel hours and trained on a more realistic mixture of the ways that real EEG data actually gets corrupted. The result is a more flexible model, able to operate across a far wider range of real-world conditions that yields  reconstruction quality better than, or equal to ZUNA1 on inputs ranging from 0.5 to 30 seconds .

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 random full-channel dropout ZUNA1 was trained on.

ZUNA Model Architecture

ZUNA1.1 maintains the 380M parameter size and transformer-based diffusion autoencoder architecture of the previous model. 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. We have improved training stability with Sandwich-norm and QK-norm throughout the architecture.  For the full architecture details, see the ZUNA1 technical paper and ZUNA1.1 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.0 s) to long (10–30 s), with the middle 1.0–10 s range more heavily represented since it is the most common operating point. This means ZUNA1.1 can serve both a 0.5 s trial snippet and a 30 s continuous stretch, with no retraining or reconfiguration.

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 that EEG gets corrupted. 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, including many real-world EEG failure modes.

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. Not all dropout schemes are equally learnable

Tracking reconstruction error separately for each scheme during training was very informative and  directly shaped the training curriculum.

We trained in three stages. For the first 250k steps (Stage 1) we sampled all eight schemes uniformly. By around 250k steps the evaluation curves for several schemes had visibly plateaued, some because the model had essentially solved them, others because they were too difficult. For steps 250k–440k (Stage 2) we adjusted the difficulty and reweighted the mixture, including raising the overall dropout probability from 0.9 to 0.99, since the model had already mastered the trivial task of reconstructing tokens shown to the encoder. Finally, from 440k to 580k steps (Stage 3), we dropped the four layout schemes entirely and trained on the four structure schemes at uniform weight (the four reported in the section above).

In the figure below, it is evident that the four layout schemes are dramatically harder than the structural ones and saturate early at a high loss. These schemes may irreversibly remove the information needed for a good reconstruction. The consumer-headset scheme, for example, can require reconstructing a full 256-channel montage from as few as four electrodes. The structure schemes, by contrast, kept improving throughout, so we concentrated  compute there. Since the layout schemes were trained on for the majority of the run (440k steps), ZUNA1.1 retains strong montage-upsampling and brain-region reconstruction performance at deployment.

Median NMSE on an evaluation set for each dropout scheme throughout training. Lower is better. Vertical dashed lines indicate where dropout schemes were amended to improve learning.

4. 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 discarded. 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 evaluation dataset, ZUNA1.1 reaches comparable or better reconstruction NMSE as ZUNA1, and both models clearly outperform classical spline interpolation. In other words, ZUNA1.1 maintains reconstruction fidelity of the original model while extending its applicability to a much wider range of real-world conditions including variable input lengths and preprocessing as well as more realistic patterns of signal corruption.

Performance of ZUNA1.1 vs ZUNA1 and MNE Spherical Spline interpolation on a range of reconstruction evaluation datasets. To compare with ZUNA1, evaluation datasets only contained  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.

How much missing data can ZUNA1.1 actually handle?

The most practical question for anyone using the model is: “How much of my EEG data can I fill in with ZUNA and still trust the results?” Because ZUNA1.1 was trained on four distinct structural dropout patterns, we can answer that separately for each one. 

What matters is not how much signal is gone, but how much redundant information survives around each gap. ZUNA1.1 is more robust to higher dropout rates when the dropout schemes are more uncorrelated, either temporally, spatially or both. Visually inspected samples from the model closely matched ground truth dropped data at NMSE ~0.5. Based on this observation, we provide a rough deployment guide:

Reconstruction NMSE on masked tokens as a function of the percentage of tokens dropped, for each of the four structural dropout schemes. Lower is better.

We also tested the performance for a fixed 15% dropout rate with a varied sample duration across the full 0.5–30s range the model supports. We found that ZUNA1.1 performs comparably well across the full possible time range, with slightly reduced performance for very short time segments. 

Reconstruction NMSE versus sample duration with token dropout fixed at 15%, for the four structural dropout schemes.

Finally, we compared the performance of ZUNA1.1 to spherical spline interpolation, MNEs default way of repairing bad channels. ZUNA1.1 outperforms spline interpolation in all three cases. One central reason for this is that ZUNA1.1 can draw on observed information from surrounding channels at the same time-point and on observed signal from the same channel at earlier or later time-points, which is exactly what you need when a channel or group of channels degrades for a short period.

ZUNA1.1 vs MNE spherical-spline reconstruction performance versus dropout percentage on a mixed evaluation dataset, for the three applicable structural dropout schemes.

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.