PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptive Anisotropic Attention for Axis-Structured Signals

Mahir Jain, Parshva Runwal, Aditya Ray Mishra, Arvasu Kulkarni, Sandeep Singh, Siddharth Panwar

arXiv:2609.08788v1cs.LGcs.AI

TL;DR

Dense all-to-all attention may mismatch the electrode–time structure of low-SNR EEG. The paper introduces AXON, which mixes temporal and spatial axis-restricted attention with learned gates, and reports improved mean balanced accuracy across six EEG tasks under linear probing and full fine-tuning, with axis factorization also helping audio.

  • Problem

    Dense attention may mismatch structured low-SNR EEG because it permits broad interactions instead of prioritizing electrode and time axes.

  • Method

    AXON replaces dense attention with temporal within-electrode and spatial same-time paths, combining them through token-conditioned convex gates.

  • Results

    AXON improves mean balanced accuracy over dense and parameter-matched baselines across six EEG tasks under linear probing and full fine-tuning, while factorization also improves audio results.

  • Takeaways & Limitations

    Aligning attention with structured-signal axes and learning the temporal–spatial balance supports transfer across EEG and audio spectrograms.

  • Takeaways & Limitations

    The evidence for removing the global path is empirical rather than a formal information-theoretic proof, and full-resolution audio and substantially different EEG populations remain unevaluated.

Abstract

from arXiv · show

Dense self-attention treats all token pairs as equally plausible before learning, an interaction-isotropic prior that can be mismatched to structured signals. For structured, low signal-to-noise ratio (SNR) signals such as EEG, dependencies are organized along the electrode and time axes, and this uniform prior exposes each token to many irrelevant interactions. We introduce Adaptive Anisotropic Attention (AAA), which splits attention into two paths: a temporal path, where each token attends to the tokens of its own electrode across time, and a spatial path, where it attends to the tokens of the other electrodes at the same time step. A small gate predicts, for every token, a convex combination of the two path outputs: two non-negative weights that sum to one. On six EEG downstream tasks, the resulting model, AXON (AXis-factorized Operator Network), improves mean balanced accuracy over a dense baseline under both linear probing and full fine-tuning. We show that both paths (temporal and spatial) are necessary and that the weighted sum beats a hard choice of one path; most of the benefit comes from the gate learning a different temporal/spatial balance at each layer of the network. Controlled audio spectrogram experiments show that axis factorization transfers beyond EEG. These results suggest that aligning attention with the natural axes of structured signals provides a useful inductive bias.

1 Introduction

Dense attention may mismatch the electrode–time structure of low-SNR EEG by allowing broad interactions that dilute localized signals. AAA replaces this uniform prior with parallel temporal and spatial paths whose balance is learned per token and layer, and the approach also helps audio spectrograms.

  • Motivation: Dense attention connects all tokens without distinguishing electrode–time or frequency–time axes, which may mismatch structured signals such as EEG.In masked autoencoding, this may encourage reconstruction from broad averages rather than localized or axis-specific signals.
  • Contribution: AAA replaces dense encoder attention with temporal attention within each electrode and spatial attention across electrodes at the same time step.A gate combines the two paths through a convex mixture that can vary across tokens and layers.
  • Evaluation: AXON evaluates this learned axis mixture on six EEG tasks under linear probing and full fine-tuning against dense baselines, including a parameter-matched model.The study also analyzes learned axis mixtures and task-dependent temporal context.
  • Scope: Axis factorization helps both EEG and audio spectrograms, indicating transfer beyond the original modality.The audio result is presented as a second axis-structured modality.
  • Relation to prior work: Related factorized-attention designs set axis order by hand and share it across tokens and layers, whereas AXON learns the balance per token and layer for low-SNR structured signals.The related-work passage also states that prior designs were not built for or tested on EEG.

2 Method: Adaptive Anisotropic Attention

