Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning

2026-05-25

AI Summary

This work addresses the unreliable uncertainty estimation in continuous-time attention mechanisms by proposing a novel architecture inspired by the neural circuitry of *Caenorhabditis elegans*. The method models attention logits as input-dependent Ornstein–Uhlenbeck stochastic differential equations and propagates stochasticity through a logistic-normal distribution to yield probabilistic outputs. Innovatively integrating neuromorphic computing with uncertainty quantification, it employs a dual-objective loss function to jointly model epistemic and aleatoric uncertainties. While preserving neuron-level interpretability, the approach achieves accuracy on par with state-of-the-art baselines across diverse tasks—including irregular function approximation, multivariate regression, long-horizon forecasting, and autonomous lane-keeping—while producing well-calibrated uncertainty estimates.

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.

Problem

• uncertainty quantification
• continuous-time representation learning
• attention mechanisms
• probabilistic modeling
• stochasticity

Innovation

• Neuronal Stochastic Attention Circuit
• Ornstein-Uhlenbeck SDE
• probabilistic attention
• uncertainty quantification
• C.elegans NCPs