Add AXON: axis-factorized EEG foundation model - #1182

Codecov Report

❌ Patch coverage is 97.24138% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.28%. Comparing base (cb1ca88) to head (30294f8).
⚠️ Report is 1 commits behind head on master.

@@            Coverage Diff             @@
##           master    #1182      +/-   ##
==========================================
+ Coverage   89.23%   89.28%   +0.05%     
==========================================
  Files         158      159       +1     
  Lines       19604    19746     +142     
==========================================
+ Hits        17493    17631     +138     
- Misses       2111     2115       +4     

Hi @bruAristimunha,

CI status: 12 of 14 checks pass. The two failing jobs appear unrelated to the AXON changes.

Would you be able to re-run the two failed jobs?

On the Windows failure: AXON's three generic Hugging Face tests each save the full 118M-parameter model (about 475 MB) to tmp_path, so it adds roughly 1.4 GB to the runner's disk during the test session. If that is a concern, I would be happy to add a smaller AXON configuration for those tests, in whatever form you prefer.

Thank you.

Integration gate (braindecode maintainers)

Target: paper Table 1 (BAC, mean ± std over 3 seeds), NeuralBench replication of the public MannasAI/axon-eeg checkpoint (epoch 10; Table 10 shows this is the validation-selected epoch behind Table 1). Protocol: AdamW, the paper's LR schedule, 30 epochs, seeds 33/87/145, LP trains only final_layer; FT uses the HF-card recipe (encoder LR 0.1× head LR, 3 head-only epochs, axis_gate/scale_gate frozen). Gate: within 5% of the paper, or within the paper's own spread if that is larger. Parity to an independent fp64 reimplementation is exact (logits ≤1.2e-16).

Splits and recording mix (updated 2026-10-06): the subject splits now follow the EEG-FM-Bench dataset builders (xw1216/EEG-FM-Bench @325398d7), which match Table 8 on all six tasks. adftd uses the deterministic label-balanced split (64/11/13 subjects). hmc uses the seed-42 subject split (103/24/24). motor uses subjects 1–69 / 70–88 / 89–109, all 109 subjects, with a common average reference.

Status: not yet within the gate. Escalated to the maintainers for a decision; no change is requested in this PR. The model code and checkpoint are not the issue: every cell is at most one dataset-protocol detail away from the paper. adftd depends heavily on which subjects land in the test set: on the same checkpoint, LP reads 0.434 on one 13-subject split and 0.658 on another. The HF card's LP adftd (0.492) also differs from Table 1 (0.538).

Hi @bruAristimunha, thank you for the review. Below are the splits:

1. Subject splits

These come from the processed datasets the Table 1 evaluation read. The splits are subject-disjoint and fixed: all three downstream seeds use the same split, and only the training seed changes.

motor_mv_img (PhysioNet EEG Motor Movement/Imagery, all 109 subjects): train = subjects 1–69, validation = 70–88, test = 89–109.

2. Run selection behind motor_mv_img

• Motor imagery runs only: R04, R06, R08, R10, R12, R14. Execution runs and baselines are excluded.
• Events T1/T2 only (T0 rest is excluded), giving 4 classes: R04/R08/R12 → left fist / right fist; R06/R10/R14 → both fists / both feet.
• One 4 s window per cue, starting at cue onset (90 windows per subject).

Preprocessing for all datasets: 0.1 Hz high-pass (plus a 100 Hz low-pass when the native rate exceeds 200 Hz), mains notch where the data are not already notched (60 Hz for motor, 50 Hz for hmc), resampling to 200 Hz in µV, then the per-window, per-channel z-score inside the model. Electrode positions are MNE standard_1020 positions in the head frame, in cm, which is what chs_info[i]["loc"] gives in the braindecode port.

3. The adftd difference (0.538 vs 0.492)

This difference does not come from the split; both numbers use the split above. It comes from the downstream epoch-selection rule. For each task, Table 1 reports the better, by test BAC, of two candidate epochs: the best-validation-BAC epoch and the lowest-validation-loss epoch. The model card selects on validation BAC only. Per task:

For a replication that selects on validation, the validation-only rows are the matching targets. We will state the selection rule explicitly in the next arXiv version.

One NeuralBench detail that may matter

If the replication feeds AXON NeuralBench's channel_positions: for datasets without layout_or_montage_name, these come from each recording's stored locations and are not always in the head frame. For Miltiadous2023Dice they match MNE's raw standard_1020 coordinates, about 5 cm from the head frame. On NeuralBench's dementia task this changed AXON's fine-tuned test BAC from 41.3 ± 4.0 to 48.3 ± 0.1 (3 seeds) when we used head-frame positions looked up from the channel names instead.

Hi @bruAristimunha, thank you for the update, and for #1251 on the Windows disk issue.

Two points in the summary may help the maintainers' decision. Both use data already in this thread.

1. hmc split. The current hmc cell uses the builder's seed-42 split over 151 subjects (103 / 24 / 24). The Table 1 evaluation read processed hmc data with 148 subjects (IDs 14, 26, 52, 53, 64 and 135 are absent), split 102 / 23 / 23; the exact lists are in my comment of 5 October. LP hmc 0.693 is therefore measured on a different test set from Table 1's 0.650.

2. adftd: 0.492 (model card) vs 0.538 (Table 1). These two numbers come from the same split and the same three runs. They differ only in the epoch-selection rule (section 3 of the same comment). adftd LP is very sensitive to the epoch: in our runs, test BAC ranges from 0.33 to 0.68 across the 30 epochs, with 11 validation and 13 test subjects. Per seed, the best-validation-BAC epoch gives 0.474 / 0.491 / 0.511 on test, and the lowest-validation-loss epoch gives 0.545 / 0.534 / 0.534; Table 1 kept the latter.

The adftd split (64 / 11 / 13, label-balanced) and the motor setup (subjects 1–69 / 70–88 / 89–109, common average reference) are the ones we used.

The branch now conflicts with master in test_foundation_models.py and test_return_features.py. I am happy to resolve it whenever that is convenient for you.