AXON replaces global attention with temporal and spatial axis-restricted operators, then mixes their outputs using learned gates. Its token-conditioned design can vary by content, position, layer, and sample, while additional temporal windows refine the temporal path without adding a global branch.

  • Input and baseline: Tokens lie on a channel–time grid, and dense attention processes all CT tokens with shared projections and (CT)^2 token pairs per layer.The pretraining grid contains 21 electrodes and 11 temporal patches, or 231 tokens.
  • Axis-factorized attention: AXON replaces dense attention with temporal attention within a channel and spatial attention across channels at one time step.The two operators use separate projections and jointly cover both axes of the token grid.
  • Axis mixture: Each AXON block computes yi = λ T(x)i + (1 − λ) S(x)i, with a convex temporal–spatial mixture.AXON-Fixed uses one sigmoid-transformed scalar per layer, shared across tokens.
  • Connectivity: Although each layer restricts interactions to one axis at a time, any two tokens on a full channel–time grid connect after two layers.A dense global branch is described as an alternative, but AXON does not use it.
  • Token-conditioned gate: The token-conditioned gate computes separate mixing weights for every token on every forward pass, allowing variation with content, position, layer, and input sample.The gate uses an MLP applied with stop-gradient input and produces the soft-gated block output.
  • Temporal windows: AXON uses two temporal windows: a short window of five consecutive patches and a long window covering the visible window, with one mixing weight per layer.This temporal-scale mixture is separate from the token-level temporal/spatial axis gate, and the final model has no global path.

3 Experiments

The experiments pretrain encoders with masked autoencoding and evaluate them on six subject-disjoint EEG tasks under linear probing and full fine-tuning. AXON outperforms dense and parameter-matched baselines, while controls support the importance of axis-specific learned mixing.

  • Pretraining: Encoders are pretrained with masked autoencoding on pooled TUH-EEG, I-CARE, and internal EEG data, using 10-second windows and 55% token masking.The encoder processes visible tokens and a decoder reconstructs masked patches.
  • Evaluation: Evaluation covers six EEG datasets spanning motor imagery, workload, sleep staging, seizure detection, and dementia diagnosis using balanced accuracy.Results use subject-disjoint splits under linear probing and full fine-tuning.
  • Main results: 0.579 Mean LP and 0.656 Mean FT are AXON’s highest scores in Table 1, improving over Dense by 2.9 and 5.0 balanced-accuracy points.These are the paper’s reported mean linear-probing and full-fine-tuning results.
  • Controls: AXON’s gains persist over parameter-matched Dense-L by 2.4 and 2.3 points and over Divided-ST by 1.9 and 2.8 points.The controls also report that restricting spatial neighbors lowers Mean LP to 0.510 and 2D rotary embeddings add only +0.008 to Dense.
  • Task-level results: +10.2 points on motor LP and +10.0 points on adftd FT are the largest task-level gains over Dense.The gains are reported as balanced-accuracy improvements.
  • External models: AXON has the highest mean scores among six released EEG foundation models: 0.579 Mean LP and 0.656 Mean FT versus REVE’s 0.527 and 0.614.The comparison uses the same six-task protocol.

4 Audio Spectrogram Experiments

The audio experiment applies axis factorization to time–frequency spectrograms, using temporal attention across time within frequency bands and frequency attention across bands within time frames. Factorized variants improve over Dense, although Token+Global is strongest in most tested audio metrics.

  • Setup: The audio model treats spectrograms as time–frequency grids and maps AAA’s temporal and spatial paths to time and frequency attention.The temporal path attends within a frequency band, while the frequency path attends within a time frame.
  • Results: All factorized variants improve over Dense on AudioSet FT and SpeechCommands LP.The comparison uses the AudioMAE pretraining recipe.
  • Results: Token+Global achieves the highest scores on three of four audio metrics, unlike the EEG setting.Its advantage on AudioSet FT and SpeechCommands LP persists at 18K, 200K, and 2M clips.

5 Conclusion

