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

AXON is part of braindecode. Until the next stable braindecode release, install a development release:

pip install --pre "braindecode[hub]"

import mne
from braindecode.models import AXON
from braindecode.util import resolve_montage_name

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

model = AXON.from_pretrained(
    "MannasAI/axon-eeg",
    chs_info=raw.info["chs"],   # any montage, any channel order
    n_outputs=4,                # the head is not pretrained
    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

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.

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},
}