Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Waleed Razzaq 1 Yun-Bo Zhao 1 2
arXiv:2605.26061v1 [cs.LG] 25 May 2026
Abstract
loss. Two types of uncertainty exist: (1) aleatoric uncertainty, which is induced by inherent noise in the data, and (2) epistemic uncertainty, which arises from model uncertainty. While post-hoc methods such as Gaussian maximumlikelihood estimation (GMLE) (Kendall and Gal, 2017), Monte-Carlo Dropout (MCD) (Gal and Ghahramani, 2016), Deep-Ensembles (DE) (Lakshminarayanan et al., 2017), and Deep Evidential Regression (DER) (Amini et al., 2020) exist, these frameworks lack integration with the temporal dynamics of CT architectures.
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 biologicallyinspired 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 wellcalibrated uncertainty estimates while being interpretable at the neuronal cell level.
Scaled-Dot-Product-Attention (SDPA) mechanisms have substantially advanced sequence modeling by enabling dynamic weighting of temporal dependencies (Vaswani et al., 2017). However, the inherently discrete nature of attention computation is not well suited for CT domain. Neuronal Attention Circuit (NAC) (Razzaq et al., 2025) address this limitation by rethinking attention logit computation as the solution to a linear ODE modulated by input-dependent, nonlinear interlinked gates derived from repurposing the wiring mechanism of C.elegans nematode’s Neuronal Circuit Policies (NCPs) (Lechner et al., 2018). While NAC provides solver-free scalability, its deterministic formulation prevents it from providing uncertainty estimates. In this work, we extend the foundational principles of NAC by reframing attention logit computation as the solution to an Ornstein–Uhlenbeck stochastic differential equation (OU-SDE) (Maller et al., 2009), where the mean (ϕ), meanreversion rate (κ), and diffusion term (ψ) are modulated by input-dependent, nonlinear, interlinked gates derived from repurposed NCPs. We refer to this formulation as the Neuronal Stochastic Attention Circuit (NSAC). NSAC yields a Gaussian distribution over temporal attention logits and propagates stochasticity through a logistic-normal distribution over attention weights to compute attention scores and probabilistic outputs. NSAC learns quantification of both aleatoric and epistemic uncertainty via a two-term training objective combining a Gaussian negative log-likelihood (NLL) with an epistemic separation regularizer. We implement the NSAC in a diverse set of learning tasks, including: (i) irregular CT function approximation; (ii) multivariate regression; (iii) multivariate long-range forecasting; (iv) industrial prognostics; and (v) the lane-keeping of autonomous vehicles. NSAC remains competitive in terms of accuracy and produces reasonably well-calibrated uncertainty tubes against several post-hoc UQ baselines and a SDE-based
1. Introduction Representation learning in Continuous-time (CT) domain with reliable uncertainty quantification (UQ) is essential for high-stakes, safety-critical systems such as robotics (Lechner et al., 2020; Chahine et al., 2023), medical (Zhang et al., 2024) and Industry 4.0 (Razzaq and Zhao, 2025a), where failures can lead to significant human or economic 1
Department of Automation, University of Science & Technology of China, Hefei, China 2 Institute of Artificial Intelligence, Hefei Comprehensive National Science Center. Correspondence to: Yun-Bo Zhao <[email protected]>, Waleed Razzaq <[email protected]>. Preprint. May 26, 2026.
1
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
architecture.
incorporated latent continuous dynamics into time-aware attention mechanisms (Chen et al., 2023), and parameterized attention matrices directly through partial differential equations (PDEs) (Zhang and Zhou, 2025). NAC (Razzaq et al., 2025) interprets attention logits as the solution of a linear ODE modulated by an input-dependent nonlinear gating function inspired by NCPs (Lechner et al., 2018). These formulations improve temporal expressivity but remain fundamentally deterministic, providing no explicit mechanism for calibrated uncertainty propagation through attention itself.
To summarize, our contributions are as follows: 1. To the best of our knowledge, we are the first to directly incorporate stochasticity into CT attention mechanism, enabling principled probabilistic CT representation learning. 2. We derive a two-term objective function that regularizes epistemic uncertainty and improves the calibration of uncertainty estimates.
Neural SDE: Neural SDE models incorporate stochasticity directly into latent dynamics. SDE-Net (Kong et al., 2020) decomposes hidden evolution drift and diffusion components to capture predictive uncertainty, whereas Latent SDE (Li et al., 2020) extends this framework to variational CT latent trajectory modeling. These methods provide stochastic representations but require numerical SDE solvers such as Euler-Maruyama and adjoint-based optimization, introducing substantial computational overhead and discretization error. More importantly, stochasticity is confined to the hidden state rather than attention generation, leaving uncertainty disconnected from the mechanism governing token interaction.
3. We validate the proposed method across diverse representation tasks: (i) irregular CT function approximation; (ii) multivariate regression; (iii) multivariate long-range forecasting; (iv) industry 4.0 prognostics; and (v) lane-keeping of autonomous vehicles. The rest of the paper is organized as follows. Section 2 reviews the related literature and highlights the positioning of the proposed method. Section 3 provides comprehensive methodological details. Section 4 presents a comprehensive evaluation across a range of representation tasks. Section 5 discusses the results. Finally, Section 6 concludes the paper.
Positioning of our work: Our work addresses these gaps by unifying CT attention dynamics and uncertainty modeling within a single formulation. Specifically, the NSAC generalizes the NAC by modeling attention-logit evolution as an OU-SDE whose closed-form margin admits Gaussian distribution without requiring numerical simulation. This makes uncertainty an intrinsic architectural property of attention generation rather than a post-hoc method, enabling joint modeling of aleatoric and epistemic through MC forward passes while preserving solver-free computational efficiency.
2. Related Works Post-hoc Methods: A large class of UQ methods augments otherwise deterministic predictors with a post-hoc mechanism for predictive variance after feature extraction. Gaussian maximum-likelihood estimation (GMLE) (Kendall and Gal, 2017) introduces heteroscedastic output variance but captures only aleatoric uncertainty. Bayesian approximation, such as Monte-Carlo Dropout (MCD) (Gal and Ghahramani, 2016), estimates uncertainty through stochastic inference-time masking, coupling uncertainty estimation to a regularization mechanism rather than explicitly modeling predictive stochasticity. Deep Ensembles (DE) (Lakshminarayanan et al., 2017) improve calibration through aggregation over independently trained predictors but require linearly increasing memory and computational cost. Deep Evidential Regression (DER) (Amini et al., 2020) instead parameterizes predictive evidence through a Normal-InverseGamma (NIG) prior, enabling closed-form uncertainty estimates; however, its evidential output is still derived from deterministic latent representations. Across these approaches, uncertainty is imposed after representation learning, leaving the internal attention or feature-composition dynamics unchanged.
3. Neuronal Stochastic Attention Circuit (NSAC) We propose to view the computation of attention logits as a stochastic Gaussian distribution, interpreting them as the solution to an OU-SDE, modulated by input-dependent, nonlinear, interlinked gates derived from repurposed C.elegans nematode NCP’s wiring mechanism:
drift
z }| { dat = fκ ([q; k], t, θκ ) ·(fϕ ([q; k], t, θϕ ) −at ) dt + | {z } | {z }
CT Attention: CT attention mechanisms address the temporal adaptivity of the SDPA by replacing the discrete layerwise transformation with a dynamical systems formulation. Prior work has modeled joint hidden-state and attention evolution through neural ODEs (Chien and Chen, 2021),
κ(u)>0
ϕ(u)
diffusion
z }| { fψ ([q; k], t, θψ ) dWt | {z } ψ(u)>0
2
(1)
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Here, u = [q; k] is the sparse Top-K concatenated input, κ(u) denotes the learnable mean-reversion rate head with parameters θκ , ϕ(u) represents the learnable long-term mean attention head;, and ψ(u) > 0 head controls the Brownian fluctuations (Wt ) in the logit trajectory. Each head is parameterized by a shared backbone network as in NAC (Razzaq et al., 2025). We refer to this formulation as the Neuronal Stochastic Attention Circuit (NSAC). It enables the logits to mirror adaptive temporal dynamics found in C.elegans nematode while simultaneously producing probabilistic outputs.
curation; (ii) handling time (t); (iii) gating structures; (iv) stochastic attention weights; (v) stochastic attention scores; (vi) scaling to multi-head; (vii) probabilistic output. An illustration of the internal architecture of NSAC layer is provided in Figure 1. Input Curation: To mitigate the quadratic memory complexity O(n2 ) of attention, we adopt the input curation strategy from NAC, which uses a hybrid Top-K selection based on square-root partitioning (Child et al., 2019). The key sequence is divided into blocks, and queries are first scored against block centroids to identify high-salience candidate regions. Fine-grained scoring is then applied only within these regions, yielding a reduced set of concatenated query√ key pairs and lowering the overall complexity to O(n n).
Motivation behind this formulation: The choice of the OU formulation is motivated by its established role in modeling mean-reverting dynamics in biological systems, particularly neuronal membrane potential fluctuations (Ricciardi and Sacerdote, 1979; Jahn et al., 2011; Ditlevsen and Lansky, 2005) and synaptic noise (Fellous et al., 2003). The neural circuits of C. elegans exhibit similar equilibrium-seeking dynamics (Iwasaki and Gomi, 2004). This behavior also appears to extend across scales, as attention-modulated circuits in primates show comparable mean-reverting dynamics (Feng, 2012). Furthermore, a key advantage of the OU formulation is its closed-form solution, which avoids the discretization errors associated with numerical SDE solvers such as the Euler–Maruyama method.
Handling Time (t): The NSAC adapts its temporal dynamics to the task following (Hasani et al., 2022) by constructing a per-head internal normalized time vector t = σ(−ta tsample + tb ), where ta and tb are learnable affine parameters and tsample is derived from the input sample. When the task contains meaningful temporal structure, tsample introduces input-dependent modulation of the dynamics. When temporal information is not meaningful, tsample is set to 1, and t reduces to a learned vector through σ(−ta + tb ). In both cases, the σ ensures t ∈ [0, 1], supporting stable stochastic dynamics modeling.
Forward- pass of the NSAC using Closed-form Solution: To obtain the attention logit state (at ), the NSAC exploits the forward pass through its closed-form solution. Although the gates are time-varying and input-dependent, they evolve on discrete internal updates. Consequently, from the perspective of CT dynamics, the gates remain constant within each interval and change only at step boundaries. This separation of time scales induces a piecewise-constant (locally frozen) parameterization (John, 1952), reducing the dynamics to a system with constant coefficients over each interval. As a result, the forward pass admits an efficient closed-form update while retaining temporal variability across steps. Given an initial condition a0 , the closed-form solution for the NSAC can be written as follows: Z t −κt −κt at = a0 e + ϕ(1 − e ) + ψ e−κ(t−s) dWs (2)
Gating Structures: To project the q, k, and v vectors, we adopt the approach of NAC, employing a sparse sensory gate derived from repurposed NCP wiring. This design produces structured, context-aware representations rather than collapsing inputs through a linear layer, thereby preserving locality and modularity and enhancing information routing. Intuitively, these functions like biological receptors, converting external stimuli into organized signals that feed downstream circuits (Razzaq et al., 2025). To modulate stochastic dynamics, we modified the sparse backbone network to Figure 1. Internal architecture produce three output heads of NSAC. producing κ, ϕ, and ψ and integrate neuronal activity over time to generate outputs that regulate signal flow. To satisfy the stability requirements,
0
The expected value and variance can be expressed as: E[at ] = ϕ + (a0 − ϕ)e−κt ψ2 1 − e−2κt Var[at ] = 2κ
(3)
The derivation is available in Appendix A.1. 3.1. Designing NSAC as Neural Network Layer We now provide the design steps to convert the NSAC into a neural network layer. It consists of 7 steps: (i) input 3
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Algorithm 1 Training Algorithm for NSAC
we enforce positivity on κ and ψ using a softplus and apply tanh nonlinearity to ϕ to allow passing negative logits. Rather than learning these parameters independently, they share a common representation, which offers several advantages: (i) faster convergence during training and (ii) improved modeling of temporal and structural dependencies. 2 The stochastic attention logit mean (µou ) and variance (σou ) are updated using Eqn. 3 under initial condition a0 = 0:
Require: Data (xID , yID ), NSAC circuit fθ , optimizer O, MC samples Nmc , OOD param. (µpert , σpert ) Ensure: Total loss: Ltotal Fake OOD samples: xOOD ← xID + N (µpert , σpert ) for n = 1 to Nmc do (n) (n) ID: (µout∼ID , σout∼ID ) ← fθ (xID ) (n) (n) OOD: (µout∼OOD , σout∼OOD ) ← fθ (xOOD ) end for Compute losses: Lnll ← E[NLL(yID , µID , σID )] Lreg ← E[Reg(Var[µID ], Var[µOOD ]] Total loss: Ltotal ← Lnll + λLreg Update: θ ← θ − ∇θ Ltotal using O return Ltotal
µou = ϕ(1 − e−κt ) 2 σou =
(4)
ψ2 1 − e−2κt 2κ
Stochastic Weights: The stochastic attention logits are modeled as a Gaussian distribution via a Brownian motion realization per head: at = µou + σou ϵ, where ϵ ∼ N (0, 1). The logits are then normalized using αt = softmax(at ), resulting in a logistic-normal distribution over the attention weights (Atchison and Shen, 1980). This formulation preserves OU-driven stochastic dynamics while remaining fully differentiable.
in-distribution (ID) data samples (DID ): h 1 (y − µout )2 i E(x,y)∼DID log(2π) + 2sout + 2 e2sout (7) To encourage separation between ID and out-of-distribution (OOD), fake OOD samples are generated by Gaussian perturbations: Lnll =
Stochastic Attention Scores: The attention scores are then computed by multiplying the αt with the value vector vt : NSAC(q, k, v) = αt · vt
xOOD = xID + ξ,
(5)
2 ξ ∼ N (µpert , σpert )
(8)
For each input, N stochastic forward passes through stochas(n) tic weights produce predictive means {µout }N n=1 . The variance across these stochastic predictions is used as a measure of model-induced epistemic uncertainty. Compared with ID inputs, the regularization term promotes greater uncertainty for perturbed inputs: ! (n) Ex∼DID Var[µout∼ID ] Lreg = log 1 + (9) (n) Ex∼DOOD Var[µout∼OOD ] + ϵ
Scalling to Multihead: To scale this mechanism to multihead attention, we divide the input sequence into H independent subspaces (heads) of dimension dmodel /H, yielding query, key, and value tensors (q (h) , k (h) , v (h) ) for h ∈ (h) {1, . . . , H}. For each head, the pairwise mean (µou ) and 2(h) variance (σou ) are computed using Eqn. 4, followed by the logistic-normal distribution to calculate the stochastic (h) attention weights αt , which are then multiplied with the value vector v (h) , producing head-specific attention scores. Finally, these scores are concatenated and linearly projected into twice the dmodel dimension.
where ϵ ensures numerical stability. The final training objective is: L = Lnll + λ Lreg (10)
Probabilistic Output: A linear projection splits the learned representation into two dmodel dimenesion streams, producing a predictive mean (µout ) and log-standard- deviation (sout ). These define a heteroscedastic Gaussian predictive distribution, allowing the model to produce probabilistic outputs.
with λ controlling the strength of the regularizer. This objective encourages accurate predictions on DID data while increasing epistemic uncertainty under DOOD , improving uncertainty separation and robustness.
4. Evaluation p(y | x) = N
2 µout , σout
,
(sout )
σout = e
.
(6)
We evaluate NSAC across diverse set of learning domains including: (i) irregular CT function approximation; (ii) multivariate regression; (iii) multivariate long-range forecasting (iv) industry 4.0 prognostics; and (iv) lane-keeping of autonomous-vehicles. Additional experiments including: (i) detailed ablation analysis on key parameters; (ii) closedloop analysis of AVs; and scalability analysis are provided in Appendix C.
3.2. Learning Uncertainty through Stochasticity NSAC learns predictive uncertainty through a two-term objective combining predictive accuracy and epistemic uncertainty separation. Predictive accuracy is enforced via the Gaussian negative log-likelihood (LNLL ) evaluated on 4
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Figure 2. Spiral trajectory with uncertainty estimates: (A) NSAC; (B) GMLE; (C) DE; (D) MCD; (E) DER; and (F) SDE-Net. NSAC and DER provide more compact uncertainty tubes with smooth mean accuracy.
Baselines: We compare NSAC against several state-of-theart UQ techniques: (i) GMLE (Kendall and Gal, 2017); (ii) MCD (Gal and Ghahramani, 2016); (iii) DE (Lakshminarayanan et al., 2017); (iv) DER (Amini et al., 2020); and (v) SDE-Net (Kong et al., 2020). For a fair comparison, the penultimate layer of all post-hoc baselines uses NAC with the same model dimensions (dmodel ), attention head (numheads ), sparsity, and Top-K as NSAC.
quantitative results are reported in Table 1. NSAC achieved the strongest predictive performance, with the lowest MSE of 0.0002±0.0001, followed by DE (0.0071 ± 0.0127) and MCD (0.0094±0.0084). It also achieved the best CRPS (0.0095 ± 0.0021), outperforming DE (0.0180 ± 0.0099), MCD (0.0616±0.0090), and SDE-Net (0.8119±0.0185), confirming its superior accuracy and distributional sharpness. In NLL, DE ranked first (-2.5632±0.3694), followed closely by NSAC (2.3483±0.0536). The DER exhibited severe distributional failure, with the CRPS degrading to 0.9801±0.05754 despite a favorable NLL (-2.5026±0.1637). With respect to calibration, DER and SDE-Net achieved the lowest ECE (0.0431±0.0443 and 0.0475±0.0340, respectively), whereas NSAC achieved moderate calibration (0.1631±0.0435). The perfect AUROC score of the GMLE and DER likely reflects monotonic variance separation rather than calibrated uncertainty estimation; the NSAC achieved a moderate AUROC of 0.5736±0.0261.
Test Setup and Metrics: We train each model with five random seeds and report the mean and standard deviation of the Mean-Squared-Error (MSE) for predictive accuracy, Negative Log-likelihood (NLL) for both prediction error and uncertainty estimation, Continuous Ranked Probability Score (CRPS) for probabilistic forecast quality, the Expected Calibration Error (ECE) for calibration of uncertainty estimates, and Area under the Receiver Operating Characteristic Curve (AUROC) for ranking quality. CRPS measures the accuracy of the predicted distribution, while ECE captures how well predicted probabilities are calibrated. Ideally, both should be low; in practice, a balance is necessary. Sharper distributions can improve CRPS but may worsen ECE, whereas broader distributions tend to improve calibration at the expense of precision. AUROC measures how well the model can distinguish higher from lower target values based on its prediction scores, reflecting its ability to correctly rank samples rather than quantify calibration or exact accuracy. More detailed definitions of metrics are provided in Appendix B.1 and data descriptions, preprocessing, and neural network details are available in Appendix B.2. Table 1 reports all the quantitative results.
Ablation Decomposition. We also perform an ablation study on spirals to examine the effect of regularization on uncertainty decomposition. The qualitative visualization of the decomposition is presented in Figure 3, and the quantitative metrics (MSE & MAE) are presented in Table 2. This confirms the functional role of Lreg . With Lreg , the NSAC successfully decomposes uncertainty into distinct aleatoric and epistemic components. Without Lreg , epistemic uncertainty is suppressed while aleatoric uncertainty is retained, and predictive accuracy is preserved with near-identical MSEs and MAEs across interpolation (MSE: 0.0002, MAE: 0.0143 vs 0.0117) and extrapolation (MSE: 0.0017 vs. 0.0012, MAE: 0.0304 vs. 0.0272) regimes.
4.1. Irregular CT Function Approximation In the first experiment, we evaluate NSAC’s ability to approximate irregular CT function with uncertainty estimates. We curated the irregular spiral dataset following the procedure in (Chen et al., 2023), which generated 300 2D spiral trajectories, each sampled at 150 equally spaced time points. To introduce irregularity, we randomly selected 50 points from each trajectory without replacement. We then visualize the resulting interpolation and extrapolation. The qualitative results with uncertainty estimates are shown in Figure 2 and
4.2. Multivariate Regression In the second experiment, we evaluate all the methods on a multivariate regression task. Specifically, we utilized two well-known regression benchmark datasets: (i) Boston Housing; and (ii) Kin8nm. Prior to training, both datasets were normalized using MinmaxScalar. Boston: The Boston Housing dataset has 506 samples 5
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning Table 1. Quantitative analysis of all methods. ETTm1
J.Climate
XJTU-SY
PRONOSTIA
Industry 4.0 HUST
Udacity
Autonomous Vehicles CarRacing
0.0002±0.0001 -2.3483±0.0536 0.0095±0.0021 0.1631±0.0435 0.5736±0.0261
0.0301±0.0061 -0.2665±0.1597 0.0940±0.0093 0.3974±0.0334 0.9923±0.0116
0.0327±0.0004 -0.3870±0.0596 0.1029±0.0007 0.3085±0.0084 0.9994±0.0005
0.0199±0.0042 -1.3816±0.3152 0.1037±0.0202 0.3970±0.1024 0.7533±0.0662
0.1675±0.0627 -1.9114±0.2946 0.1422±0.0977 0.3179±0.0609 0.7750±0.1088
0.0048±0.0018 -0.0219±0.0012 0.0261±0.0049 0.1404±0.0098 0.3436±0.1874
0.0294±0.0123 -1.7521±0.1832 0.0310±0.0032 0.1337±0.0106 0.4418±0.0455
0.0033±0.0008 -1.5889±0.1378 0.0231±0.0031 0.0841±0.0104 0.4191±0.0868
0.0249±0.0068 -0.4122±0.1679 0.0814±0.0075 0.1982±0.0335 0.9984±0.0032
0.0154±0.0064 -1.6450±0.2422 0.7271±0.2619 0.0763±0.1265 0.6763±0.1342
MSE (↓) NLL (↓) CRPS (↓) ECE (↓) AUROC (↑)
0.0034±0.0036 -2.1598±0.3487 0.3003±0.0063 0.4727±0.0864 1.0±0.0
0.0304±0.0063 -0.2592±0.1759 0.0937±0.0091 0.4084±0.0315 1.0±0.0
0.0327±0.0004 -0.3893±0.0593 0.1031±0.0008 0.5144±0.0055 1.0±0.0
0.0078±0.0032 -0.9513±0.0313 0.0755±0.0043 0.4457±0.0378 1.0±0.0
0.8878±0.0545 0.9812±0.5706 0.8181±0.0156 0.6977±0.0580 1.0±0.0
0.0221±0.0022 0.2280±0.0988 0.1353±0.0123 0.8824±0.0201 1.0±0.0
0.0588±0.0274 -1.4809±0.2681 0.0981±0.0077 0.9066±0.0068 1.0±0.0
0.0079±0.0034 -1.1007±0.2740 0.0964±0.0079 0.8609±0.0403 1.0±0.0
0.0251±0.0068 -0.4129±0.1664 0.0816±0.0075 0.2060±0.0340 1.0±0.0
0.1012±0.1045 -0.4716±0.0933 0.2087±0.0597 0.7531±0.1785 1.0±0.0
MSE (↓) NLL (↓) CRPS (↓) ECE (↓) AUROC (↑)
0.0071±0.0127 -2.5632±0.3694 0.0180±0.0099 0.4285±0.1073 0.5458±0.0244
0.0303±0.0061 -0.2633±0.1636 0.0945±0.0095 0.4012±0.0388 0.9895±0.0159
0.0327±0.0005 -0.3905±0.0577 0.1032±0.0008 0.5025±0.0055 0.9996±0.0002
0.0059±0.0025 -0.8922±0.0712 0.0494±0.0050 0.2717±0.0648 0.4209±0.0493
0.6881±0.2989 2.3523±1.7854 0.5005±0.1006 0.7919±0.0632 0.6156±0.0522
0.0139±0.0023 0.3432±0.0877 0.0476±0.0081 0.5832±0.1622 0.2577±0.1186
0.0536±0.0372 -1.6891±0.2163 0.0376±0.0139 0.5451±0.2481 0.4546±0.1402
0.0034±0.0005 0.1531±0.2809 0.0217±0.0027 0.5916±0.1966 0.3814±0.0784
0.0251±0.0069 -0.4140±0.1662 0.0817±0.0076 0.2015±0.0339 0.9991±0.0018
0.0362±0.0191 0.1456±0.0278 0.0809±0.0278 0.5983±0.1189 0.5484±0.1366
MSE (↓) NLL (↓) CRPS (↓) ECE (↓) AUROC (↑)
0.0094±0.0084 1.7845±0.1469 0.0616±0.0090 0.1977±0.0417 0.8169±0.0225
0.0418±0.0058 -0.2618±0.1704 0.1168±0.0073 0.3613±0.0345 0.8058±0.0434
0.0327±0.0005 -0.3889±0.0598 0.1033±0.0008 0.5056±0.0071 0.9994±0.0002
0.0070±0.0040 -0.7977±0.1874 0.0525±0.0096 0.2810±0.0727 0.3620±0.0902
0.6464±0.1363 6.7464±6.8760 0.5242±0.0746 0.8318±0.0620 0.7483±0.0531
0.0145±0.0087 0.0394±0.0028 0.0458±0.0131 0.3476±0.0597 0.2260±0.0992
0.0670±0.0276 -0.6656±0.1664 0.0433±0.0094 0.4348±0.1056 0.1926±0.0513
0.0082±0.0070 0.5466±0.2362 0.0277±0.0090 0.4489±0.1319 0.4319±0.0748
0.0252±0.0068 -0.4053±0.1577 0.0820±0.0074 0.2013±0.0315 0.9960±0.0064
0.0102±0.0034 0.4916±0.5748 0.0354±0.0063 0.4249±0.0797 0.5336±0.1266
MSE (↓) NLL (↓) CRPS (↓) ECE (↓) AUROC (↑)
0.1694±0.0033 -2.5026±0.1637 0.9803±0.05754 0.0431±0.0443 1.0±0.0
0.0310±0.0059 -0.2425±0.1352 0.2111±0.0458 0.1208±0.0296 1.0±0.0
0.0328±0.0006 -0.3878±0.0123 0.8011±0.1905 0.0001±0.0000 1.0±0.0
0.0154±0.0028 -0.9593±0.2032 0.9438±0.3740 0.0012±0.0003 1.0±0.0
1.0938±0.0177 3.1940±0.1607 0.9357±0.2642 0.0113±0.0030 1.0±0.0
0.0331±0.0004 -2.2316±0.5052 0.1099±0.0055 0.8621±0.0179 1.0±0.0
0.0252±0.0001 -2.1217±1.3504 0.6716±0.0978 0.1484±0.1436 1.0±0.0
0.0296±0.0001 -2.2585±1.0695 0.0814±0.0070 0.3441±0.1347 1.0±0.0
0.0253±0.0068 -0.6008±0.0965 0.7994±0.0004 0.0125±0.0024 1.0±0.0
0.0913±0.0062 -1.7055±1.7941 0.7217±0.0319 0.0582±0.1119 1.0±0.0
MSE (↓) NLL (↓) CRPS (↓) ECE (↓) AUROC (↑)
0.0242±0.0190 -2.0640±0.2909 0.8119±0.0185 0.0475±0.0340 0.6923±0.0265
0.0348±0.0071 -0.1729±0.1785 0.4041±0.0020 0.0632±0.0086 0.9338±0.0254
0.0328±0.0005 -0.3822±0.0578 0.4048±0.0020 0.0673±0.0007 0.9989±0.0004
0.0199±0.0035 -1.5697±0.2236 0.7437±0.0129 0.0285±0.0038 0.8664±0.0305
1.6981±0.2651 -1.1297±0.48276 0.8338±0.0743 0.5171±0.0396 0.9020±0.0297
0.0068±0.0023 0.9260±0.1421 0.9017±0.0803 0.0129±0.0066 0.2249±0.0627
0.0239±0.0032 -1.1261±0.1385 0.8534±0.0731 0.0159±0.0034 0.4274±0.0388
0.0170±0.0017 -1.9123±0.2048 0.9478±0.0880 0.0204±0.0034 0.3933±0.0650
0.0255±0.0071 -0.3875±0.1544 0.3771±0.0143 0.0424±0.0119 1.0±0.0
0.0550±0.0213 -1.7786±0.7054 0.9114±0.0745 0.0803±0.0326 0.6624±0.225
SD E-
et
M
G
LE
N SA C
MSE (↓) NLL (↓) CRPS (↓) ECE (↓) AUROC (↑)
M
Kin8nm
E
Long-range Forecasting
Boston
D
Multivariate Regression
C D
Spiral
D ER
Metrics
N
Model
Note: (↑) higher is better; (↓) lower is better. Table 2. Ablation decomposition results Metric
Aleatoric Lreg
W/O Lreg
Epistemic Lreg
W/O Lreg
Task = Interpolation MSE (↓) MAE (↓)
Figure 3. Ablation Decomposition of NSAC uncertainty estimates: (A) Aleatoric uncertainty with regularizer; (B) Aleatoric uncertainty without regularizer; (C) Epistemic uncertainty with regularizer; (D) Epistemic uncertainty without regularizer.
0.0002 0.0143
0.0002 0.0117
0.0002 0.0119
0.0002 0.0058
Task = Extrapolation MSE (↓) MAE (↓)
of homes in Boston suburbs, each with 13 features, such as crime rate, average rooms, and accessibility to highways, with the target being the median home value. NSAC achieves the lowest MSE (0.0301±0.0061), outperforming DE (0.0303±0.0061) and the GMLE (0.0304±0.0063), while improving over MCD (0.0418±0.0058) and SDENet (0.0348±0.0071). The DER attains a competitive MSE but exhibits severe distributional distortion, with the CRPS increasing to 0.211±0.0458, which is double those of the NSAC (0.0940±0.0093) and DE (0.0945±0.0095). NLL values are tightly clustered for NSAC (-0.2665±0.1597), DE (-0.2633±0.1636), and MCD (-0.2618±0.1785). In terms of calibration, SDE-Net achieves the lowest ECE (0.0632±0.0086), followed by DER (0.1208±0.0296), while NSAC (0.3974±0.0334) is comparable to DE (0.4012±0.0388). NSAC attains an AUROC of 0.9923±0.0116.
0.0017 0.0304
0.0012 0.0272
0.0018 0.0300
0.0012 0.0237
Kin8nm: The Kin8nm dataset contains a total of 8192 samples, where each sample contains 8 numerical input features representing physical measurements and a single continuous target variable. In this test, the prediction accuracy is nearly saturated, with all methods achieving MSEs between 0.0327 and 0.0328, making the distributional quality the primary differentiator. NLL is also effectively indistinguishable for NSAC (-0.3870±0.0596), GMLE (-0.3893±0.0593), DE (-0.3905±0.0577), MCD (-0.3889±0.0598), and DER (-0.3878±0.0123). NSAC achieves a CRPS of 0.1029±0.0007, matching those of the GMLE, DE, and MCD, whereas DER and SDE-Net substantially decrease despite having comparable MSE. The near-zero ECE of the DER likely reflects evidential overfitting rather than genuine calibration, which is consistent with its elevated CRPS variance. NSAC has an ECE of 6
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
0.3085±0.0084 and an AUROC of 0.9994±0.0005, matching DE and MCD.
(48 look-back ) to forecast the following 4 hours (24 horizon). NSAC achieves the lowest MSE of 0.1675±0.0627, substantially outperforming MCD (0.6464±0.1363), DE (0.6881±0.2989), DER (1.0938±0.0477), and SDE-Net (1.6981±0.2651). It also achieves the most favorable NLL of -1.9114±0.29 and the lowest CRPS of 0.1422±0.0977, while all the baselines exhibit severely degraded density estimation. Furthermore, NSAC records a competitive AUROC of 0.7750±0.1088, outperforming DE (0.6156±0.0522) and MCD (0.7483±0.0531). Figure 4(B) shows that NSAC produces more accurate forecasting trajectories than competing methods. 4.4. Industry 4.0 In this fourth experiment, we evaluate the ability of the NSAC to model physical dynamics through the task of bearing degradation estimation. Bearing degradation prediction is a classical problem in industrial engineering and is typically addressed using data-driven approaches. However, such methods often suffer from limited generalizability beyond their training regimes, while their large model sizes make them impractical for deployment in resourceconstrained environments. The objective of this experiment is therefore to learn degradation dynamics from a single operating condition and assess zero-shot generalization to entirely unseen datasets. It also aims to assess the distributional generalization of uncertainty estimates and prediction accuracy across datasets.
Figure 4. Visualization of forecast projections with uncertainty estimates: (A) ETTm1; and (B) Jena-Climate.
4.3. Multivariate Long-range Forecasting In the third experiment, we evaluate the multivariate longrange forecasting capability of NSAC. Two widely used benchmark datasets are utilized: (i) ETTm1; and (ii) JenaClimate. Prior to training, both datasets were transformed using MinMaxScaler, and predictions were inversetransformed back for post-training evaluation. Each dataset was split into training (80%), validation (10%), and test (10%) folds. Figure 4 illustrates the forecasting projections with uncertainty tubes.
We conduct experiments on three widely used benchmark datasets: (i) XJTU-SY (Wang et al., 2018), (ii) PRONOSTIA (Nectoux et al., 2012), and (iii) HUST (Thuan and Hong, 2023). Training is performed using the first three bearing from the XJTU-SY dataset, while evaluation is conducted in a zero-shot manner on the remaining datasets. Prior to training, the raw vibration signals are transformed into physically meaningful features following (Razzaq and Zhao, 2025a), and nonlinear degradation labels are generated using the procedure described in (Razzaq and Zhao, 2025b). This formulation reduces the problem to a longhorizon regression task. Figure 5 presents the qualitative results.
ETTm1: The ETTm1 dataset consists of 7 time-dependent features sampled at 15-minute intervals. For forecasting, the model was conditioned on the preceding 12 hours (48 lookback) to predict the subsequent 6 hours (24 horizons). Posthoc baselines achieve stronger predictive accuracy, with DE attaining the lowest MSE of 0.0059±0.0025, followed by MCD (0.0070±0.0040) and GMLE (0.0078±0.0032). In contrast, both NSAC and SDE-Net record an MSE of 0.0199, indicating a predictive trade-off associated with stochastic inductive bias. However, the NSAC achieves the best NLL of -1.3816±0.3152, outperforming all the baselines. GMLE and DE attain the lowest CRPS values of 0.0755 and 0.0494, respectively, while NSAC achieves 0.1037±0.0202. Similar to Kin8nm test, the DER’s nearzero ECE of 0.0012±0.0003 likely reflects evidential overfitting, whereas the NSAC records an ECE of 0.3970±0.1024 alongside a competitive AUROC of 0.7533±0.0662, outperforming DE (0.4209±0.0493) and MCD (0.3620±0.0902). Figure 4(A) indicates that NSAC forecasting trajectories align more closely with the ground truth..
XJTU-SY (in-distribution): NSAC achieves the lowest MSE of 0.0048±0.0018, outperforming SDENet (0.0068±0.0023), MCD (0.0145± 0.0087), DE (0.0139±0.0023), GMLE (0.0221±0.0022), and DER (0.0331±0.0004). Although the DER attains the most favorable NLL of -2.2316±0.5052, the NSAC achieves a competitive value of -0.0219±0.0012, surpassing those of the GMLE, DE, MCD, and SDE-Net. NSAC further records the lowest CRPS of 0.0261±0.0049 and an ECE of 0.1404±0.0098. Figure 5(A) confirms NSAC’s superior probabilistic accuracy and calibration relative to those of all
Jena-Climate: The Jena-Climate dataset contains 14 atmospheric variables recorded in Jena, Germany, at 10-minute intervals. The model was trained using the past 8 hours 7
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Figure 5. Visualization of degradation trajectory with uncertainty estimates. (A) XJTU-SY; (B) PRONOSTIA; and (C) HUST.
the baselines on XJTU-SY.
corresponding steering angles are recorded. This produces a dataset of 15647 RGB images of size 320 × 160 × 3 with corresponding steering angles. We adopt the preprocessing pipeline and architecture introduced in (Shibuya, 2017), replacing the latent dense layers with NSAC or baseline. The quantitative results are summarized in Table 1. All the models converge to nearly identical predictive accuracies, with MSE values ranging from 0.0249 to 0.0255. DER achieves the lowest NLL of -0.6008±0.0965, followed by DE (-0.4140±0.1662), the GMLE (-0.4129±0.1664), and the NSAC (-0.4122±0.1679). NSAC attains the lowest CRPS of 0.0814±0.0075, marginally outperforming the GMLE (0.0816±0.0075) and DE (0.0817±0.0076) and substantially surpassing the DER (0.7994±0.0004). DER achieves the lowest ECE of 0.0125±0.0024, followed by SDE-Net (0.0424±0.0119), while NSAC records an ECE of 0.1982±0.0335 and AUROC of 0.9984±0.0032.
PRONOSTIA (out-of-distribution): Under a distributional shift, the NSAC achieves the lowest CRPS (0.0310±0.0032) and ECE (0.1337±0.0106), demonstrating strong generalizability in terms of probabilistic estimation. Although DER achieves the lowest MSE (0.0252±0.0001) and NLL (-2.1217±1.3504), its substantially higher CRPS (0.6716±0.0978) suggests significant distributional distortion and poor uncertainty quality. Figure 5(B) illustrates the degradation trajectory with uncertainty tubes for PRONOSTIA. HUST (out-of-distribution): NSAC achieves the lowest MSE (0.0033±0.0008) and ECE (0.0841±0.0104). DE records the lowest CRPS (0.0217±0.0027), marginally outperforming NSAC (0.0231±0.0031). DER again achieves the most favorable NLL (-2.2585±1.06), followed by NSAC (-1.5889±0.1378). Figure 5(C) illustrates the degradation trajectory with uncertainty tubes for HUST.
OpenAI-CarRacing: For OpenAI CarRacing, the training dataset is generated using a PPO agent trained for two million steps to autonomously drive the vehicle while recording observations and corresponding control actions, including the steering angle, throttle, and brake. This results in a dataset of 48102 grayscale images of size 84 × 84 paired with control actions. We develop an end-to-end neural architecture consisting of convolutional layers for spatial feature extraction followed by NSAC or baseline methods for action prediction in a regression setting. The quantitative results are reported in Table 1. NSAC achieves the most favorable NLL of -1.6450±0.2422 and a competitive ECE of 0.0763±0.1265, with an MSE of 0.0154±0.0064. The MCD attains the lowest MSE of 0.0102±0.0034 and the lowest CRPS of 0.0354±0.0060. A detailed closed-loop analysis with trajectory analysis is presented in Appendix C.2.
4.5. Autonomous Vehicle In the fifth experiment, we evaluate the ability of NSAC to function as a step-wise controller for autonomous vehicles. To perform this evaluation, we utilize two widely used autonomous driving simulation environments: (i) the Udacity Self-Driving Car simulator and (ii) OpenAI CarRacing. The objective is to develop an end-to-end controller that takes an input image stream and predicts steering actions based on the observed scene. Specifically, the task requires learning the relationship between the road horizon and the corresponding steering commands. Udacity Simulator: For the Udacity simulator, the training dataset is generated by manually driving the vehicle for 30 minutes while the front-mounted camera input and 8
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Figure 6. Intuitive visualization of NSAC cell activity during driving. (A) Test drive map; (B) Sensory neurons output (Centroid) for q, k, v projections; (C) Backbone activity with interneurons (Centroid), Command neurons (Centroid) and OU coefficients (κ, ϕ, and ψ) outputs (D) Attention logit (at ) computation. Internal plots are actual neuron activities.
4.6. NSAC Carries Interpretability
6. Conclusion
Interpretability refers to explaining AI decisions in a humanunderstandable manner through the analysis of cell-level potentials. We conduct a cell-level analysis of the NSAC on the OpenAI-CarRacing environment using a the test track shown in Figure 6(A). Specifically, we trace sensory projections (q, k, v) through the backbone, from inter-command neurons to the output heads κ, ϕ, and ψ. The cell potentials at each point along the test track. The output head ϕ varies within [-0.2, 0.2], with values below -0.10 associated with the left steering turn, values above 0.10 associated with the right steering turn, and intermediate values corresponding to straight. Similarly, κ remains within [0.65, 0.8], where values below 0.7 indicate left turns and values above 0.75 indicate right turns. The head ψ lies within [0.625, 0.750], with values below 0.685 corresponding to left turns and values above 0.715 corresponding to right turns. The attention logits vary within [-0.2, 0.2], where values above 0.1 indicate right steering, values below -0.1 indicate left steering, and intermediate values correspond to straight driving.
We introduce NSAC, a novel biologically- inspired CT stochastic attention architecture that reformulates attention logits as the solution of an Ornstein–Uhlenbeck SDE modulated by nonlinear, interlinked gates derived from the C.elegans NCP wiring mechanism. NSAC admits a closed-form forward pass, avoiding expensive numerical SDE solvers while enabling principled stochasticity within the attention logits. Logistic-normal distribution normalization propagates this stochasticity over attention weights to compute attention scores and heteroscedastic probabilistic outputs. A two-term objective that combines likelihood estimation with an epistemic-separation regularizer enables the model to learn accurate and aleatoric and epistemic uncertainty estimates. Empirically, we evaluate the NSAC across a diverse set of learning tasks and compare it with several state-of-the-art UQ baselines. We observe that the NSAC performs competitively in terms of accuracy and produces reasonably well-calibrated uncertainty tubes while being interpretable at the neuronal cell level. NSAC provides a model-integrated approach for incorporating uncertainty in CT attention architectures, which is particularly relevant in settings that require temporal continuity and tight architectural integration. NSAC is not intended as a replacement for established uncertainty quantification methods, but rather as a complementary mechanism that enables uncertainty to be modeled directly within the architecture without introducing instability.
5. Discussion NSAC achieved strong CRPS and NLL across benchmarks, indicating reliable probabilistic output rather than improved point prediction alone. The moderate ECE observed in some settings is consistent with the stochastic inductive bias of the OU formulation, where locally sharper uncertainty estimates do not always translate into exact marginal coverage. The perfect AUROC for GMLE and DER is misleading as their variances are deterministically induced and monotonically related to input features, so median-based binarization yields near-trivial separability without reflecting calibration quality. In contrast, NSAC produces variance that responds nonmonotonically to input dynamics, yielding more informative but less separable uncertainty ranking. DER further shows evidential collapse, with near-zero ECE but degraded CRPS on spiral and ETTm1, indicating overconfident predictions.
Reproducibility Statement The code for reproducibility is available at https://github.com/itxwaleedrazzaq/ neuronal_stochastic_attention_circuit.
9
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
References
Yuishi Iwasaki and Sohei Gomi. Stochastic formulation for a partial neural circuit of c. elegans. Bulletin of mathematical biology, 66(4):727–743, 2004.
Alexander Amini, Wilko Schwarting, Ava Soleimany, and Daniela Rus. Deep evidential regression. Advances in neural information processing systems, 33:14927–14937, 2020.
Patrick Jahn, Rune W Berg, Jørn Hounsgaard, and Susanne Ditlevsen. Motoneuron membrane potentials follow a time inhomogeneous jump diffusion process. Journal of computational neuroscience, 31(3):563–579, 2011.
Jhon Atchison and Sheng M Shen. Logistic-normal distributions: Some properties and uses. Biometrika, 67(2): 261–272, 1980.
Fritz John. On integration of parabolic equations by difference methods: I. linear and quasi-linear equations for the infinite interval. Communications on Pure and Applied Mathematics, 5(2):155–211, 1952.
Makram Chahine, Ramin Hasani, Patrick Kao, Aaron Ray, Ryan Shubert, Mathias Lechner, Alexander Amini, and Daniela Rus. Robust flight navigation out of distribution with liquid neural networks. Science Robotics, 8(77): eadc8892, 2023.
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Yuqi Chen, Kan Ren, Yansen Wang, Yuchen Fang, Weiwei Sun, and Dongsheng Li. Contiformer: Continuous-time transformer for irregular time series modeling. Advances in Neural Information Processing Systems, 36:47143– 47175, 2023.
Lingkai Kong, Jimeng Sun, and Chao Zhang. Sde-net: Equipping deep neural networks with uncertainty estimates. arXiv preprint arXiv:2008.10546, 2020. Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Jen-Tzung Chien and Yi-Hsiang Chen. Continuous-time attention for sequential learning. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 7116–7124, 2021.
Mathias Lechner, Ramin M Hasani, and Radu Grosu. Neuronal circuit policies. arXiv preprint arXiv:1803.08554, 2018.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
Mathias Lechner, Ramin Hasani, Alexander Amini, Thomas A Henzinger, Daniela Rus, and Radu Grosu. Neural circuit policies enabling auditable autonomy. Nature Machine Intelligence, 2(10):642–652, 2020.
Susanne Ditlevsen and Petr Lansky. Estimation of the input parameters in the ornstein-uhlenbeck neuronal model. Physical Review E—Statistical, Nonlinear, and Soft Matter Physics, 71(1):011907, 2005.
Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International conference on artificial intelligence and statistics, pages 3870–3882. PMLR, 2020.
J-M Fellous, Michael Rudolph, Alain Destexhe, and Terrence J Sejnowski. Synaptic background noise controls the input/output characteristics of single cells in an in vitro model of in vivo activity. Neuroscience, 122(3): 811–829, 2003.
Ross A Maller, Gernot Müller, and Alex Szimayer. Ornstein– uhlenbeck processes and extensions. Handbook of financial time series, pages 421–437, 2009.
Samuel F Feng. Extensions and applications of stochastic accumulator models in attention and decision making. PhD thesis, Princeton University, 2012.
Patrick Nectoux, Rafael Gouriveau, Kamal Medjaher, Emmanuel Ramasso, Brigitte Chebel-Morello, Noureddine Zerhouni, and Christophe Varnier. Pronostia: An experimental platform for bearings accelerated degradation tests. In IEEE International Conference on Prognostics and Health Management, PHM’12., pages 1–8. IEEE Catalog Number: CPF12PHM-CDR, 2012.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050–1059. PMLR, 2016. Ramin Hasani, Mathias Lechner, Alexander Amini, Lucas Liebenwein, Aaron Ray, Max Tschaikowski, Gerald Teschl, and Daniela Rus. Closed-form continuous-time neural networks. Nature Machine Intelligence, 4(11): 992–1003, 2022.
Waleed Razzaq and Yun-Bo Zhao. Carle: a hybrid deepshallow learning framework for robust and explainable rul estimation of rolling element bearings. Soft Computing, 29(23):6269–6292, 2025a. 10
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
ϕ(u), and ψ(u) are treated as piecewise constant within each discrete update interval (locally frozen coefficients (John, 1952)). Under this assumption, the system reduces to a linear SDE with constant coefficients on each interval, admitting an exact analytical solution. Let the attention logit at evolve according to
Waleed Razzaq and Yun-Bo Zhao. Developing distanceaware uncertainty quantification methods in physicsguided neural networks for reliable bearing health prediction, 2025b. URL https://arxiv.org/abs/ 2512.08499. Waleed Razzaq, Izis Kanjaraway, and Yun-Bo Zhao. Neuronal attention circuit (nac) for representation learning. arXiv preprint arXiv:2512.10282, 2025.
dat = κ(ϕ − at ) dt + ψ dWt ,
(11)
Rewriting it in standard linear form,
Luigi M Ricciardi and Laura Sacerdote. The ornsteinuhlenbeck process as a model for neuronal activity: I. mean and variance of the firing time. Biological cybernetics, 35(1):1–9, 1979.
dat + κat dt = κϕ dt + ψ dWt .
(12)
Multiplying both sides by the integrating factor eκt and applying Itô’s product rule yields d eκt at = κϕeκt dt + ψeκt dWt . (13)
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference on computer vision, pages 618– 626, 2017.
Integrating from 0 to t, Z t
κt
e at − a0 = κϕ
Naoki Shibuya. Car behavioral cloning, 2017. URL https://github.com/naokishibuya/ car-behavioral-cloning. Accessed: 202510-05.
e
κs
Z t ds + ψ
0
eκs dWs .
(14)
0
Evaluating the deterministic integral: Z t κϕ
Nguyen Duc Thuan and Hoang Si Hong. Hust bearing: a practical dataset for ball bearing fault diagnosis. BMC research notes, 16(1):138, 2023.
eκs ds = ϕ eκt − 1 ,
(15)
0
and dividing through by eκt gives the closed-form solution
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
at = ϕ + (a0 − ϕ)e−κt + ψ
Z t
e−κ(t−s) dWs .
(16)
0
Taking expectations and using E[dWt ] = 0,
Biao Wang, Yaguo Lei, Naipeng Li, et al. Xjtu-sy bearing datasets. GitHub, GitHub Repository, 2018.
E[at ] = ϕ + (a0 − ϕ)e−κt .
Xi Zhang, Yuan Pu, Yuki Kawamura, Andrew Loza, Yoshua Bengio, Dennis L Shung, and Alexander Tong. Trajectory flow matching with applications to clinical time series modelling. Advances in Neural Information Processing Systems, 37:107198–107224, 2024.
(17)
The stochastic integral is the sole source of randomness. Applying Itô isometry, Var(at ) = ψ
2
Z t
e−2κ(t−s) ds =
0
Yukun Zhang and Xueqing Zhou. Continuous-time attention: Pde-guided mechanisms for long-sequence transformers. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, pages 21654–21674, 2025.
ψ2 1 − e−2κt . (18) 2κ
Thus the variance increases monotonically from 0 and saturates at ψ2 lim Var(at ) = . (19) t→∞ 2κ The mean reverts exponentially from the initial value a0 toward the long-term mean ϕ at rate κ, while the variance approaches a stationary limit determined by the noise scale ψ and the reversion strength κ.
Appendix A. Derivations & Proofs A.1. Derivation of Closed-form Solution
B. Evaluation Details
We derive the closed-form solution of NSAC. Throughout this derivation, the input-dependent gate parameters κ(u),
In this section, we provide comprehensive definitions of evaluation metrics used to assess the model. 11
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
a confidence score ci = 2Φ(zi ) − 1, where Φ is the standard normal cumulative distribution function (CDF). Confidence scores are partitioned into B uniform bins (10 bins in our implementation), and for each bin we calculate the empirical coverage, defined as the fraction of samples with zi below the bin midpoint. The ECE is then the weighted average absolute difference between predicted confidence and empirical coverage across bins.
B.1. Evaluation Metrics B.1.1. M EAN S QUARED E RROR (MSE) MSE is a widely used metric for evaluating regression models, measuring the average squared difference between predicted values ŷi and true targets yi : N
MSE =
1 X (ŷi − yi )2 , N i=1
(20)
B
ECE =
where N is the total number of samples. Lower MSE values indicate that predictions are closer to the true targets.
where, N be the total number of samples, and Sk the set of samples in bin k. Lower ECE indicates better alignment between predicted uncertainties and observed errors, enabling reliable uncertainty intervals.
NLL is a strictly proper scoring rule that evaluates the quality of probabilistic regression by measuring how much probability mass the predicted Gaussian distribution assigns to the true target. For a predicted mean µi , log-standard deviation si = log σi , and true target yi , the NLL is defined as N 1 X 1 (yi − µi )2 log(2π) + si + N i=1 2 2e2si
(23)
k=1
B.1.2. N EGATIVE -L OG L IKELIHOOD (NLL)
NLL =
1 X |p̂k − ĉk | · |Sk | N
B.1.5. A REA UNDER THE R ECEIVER O PERATING C HARACTERISTIC C URVE (AUROC) AUROC is a standard model discriminative quality metrix for binary classifier but we adopt it to evaluate the probabilistic predictions in a continuous regression setting. Since AUROC requires binary labels, we transform continuous targets into binary indicators by thresholding at the empirical median of the target distribution i.e., ỹi = I(yi > m), where m = median(y). For each predictive Gaussian distribution parameterized by mean µi and standard deviation σi , we compute the probability that the target exceeds the median: m − µi pi = 1 − Φ , (24) σi where Φ denotes the cumulative distribution function (CDF). These probabilities are treated as a confidence score for the positive class, then AUROC is computed by comparing the ranking induced by pi the probabilities against the binary labels ỹi : 1 XX AUROC = I(pi > pj ), (25) |P||N |
(21)
Where N is the total number of samples. Intuitively, NLL jointly evaluates the predictive accuracy and uncertainty estimates by penalizing both inaccurate prediction and poorly estimated predictive variance. Lower NLL values indicate that the predicted distribution assign higher likelihood to the observed target, reflecting better probabilistic modeling. B.1.3. C ONTINUOUS R ANKED P ROBABILITY S CORE (CRPS) CRPSis a strictly proper scoring rule that assesses both calibration and sharpness. For a predicted cumulative distribution function (CDF) Φi and true target yi , the CRPS is defined as Z ∞ 2 CRPS(Φi , yi ) = Φi (t) − 1(t ≥ yi ) dt, (22)
i∈P j∈N
Where P and N denote the sets of positive and negative samples, respectively. Higher AUROC indicates stronger discriminative capability.
−∞
where 1(·) is the indicator function. Intuitively, CRPS measures the squared distance between the predicted CDF and the empirical step function at the observation. Lower CRPS values indicate more accurate and well-calibrated probabilistic forecasts, as the metric jointly rewards correct central tendency and uncertainty estimation.
B.2. Experiment Details In this section, we briefly overview the methodology of each experiments including dataset details, preprocessing and neural network architecture utilized to carry the experiments. The Table 3 presents the neural networks architecture utilized during the training/testing.
B.1.4. E XPECTED C ALIBRATION E RROR (ECE) ECE is commonly used as probabilistic classification metric but we adopt it for continuous regression prediction. For each prediction (µi , σi ) and true target yi , we compute the standardized absolute error zi = |yi − µi |/σi and map it to
B.2.1. M ULTIVARIATE R EGRESSION Dataset Explanation: The Boston Housing dataset comprises 506 samples of residential areas in the Boston suburbs, 12
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning Table 3. Summary of Key Hyperparameters of All Experiments
with each sample described by 13 features, including crime rate, average number of rooms per dwelling, and accessibility to highways. The target variable is the median value of owner-occupied homes. The Kin8nm dataset contains 8,192 samples, where each instance consists of 8 continuous input features derived from physical measurements and a single continuous target variable. It is a synthetic benchmark dataset designed to evaluate regression models on smooth but nonlinear input-output relationships. Preprocess: Prior to training both dataset were normalized using MinMaxScaler function. Neural Network: A single NSAC layer with a model dimension of 64, 16 heads, Top-K of 8, and sparsity of 0.5 was used for both tests, with all baselines configured with the same specifications.
Param.
Exp. 1 to 4
Udacity
CarRacing
Conv-layers NSAC/NAC Sparsity OOD params. (µpert ,σpert ) Dense MC-steps Opt. LR Batch Epochs
– 64-d, 16h 0.5 (0,5.0) – 5 AdamW 0.001 64 100
5×2D(24–64@5–3, ELU) 100-d, 20h 0.5 (0,5.0) – 5 AdamW 0.0001 40 10
3×TD-2D(32–640@3–5) 64-d, 16h 0.5 (0,5.0) – 5 AdamW 0.0001 32 20
Note: TD = TimeDistributed; 1D/2D = Conv1D/2D; d = model dimension; h = attention heads. Baseline clarification: For Deep-Ensemble we use ×3 models and for MC-Dropout, we placed dropout layer as penultimate layer with drop rate=0.1 with MC steps same as NSAC.
was used for both tests, with all baselines configured with the same specifications.
B.2.2. M ULTIVARIATE L ONG - RANGE FORECASTING Dataset Explanation: We utilized two well-known longrange modeling datasets: (i) ETTm1; and (ii) Jena-Climate. ETTm1 consists of 7 features and is used for long-term oil temperature forecasting. It provides data at 15-minute intervals, capturing real-world temporal patterns for timeseries analysis. Jena-Climate contains 14 features capturing various atmospheric and environmental measurements. It provides data at 10-minute intervals for modeling and forecasting climate-related time series. Preprocess: Prior to training both dataset were normalized using MinMaxScaler function. Neural Network: A single NSAC layer with a model dimension of 64, 16 heads, a top-k of 8, and a sparsity of 0.5 was used for both tests, with all baselines configured with the same specifications.
B.2.4. AUTONOMOUS V EHICLE Data Curation: We collected training data for both the Udacity and OpenAI-CarRacing environments by manually playing each game for 30 minutes, recording images and the corresponding steering values. Preprocessing: For Udacity, we applied the preprocessing steps proposed in (Shibuya, 2017), while for OpenAICarRacing, we did not apply any preprocessing. Neural Network: Both neural networks are trained endto-end and consist of a series of convolutional heads that extract spatial features, which are then passed to NSAC or baseline models for temporal modeling. Detailed network architecture and all hyperparameters are provided in Table 3.
C. Additional Experiments B.2.3. I NDUSTRY 4.0
In this section, we provide additional experiments including: (i) Ablation study of key-hyperparameters; (ii) Closed-loop trajectory analysis of AVs; and (iii) Scalability & Efficiency analysis.
Dataset Explanation: We utilized three well-known bearing degradation datasets to conduct our experiments: XJTUSY, PRONOSTIA, and HUST. Each dataset contains 1D vibrational signals captured from at least two sensors under three distinct operating conditions. To train the neural network, we used the first three bearings from operating condition 1 of XJTU-SY and evaluated the model in zeroshot manner on PRONOSTIA and HUST, also using first operating condition. This setup allows us to assess OOD calibration and zero-shot testing performance of NSAC. Preprocessing: The raw vibrational signals needs to be converted into a meaningful set of features to capture degradation patterns. We followed the feature engineering approach proposed in (Razzaq and Zhao, 2025a), which employs a wavelet transform to extract seven physically interpretable features, including energy, dominant frequency, skewness, and others. Neural Network: A single NSAC layer with a model dimension of 64, 16 heads, a top-k of 8, and a sparsity of 0.5
C.1. Ablation Analysis on Key Parameters In this experiment, we analyze the effects of key NSAC hyperparameters: (i) the number of attention heads (1–32); (ii) sparsity (0.1–0.9); (iii) Top-K selection (2–32) across sequence lengths (100–50000); (iv) MC samples (1–20); effects of (v) µpert (0–5); and (vi) σpert (0–15). We conduct this evaluation on a modified irregular spiral dataset while fixing the dmodel to 64. Figure 7 presents the quantitative results of this study. For the first two analyses (Figure 7(A,B)), we report both MSE and MAE under the interpolation and extrapolation regimes. Both the number of heads and sparsity exhibit non-monotonic behavior. Increasing the number of heads improves performance, with 32 heads achieving the lowest 13
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
indicating that it most consistently remains aligned with the center of the road. Figure 8(B) presents the sequence of actions produced by each model. Notably, the ground-truth trajectory itself exhibits distortions, which propagate into the learned baselines. In contrast, the NSAC, implemented with the penultimate NSAC layer, exhibits a self-regulating behavior that filters out such distortions and maintains close adherence to the road structure. Although SDE-Net is able to complete the lap, its predicted actions are severely distorted, indicating instability in sequential control. In Figure 8(C), we further analyze saliency maps using Grad-CAM (Selvaraju et al., 2017) to identify the regions of the input that influence action decisions. The NSAC consistently focuses on the full road curvature when actions are selected, suggesting that causally meaningful attention is aligned with the underlying driving structure. In contrast, other baselines exhibit less consistent and often scattered attention patterns across irrelevant regions of the environment. SDE-Net in particular appears to exhibit weak and unreliable visual grounding during decision-making, which explains its degraded action quality.
Figure 7. Visualization of impact of key hyperparameters of NSAC. (A) No. of attention heads; (B) Sparisity level; (C) Sequence length vs. Top-K selection; (D) No. of MC-Samples; (E) OOD mean (µpert ); and (F) OOD std (σpert ).
We additionally evaluate robustness to environmental perturbations by modifying the color scheme of the simulation environment while preserving the underlying road trajectory. This introduces visual noise without altering the control task. We conduct 10 runs with 10 trials each and report the mean success rate and standard deviation. All the models perform reasonably well under this setting. NSAC achieves an average success rate of approximately 80%, which is comparable to that of GMLE and SDE-Net, while all the remaining baselines exceed a 60% success rate.
MSE and MAE across both regimes. Similarly, although 90% sparsity achieves the lowest MSE, excessive sparsity may reduce representational capacity. Sparsity levels between 0.2 and 0.7 yield nearly identical performance, with 0.5 providing the most balanced trade-off. Figure 7(C) shows the effect of Top-K selection across varying sequence lengths under the interpolation regime. Among all the settings, K=8 consistently achieves the lowest error while maintaining a moderate computational cost. For the last three experiments, we report the CRPS. Increasing the number of MC samples improves performance, as shown in Figure 7(D), although the gains diminish and performance saturates beyond 15 samples. Both µpert (Figure 7(E)) and σpert (Figure 7(F)) exhibit non-monotonic effects, with the best performance achieved at µpert = 4 and σpert = 5. However, these settings require careful tuning. Therefore, we choose µpert = 0 and σpert = 5 as the default configuration and use them for all experiments.
Table 4. Run-Time & Memory Results Model
Run-Time (s)
Throughput (seq/s)
Peak Memory (MB)
Params.
GMLE DE MCD DER
14.47±0.21 14.02±0.18 14.09±0.23 14.21±0.23
0.07 0.07 0.07 0.07
75.78 75.98 75.89 76.00
36218 36218 36218 36416
NSAC
14.14±0.46
0.07
75.94
37278
C.3. Scalability & Efficiency Analysis C.2. Closed-loop Analysis of AVs
We evaluate the efficiency and scalability of NSAC’s solverfree forward-pass update on configurations used for spiral with synthetic inputs (batch size 1, sequence length of 1000). Experiments are conducted on Google Colab T4-GPU and compared against all post-hoc methods using solver-free NAC as the penultimate layer for fair comparison. We report the mean and standard deviation over five forward passes for runtime, throughput, and parameter count (see Table 4). NSAC achieves a mean run-time comparable to the
Beyond the predictive metrics, we further conduct a closedloop evaluation on the OpenAI CarRacing environment using a controlled test lap configuration, as shown in Figure 8(A). This setting enables a qualitative analysis of steering behavior under sequential decision-making. The figure illustrates the trajectories generated by each model, along with the corresponding position error relative to the road centerline. The NSAC achieves the lowest positional error, 14
Neuronal Stochastic Attention Circuit (NSAC) for Probabilistic Representation Learning
Figure 8. Closed-loop driving analysis on OpenAI-CarRacing. (A) Test drive map; (B) Trajectory followed by positional error to the center of the road; (C) Individual model actions at each step on the map; (D) Visual saliency maps in the highlighted region; (E) Noise test.
strongest baseline (14.13 s) with a throughput of 0.07 and 75.94 MB peak memory. It incurs slightly higher parameter count than baselines (37288 vs. 36218) due to an additional backbone head.
15