A Transformer Learns Wavefunction Evolution: Probability Normalization Built Into the Output, 22× Faster Than Crank–Nicolson at Long Horizons
Synopsis
Wavefunction evolution in a one-dimensional infinite square well is now predicted directly by a transformer, with probability normalization imposed as a hard architectural constraint on the output rather than a soft loss penalty. An unconstrained multilayer perceptron left total probability between 1.039418 and 1.045530, away from unity; the constrained model held the maximum probability deviation to 2.98×10⁻⁷ across ground, first-excited, Gaussian wavepacket and superposition initial states, with mean fidelity 0.999949. At matched spatial resolution, one forward pass takes roughly 0.0006 to 0.001 s, reaching a 22.35× speedup over sequential Crank–Nicolson at a 195-step horizon.
Interpretation
A short history of a quantum wavefunction now maps directly to the complex wavefunction at the next time step, and the prediction satisfies unit total probability by construction. An unconstrained multilayer perceptron could fit wavefunction evolution but produced total probabilities between 1.039418 and 1.045530, showing that minimizing prediction error alone does not enforce normalization. On the standard dataset the multi-token probability-preserving transformer held maximum total-probability deviation near 1.79×10⁻⁷, with mean squared error 3.46796×10⁻⁴ and relative L2 error 2.64688×10⁻².
The same architecture was applied to four initial-condition families, ground states, first-excited states, Gaussian wavepackets and superposition states, and kept the probability constraint on unseen trajectories. Training and test sets were split at the trajectory level, so temporal samples from one trajectory never appear on both sides, making the evaluation trajectory-level generalization rather than interpolation within a trajectory. Across 60 Crank–Nicolson reference trajectories and 1.2 million space-time samples, held-out trajectories gave mean squared error 2.5284×10⁻⁵, relative L2 error 7.1469×10⁻³, minimum fidelity 0.998270 and maximum probability deviation 2.98×10⁻⁷.
The time needed for one prediction does not grow with the prediction horizon, while the sequential reference solver's runtime rises linearly with it. Crank–Nicolson must advance step by step and compute intermediate states, with runtime growing from 0.000673 s at 10 steps to 0.020106 s at 195 steps. On the same CPU with 100 spatial points and time step 0.001, single-state inference stayed between 0.000609 and 0.001038 s, and the measured speedup reached 22.35× at 195 steps, averaged over 1000 independent runs.
Perspective
The demonstrated scope is a one-dimensional infinite square well in natural units, with 100 spatial points, time step 0.001 and 200 stored temporal states per trajectory, using Crank–Nicolson references. The intended readers are researchers and engineers working on quantum-dynamics numerics and learned surrogate models. Stated next directions are extended temporal rollouts, derivative-aware Sobolev training and extension to higher-dimensional quantum systems.
Prediction accuracy varies across initial-condition families, and one Gaussian wavepacket trajectory, gaussian 18, stands out with mean RMSE 0.016202 and maximum RMSE 0.029289, pulling the aggregate mean error upward. Maximum relative energy error is 2.4875×10⁻² and mean relative energy error is 1.1675×10⁻³, so energy stability is weaker than probability conservation. The spatiotemporal error heatmap shows vertical stripes, indicating errors depend mainly on spatial location rather than accumulating over time.