AXON aligns attention with EEG’s electrode and time axes through learned soft mixing, improving mean balanced accuracy across six tasks under linear probing and fine-tuning. The evaluation also compares external models and documents the pretraining and data setup, including internal clinical sources and variable electrode configurations.

  • Conclusion: AXON improves mean balanced accuracy over dense and parameter-matched dense baselines across six subject-disjoint EEG tasks under linear probing and full fine-tuning.The evaluation uses the same six tasks, subject-disjoint splits, and balanced-accuracy metric for the compared models.
  • Conclusion: Removing either attention axis or using hard routing causes larger performance drops than replacing token-level gates with layer means.This supports the contribution of both temporal and spatial paths and of soft learned mixing.
  • Conclusion: Motor imagery shows the largest gap over external models at +19% LP, while REVE retains advantages on workload and siena FT.The motor-imagery comparison is consistent with preserving localized left/right sensorimotor structure, but does not identify which features cause the gain.
  • Data and pretraining: The pretraining corpus pools public TUH-EEG and I-CARE data with two internal clinical EEG collections that are de-identified but not publicly released.The internal sources were collected under institutional ethics approval; their restricted availability limits direct reproducibility of the full corpus.
  • Data and pretraining: The data span 16- to 256-channel systems, so AXON retains recordings resolving to available international 10–20 positions and uses C = 21 as the representative channel count.Not every recording contains all standard 10–20 electrodes, and the model handles missing channels through tokenization and attention masking.
  • Data and pretraining: Masked autoencoding pretrains AXON and the dense baseline on the same schedule and corpus, with the encoder attention operator as the stated difference.The training setup uses a lightweight decoder and an auxiliary pooled-attention reconstruction loss.

C.2 Parameter count and compute

AXON’s axis-factorized attention adds parameters through separate temporal and spatial projections, so comparisons with dense attention must distinguish capacity from inductive bias. The paper reports controls and gate analyses intended to separate these effects while evaluating six subject-disjoint EEG tasks.

  • Parameter count and compute: 26M parameters (22.5%) are added when one dense attention operator is replaced by two full-width temporal and spatial operators.AXON is therefore not parameter-matched to the original dense baseline.
  • Parameter count and compute: AXON uses separate QKV and output projections for temporal and spatial paths, doubling attention-projection parameters relative to one dense attention.The authors state that separate projections support axis-specific specialization but increase projection cost.
  • Parameter count and compute: 142.5M parameters (within 0.5% of AXON’s 141.9M) are used by the parameter-matched dense baseline, which changes only hidden dimension.The matched baseline uses the same pretraining data, optimizer, epochs, and mask ratio.
  • Parameter count and compute: 2.4 points on Mean LP and 2.3 points on Mean FT separate AXON from the parameter-matched dense baseline.The reported relative gaps are +4.3% for Mean LP and +3.6% for Mean FT.
  • Parameter count and compute: ∼7× fewer pairwise attention interactions occur per layer with axis factorization, but the computational advantage is attributed to the structured prior rather than raw speed.For the reported sequence length and hidden dimension, projection costs dominate total FLOPs.
  • Parameter count and compute: 0.932 is the mean gate entropy ratio, while layer-mean gates reduce performance by only −2.1% and shuffled gates by −1.3%.The gate is neither collapsed nor trivially uniform; the dominant learned structure is a depth-dependent soft schedule.

D.2 Training budget

AXON’s advantage persists across matched pretraining budgets rather than appearing only as a faster-learning effect. Its axis-restricted attention also retains global reach because any channel-time token pair connects within two stacked layers.

  • Training budget: AXON leads Dense by +0.025–0.028 Mean LP across every tested pretraining budget.At the final epoch, AXON reaches 0.568 versus 0.543 for Dense.
  • Training budget: AXON’s ranking is stable from epochs 5–50, and Dense never catches up.Validation-best checkpoints occur at epoch 10 for AXON and epoch 9 for Dense.
  • Global reach: Two stacked axis-attention layers allow any token to receive information from any other token on the full channel-time grid.The graph has diameter at most two because an arbitrary pair can connect through an intermediate token sharing one channel or time index.
  • Global reach: Axis-restricted attention reduces one-hop interactions relative to the dense graph while preserving global reach through composition across layers.The comparison is between axis-group edges and the dense all-to-all graph.

