ZUNA: Flexible EEG Superresolution with Position-Aware Diffusion Autoencoders

TL;DR

ZUNA is a large, flexible diffusion autoencoder that improves EEG channel infilling and superresolution across arbitrary electrode configurations, outperforming traditional interpolation methods and generalizing well across datasets.

Contribution

The paper introduces ZUNA, a novel 380M-parameter diffusion autoencoder for EEG superresolution that handles arbitrary electrode positions and outperforms existing methods.

Findings

ZUNA significantly outperforms spherical-spline interpolation at high dropout rates.

ZUNA generalizes across diverse EEG datasets and electrode configurations.

The model is computationally practical for real-world EEG analysis.

Abstract

We present \texttt{ZUNA}, a 380M-parameter masked diffusion autoencoder trained to perform masked channel infilling and superresolution for arbitrary electrode numbers and positions in EEG signals. The \texttt{ZUNA} architecture tokenizes multichannel EEG into short temporal windows and injects spatiotemporal structure via a 4D rotary positional encoding over (x,y,z,t), enabling inference on arbitrary channel subsets and positions. We train ZUNA on an aggregated and harmonized corpus spanning 208 public datasets containing approximately 2 million channel-hours using a combined reconstruction and heavy channel-dropout objective. We show that \texttt{ZUNA} substantially improves over ubiquitous spherical-spline interpolation methods, with the gap widening at higher dropout rates. Crucially, compared to other deep learning methods in this space, \texttt{ZUNA}'s performance \emph{generalizes} across datasets and channel positions allowing it to be applied directly to novel datasets and problems. Despite its generative capabilities, \texttt{ZUNA} remains computationally practical for deployment. We release Apache-2.0 weights and an MNE-compatible preprocessing/inference stack to encourage reproducible comparisons and downstream use in EEG analysis pipelines.