Probability-Preserving Transformer for the Time-Dependent Schrödinger Equation

Mushtaq Ali; Muzamil Tariq; Niaz Ali Khan

cs.LGAug 15, 2026

Probability-Preserving Transformer for the Time-Dependent Schrödinger Equation

Authors: Mushtaq Ali, Muzamil Tariq, Niaz Ali Khan

Organizations: Department of Physics, Division of Science and Technology, University of Education Lahore, Pakistan · Department of Physics, Xiamen University, Xiamen 361005, China

Abstract

Solving the time-dependent Schrödinger equation (TDSE) via traditional numerical methods is computationally intensive. Transformer models offer a compelling alternative, but standard implementations rely on soft constraints that cannot rigorously guarantee probability conservation. Here, we introduce a Transformer architecture that enforces TDSE probability conservation as a hard constraint. The design intrinsically ensures unitarity across temporal evolution without requiring repeated retraining. Our empirical findings show that this hard-constraint approach is not only physically exact but also computationally superior to conventional soft-constraint methods.

Figures & tables

Figure 1:: Architectural schematic of the proposed multi-token PPT. The model comprises input time embeddings, a multi-head self-attention encoder, a linear output projection, and a final hard-normalization layer designed to preserve unit probability.


Test	Total Probability ( $P$ )

Sample 0	1.039418
Sample 1	1.041417
Sample 2	1.043455
Sample 3	1.045530

Table 1: Violation of the probability conservation in the MLP Technique.

Figure 2:: Mean squared error (MSE) of the MLP as a function of training epochs, shown over $1,000$ iterations.

Performance Metric	CNM	Multi-token PPT
Minimum Total Probability ( $P_{\min}$ )	$1.000000000000$	$0.999999821186$
Maximum Total Probability ( $P_{\max}$ )	$1.000000000006$	$1.000000238419$
Mean Total Probability ( $\bar{P}$ )	$1.000000000000$	$1.000000000000$
Probability Deviation / Std. Dev. ( $\sigma$ )	$5.4800\times 10^{-14}$	$1.0324\times 10^{-7}$
Mean Squared Error (MSE)	$7.1230\times 10^{-9}$	$5.8011\times 10^{-4}$
Root Mean Squared Error (RMSE)	$8.4400\times 10^{-5}$	$2.4086\times 10^{-2}$

Table 2: Performance comparison between the classical CNM and the proposed multi-token PPT architecture.

Panel of Figure 3 — Figure 3:: Comparison of the real ( $\text{Re}(\psi)$ , top) and imaginary ( $\text{Im}(\psi)$ , bottom) components across the spatial domain.

Figure 4:: Pointwise wavefunction error across the spatial domain grid points.


Prediction Horizon ( $N$ )	Crank–Nicolson Sequential (s)	PPT Batched (s)	PPT Constant-Time (s)	Speedup

10	0.000612	0.000897	0.000579	1.06 $\times$
20	0.003159	0.001589	0.000636	4.96 $\times$
50	0.003158	0.002209	0.000561	5.63 $\times$
100	0.009180	0.003832	0.000589	15.58 $\times$

Table 3: Comparative computational performance of the classical CNM and the proposed multi-token PPT.

Panel of Figure 5 — Figure 5:: Computational scaling and relative performance speedup of the proposed PPT framework versus the sequential Crank–Nicolson solver.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 540 papers

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs
Mar 3, 2026Pengyu Lai, Yixiao Chen, Dewu Yang +3PDE Surrogate Modeling Neural Operators
School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China
Physics Transformer: Tailoring Transformer for General PDE Prediction
Jul 27, 2026Guoze Sun, Rui Zhang, Jiankai Tang +4Transformer PDE Surrogate Modeling
Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China · School of Mechanics and Engineering Science, Peking University, Beijing, China
Efficient Generative Transformer Operators For Million-Point PDEs
Dec 4, 2025Armand Kassaï Koupaï, Lise Le Boudec, Patrick GallinariPDE Surrogate Modeling Transformer Neural Operator
Sorbonne Université, CNRS, ISIR, 75005 Paris, France · Criteo AI Lab, Paris, France

Mar 3, 2026J/K move · Enter open · S save

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs

Pengyu Lai, Yixiao Chen, Dewu Yang, Rui Wang+2

School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China

PDE Surrogate Modeling Neural Operators Physics-Informed Neural Operators Transformer Attention Attention Mechanisms

Open paper AskPDF arXiv

Partial differential equations (PDEs) are fundamental for modeling complex physical systems, yet classical numerical solvers face prohibitive computational costs in high-dimensional and multi-scale regimes. While Transformer-based neural operators have emerged as powerful data-driven alternatives, they conventionally treat all discretized spatial points as uniform, independent tokens. This monolithic approach ignores the intrinsic scale separation of physical fields, applying computationally prohibitive global attention that redundantly mixes smooth large-scale dynamics with high-frequency fluctuations. Rethinking Transformers through the lens of complex dynamics, we propose DynFormer, a novel dynamics-informed neural operator. Rather than applying a uniform attention mechanism across all scales, DynFormer explicitly assigns specialized network modules to distinct physical scales. It leverages a Spectral Embedding to isolate low-frequency modes, enabling a Kronecker-structured attention mechanism to efficiently capture large-scale global interactions with reduced complexity. Concurrently, we introduce a Local-Global-Mixing transformation. This module utilizes nonlinear multiplicative frequency mixing to implicitly reconstruct the small-scale, fast-varying turbulent cascades that are slaved to the macroscopic state, without incurring the cost of global attention. Integrating these modules into a hybrid evolutionary architecture ensures robust long-term temporal stability. Extensive memory-aligned evaluations across four PDE benchmarks demonstrate that DynFormer achieves up to a 95% reduction in relative error compared to state-of-the-art baselines, while significantly reducing GPU memory consumption. Our results establish that embedding first-principles physical dynamics into Transformer architectures yields a highly scalable, theoretically grounded blueprint for PDE surrogate modeling.

Abstract

Transformer

Explore similar work

From Complex Dynamics to DynFormer: Rethinking Transformers for PDEs

Physics Transformer: Tailoring Transformer for General PDE Prediction

Efficient Generative Transformer Operators For Million-Point PDEs