Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
For practitioners in probabilistic representation learning and uncertainty quantification, this work provides a novel attention mechanism that yields interpretable uncertainty estimates, though its empirical gains over baselines are modest.
The paper introduces NSAC, a biologically-inspired continuous-time attention architecture that models attention logits as solutions to an Ornstein-Uhlenbeck process, enabling probabilistic outputs and joint aleatoric/epistemic uncertainty quantification. It achieves competitive accuracy and well-calibrated uncertainty across diverse tasks including function approximation, regression, forecasting, and autonomous driving.
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.