ZUNA1.1: A more flexible EEG foundation model for Denoising and Super-resolution

# ZUNA1.1: A more flexible EEG foundation model for Denoising and Super-resolution
Source: [https://arxiv.org/html/2607.27308](https://arxiv.org/html/2607.27308)
###### Abstract

We introduceZUNA1\.1, a 380M\-parameter diffusion autoencoder for flexible EEG signal reconstruction\.ZUNA1\.1is capable of reconstructing variable length sequences of up to 30 s, with an arbitrary number of EEG channels at arbitrary scalp locations, and can reconstruct arbitrary temporal intervals within channels in addition to reconstructing entire channels\. We demonstrate thatZUNA1\.1performs at least on par with our earlierZUNA1model, while being far more flexible and capable of handling a wide range of reconstruction tasks\.ZUNA1\.1continues to substantially outperform standard EEG denoising and reconstruction methods such as spherical spline interpolation, which is ubiquitously deployed in the MNE package\. TheZUNA1\.1model is released open source under the permissive Apache 2\.0 license\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/region_occlusion_brain_region_dataset_average_light.png)Figure 1:Per\-region reconstruction accuracy \(topographic\)\.Average NMSE by scalp region forZUNA1\.1,ZUNA1\(labeled with theΔ\\Deltarelative toZUNA1\.1\), and spherical\-spline, averaged over the four datasets\. All electrodes in one region are deleted and reconstructed from the remaining seven regions\. Lower \(greener\) is better\.## IIntroduction

