Adaptive attention improves signal analysis along time and space axes

Adaptive Anisotropic Attention for Axis-Structured Signals

Machine LearningArtificial Intelligence

Summary

Signals like brain activity recordings (EEG) have natural patterns across time and different sensors, but usual attention methods treat all parts equally, making it hard to focus on the important parts. The authors introduced a new attention method that separately considers changes over time and differences across sensors, then learns how much to weigh each part. Their approach helps models better understand EEG signals and also works on other structured data like audio spectrograms. This shows that designing attention to match the natural structure of signals can lead to better results.

What this means in practice

  • For medical device developers: Build brain signal analysis tools that better separate time and electrode relationships to improve diagnostic accuracy in EEG-based applications.$Commercial implications: This enables more reliable EEG interpreting devices that can be marketed for clinical and neurotechnology use by exploiting structured attention.
  • For audio signal engineers: Enhance audio analysis and processing systems by applying separate attention along frequency and time axes for improved performance on spectrogram data.

Authors

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

Abstract

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.