Source-linked AI summary
Physics-Constrained Neural Flow Maps for Long-Horizon Prediction of Spin Dynamics
Haoen Feng, Shenglan Yuan, Shirong Lin
TL;DR
Long-horizon prediction of current-driven magnetization is challenged by fine-step integration and preservation of the unit-sphere constraint. The paper introduces a physics-constrained neural flow map for finite-time, recursive propagation, achieving 0.00425 in-domain RMSE with norm drift at 10^-7 while outperforming an adapted LSTM in-domain. Its main remaining limitation is weaker dynamical extrapolation under unseen stronger drive.
Problem
Long-horizon magnetization surrogates must control accumulated radial drift and tangential phase error while respecting the unit-sphere state constraint.
Method
A parameter-conditioned neural flow map predicts future magnetization from the current state, torque strength, and time span using tangential updates retracted to the unit sphere.
Results
0.00425 in-domain RMSE with 10^-7 norm drift was achieved beyond the training horizon, while Flow was more accurate and geometrically stable in-domain than an adapted LSTM.
Takeaways & Limitations
The geometry-preserving propagator provides a physically admissible surrogate for long-horizon magnetization dynamics without repeated fine-step integration.
Takeaways & Limitations
Mean RMSE rises from 0.00425 in-domain to 0.0259 under unseen stronger drive, indicating limited dynamical extrapolation despite geometric preservation.
Abstract
from arXiv · showhide
Conventional simulation of current-driven magnetization relies on fine-step integration of the spin-transfer-torque Landau--Lifshitz--Gilbert equation, creating a computational bottleneck in parameter sweeps and control searches. In this work, we propose a physics-constrained neural flow map that learns finite-time dynamics directly on the unit sphere. The model maps the current magnetization, spin-torque strength, and requested time span to a future state in a single forward pass. Tangent-space projection and spherical retraction preserve unit magnetization during recursive, composition-consistent rollout. We validate the framework on single-spin trajectories under in-domain torques and previously unseen but stronger drive. Beyond the training horizon, it achieves an in-domain root mean square error of $0.00425$ with norm drift at the $10^{-7}$ level. The flow outperforms an adapted Long Short-Term Memory (LSTM) in in-domain accuracy and geometric stability, although the LSTM retains slightly lower out-of-distribution state error. The resulting geometry-preserving propagator reduces reliance on fine-step integration and enables physically admissible long-horizon prediction.
1. Introduction
The introduction motivates efficient learning of current-driven spin dynamics governed by the STT-extended LLG equation. It identifies unit-sphere preservation and long-horizon stability as key challenges and presents a parameter-conditioned, physics-constrained neural flow map as the proposed solution.
- Motivation: Current-driven magnetization reversal is governed by the Landau–Lifshitz–Gilbert equation extended with spin-transfer torque from spin-polarized currents.The extension models angular-momentum transfer to the free-layer magnetic moment.
- Related work: Learned propagators can amortize repeated integrations and reduce online computational cost, but long autoregressive rollouts and unseen parameters remain challenges.Prior approaches include latent-space, kernel-based, low-rank, and LSTM representations.
- Problem: The single-macrospin state lies on the unit sphere S2, so unconstrained three-dimensional neural outputs can violate the physical condition ∥m∥2 = 1.The introduction identifies norm preservation as a central constraint because small one-step errors can still affect stability.
- Contribution: The proposed flow map maps magnetization, normalized spin-transfer-torque strength, and requested time span directly to a future state.A residual decoder predicts a tangential increment, retracts it to the unit sphere, and uses multi-span training with composition consistency across temporal scales.
2. Physical Problem and Reference Dynamics
The reference dynamics model magnetization on the unit sphere using a macrospin spin-transfer-torque equation with field-induced precession, damping, and current-driven contributions. Reference trajectories use a four-stage frozen-effective-field propagator with normalization after each fine step, defining the discrete flow used for training targets and RMSE evaluation.
- Physical model: The state is the unit magnetization m = (mx, my, mz)T ∈ S2, with time expressed in nanoseconds.The continuous model preserves the unit sphere because its cross-product terms are orthogonal to m whenever ∥m∥2 = 1.
- Physical model: The governing equation combines field-induced precession and damping with damping-like and field-like current-driven spin-transfer torque.The torque strength aJH is dimensionless, normalized by the characteristic field H0; its magnitude measures relative torque strength and its sign encodes current polarity and action relative to n.
- Numerical reference: The implementation uses hx = 0.02, hz = 0.5, α = 0.02, and a fine reference step of Δt = 0.045 ns.Each neural flow-map call spans ten fine steps, advancing the state by 0.45 ns during closed-loop evaluation.
- Numerical reference: Reference trajectories use a four-stage explicit propagator that freezes heff at the start of each step and normalizes the state afterward.This differs from standard RK4, which recomputes the complete right-hand side at every intermediate stage.
- Numerical reference: Training data and online targets follow the same discrete propagation path, and reported RMSE values are defined relative to this reference flow.The pregenerated data and online targets coincide elementwise within each training batch.
3. Physics-Constrained Neural Flow Map
The model learns finite-time magnetization evolution from the current state, torque strength, and requested span, using span features to represent variable-duration maps. Tangent-space projection and spherical retraction enforce unit norm, while composition training supports recursive propagation without guaranteeing low state error.
- Input and span encoding: The input encodes current magnetization, standardized torque strength, and requested span through eight features, including d, d^2, sin(πd), and cos(πd).Training samples δ uniformly from 1 to 10, corresponding to Δt ∈ [0.045, 0.45] ns; closed-loop evaluation uses δ = 10.
- Network architecture: An 8 → 192 SiLU encoder, five width-192 residual blocks, and a three-dimensional output increment form a 374,787-parameter network.Each residual block contains LayerNorm, two linear transformations, and SiLU activations.
- Geometric constraint: The raw increment is projected onto the tangent plane at the normalized base point and retracted to the sphere, enforcing unit norm by construction.This removes radial drift without requiring the network to learn the unit-norm constraint or balance a norm penalty against state accuracy.
- Composition constraint: The composition term compares direct and decomposed predictions for sampled spans with δ ≥ 2, within an autonomous constant-torque slice.It is an empirical one-step-plus-remaining-steps relation, not a proof of a strict semigroup property under arbitrary time partitions.
- Training objective: Because the duplicate target is elementwise identical to the data target, Ldup ≡ Ldata and the so-called RK4 term acts as supervised reweighting rather than independent physical consistency.The main configuration uses total coefficients 2.0 on Ldata and 0.1 on Lcomp.
4. Training and Evaluation Protocol
The protocol partitions complete parameter trajectories into training, validation, test, and strictly separated stronger-drive OOD sets, then evaluates long-horizon closed-loop rollouts without post-boundary reference correction. Flow and LSTM models use matched splits and evaluation windows, with checkpoints selected solely by validation closed-loop RMSE across five independent runs.
- Data partitioning: Complete parameter trajectories are partitioned into training, validation, and test sets in a 10:1:1 ratio, preventing time-slice leakage across subsets.The in-domain grid spans aJH ∈[−0.08, −0.03], while the separate OOD set uses aJH ∈[−0.10, −0.08001].
- Model selection: Checkpoints are selected exclusively by full closed-loop RMSE on the future validation window, with test and OOD sets held out until evaluation.The main results aggregate five independent training runs sharing the data partition and protocol but differing in initialization and minibatch sampling.
- Closed-loop evaluation: 222 neural steps of ∆t = 0.45 ns extend evaluation from the 99.99 ns boundary to 199.89 ns without future ground truth or reference-integrator calls.In-domain testing stays within the observed drive-strength range, whereas the stronger-drive OOD setting adds one-sided parameter extrapolation.
- Closed-loop evaluation: Every future prediction is fed back as the next input, so evaluation uses closed-loop rollout rather than teacher forcing or sliding restarts.Each call uses the largest locally observed training span, 0.45 ns; extrapolation concerns successive compositions and cumulative absolute time.
- Baseline comparison: The LSTM and Flow models share parameter-trajectory splits, training windows, future evaluation windows, and five independent runs, but the LSTM warms up on the observed sequence before generating 222 autoregressive steps.Flow reads only the current state at the boundary, while both checkpoints are selected using validation closed-loop RMSE over 99.99–199.89 ns.
5. Model Validation and Generalization
The Flow model maintains physically admissible unit-norm states during long-horizon rollouts while achieving substantially lower in-domain error than LSTM. Stronger-drive and extended extrapolation increase dynamical error, but spherical retraction preserves geometric stability.
- Long-horizon validation: 0.00425±0.00126 in-domain RMSE rises to 0.0259 ± 0.00656 under stronger-drive evaluation, while every model completes 222 closed-loop steps without numerical divergence.The stronger drive makes dynamical extrapolation more difficult.
- Flow–LSTM comparison: 0.004247 ± 0.001258 RMSE for Flow versus 0.021673 ± 0.002875 for LSTM gives Flow an approximately 80.4% in-domain error reduction.Under stronger-drive OOD evaluation, Flow records 0.025859 ± 0.006555 RMSE versus 0.021570 ± 0.005417 for LSTM, so LSTM has approximately 16.6% lower state error.
- Geometric stability: 1.788 × 10−7 in-domain and 1.192×10−7 OOD are Flow’s mean maximum norm drifts, compared with LSTM values of (5.577±1.143)×10−2 and (2.391±1.606)×10−2.Spherical retraction removes radial error from Flow’s closed-loop propagation, whereas LSTM has no unit-sphere guarantee.
- Parameter extrapolation: 0.0772 representative-trajectory RMSE at the unseen aJH = −0.10 drive accompanies ∥bm∥2 −1 at the 10−7 level, separating dynamical accuracy from geometric stability.The rollout initially reproduces the decay trend but later converges to a shifted asymptotic direction in mx and my.
- Ablation study: 0.006599 in-domain and 0.030187 OOD RMSE for the full ablation model increase to 0.007005 and 0.031224 when composition consistency is removed.The ablation uses shared initialization and minibatch order, isolating the structural change.
- Ablation study: 5.065 × 10−3 in-domain maximum norm drift for the raw-output variant is approximately 2.83 × 104 times the full model’s value, while its OOD RMSE reaches 0.036993.The exponential map preserves unit norm but produces higher closed-loop RMSE than spherical retraction; single-run ablation differences should be interpreted cautiously.
6. Conclusion and Challenges
The study presents a geometry-preserving neural flow map for recursive, long-horizon magnetization prediction, with accurate in-domain dynamics and near-machine-precision norm preservation. Its main challenges are weaker dynamical extrapolation under stronger unseen drives and sensitivity to insufficient temporal coverage.
- Conclusion: The flow map directly predicts finite-time magnetization states on the unit sphere and applies the same propagator recursively for long-horizon prediction.It uses current magnetization, spin-torque strength, and requested time span without future reference states.
- Future work: Future work will target torque-parameter generalization, multiple initial states, time-dependent currents, and matched comparisons of alternative learned propagators.The study also recommends sampling near parameter boundaries and over long transients.
- Comparison: The Flow is more accurate in-domain than the adapted LSTM and preserves near-machine-precision norm constraints, whereas the LSTM has slightly lower RMSE under unseen, stronger drive.The LSTM aggregates observation history through its hidden state, while Flow carries no hidden history.
- Challenges: Tangent projection and spherical retraction nearly eliminate radial drift, while removing the geometric layer increases norm error by four orders of magnitude.Tangential error can still shift oscillation phase and asymptotic direction under unseen, stronger STT drive.
- Challenges: 0.0259 OOD mean RMSE follows 0.00425 in-domain RMSE when normalized torque strength leaves the training range, despite exact unit-sphere preservation.The representative trajectory converges to a shifted asymptotic direction in mx and my.
- Challenges: 0.0332 representative-trajectory RMSE results from extending Flow training horizon from approximately 50 to 100 ns, down from 0.2373.Multi-span supervision and composition consistency do not eliminate phase accumulation caused by insufficient temporal coverage.
Declaration of competing interest
The authors report no potential or pertinent competing interests associated with this work.
- No author disclosed any potential or pertinent conflicts associated with this work.