EEG is the most accessible modality for measuring brain activity, valued in clinical, research, and consumer settings for its simple setup and low cost\. Although the skull blurs the signal at each electrode to approximately an average of over a billion neurons, scalp recordings still carry substantial decodable information\. Prior work has classified sleep stages\(Suprataket al\.,[2017](https://arxiv.org/html/2607.27308#bib.bib7); Liet al\.,[2022](https://arxiv.org/html/2607.27308#bib.bib272)\), decoded emotional states\(Lianget al\.,[2019](https://arxiv.org/html/2607.27308#bib.bib283); Duanet al\.,[2020](https://arxiv.org/html/2607.27308#bib.bib273)\), and tracked attentional focus\(Suet al\.,[2022](https://arxiv.org/html/2607.27308#bib.bib274); Kaushiket al\.,[2022](https://arxiv.org/html/2607.27308#bib.bib275)\)\. This success has driven adoption beyond the lab\. A growing range of consumer\-grade headsets now makes EEG affordable and easy to deploy outside research settings\. However, consumer and inexpensive lab\-grade headsets often have a much smaller number of channels and poorer data quality than full lab\-grade devices\. Moreover, all EEG data collection methods often suffer from occasional corruption, artifacts, and general low quality recordings\. For instance, electrodes can lose contact or their signal can weaken, motion artifacts can disrupt entire groups of channels at a time, and general noise remains prevalent\.

To address these limitations, we recently introducedZUNA1\(Warneret al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib8)\), a 380M\-parameter masked diffusion autoencoder that repairs bad channels and upsamples consumer\-grade recordings to as many as 256 channels, matching the highest quality lab\-grade systems\.ZUNA1tokenizes 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\. It substantially outperformed spherical\-spline interpolation\(Perrinet al\.,[1989](https://arxiv.org/html/2607.27308#bib.bib276)\)for channel reconstruction across heterogeneous datasets\. Our thinking was that much of the EEG signal lost due to various sources of noise was contained redundantly across the channel montage as a whole, however, it is encoded nonlinearly and is thus inaccessible to conventional spline interpolation methods\. Instead, by training EEG foundation models, which can learn generalizable patterns across millions of channel\-hours of EEG data, these foundation models would be able to provide a similar boost to that which foundation models have produced in other modalities such as text, images, and audio\.

In addition to ZUNA1, a number of other EEG foundation models have been developed in recent years\(Kuruppuet al\.,[2025](https://arxiv.org/html/2607.27308#bib.bib277)\), differing mainly in their self\-supervised objective\. Contrastive models such as BENDR\(Kostaset al\.,[2021](https://arxiv.org/html/2607.27308#bib.bib260)\)and SleepFM\(Thapa and others,[2024](https://arxiv.org/html/2607.27308#bib.bib261)\), and a large family of masked\-reconstruction models including BIOT\(Yanget al\.,[2023](https://arxiv.org/html/2607.27308#bib.bib278)\), LaBraM\(Jianget al\.,[2024b](https://arxiv.org/html/2607.27308#bib.bib279)\), CBraMod\(Wang and others,[2024b](https://arxiv.org/html/2607.27308#bib.bib262)\), Brant\(Zhanget al\.,[2023](https://arxiv.org/html/2607.27308#bib.bib263)\), BrainWave\(Yuanet al\.,[2024](https://arxiv.org/html/2607.27308#bib.bib280)\), LUNA\(Doner and others,[2025](https://arxiv.org/html/2607.27308#bib.bib264)\), and REVE\(El Ouahidi and others,[2025](https://arxiv.org/html/2607.27308#bib.bib267)\), dominate, alongside autoregressive\(Jianget al\.,[2024a](https://arxiv.org/html/2607.27308#bib.bib281); Cuiet al\.,[2024](https://arxiv.org/html/2607.27308#bib.bib282)\)and latent\-predictive \(JEPA\)\(Wang and others,[2024a](https://arxiv.org/html/2607.27308#bib.bib269); Panchavati and others,[2026](https://arxiv.org/html/2607.27308#bib.bib270)\)approaches\.ZUNAdiffers from this body of work in three respects\. First, unlike all but a handful of recent efforts, it is trained at scale \(∼3\.5\{\\sim\}3\.5M channel\-hours\)\. Second, it uses a position\-conditioned*diffusion*architecture rather than a discriminative encoder\. Third, and most importantly, nearly all of these models are optimized to learn a general\-purpose latent space which can then be fine\-tuned or linearly probed for downstream*classification*, whereasZUNA’s training objective is the signal*reconstruction*itself\. This makesZUNAdirectly suited to flexibly infilling missing or corrupted EEG, the problem we target, rather than producing features that a separate model must then decode\.

However, deploying ZUNA1 to real workflows revealed two practical limitations\. Firstly, the model operated only on fixed 5\-second windows, with no support for longer or shorter sub\-segments\. This is clearly a major limitation since almost all naturally occurring EEG data is not exactly five seconds in duration, and chopping existing data into the 5\-second chunks often causes unnatural boundary artifacts\. Secondly, the original ZUNA1 training used a single dropout scheme that randomly sub\-selected channels\. The problem of the random dropout scheme was that it does not necessarily match the patterns of missing channels in real world cases\. In real\-world recordings, an individual channel completely dropping and remaining unavailable for the entire duration is rare\. A much more common pattern is bursts of errors, often individual or correlated across many channels due to temporary electrode glitches or motion artifacts\. Moreover, consumer headsets impose specific sparse layouts of channels that tend to fail in spatially or temporally correlated bursts\.

To address these limitations, we introduceZUNA1\.1, which retains the architecture ofZUNA1and addresses each of the limitations above through the following contributions:

- •*Variable\-length training*:ZUNA1\.1trains on randomly cropped windows from0\.50\.5to3030seconds, packed into batches via flex attention with sample\-aware masks for efficient training and inference\.
- •*Implicit augmentation*:ZUNA1\.1sees each underlying signal under multiple transformations across exposures: random temporal crops combined with per\-window z\-score normalization, and two pre\-computed filter variants of all data \(highpass\+notch and bandpass\)\. This provides substantial implicit data augmentation that lets us train for many more passes over the corpus without overfitting\.
- •*Quality\-aware preprocessing*: each recording is paired with a per\-channel, per\-second quality matrix computed from amplitude extremes, segment\-level standard deviation, and flatness/clipping detectors, with thresholds applied at load time rather than at preprocessing time, so the same on\-disk corpus can be re\-thresholded for any desired stringency without re\-preprocessing\.
- •*Diverse dropout mixture*: training uses a mixture of eight different dropout schemes spanning uncorrelated, contiguous, anatomical, montage\-density, and consumer\-headset patterns, broadening the range of realistic infilling scenarios the model handles without retraining\.
- •*Expanded corpus and training scale*: quality\-aware loading and variable\-length training together let us recover signal from partially noisy channels and ingest pre\-epoched datasets the original pipeline could not use, growing the training corpus from approximately 2M channel\-hours\(Warneret al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib8)\)to approximately 3\.5M channel\-hours\. In addition to the increased training corpus, ZUNA1\.1 was trained for substantially more steps\.

As withZUNA1, we release theZUNA1\.1weights under the Apache License 2\.0 on Hugging Face, together with inference and preprocessing code on GitHub\. The package is available viapip install zuna\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/zuna11_eeg_architecture_light.png)Figure 2:ZUNA1\.1architecture\.A transformer\-based diffusion autoencoder\. The main modifications to theZUNA1architecture are adding sandwich norm and QK\-norm to improve training stability\. This allowed us to remove Adaptive Loss Weighting from the decoder RF loss term\.
## IIBackground and Architecture

ZUNA1\.1is built on a slightly modified version of the ZUNA1 model architecture and we refer readers to our original release\(Warneret al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib8)\)for a full overview of the model specifications\.ZUNA1\.1is a transformer\-based encoder–decoder*diffusion autoencoder*for masked channel reconstruction\. The encoder produces latent representations that are injected into the decoder\. The decoder is trained with a rectified\-flow loss\. The encoder latent is regularized with an auxiliary MMD loss\. In an effort to improve the quality of latent representations for downstream tasks, we experimented with reducing the weight on the MMD regularization term and discuss these experiments further below\. To improve training stability, we added sandwich normalization\(Dinget al\.,[2021](https://arxiv.org/html/2607.27308#bib.bib1)\)and Query\-Key normalization\(Henryet al\.,[2020](https://arxiv.org/html/2607.27308#bib.bib5)\)to each layer in the encoder and decoder of theZUNA1architecture, shown in Figure[2](https://arxiv.org/html/2607.27308#S1.F2)\. These architectural improvements allowed us to remove adaptive loss weighting\(Genget al\.,[2024](https://arxiv.org/html/2607.27308#bib.bib6)\)from the decoder RF\-loss without reintroducing training instabilities\.

ZUNA1\.1retains the same discretized positional encoding scheme asZUNA1implemented with 4D\-RoPE but extends training support to segment lengths from0\.50\.5to3030seconds\. This encoding scheme allows the model to process arbitrary sequence lengths \(up to 30 s\) as well as an arbitrary number of channels at arbitrary scalp locations\. The total parameter count is unchanged at 380M\. In early ablations, we explored scaling the model size to 1\.2B parameters, and found that scaling up did not meaningfully improve reconstruction or performance on downstream tasks\.

## IIIMethod

TABLE I:Channel\-dropout schemes and their sampling weights across the three training stages\.Stage 1 \(0–250250k steps\) samples all eight schemes with the1/81/8weights shown; Stage 2 \(250250–440440k steps\) samples all eight schemes with the \(non\-uniform\) weights shown; Stage 3 \(440440–580580k steps\) drops the four*layout*schemes \(bottom four, whose loss had plateaued\) and samples the four*structure*schemes \(top four\) uniformly\. Within any selected scheme, dropout is applied with probabilitypdrop=0\.9p\_\{\\text\{drop\}\}=0\.9\.### III\-AVariable\-length sequence training

InZUNA1, every training sample was a fixed five\-second segment\.ZUNA1\.1instead samples a window lengthwwuniformly at random per training example, snapped to a multiple of theτ=0\.125​s\\tau=0\.125\\,\\text\{s\}coarse\-time unit\. Window lengths are drawn from a piecewise\-uniform distribution over four bins:

This distribution focuses on the11–1010s range, which we found to be the most common operating regime in downstream applications, while preserving meaningful coverage of both very short windows \(which exercise the model’s short\-context behavior\) and very long windows \(which let the model learn long\-range structure\)\.

A consequence of variable\-length training is that the per\-sample token count varies dramatically\. We extend the sample\-packing strategy fromZUNA1to handle this efficiently\. Random\-length windows are packed contiguously into a batch and padded to fit a target packed sequence length\. Flex attention is then used to compute a sample\-aware attention mask that prevents tokens from one sample attending to tokens of another\. This approach allows a compiled model to efficiently process a large batch of tokens from multiple samples in parallel, filling GPU memory and maximizing utilization\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/dropout_schemes_schematic_light.png)Figure 3:The four channel\-dropout schemes retained throughout training\.Each panel shows one training sample: channels run along the vertical axis and coarse\-time windows \(τ=0\.125\\tau=0\.125s\) along the horizontal axis\. Observed tokens \(unshaded traces\) are passed to the encoder; the shaded spans are set to zero in the encoder and reconstructed by the decoder based on observed signals\.full\-channel\-randomremoves entire channels;full\-time\-pt\-randomremoves short time spans across all channels;correlated\-channel\-timeremoves those spans from a subset of channels;random\-uniformscatters missing tokens independently\. These are the four*structural*dropout schemes retained for the full training run \([section˜III\-D](https://arxiv.org/html/2607.27308#S3.SS4)\); the four*layout*schemes, phased out after440440k steps, are not shown\.
### III\-BQuality\-aware preprocessing

Previously,ZUNA1’s preprocessing pipeline operated at the recording level such that a channel which exceeded noise thresholds for any portion of a recording was zeroed for the whole recording, and epochs with too many bad channels were discarded entirely\. Instead,ZUNA1\.1’s pipeline stores the full continuous recording on disk and computes a per\-channel, per\-second*quality matrix*Q∈\[0,1\]C×SQ\\in\[0,1\]^\{C\\times S\}, whereCCis the number of channels andSSis the number of one\-second segments in the recording\. The quality score combines three markers, each comparing a segment against the statistics of its own channel \(or, for flatnessqflatq^\{\\text\{flat\}\}, of the whole recording\)\. Letσc,s\\sigma\_\{c,s\}denote the standard deviation of channelccin segmentss,σ~c\\tilde\{\\sigma\}\_\{c\}the median ofσc,s\\sigma\_\{c,s\}over segments, andρc,s\\rho\_\{c,s\}the segment’s peak\-to\-peak amplitude relative to its channel\-median peak\-to\-peak\. The three scores are

qc,svar\\displaystyle q^\{\\text\{var\}\}\_\{c,s\}=exp⁡\(−12​\(log⁡\(σc,s/σ~c\)0\.7\)2\),\\displaystyle=\\exp\\\!\\left\(\-\\tfrac\{1\}\{2\}\\left\(\\tfrac\{\\log\(\\sigma\_\{c,s\}/\\tilde\{\\sigma\}\_\{c\}\)\}\{0\.7\}\\right\)^\{\\\!2\}\\right\),qc,sptp\\displaystyle q^\{\\text\{ptp\}\}\_\{c,s\}=\(1\+e2​\(ρc,s−3\.5\)\)−1,\\displaystyle=\\left\(1\+e^\{\\,2\(\\rho\_\{c,s\}\-3\.5\)\}\\right\)^\{\-1\},qc,sflat\\displaystyle q^\{\\text\{flat\}\}\_\{c,s\}=𝟙​\[σc,s≥0\.05⋅medianc⁡σ~c\],\\displaystyle=\\mathbbm\{1\}\\\!\\left\[\\sigma\_\{c,s\}\\geq 0\.05\\cdot\\operatorname\{median\}\_\{c\}\\tilde\{\\sigma\}\_\{c\}\\right\],and the final quality is their element\-wise minimum,Qc,s=min⁡\(qc,svar,qc,sptp,qc,sflat\)Q\_\{c,s\}=\\min\(q^\{\\text\{var\}\}\_\{c,s\},\\,q^\{\\text\{ptp\}\}\_\{c,s\},\\,q^\{\\text\{flat\}\}\_\{c,s\}\)\. The variance score penalizes both noise bursts and flat stretches \(deviations of the segment’s variance from the channel’s typical variance, measured in log\-space\); the peak\-to\-peak score falls off steeply for large transient artifacts; and the flatness score hard\-zeroes dead channels, including reference channels that are flat by construction\.

The quality matrix is saved alongside the recording as a separate mmap file, and quality thresholds are applied at*load time*\. The data loader takes two thresholds,min\_quality\_anyandmin\_quality\_mean: a channel is included in the current training window if every one\-second segment overlapping the window has quality≥\\geqmin\_quality\_any*and*the mean quality across those segments is≥\\geqmin\_quality\_mean\(set to0\.10\.1and0\.30\.3respectively for the released model\)\. This means that a channel which is noisy for part of a recording can still contribute usable training data drawn from the clean portions, and the same on\-disk corpus can be re\-thresholded for any desired quality stringency without rerunning preprocessing\.

The recording is stored as a memory\-mappedfloat32array of shape\(C,T\)\(C,T\), whereTTis the number of samples in the recording at256​Hz256\\,\\text\{Hz\}\. Pre\-epoched recordings \(which the originalZUNA1pipeline could not ingest\) are stored as\(E,C,Te\)\(E,C,T\_\{e\}\)whereEEis the number of epochs andTeT\_\{e\}is the number of samples per epoch\. Per\-channel z\-score normalization parameters used at preprocessing time are stored in the recording’s metadata so that the operation can be reverted if needed\.

### III\-CMultiple filter variants

EEG analyses use a wide variety of upstream filtering conventions, and there is no universally agreed default\. To allowZUNA1\.1to interoperate with users’ existing pipelines, we precompute and save two filter variants of every recording:

- •Notch variant: A0\.01​Hz0\.01\\,\\text\{Hz\}highpass plus zero\-phase notch filters at powerline frequencies and their harmonics \(5050,6060,100100,120120,150150,180180,200200, and240​Hz240\\,\\text\{Hz\}, restricted to frequencies below Nyquist\)\. Since our corpus mixes recording sites on50​Hz50\\,\\text\{Hz\}and60​Hz60\\,\\text\{Hz\}mains power, both families are removed\. This variant preserves the broadband signal, including slow drifts and high\-frequency content, apart from line noise\.
- •Bandpass variant: A0\.10\.1–45​Hz45\\,\\text\{Hz\}bandpass filter applied independently to the raw recording\. Suitable for users whose downstream analyses operate in the conventional clinical EEG band\.

Both variants are sampled with equal probability during training: each exposure of a recording samples one of the two, which contributes to the implicit augmentation described in[section˜III\-A](https://arxiv.org/html/2607.27308#S3.SS1)\. Users can select a variant at inference time to match their preferred preprocessing convention\. Filtering is applied with EEGLAB\-compatible boundary handling \(skip\_by\_annotation\) so that boundary annotations from re\-concatenated cleaned epochs do not introduce filter ringing across discontinuities\.

### III\-DDiverse channel dropout mixture

Perhaps the most consequential advancement is the dropout mixture used at training time\. The originalZUNA1model was trained exclusively with a single dropout scheme that randomly removed whole channels with uniform probability\. However, this did not match many of the patterns of degradation or noise encountered in practice\. For instance, consumer EEG devices have very specific sparse channel layouts, motion artifacts tend to corrupt anatomically clustered channels rather than randomly distributed channels, and electrode dropout is often correlated in time as well as in space\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/median_NMSE_do_vs_step_newB.png)Figure 4:Reconstruction performance for each dropout scheme throughout training\.Median NMSE \(lower is better\) on the frozen evaluation dataset for each dropout scheme\. See Table[I](https://arxiv.org/html/2607.27308#S3.T1)for details on how the dropout schemes were adjusted during training\.ForZUNA1\.1, we instead trained on a diverse mixture of eight dropout schemes, summarized in Table[I](https://arxiv.org/html/2607.27308#S3.T1), which fall into two families: four that control the dropout*structure*and four that control the*layout*of the retained channels\. During training, we used a three\-stage curriculum\. See Table[I](https://arxiv.org/html/2607.27308#S3.T1)for details\. For the first250250k steps we sample all eight schemes using uniform weights\. From250250k \-440440k steps, observing that evaluation metrics had plateaued for many of the dropout schemes \(cf\. Fig\.[4](https://arxiv.org/html/2607.27308#S3.F4)\) because they were too hard or already learned, we changed the parameters of individual dropout schemes to adjust their difficulty level and adjusted the weight of each scheme as reflected in Stage 2\. This included increasing overall dropout probability from0\.90\.9to0\.990\.99because the model had perfectly learned the easy task of reconstructing non\-dropped tokens\. Finally, in Stage 3, from440440k \-580580k steps, we retained only the 4 structure dropout schemes at uniform weighting\. We found the four layout schemes were substantially more challenging for the model to learn than the structural schemes\. For instance, the consumer\-headset scheme can require reconstructing a full 256\-channel montage from as few as four electrodes\.

The two families of dropout schemes play complementary roles\. The four*structure*schemes control the*pattern*of dropout, from fully unstructured token dropout through to spatio\-temporally correlated dropout that mimics motion artifacts and transient electrode failures\. The four*layout*schemes instead control which channels are retained\. The montage and consumer\-headset schemes select channels by matching each target electrode position to the nearest recorded electrode in 3D space, so they apply to any input montage, while the brain\-region scheme partitions the recorded electrodes themselves by scalp coordinates\. Together, the eight schemes expose the model to a wide range of infilling scenarios we target at deployment: upsampling a sparse consumer headset to a dense research montage, repairing an anatomically\-clustered block of bad channels, and filling temporal gaps from transmission dropouts\. Because the layout schemes proved hardest and appeared to quickly saturate at a high loss, we concentrated the final training stretch on the four structure schemes which were amenable to further progress\. However, since we trained on the layout schemes for the majority \(440440k\) of its steps,ZUNA1\.1retains strong montage\-upsampling and region\-reconstruction performance at deployment\.

It should be noted that by combining the random crop dropout schemes and the per\-window z\-score normalization, we achieve substantial*implicit data augmentation*\. Each time the same recording is loaded, a different temporal crop is randomly sampled, and the per\-channel mean and standard deviation used for normalization are computed locally on that crop\. The resulting numerical input therefore differs across exposures even when the underlying signal is identical, which broadens the input distribution the model sees and reduces effective overfitting\. This effect is essentially free, requiring no explicit augmentation pipeline\. Supporting this, we trained for 6 epochs without noticing any signs of overfitting in validation loss\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/zuna11_reconstruction_dropout_curves_light.png)Figure 5:Reconstruction accuracy as channels drop\.Reconstruction NMSE versus channel\-dropout rate on four held\-out datasets \(ANPHY\-Sleep, BerlinBCI, BCI2000, AAD\), comparingZUNA1\.1,ZUNA1, and MNE spherical\-spline interpolation; lower is better\.ZUNA1\.1matches or improves onZUNA1, and both clearly outperform spline interpolation, with the gap widening as more channels are removed\. Evaluation restricted to55s samples for comparability withZUNA1\.![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/region_occlusion_region_error_bars_dataset_average_light.png)Figure 6:Per\-region reconstruction error\.The region\-occlusion evaluation of[fig\.˜1](https://arxiv.org/html/2607.27308#S0.F1)shown as grouped bars: average NMSE across the four datasets with error bars giving the propagated standard deviation\. Lower is better\.
### III\-ETraining details

ZUNA1\.1was trained for 580k steps on approximately 3\.5 million channel\-hours of processed and cleaned EEG data\. Since predicting the non\-dropped tokens is a fairly trivial identity operation, which the model learns extremely quickly, we downweight the loss contribution of the non\-dropped tokens to 10% of the weight for dropped tokens, concentrating the objective on the infilling task while still lightly supervising the retained signal\.\. We used the standard weight decay of 0\.1 on all layers except for embedding and norm layers\.

ZUNA1\.1is trained with AdamW \(β1=0\.9\\beta\_\{1\}=0\.9,β2=0\.95\\beta\_\{2\}=0\.95\) with a weight decay of0\.10\.1applied to the linear, MLP, and attention\-projection weights but not to the embedding or normalization parameters\. The learning rate is cosine\-annealed from a peak of5×10−45\\times 10^\{\-4\}to9×10−59\\times 10^\{\-5\}over580580k steps, following a1,0001\{,\}000\-step linear warmup\. The model trains inbfloat16with normalization layers and the residual stream kept infloat32for stability, and is compiled end\-to\-end withtorch\.compile\. We additionally maintain an exponential moving average \(EMA\) version of the model weights with decay0\.99990\.9999per step, which is used for evaluation and inference\. We ran data\-parallel training across 6 nodes with a per\-rank micro\-batch of 22,000 packed tokens and 2 gradient\-accumulation steps; giving a global batch size of roughly22M tokens per optimizer step\.

Training at this scale on heterogeneous public EEG surfaced stability issues worth recording\. Rare extreme samples produced sudden gradient\-norm spikes followed by loss excursions, and in a handful of cases NaN losses\. We addressed these with three safeguards\. First, optimizer updates are skipped entirely whenever a non\-finite loss is detected\. Secondly, we add an epsilon to the z\-score denominator to protect against zero\-variance \(flat or clipped\) channel segments\. Thirdly, we conducted a corpus scan and removed a small set of recordings whose data and quality matrices contained NaNs\. The removed data comprised approximately0\.3%0\.3\\%of our total dataset, and we traced the nans to conversion artifacts in three sleep datasets\. With architectural improvements and these mitigations in place, we were able to trainZUNA1\.1for4×4\\timesas many steps asZUNA1with no signs of instability\.

## IVResults

### IV\-AReconstruction

To compare to previous work in\(Warneret al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib8)\), we evaluateZUNA1\.1on four carefully chosen datasets that represent a diverse set of tasks and channel counts, ranging from sleep to motor control, each with 5 second samples thatZUNA1can operate on\. Note that 5 second samples are also at the heart ofZUNA1\.1’s training distribution\. The first experiment repeats Figure 4 of the originalZUNA1paper, where we test the sensitivity of the reconstruction methods with respect to the proportion of deleted channels, progressively deleting more channels and giving the models less information to reconstruct channels\. In this experiment, deleted channels are chosen uniformly at random, and are deleted for all time in the sample \(matching the dropout training scheme ofZUNA1\)\. We find in Figure[5](https://arxiv.org/html/2607.27308#S3.F5)that, on average, the performance difference betweenZUNA1\.1andZUNA1is relatively small, whereZUNA1\.1outperformsZUNA1, on average but not uniformly across all datasets\. However, we unequivocally find that both methods are substantially better than spherical\-spline interpolation\. This performance comes in conjunction withZUNA1\.1being trained on a wider range of sample durations with more diverse dropout schemes, which represent more realistic experimental setups such thatZUNA1\.1has much greater range and usability for real\-world EEG signals\.

Since spherical spline interpolation heavily leverages proximal channels to predict missing signals, it can maintain reasonable performance predicting channels when signal is retained in nearby channels\. To stress test the channel prediction methods, we stratified channels into eight regions and dropped all channels from each region, predicting them from the remaining seven regions\. Figures[6](https://arxiv.org/html/2607.27308#S3.F6)and[1](https://arxiv.org/html/2607.27308#S0.F1)illustrate thatZUNA1\.1performs at least as well asZUNA1and both outperform the spherical spline method on this reconstruction task\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/nmse_do_vs_pct_dropout_linear_light.png)\(a\)NMSE vs token dropout percentage\.
![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/nmse_do_vs_duration_fixed0.15_linear_light.png)\(b\)NMSE vs sample duration with token dropout fixed at15%15\\%\.

Figure 7:Reconstruction performance for 4 structural dropout schemeson masked tokens in Evaluation dataset\. NMSE \(lower is better\)\. Description in Figure[3](https://arxiv.org/html/2607.27308#S3.F3)As discussed above, two major areas of advancement forZUNA1\.1relative toZUNA1are the ability to process more varied sample durations \(from 0\.5 \- 30 seconds\) and the ability to predict missing data from an assortment of realistic structural dropout patterns\. Figure[7](https://arxiv.org/html/2607.27308#S4.F7)illustrates the model’s reconstruction accuracy for the four different structural dropout schemes retained throughout training for an evaluation dataset as we vary what percentage of tokens were masked in panel a, and the duration of the sample with a fixed15%15\\%of tokens masked in panel b\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/zuna_vs_mne_full-channel-random-dropout-train_linear_light.png)\(a\)Full Channel Random Dropout
![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/zuna_vs_mne_random-uniform-dropout_linear_light.png)\(b\)Random\-Uniform Dropout
![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/zuna_vs_mne_correlated-channel-time-dropout_linear_light.png)\(c\)Spatio\-Temporally Correlated Dropout

Figure 8:ZUNA1\.1vs MNE Spherical Spline\- reconstruction performance vs dropout percent on mixed evaluation dataset for structural dropout schemesFinally, one can apply the spherical spline interpolation method to impute missing data in three of the four structural dropout schemes, using only information from observed channels at a given time\-point to predict missing channels at that same time\-point\.ZUNA1\.1outperforms spline interpolation in these use\-cases as well, as seen in Figure[8](https://arxiv.org/html/2607.27308#S4.F8)\.ZUNA1\.1can use observed information from surrounding channels at the same time\-point as well as observed signals from the same channel at previous or later time\-points if signal degrades in a channel or group of channels for a short period of time\.

## VDiscussion

ZUNA1\.1builds upon and substantially extends our originalZUNA1model, making it far more flexible and ultimately usable in real\-world situations by EEG scientists and practitioners\. The new dropout training schemes we introduced into the training ofZUNA1\.1significantly expand its range and ability to process and reconstruct realistic “in\-the\-wild” EEG data from a variety of sources\. The ability ofZUNA1\.1to handle variable length sequences \(up to 30 seconds\) as well as reconstruct arbitrary time windows within a channel are crucial here\. Beyond the dropout schemes, we also substantially improve the size and quality of our training dataset, and train for4×4\\timeslonger on75%75\\%more data than the originalZUNA1model\. Architecture and training recipe improvements enabled us to train stably far beyond the original model, while computational efficiency improvements meant that training required only roughly the same level of computational resources\.

WhileZUNA1\.1substantially improves over its predecessor in its range and versatility, on the specific tasks which are suitable forZUNA1, the two models are roughly comparable, withZUNA1\.1achieving only a small advantage on average\. Both models still substantially outperform non\-neural reconstruction methods such as MNE’s spherical spline method\. SinceZUNA1\.1is trained for substantially longer on a much greater quantity \(and higher quality\) of data, this is somewhat surprising\. One possibility is that theZUNA1\.1model is generally superior but that due to having to handle a much wider range of data, the model has to spread its competence over a much wider distribution, and hence necessarily dilute its capabilities in the specific kind of task which the originalZUNA1model has specialized in\. This implies that the model is primarily limited by capacity\. However, preliminary experiments with naively scaling up the model architecture from 380M to approximately 1\.2B parameters did not show dramatically improved performance versus the smaller model, which led toZUNA1\.1being trained at the same size as the originalZUNA1model\. This implies that parameter count alone is not the primary bottleneck for the model’s capabilities\.

Additional possibilities are that the diffusion encoder\-decoder architecture ofZUNA1\.1is itself the bottleneck\. It is possible that the architectural scheme of passing the core information through the encoder bottleneck in ‘registers’ imposes a heavy constraint on the ultimate reconstruction capabilities of the model, although expanding this bottleneck naively does not immediately result in better reconstructions\. The encoder\-decoder architecture itself could be a bottleneck\. Another possibility is that the diffusion training scheme and especially the chosen noise schedule reduced the potential SNR of training causing an effective plateau in performance\.

Finally, it is of course possible, although we think it is unlikely, that even at 380M parameters, the model has already exhausted all of the signal present in the data which could be used for reconstruction\. If this were true, it would mean that the irreducible noise floor in the EEG data itself is relatively high while the information redundancy across channels and time is fairly low\. Some evidence that indirectly supports this is that we observed a similarly unexpected phenomenon with our initial set of dropout schemes, where a few of them appeared extremely challenging to learn such that the model made effectively little progress over the course of training\. These were the intuitively harder dropout schemes such as the consumer or standard montage channel dropout\. The lack of improvement on these schemes seems to imply that good reconstructions in these cases may be largely impossible in that the dropout scheme irreversibly removes the information required to make good predictions\.

A further point is that during development ofZUNA1\.1, we observed a growing dissociation between reconstructive accuracy and the utility of the learned representations for downstream tasks\. We discuss these explorations in appendix[appendix˜C](https://arxiv.org/html/2607.27308#A3)\. While uncertain, we hypothesize that a substantial fraction of this dissociation may be due to the diffusion objective and architecture, which in other modalities is known to produce less clearly generalizable representations than more standard autoregressive objectives\(Heet al\.,[2021](https://arxiv.org/html/2607.27308#bib.bib2)\)\. More generally, for the goal of constructing powerful representations for downstream classifiers, pure reconstruction may not be the optimal objective in general, and latent\-space objectives such as contrastive\(Oordet al\.,[2018](https://arxiv.org/html/2607.27308#bib.bib286); Heet al\.,[2019](https://arxiv.org/html/2607.27308#bib.bib287); Chenet al\.,[2020](https://arxiv.org/html/2607.27308#bib.bib288); Radfordet al\.,[2021](https://arxiv.org/html/2607.27308#bib.bib289)\)and JEPA\-like\(LeCun and others,[2022](https://arxiv.org/html/2607.27308#bib.bib284); Assranet al\.,[2023](https://arxiv.org/html/2607.27308#bib.bib285)\)losses may be preferred\. In future work, we aim to focus more on developing general encoder models of EEGs with representations well\-suited to downstream tasks\. This is ultimately a necessary step to establish the latent space which can then be joined with pretrained language or image models to train effective EEG\-to\-X models\. This will likely require moving beyond regularized reconstruction as an objective, and developing novel architectures to move beyondZUNA1\.1’s diffusion autoencoder base\.

## Acknowledgements

We would like to thank Paul White, Danny Martinelli, and Kristina Zhao for their assistance with the release, and Nathan Kolbas for integratingZUNA1\.1into Zyphra Cloud\.

## References

- M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. Ballas \(2023\)Self\-supervised learning from images with a joint\-embedding predictive architecture\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 15619–15629\.Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- H\. Banville, S\. d’Ascoli, S\. Dahan, J\. Rapin, M\. Careil, Y\. Benchetrit, J\. Lévy, S\. Panchavati, A\. Ratouchniak, E\. Cascardi,et al\.\(2026\)NeuralBench: a unifying framework to benchmark neuroai models\.arXiv preprint arXiv:2605\.08495\.Cited by:[Appendix C](https://arxiv.org/html/2607.27308#A3.p1.1)\.
- T\. Chen, S\. Kornblith, M\. Norouzi, and G\. Hinton \(2020\)A simple framework for contrastive learning of visual representations\.InInternational conference on machine learning,pp\. 1597–1607\.Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- W\. Cui, W\. Jeong, P\. Thölke, T\. Medani, K\. Jerbi, A\. A\. Joshi, and R\. M\. Leahy \(2024\)Neuro\-gpt: towards a foundation model for eeg\.In2024 IEEE International Symposium on Biomedical Imaging \(ISBI\),pp\. 1–5\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- M\. Ding, Z\. Yang, W\. Hong, W\. Zheng, C\. Zhou, D\. Yin, J\. Lin, X\. Zou, Z\. Shao, H\. Yang, and J\. Tang \(2021\)CogView: mastering text\-to\-image generation via transformers\.External Links:2105\.13290,[Link](https://arxiv.org/abs/2105.13290)Cited by:[§II](https://arxiv.org/html/2607.27308#S2.p1.1)\.
- Doneret al\.\(2025\)LUNA: a latent unified network architecture for topology\-agnostic eeg modeling\.External Links:2510\.22257Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- L\. Duan, H\. Duan, Y\. Qiao, S\. Sha, S\. Qi, X\. Zhang, J\. Huang, X\. Huang, and C\. Wang \(2020\)Machine learning approaches for mdd detection and emotion decoding using eeg signals\.Frontiers in Human Neuroscience14,pp\. 284\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p1.1)\.
- Y\. El Ouahidiet al\.\(2025\)REVE: representation for eeg with versatile embeddings\.External Links:2510\.21585Cited by:[Appendix C](https://arxiv.org/html/2607.27308#A3.p3.1),[Appendix C](https://arxiv.org/html/2607.27308#A3.p4.1),[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- Z\. Geng, A\. Pokle, W\. Luo, J\. Lin, and J\. Z\. Kolter \(2024\)Consistency models made easy\.External Links:2406\.14548,[Link](https://arxiv.org/abs/2406.14548)Cited by:[§II](https://arxiv.org/html/2607.27308#S2.p1.1)\.
- K\. He, X\. Chen, S\. Xie, Y\. Li, P\. Dollár, and R\. Girshick \(2021\)Masked autoencoders are scalable vision learners\.External Links:2111\.06377,[Link](https://arxiv.org/abs/2111.06377)Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- K\. He, H\. Fan, Y\. Wu, S\. Xie, and R\. Girshick \(2019\)Momentum contrast for unsupervised visual representation learning\.arXiv preprint arXiv:1911\.05722\.Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- A\. Henry, P\. R\. Dachapally, S\. Pawar, and Y\. Chen \(2020\)Query\-key normalization for transformers\.External Links:2010\.04245,[Link](https://arxiv.org/abs/2010.04245)Cited by:[§II](https://arxiv.org/html/2607.27308#S2.p1.1)\.
- W\. Jiang, Y\. Wang, B\. Lu, and D\. Li \(2024a\)NeuroLM: a universal multi\-task foundation model for bridging the gap between language and eeg signals\.arXiv preprint arXiv:2409\.00101\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- W\. Jiang, L\. Zhao, and B\. Lu \(2024b\)Large brain model for learning generic representations with tremendous eeg data in bci\.arXiv preprint arXiv:2405\.18765\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- P\. Kaushik, A\. Moye, M\. v\. Vugt, and P\. P\. Roy \(2022\)Decoding the cognitive states of attention and distraction in a real\-life setting using eeg\.Scientific Reports12\(1\),pp\. 20649\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p1.1)\.
- K\. Kontras, T\. Osselaer, S\. G\. Mouslech, A\. Karaiskou, G\. Gagliardi, T\. Strypsteen, M\. H\. Badiei, A\. Rani, M\. Vanmarcke, M\. Bhagubai,et al\.\(2026\)NeuroAtlas: benchmarking foundation models for clinical eeg and brain\-computer interfaces\.arXiv preprint arXiv:2605\.14698\.Cited by:[Appendix C](https://arxiv.org/html/2607.27308#A3.p1.1)\.
- D\. Kostas, S\. Aroca\-Ouellette, and F\. Rudzicz \(2021\)BENDR: using transformers and a contrastive self\-supervised learning task to learn from massive amounts of eeg data\.External Links:2101\.12037Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- G\. Kuruppu, N\. Wagh, V\. Kremen, S\. Pati, G\. Worrell, and Y\. Varatharajah \(2025\)Eeg foundation models: a critical review of current progress and future directions\.arXiv preprint arXiv:2507\.11783\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- Y\. LeCunet al\.\(2022\)A path towards autonomous machine intelligence version 0\.9\. 2, 2022\-06\-27\.Open Review62\(1\),pp\. 1–62\.Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- C\. Li, Y\. Qi, X\. Ding, J\. Zhao, T\. Sang, and M\. Lee \(2022\)A deep learning method approach for sleep stage classification with eeg spectrogram\.International journal of environmental research and public health19\(10\),pp\. 6322\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p1.1)\.
- Z\. Liang, S\. Oba, and S\. Ishii \(2019\)An unsupervised eeg decoding system for human emotion recognition\.Neural Networks116,pp\. 257–268\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p1.1)\.
- A\. v\. d\. Oord, Y\. Li, and O\. Vinyals \(2018\)Representation learning with contrastive predictive coding\.arXiv preprint arXiv:1807\.03748\.Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- Panchavatiet al\.\(2026\)Laya: a latent\-predictive eeg foundation model\.External Links:2603\.16281Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- F\. Perrin, J\. Pernier, O\. Bertrand, and J\.F\. Echallier \(1989\)Spherical splines for scalp potential and current density mapping\.Electroencephalography and Clinical Neurophysiology72\(2\),pp\. 184–187\.External Links:ISSN 0013\-4694,[Document](https://dx.doi.org/https%3A//doi.org/10.1016/0013-4694%2889%2990180-6),[Link](https://www.sciencedirect.com/science/article/pii/0013469489901806)Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p2.1)\.
- A\. Radford, J\. W\. Kim, C\. Hallacy, A\. Ramesh, G\. Goh, S\. Agarwal, G\. Sastry, A\. Askell, P\. Mishkin, J\. Clark,et al\.\(2021\)Learning transferable visual models from natural language supervision\.InInternational conference on machine learning,pp\. 8748–8763\.Cited by:[§V](https://arxiv.org/html/2607.27308#S5.p5.1)\.
- E\. Su, S\. Cai, L\. Xie, H\. Li, and T\. Schultz \(2022\)STAnet: a spatiotemporal attention network for decoding auditory spatial attention from eeg\.IEEE Transactions on Biomedical Engineering69\(7\),pp\. 2233–2242\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p1.1)\.
- A\. Supratak, H\. Dong, C\. Wu, and Y\. Guo \(2017\)DeepSleepNet: a model for automatic sleep stage scoring based on raw single\-channel eeg\.IEEE Transactions on Neural Systems and Rehabilitation Engineering25\(11\),pp\. 1998–2008\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p1.1)\.
- R\. Thapaet al\.\(2024\)SleepFM: multi\-modal representation learning for sleep across brain activity, ecg and respiratory signals\.External Links:2405\.17766Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- G\. Wanget al\.\(2024a\)EEGPT: pretrained transformer for universal and reliable representation of eeg signals\.InAdvances in Neural Information Processing Systems,Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- J\. Wanget al\.\(2024b\)CBraMod: a criss\-cross brain foundation model for eeg decoding\.External Links:2412\.07236Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- C\. Warner, J\. Mago, J\. Huml, M\. Osman, and B\. Millidge \(2026\)ZUNA: flexible eeg superresolution with position\-aware diffusion autoencoders\.arXiv preprint arXiv:2602\.18478\.External Links:2602\.18478,[Link](https://arxiv.org/abs/2602.18478)Cited by:[5th item](https://arxiv.org/html/2607.27308#S1.I1.i5.p1.1),[§I](https://arxiv.org/html/2607.27308#S1.p2.1),[§II](https://arxiv.org/html/2607.27308#S2.p1.1),[§IV\-A](https://arxiv.org/html/2607.27308#S4.SS1.p1.1)\.
- W\. Xiong, J\. Li, J\. Li, K\. Zhu, and C\. Jiang \(2026\)EEG\-fm\-bench: a comprehensive benchmark for the systematic evaluation and diagnostic analyses of eeg foundation models\.External Links:2508\.17742,[Link](https://arxiv.org/abs/2508.17742)Cited by:[Appendix C](https://arxiv.org/html/2607.27308#A3.p1.1)\.
- C\. Yang, M\. Westover, and J\. Sun \(2023\)Biot: biosignal transformer for cross\-data learning in the wild\.Advances in Neural Information Processing Systems36,pp\. 78240–78260\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- Z\. Yuan, F\. Shen, M\. Li, Y\. Yu, C\. Tan, and Y\. Yang \(2024\)Brainwave: a brain signal foundation model for clinical applications\.arXiv preprint arXiv:2402\.10251\.Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.
- D\. Zhang, Z\. Yuan, Y\. Yang,et al\.\(2023\)Brant: foundation model for intracranial neural signal\.InAdvances in Neural Information Processing Systems,Cited by:[§I](https://arxiv.org/html/2607.27308#S1.p3.1)\.

## Appendix ATraining hyperparameters

TABLE II:ZUNA1\.1 training configuration\.
## Appendix BReconstruction performance

TABLE III:Average NMSE for region occlusion reconstruction across datasets, methods, and target regions\.TABLE IV:Channel reconstruction NMSE on dropped channels\. Lower is better\.
## Appendix CDownstream Tasks

We evaluatedZUNA1\.1on a range of downstream classification tasks\. We probed the latent representation with a linear decoder that averages over tokens and finetunes the encoder as in NeuralBench\(Banvilleet al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib9)\)\. We note that frozen backbones with linear heads performed significantly worse than finetuned models forZUNA1\.1, which is consistent across EEG foundation models \(FMs\) as found in NeuroAtlas\(Kontraset al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib4)\)and EEG\-FM\-Bench\(Xionget al\.,[2026](https://arxiv.org/html/2607.27308#bib.bib3)\)\.

However, we note that we find no correlation between model size and classification performance on finetuned FMs, not only in NeuralBench but also when scalingZUNA1\.1up to 1\.2B parameters, which performed worse than the current 380M model\. A simple linear regression across average model ranking in classification and number of parameters yieldedR2=0\.007R^\{2\}=0\.007with no obvious trends across eight FMs, including bothZUNA1\.1\-380M\(i\.e\.ZUNA1\.1\) andZUNA1\.1\-1\.2B\. Our experiments and the NeuralBench results suggest that new architectures and training objectives may be needed, and simply adding parameters is not a panacea to representation learning\.

TABLE V:Classification accuracy\(0\-1 scale\) for the MMD loss ablation across downstream tasks\. Bold indicates the highest classification accuracies\. We find a peak early on in our accuracies, and then a sharp decay as the latent representations are increasingly specialized towards reconstruction\.TABLE VI:Reconstruction NMSEunder 50% random\-uniform dropout\. Bold indicates the lowest errors, showing that our final checkpoint is almost uniformly superior for reconstruction\.In particular, there seems to be a tension between masked autoencoders and downstream representations\.El Ouahidi and others \([2025](https://arxiv.org/html/2607.27308#bib.bib267)\)argues that the final layer of the encoder in a masked autoencoder is prone to overfitting to the reconstruction task, harming representations for downstream adaptation\. However, they provide no ablations for their regularization that attempts to mitigate this issue\. We provide experiments to show how this tension manifests in classification and reconstruction evaluations\.

In Table[V](https://arxiv.org/html/2607.27308#A3.T5)and Table[VI](https://arxiv.org/html/2607.27308#A3.T6), we show the counterintuitive relationship between classification accuracy and reconstruction error: as the number of training steps increases, reconstruction accuracyincreaseswhile classification accuracydecreases\. The training objective contains a term that rewards reconstruction \(decoder RF\-loss\) and a term that encourages latent representations to be more Gaussian distributed \(encoder MMD\-loss\); however, there is no explicit term that shapes latents to be informative or useful for downstream tasks\. At first, in the style of REVE\(El Ouahidi and others,[2025](https://arxiv.org/html/2607.27308#bib.bib267)\), we hypothesized that our loss function was not encouraging useful latent structure, and we ablated our Maximum Mean Discrepancy \(MMD\) loss that encourages latent Gaussian structure\. Deleting MMD did seem to help classification accuracy by about 2% per task on average, and motivated our switch to turning off MMD loss in our main training run at training step 400K \(i\.e\. our ”Hybrid” variation trained with MMD loss up to 400K training steps, and turned off thereafter\)\.

With or without the MMD loss, classification accuracy plateaued early in training and then decreased sharply thereafter\. Our final checkpoint was almost uniformly superior on reconstruction while becoming progressively worse at classification\. Even after turning off the MMD loss, which proved to be useful in our first ablation, the classification accuracies degraded until some of the datasets were even near\-chance on classification accuracy even though the model performed extremely well at reconstruction\.

In Figure[9](https://arxiv.org/html/2607.27308#A3.F9), we show the model rankings on EEG foundation models includingZUNA1\.1using the NeuralBench metrics\. Here, the “No MMD\-50K” checkpoint yields respectable performance on the NeuralBench datasets\. However, our results cast serious doubt on the usage of masked autoencoders trained to reconstruct missing data for downstream tasks\. If REVE were trained longer with more data, would its classification accuracy also decrease given its objective is entirely reconstructive? Without these ablations, the answer is inconclusive\. However, this seems to be the case forZUNA1\.1\.

Finally, investigating the claim that the final encoder layer was overfit to the reconstruction task, we systematically explored whether building a linear classifier head on earlier layers of the encoder improves classification performance\. Figure[10](https://arxiv.org/html/2607.27308#A3.F10)shows balanced accuracy for classifiers built on top of different layers in the encoder for a few datasets in the NerualBench benchmark\. We observed some improved classification performance when using representations from earlier encoder layers, but nothing systematic\. We include these results for completeness\.

![Refer to caption](https://arxiv.org/html/2607.27308v1/figures/new_boxplot.png)Figure 9:Model ranking on downstream tasks\.Models are ordered by their mean normalized rank, displayed on the right hand side of the boxes\. Lower is better\. ZUNA1\.1 is shown for No MMD\-50K\.TABLE VII:Classification accuracy for ZUNA1\.1 \(No MMD\) compared with NeuralBench EEG foundation models and baseline methods\.![Refer to caption](https://arxiv.org/html/2607.27308v1/classifiers_by_layer/audiovisual.png)\(a\)Audiovisual
![Refer to caption](https://arxiv.org/html/2607.27308v1/classifiers_by_layer/mental_arithmetic.png)\(b\)Mental Arithmetic
![Refer to caption](https://arxiv.org/html/2607.27308v1/classifiers_by_layer/mental_imagery.png)\(c\)Mental Imagery
![Refer to caption](https://arxiv.org/html/2607.27308v1/classifiers_by_layer/n170.png)\(d\)N170
![Refer to caption](https://arxiv.org/html/2607.27308v1/classifiers_by_layer/n400.png)\(e\)N400

Figure 10:Balanced accuracy for linear classifiers attached at different ZUNA representation layers\. We ran this experiment to probe the representational quality due to previous hypotheses in EEG literature that the final layer of masked autoencoders is prone to overfitting to the reconstruction task\. We do not see uniform evidence of this across datasets\.