Neuronal Attention Circuit (NAC) for Representation Learning

TL;DR

The paper introduces Neuronal Attention Circuit (NAC), a biologically inspired continuous-time attention mechanism that improves representation learning by modeling attention as a linear ODE with sparse, adaptive gates, offering efficiency, interpretability, and competitive performance.

Contribution

NAC is a novel CT-attention mechanism inspired by C.elegans neuronal circuits, featuring sparse gates and a subquadratic sparse interaction mechanism, with theoretical guarantees and diverse empirical applications.

Findings

NAC matches or outperforms baselines in accuracy.

NAC has intermediate runtime and memory consumption.

NAC is interpretable at the neuron cell level.

Abstract

Attention improves representation learning over RNNs, but its discrete nature limits continuous-time (CT) modeling. We introduce Neuronal Attention Circuit (NAC), a novel, biologically inspired CT-Attention mechanism that reformulates attention logit computation as the solution to a linear first-order ODE with nonlinear interlinked gates derived from repurposing C.elegans Neuronal Circuit Policies (NCPs) wiring. NAC replaces dense projections with sparse sensory gates for key-query projections and a sparse backbone network with two heads for computing content-target and learnable time-constant gates, enabling efficient adaptive dynamics. To improve efficiency and memory consumption, we implemented an adaptable subquadratic sparse Top-K pairwise concatenation mechanism that selectively curates key-query interactions. We provide rigorous theoretical guarantees, including state stability and bounded approximation errors. Empirically, we implemented NAC in diverse domains, including irregular time-series classification, lane-keeping for autonomous vehicles, and industrial prognostics. We observed that NAC matches or outperforms competing baselines in accuracy and occupies an intermediate position in runtime and memory consumption compared with several CT state-of-the-art baselines, while being interpretable at the neuron cell level.