Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
2026-05-25 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors introduce NSAC, a new way for computers to pay attention to information over continuous time, inspired by how neurons work in a simple worm's brain. They use a special math equation to add controlled randomness to the process, helping the model understand uncertainty better. Their method can predict outcomes and how unsure the model is, useful in tasks like forecasting and autonomous driving. They show NSAC works well compared to other methods and can explain its decisions in terms of artificial neurons.
Continuous-time representation learningAttention mechanismOrnstein-Uhlenbeck processNeuronal Circuit PoliciesGaussian distributionAleatoric uncertaintyEpistemic uncertaintyNegative log-likelihoodAutonomous vehiclesLong-range forecasting
Authors
Waleed Razzaq, Yun-Bo Zhao
Abstract
Reliable quantification of uncertainty estimates in continuous-time (CT) representation learning remains nascent, particularly within CT attention architectures. 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 Gaussian distribution over logits that propagates principled stochasticity through 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 and enables joint quantification of aleatoric and epistemic uncertainty. 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) the lane-keeping of autonomous vehicles. We observe that the NSAC remains competitive against several baselines in terms of accuracy and produces reasonably well-calibrated uncertainty estimates while being interpretable at the neuronal cell level.