Noise-Aware Training (NAT) Overview

• Noise-Aware Training (NAT) is a framework that explicitly models and compensates for stochastic noise during training using Neural-SDEs to enhance robustness.
• It constructs digital twins of physical devices and employs gradient-based optimization through both deterministic and stochastic components to reduce simulation–hardware gaps.
• Empirical results demonstrate that NAT-trained networks drastically reduce transfer discrepancies, achieving up to ~95% accuracy with significantly lower error margins.

Noise-Aware Training (NAT) is a broad methodological framework that systematically incorporates the explicit modeling, injection, estimation, or compensation of noise during the training of neural and hybrid dynamic systems. With roots in both statistical modeling and robust optimization, NAT methods are motivated by the need to develop models that generalize under mismatched, stochastic, or ill-characterized conditions, particularly when analog devices, physical reservoirs, or real-world data channels impose non-negligible stochasticity that impedes direct deployment of deterministic models. In neuromorphic device networks, NAT is instantiated via Neural Stochastic Differential Equations (Neural-SDEs) as differentiable digital twins of physical devices, enabling backpropagation through genuinely stochastic dynamics and closing the reality-simulation transfer gap (Manneschi et al., 2024).

1. Mathematical Formulation and Noise-Aware Modeling

In the context of neuromorphic dynamic device networks, each device ("node") is formulated as a continuous-time stochastic dynamical system governed by a stochastic differential equation,

$$d\mathbf{x}_t = \mathbf{f}(\mathbf{x}_t,\mathbf{u}_t;\theta)\,dt + \mathbf{g}(\mathbf{x}_t,\mathbf{u}_t;\theta)\,d\mathbf{W}_t,$$

where $\mathbf{x}_t$ encodes internal device states (e.g., magnetization vectors), $\mathbf{u}_t$ denotes the external driving/input signals (voltages, fields, or upstream activations), $\mathbf{f}$ and $\mathbf{g}$ are neural-network–parameterized drift and diffusion coefficients (collectively indexed by $\theta$), and $d\mathbf{W}_t$ is an increment of an $n$-dimensional Wiener process modeling physical noise sources.

Crucially, Neural-SDE surrogates capture both deterministic device response (via $\mathbf{f}$) and device-specific stochasticity (via $\mathbf{g}$). The framework accommodates colored (temporal) noise by introducing auxiliary variables $\mathbf{a}_t$ with their own time constants $\tau$, yielding coupled SDEs that embed both fast and slow noise modes. Delayed states can be embedded to represent memory effects critical for dynamical hardware platforms.

The loss for a temporal task (e.g., sequence classification, regression) is an expected value over both data samples and SDE noise realizations,

$$\mathcal{L}(\theta) = \mathbb{E}\left[\sum_t \|\hat{\mathbf{y}}_t - \mathbf{y}^*_t\|^2 \right],$$

with $\hat{\mathbf{y}}_t$ the output trajectory and $\mathbf{y}^*_t$ the target. The total risk is

$$R(\theta) = \mathbb{E}_{\text{data}}\; \mathbb{E}_{\text{noise}} [\mathcal{L}(\theta; \text{noise})],$$

where Monte Carlo samples provide unbiased stochastic gradient estimates.

2. Neural-SDE Surrogate Construction and Cascade Learning

NAT proceeds in two major stages: device-model (digital twin) training, and network-level connectivity optimization:

Stage 1: Digital Twin Fitting

• For each physical device class (e.g., spintronic oscillator, Duffing oscillator), a dataset of input–output trajectories under random excitations is collected.
• The Neural-SDE's parameterization $(\hat{\mathbf{f}},\hat{\mathbf{g}})$ is fit such that simulated trajectories match device statistics, including mean, variance, and autocovariance, capturing both average behavior and stochastic fluctuations.

Stage 2: Network Training

• An in-silico network is assembled using the trained neural-SDE surrogates.
• The network-level loss is differentiated via backpropagation-through-time (BPTT), traversing both the deterministic and stochastic computational graph.
• Only inter-device connectivity weights are optimized at this stage, dramatically reducing the need for repeated physical device measurement. This cascade decouples device modeling from network optimization and allows for scalable, gradient-based design.

3. Backpropagation and Gradient Estimation in Stochastic Dynamics

Gradient-based optimization in this stochastic, dynamical setting employs either:

