Zyphra has released ZUNA1.1, an open-source EEG foundation model licensed under Apache 2.0. The update allows researchers to process recordings of variable length, ranging from 0.5 to 30 seconds, rather than the fixed five-second segments required by the original ZUNA1.
The shift addresses a specific problem in real-world data. Sessions often vary in duration, and channels frequently drop out or become noisy mid-stream. Montages differ significantly, from four-electrode headbands to 256-channel research caps. While ZUNA1 handled fixed windows, ZUNA1.1 accepts inputs of arbitrary length.
What is ZUNA1.1?
To understand the flexibility, it helps to look at the model’s function. ZUNA1.1 is a masked diffusion autoencoder with 380 million parameters designed for scalp-EEG signals.
Given a subset of channels, the model denoises existing segments and reconstructs missing ones. It also predicts signals for novel channel positions based on physical coordinates on the scalp.
The parameter count remains unchanged from ZUNA1. The model runs on a consumer GPU and performs acceptably on CPU for many workloads. Weights are hosted on Hugging Face, while inference and preprocessing code are on GitHub. Installation requires `pip install zuna`. Zyphra also provides a free browser EEG Playground, though the tools are for research use only.
How The Architecture Works
This flexibility rests on the tokenization method.
ZUNA functions as a transformer encoder–decoder diffusion autoencoder. It slices each channel into 0.125 second segments, which equates to 32 samples at 256 Hz. Each segment becomes a continuous-valued token, serialized in channel and time order.
The positional encoding is the core idea. Each token carries a 4D rotary positional encoding over (x, y, z, t). This represents the electrode’s 3D scalp coordinate along with its coarse-time index. Because position, not array index, tells the model where a channel sits, ZUNA is channel-agnostic. It accepts any electrode layout and can generate signals at positions never recorded. This capability enables arbitrary channel upsampling by location.
The encoder compresses the signal into a latent. That latent conditions the decoder via adaptive-RMS norm. The decoder is trained with a rectified-flow objective. ZUNA1.1’s architectural changes targeted training stability, including added normalization layers.
What Changed From ZUNA1
Since the architecture stayed close, the differences come from training.
1. Variable-length inputs (0.5–30 seconds): ZUNA1.1 samples a segment length per training example, snapped to the 0.125 s token grid. Lengths are drawn across four bins, from very short to long. The middle 1.5–10 s range is oversampled, since it is the most common operating point. Because token counts vary, Zyphra packs multiple segments per batch up to a fixed budget. Flex attention with a sample-aware mask stops tokens attending across samples. One model therefore serves a 0.5 s snippet and a 30 s stretch without reconfiguration.
2. A richer mixture of reconstruction tasks: ZUNA1 trained on one dropout pattern: uniformly random whole channels. ZUNA1.1 trains on four. The first is whole-channel dropout, covering sparse montages and dead electrodes. The second removes short time stretches across every channel. The third removes those stretches from only some channels, clustering gaps in space and time. The fourth scatters missing values across individual points.
3. Quality-aware preprocessing and a bigger corpus: ZUNA1 made channel-quality calls at the whole-recording level, discarding usable signal. ZUNA1.1 instead computes a per-channel, per-second quality score, thresholded at load time. That grew the corpus from roughly 2M to roughly 3.5M channel-hours of public EEG data. Zyphra team also precomputes two filter variants per recording: a 0.1–45 Hz bandpass, and a 0.01 Hz highpass along with notch. Generalizing across preprocessing strategies is a stated goal, not a benchmarked result.
The Results
Consequently, the question is whether flexibility cost accuracy.
On held-out tasks, ZUNA1.1 reaches better or essentially the same reconstruction NMSE as ZUNA1. Both clearly outperform classical spherical-spline interpolation from MNE. For fair comparison, those evaluation sets used exactly five-second samples.
Zyphra also ran a region-based test. Electrodes from one brain region are deleted, then reconstructed from the remaining seven. That setup is more realistic than random channel dropping. ZUNA1.1 outperforms both spherical-spline and ZUNA1 there.
Interactive Explainer
To make those mechanics concrete, the demo below animates the pipeline end to end.
ZUNA1 vs ZUNA1.1
Taken together, the releases differ mostly in training, not architecture.
| Attribute | ZUNA1 | ZUNA1.1 |
|---|---|---|
| Parameters | 380M | 380M |
| Architecture | Transformer encoder–decoder diffusion autoencoder | Same, plus extra normalization layers |
| Input length | Fixed 5 s | 0.5–30 s, snapped to 0.125 s grid |
| Token | 0.125 s / 32 samples at 256 Hz | Same |
| Positional encoding | 4D RoPE over (x, y, z, t) | Same |
| Decoder objective | Rectified flow | Rectified flow |
| Dropout schemes in training | 1 (uniform random whole-channel) | 4 (channel, time, channel×time, scattered) |
| Training corpus | ~2M channel-hours | ~3.5M channel-hours |
| Quality filtering | Whole-recording level | Per-channel, per-second score at load time |
| Preprocessing variants | Single | Two (0.1–45 Hz bandpass; 0.01 Hz highpass + notch) |
| License | Apache 2.0 | Apache 2.0 |
| Reconstruction NMSE | Baseline | Equal or better |
Running It
Turning to practice, `reconstruct_fif` runs directly on `.fif` files with no `.pt` round-trip. The older four-step pipeline still ships alongside it.
from zuna import reconstruct_fif
reconstruct_fif(
input_dir="fif_in",
output_dir="fif_out",
figures_dir="figures",
gpu_device=0, # GPU id, or "" for CPU
segment_sec=5.0, # window length; default is 5.0, not the full 30 s
montage="standard_1020", # fallback, used only if the file has no positions
repair_channels=["Cz"], # channel(s) to fully reconstruct
target_channel_count=["Fz", "Pz"], # add/upsample new channels by name (or an int for auto)
bad_segments=[(5, 6), (10, 11, "C3")], # mark time spans bad (all channels, or one)
sample_steps=50, # diffusion steps; note: not "diffusion_sample_steps"
)Note the defaults. `segment_sec` is 5.0, so the 0.5–30 s range needs setting explicitly. Electrode positions are read from the file itself. The `montage` argument is only a fallback when positions are absent, and channels without 3D coordinates are dropped.
The reconstruction target is a union. It combines the file’s own MNE bad channels and `BAD_` annotations with anything requested above. Two directories are written. `full_reconstruction/` holds model output everywhere. `hybrid/` keeps the original and infills only inferred cells, plus a `_mask.npz`.
Use Cases With Examples
Because masking is now flexible, several practical patterns open up.
- Dead electrode: Mark `repair_channels=[“Cz”]` to rebuild the channel from its neighbours.
- Motion artifact in a trial: Pass `bad_segments=[(10, 11, “C3”)]` to clean one span on one channel.
- Headband upsampling: Feed four electrodes, then request extra `standard_1005` positions.
- UI-driven cleaning: Supply per-file masks via `mask_dir`, unioned into the target.
What it means
For people making things, the change is practical rather than theoretical. Researchers no longer need to chop long recordings into identical five-second chunks before processing. They can feed the model a 0.5 second snippet or a full 30 second stretch. This reduces manual prep work and makes the tool viable for messy, real-world data where channel dropouts and varying session lengths are common.