F.1 Interpreted ablation summary

The ablations attribute AXON’s gains primarily to axis factorization and learned gating, rather than parameter count, dense global mixing, or positional encoding alone. Fixed axis order, hard routing, and extra temporal-scale machinery provide weaker or task-selective benefits.

  • Core ablation conclusions: Axis factorization improves over dense attention after parameter matching, while adding a dense global branch still reduces Mean LP and Mean FT.The parameter-matched dense model remains below AXON, and AXON-withGlobal underperforms despite having more parameters.
  • Core ablation conclusions: +0.020 Mean LP is the largest additional gain from token gating over fixed factorization, while two-scale temporal routing contributes a smaller refinement.The scale-gate improvement is +0.001 Mean LP and +0.013 Mean FT over the single-scale variant.
  • Gate ablations: The learned weighted gate outperforms uniform and hard-argmax routing, whereas layer-mean, position-mean, and shuffled gates remain close to learned routing.These interventions indicate that exact token-gate alignment is not the dominant inference-time effect.
  • Task-selective variants: AXON-Decay improves short-event tasks but lowers Mean FT, so it is treated as a diagnostic task-selective extension rather than a replacement for AXON.Motor imagery improves from 0.455 to 0.487 LP and seizure detection from 0.867 to 0.894, while Mean FT changes from 0.656 to 0.633.
  • Controls: 2D RoPE gains only +0.008 Mean LP, compared with +0.029 for axis factorization, because RoPE leaves dense off-axis interactions available.The dense model continues assigning substantial attention to off-axis pairs after pretraining.

G Analysis of the Global-Path Variant

The global-path analysis shows that comparable reconstruction loss does not imply comparable downstream representations. Removing global mixing yields richer and more task-relevant features, while learned axis gates vary systematically across depth but are not fully dependent on exact token-level alignment.

  • Representation consequences: Comparable reconstruction loss coexists with worse transfer for AXON-withGlobal: Mean LP is 0.555 versus 0.579 and Mean FT is 0.619 versus 0.656.The global path changes the features formed by the encoder rather than improving reconstruction.
  • Representation consequences: Removing the global path raises effective rank on five of six datasets, with motor imagery increasing from 19.97 to 29.05.Motor imagery is also the task where AXON gains most over Dense.
  • Representation consequences: AXON and AXON-withGlobal are least CKA-similar on motor imagery, sleep staging, and workload, but nearly identical on seizure detection and dementia.The reported CKA values are 0.84, 0.86, and 0.87 for the least-similar tasks, versus above 0.96 for the most similar tasks.
  • Gate behavior: The gate is temporal-heavy early, spatial-heavy in middle layers, and temporal-heavy again late, with α = 0.745 at layer 1 and α = 0.651 at layer 19.Layers 4–7 favor the spatial path with β = 0.686–0.707.
  • Gate behavior: Intervention results show that shuffled gates lose only −1.3% versus learned gates, while layer-mean and position-mean gates remain close to learned routing.Token-level gate diversity exists, especially around layers 9 and 12, but exact alignment is not the dominant downstream source.

H.4 Analysis of the AXON-TokenGated vs AXON-Fixed gap

