Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning

Abstract

Reliable uncertainty quantification in continuous-time (CT) representation learning remains nascent, particularly within CT attention literature. We introduce the Neuronal Stochastic Attention Circuit (NSAC), a novel biologically-inspired CT attention architecture that reformulates attention logit computation as the solution of an Ornstein-Uhlenbeck stochastic differential equation modulated by input-dependent, nonlinear interlinked gates derived from repurposed C. elegans Neuronal Circuit Policies (NCPs) wiring mechanism. It induces a Gaussian distribution over logits that propagates principled stochasticity through a logistic-normal distribution over attention weights to yield probabilistic output. A two-term objective function combining Gaussian negative log-likelihood with an epistemic-separation regularizer enforces higher predictive variance under distributional shifts and enables joint quantification of aleatoric and epistemic uncertainty. Theoretically, we provide: (i) state stability bounds; (ii) closed-form guarantees; and (iii) frozen-coefficient error approximation. Empirically, we implement NSAC in a diverse set of learning tasks including: (i) irregular CT function approximation; (ii) multivariate regression; (iii) long-range forecasting; (iv) Industry 4.0; and (v) lane-keeping of autonomous vehicles. We observe that NSAC remains competitive against several baselines in terms of accuracy and produces informative uncertainty estimates while being interpretable at the neuronal cell level.