NeuroDX/axon-eeg

AXON: axis-factorized EEG foundation model

AXON (AXis-factorized Operator Network) is a transformer encoder for EEG, pretrained with masked autoencoding. Each recording window is cut into tokens, one per electrode per one-second patch, so the tokens form a grid of electrodes by time steps. Instead of dense attention over all tokens, every layer runs two attention paths in parallel: a temporal path (each token attends to its own electrode across time) and a spatial path (each token attends to all electrodes at the same time step). A small per-token gate mixes the two.

• Parameters: 118.6M (encoder 118.56M + default head). The 141.9M in the paper include the pretraining decoder and auxiliary head, which are not part of this model.
• Code: braindecode.models.AXON (Apache-2.0).
• Paper: Adaptive Anisotropic Attention for Axis-Structured Signals (arXiv:2609.08788).

How to use

Status: AXON is being contributed to braindecode (pull request pending). Until it is merged, install braindecode from the pull-request branch:

pip install "braindecode[hub] @ git+https://github.com/TODO-fork/braindecode@add-axon"

import mne
from braindecode.models import AXON

raw = mne.io.read_raw_edf("recording.edf", preload=True)
raw.set_montage("standard_1020", match_case=False)   # AXON needs 3D electrode positions
raw.resample(200)                                     # AXON expects 200 Hz

model = AXON.from_pretrained(
    "NeuroDX/axon-eeg",
    chs_info=raw.info["chs"],   # any montage, any channel order
    n_outputs=4,                # the head is always newly initialised
    n_times=800,                # window length in samples (here 4 s)
)
logits = model(x)                                        # x: (batch, n_chans, n_times), volts or µV
features = model(x, return_features=True)["features"]    # (batch, 512) pooled embedding
tokens = model.encode(x)                                 # (batch, n_chans, n_patches, 512)

Linear probing: freeze everything except model.final_layer. Fine-tuning (reference recipe): encoder learning rate 0.1× the head's, 3 head-only epochs first, and the two gates (*.axis_gate, *.scale_gate) kept frozen.

Input requirements

Sampling rate	200 Hz
Window length	at least 200 samples (1 s); patches of 1 s with a 0.9 s stride
Channels	any number and order; each needs a 3D position
Positions	chs_info[i]["loc"][:3] (MNE head coordinates). Channels without a position are looked up by name in MNE's 10-20 / 10-05 montages
Scaling	none needed: each channel is z-scored within each window inside the model; volts and microvolts give the same output

Pretraining

• Preprocessing: resample to 200 Hz, notch 50/60 Hz, band-pass 0.5–99.5 Hz, per-channel z-score, clip at ±15σ, non-overlapping 10 s windows.
• Objective: masked autoencoding, 55% spatio-temporal block masking, L1 on masked patches plus an auxiliary pooled reconstruction loss (weight 0.5); a 4-layer dense decoder (not included).
• Optimisation: AdamW (β = 0.9/0.95, weight decay 0.05), peak LR 2.4e-4 with cosine decay, batch 4,096, bfloat16, 20 epochs; this is the epoch-10 checkpoint.
• Augmentation: token dropout 0.3, electrode-position noise σ = 0.25 cm.
• During pretraining, electrode positions came from MNE's standard_1020 montage; the downstream evaluation below (and this package) use the same montage in MNE head coordinates, which are offset by about 5 cm. The reported numbers use the head-coordinate convention.

Evaluation

Six subject-disjoint tasks (motor imagery on PhysioNet MMI and BCI Competition IV-2a, STEW workload, HMC sleep staging, Siena seizure detection, ADFTD dementia), balanced accuracy, mean over 3 downstream seeds. For each downstream run the checkpoint is selected on validation subjects only.

Linear probe	motor	bcic	workload	hmc	siena	adftd
AXON	0.456	0.283	0.673	0.650	0.866	0.492

Citation

@misc{jain2026adaptiveanisotropicattentionaxisstructured,
      title={Adaptive Anisotropic Attention for Axis-Structured Signals},
      author={Mahir Jain and Parshva Runwal and Aditya Ray Mishra and Arvasu Kulkarni and Jeet Bandhu Lahiri and Sandeep Singh and Siddharth Panwar},
      year={2026},
      eprint={2609.08788},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2609.08788},
}