The TokenGated–Fixed gap is mainly a training effect: per-token routing exposes the two axis paths to more varied mixtures, even though exact token-specific weights matter little at inference. The dense model also retains off-axis interactions that AXON removes by construction, while electrode occlusion supports task-relevant spatial use.

  • Training versus inference: One weight per layer costs only −2.1% at test time, yet AXON-Fixed reaches 0.558 Mean LP versus 0.578 for AXON-TokenGated.This separates inference-time gate alignment from the larger benefit of token gating during training.
  • Training versus inference: Initializing AXON-Fixed with AXON’s learned per-layer weights still yields 0.555 Mean LP, so initialization does not explain the gap.The neutral initialization reaches 0.558, while the learned-balance initialization performs similarly or slightly worse.
  • Training versus inference: Per-token gates create varied training mixtures for the temporal and spatial paths, while shuffled gates cost only −1.3% after training.The evidence is consistent with training-time gradient diversity improving the learned paths.
  • Temporal scale routing: The scale gate routes about 79% of mass to the long temporal window, with selected layers using appreciable short-window mass.This supports broad temporal context with selective local processing rather than uniformly local attention.
  • Off-axis interactions: After pretraining, dense attention retains 30–63% off-axis mass across depth, whereas AXON assigns zero one-layer off-axis attention by construction.The dense model learns axis structure only partly and unevenly.
  • Task-relevant spatial use: Covering the left motor strip drops right-hand recall by 0.101, while covering the right strip drops left-hand recall by 0.127.The asymmetric effects match the expected contralateral motor-cortex organization.
  • Task-relevant spatial use: AXON remains ahead of dense under random electrode deletion, from 0.377 versus 0.321 intact to 0.274 versus 0.263 with 94% removed.Only the model-to-model comparison is meaningful because this test uses a different frozen-feature evaluation setup.
  • Scope boundary: Audio absolute mAP is not directly comparable with published AudioMAE results because the experiment uses lower resolution, compute, capacity, decoder depth, and augmentation.The factorized and dense variants share these constraints, preserving the within-setup comparison.

J.1 Small-scale preliminary results (18K AudioSet clips)

At the smallest tested AudioSet scale, factorized attention already improves over dense attention, and this advantage persists at 200K clips. However, Token+Global’s ranking varies across scale and metrics, so the consistent conclusion is factorized-versus-dense improvement rather than a stable global-branch ranking.

  • 18K-scale results: At 18K AudioSet clips, all factorized variants improve AudioSet FT mAP by +49–55% relative over the dense baseline.The preliminary experiment used the same four-variant design and identical optimizer and schedule.
  • 200K-scale results: At 200K clips, the factorized advantage persists across all four metrics.Token+Global wins three of four metrics but underperforms the token-gate model on SpeechCommands LP, 20.00 versus 21.08.
  • 200K-scale results: Token+Global’s SpeechCommands ordering changes across scale, leading at 18K, trailing at 200K, and approximately tying at 2M.The reported pairs are 12.64 versus 12.26, 20.00 versus 21.08, and 21.87 versus 21.96.

K Cross-Modal Gate Mechanism Analysis

Cross-modal gate interventions show that axis-specific paths and soft mixing matter in both EEG and audio, while the useful axis schedule differs by modality. EEG benefits more from layer-varying preferences, whereas audio remains comparatively robust to uniform mixing.

  • Audio evaluation: Audio gate interventions were evaluated on ESC-50 and SpeechCommands v2 using a frozen encoder and a trained linear probe.The experiment used 200K AudioSet pretraining and the same seven gate overrides as the EEG analysis.
  • Shared intervention pattern: Removing either axis or replacing soft mixing with hard selection is more damaging than averaging or shuffling gates across EEG and audio.The broad intervention pattern is shared, although exact rankings and magnitudes differ by modality.
  • Audio gate schedule: Audio’s learned layer schedule favors frequency early, is nearly balanced in middle layers, and shifts toward time late.Reported gate means include α1 = 0.422, α3–5 ≈ 0.485, α10 = 0.564, and α11 = 0.556.
  • Intervention heatmap: Audio intervention heatmaps show that single-axis routing is catastrophic, while shuffled and layer-mean gates are indistinguishable from the learned gate.This pattern mirrors the EEG intervention findings.
  • Cross-modal robustness: −5.6% uniform mixing degradation in EEG contrasts with −0.4% in audio, whose gate range is 0.14 versus 0.45 in EEG.The narrower audio range indicates greater robustness to uniform mixing than the EEG schedule exhibits.
  • Limitations: The study’s scope is bounded by reduced audio resolution, limited EEG source diversity, and empirical rather than formal support for removing global attention.The authors identify full-resolution audio testing, broader populations, and formal analysis as open directions.
Loading 2609.08788v1…