Adaptive attention improves analysis of brain and audio signals

Adaptive Anisotropic Attention for Axis-Structured Signals

Machine LearningArtificial Intelligence

Summary

Analyzing signals like brain waves (EEG) or audio can be tricky because important details happen along different lines, like time or space. The authors found that treating all parts of the signal equally can mix irrelevant information together. They created a new type of attention method that looks separately along time and space, then combines these views in a smart way. This method, tested on brain signals and audio, worked better than older approaches by matching the natural structure of the data.

self-attentionEEGanisotropictemporal axisspatial axissignal-to-noise ratiomachine learninginductive biasconvex combinationspectrogram

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.