• Pathwise gradient (reparameterization) estimators, which propagate fixed Wiener realizations through discretized SDE integrators (e.g., Euler–Maruyama),

$$\mathbf{x}_{n+1} = \mathbf{x}_n + f(\mathbf{x}_n,\mathbf{u}_n;\theta)\,\Delta t + g(\mathbf{x}_n,\mathbf{u}_n;\theta)\sqrt{\Delta t}\,\xi_n,$$

where $\xi_n \sim \mathcal{N}(0, I)$. The chain rule is applied through each discretization step for $\nabla_\theta \mathcal{L}$, accumulating eligibility traces for temporal credit assignment,

$$\frac{d}{d\theta}E(t) = \frac{\partial E}{\partial \mathbf{y}(t)} \left[ \frac{\partial \mathbf{y}}{\partial \theta} + \int_{0}^{t} \frac{\partial \mathbf{y}(t)}{\partial \mathbf{y}(s)}\frac{\partial \mathbf{y}(s)}{\partial \theta}ds \right].$$

• Adjoint/sensitivity-based approaches, which solve backward-in-time ODEs for gradients efficiently, as in [Kidger et al., 2021].

This framework supports both SDEs with Markovian and non-Markovian (long-memory) dynamics. For physical platforms with "short-memory" devices, truncating the BPTT window reduces memory cost.

4. Experimental Validation and Quantitative Assessment

The NAT methodology was validated on physical spintronic device networks benchmarked using the temporally-partial MNIST task: each digit is revealed over $T$ frames, each showing only a random subset of pixels. The network must aggregate this temporal evidence for final classification.

Key empirical findings:

• Conventional neural-ODE–trained networks, which ignore intrinsic hardware noise, achieve near-perfect accuracy in simulation, yet lose 30–50 points in accuracy ("transfer gap") on real hardware due to noise-induced errors.
• Simple augmentation with additive external Gaussian noise during ODE training reduces the gap but leaves 10–20 points discrepancy.
• NAT-trained (neural-SDE) networks close the transfer gap to under 2 points and achieve up to ~95% accuracy when only 20% of pixels are visible.
• Against a hardware-random reservoir baseline, NAT networks outperform by 20–30 points in accuracy.

These results establish that matching both deterministic and stochastic aspects of device dynamics in the digital twin is critical for successful model–hardware transfer.

5. Scalability, Applicability, and Implementation Guidelines

NAT is broadly applicable to any experimental platform where full input–output trajectories can be recorded, including spintronic, photonic, or mechanical oscillators. Once the digital twin is trained, gradients of task loss with respect to hardware-programmable parameters (e.g., interconnect weights) can be computed via standard autodiff frameworks.

Implementation recommendations include:

• Pretraining the deterministic drift (f) for stability before introducing the noise components or adversarial objectives.
• Using mini-batches of trajectory segments with random starts to address long-memory effects and de-correlate samples.
• Keeping surrogate neural net architectures (per digital twin) small to conserve adjoint memory and computation.
• Truncating BPTT windows where device memory is short.
• Selective noise injection: regularize $g(\cdot)$ so that noise is added only to the most recent state variables and not repeatedly through delay buffers.
• Contrastive/moment-matching loss terms on trajectory statistics (mean, variance, autocovariance) to prevent degeneracies in adversarial/discriminator training stages.

This best-practice regime ensures minimal training set requirements and guarantees stable, noise-matched learning.

6. Impact, Limitations, and Outlook

Noise-Aware Training with digital twin Neural-SDEs enables robust, gradient-driven design of dynamical hardware networks that are nearly free of simulation–hardware discrepancies. It is particularly crucial for platforms with device-intrinsic stochasticity, nontrivial memory, or limited analytical models.

Intrinsic limitations and open problems include:

• Computational cost of adjoint-based BPTT (though manageable for compact device surrogates).
• The challenge of capturing and fitting high-order colored noise when device noise statistics are complex.
• Extending the methodology to platforms exhibiting non-Gaussian, non-i.i.d., or heavy-tailed noise.

A plausible implication is that as experimental hardware advances and larger networks are physically realized, the practice of twin-based NAT may become the standard for hardware-software co-design in neuromorphic and non-von Neumann computing (Manneschi et al., 2024).

---

References

• "Noise-Aware Training of Neuromorphic Dynamic Device Networks" (Manneschi et al., 2024)