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
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 ( Pmin )
1.000000000000
0.999999821186
Maximum Total Probability ( Pmax )
1.000000000006
1.000000238419
Mean Total Probability ( Pˉ )
1.000000000000
1.000000000000
Probability Deviation / Std. Dev. ( σ )
5.4800×10−14
1.0324×10−7
Mean Squared Error (MSE)
7.1230×10−9
5.8011×10−4
Root Mean Squared Error (RMSE)
8.4400×10−5
2.4086×10−2
Table 2: Performance comparison between the classical CNM and the proposed multi-token PPT architecture.
Figure 3:: Comparison of the real ( Re(ψ) , top) and imaginary ( Im(ψ) , 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 ×
20
0.003159
0.001589
0.000636
4.96 ×
50
0.003158
0.002209
0.000561
5.63 ×
100
0.009180
0.003832
0.000589
15.58 ×
Table 3: Comparative computational performance of the classical CNM and the proposed multi-token PPT.
Figure 5:: Computational scaling and relative performance speedup of the proposed PPT framework versus the sequential Crank–Nicolson solver.
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.
Pengyu Lai, Yixiao Chen, Dewu Yang +3
School of Aeronautics and Astronautics, Shanghai Jiao Tong University, Shanghai, 200240, China
Transformer architectures have attracted increasing attention for solving partial differential equations (PDEs), owing to their flexibility in handling irregular discretizations and their ability to capture long-range physical dependencies. However, unlike discrete language tokens or fixed-resolution image patches, observed physical fields are finite samples of underlying infinite-dimensional functions. Consequently, effectively applying Transformers to PDEs requires a tokenizer that respects the functional nature of physical fields and constructs physically expressive tokens from arbitrary discretizations.To this end, we propose \methodname{Physics Transformer}, a function-projection-based Transformer architecture for physical field prediction. Physics Transformer treats a physical field as a continuous function and partitions its discretization into locality-preserving spatial patches. Within each patch, it dynamically learns a set of adaptive local basis functions and projects the sampled field onto these bases to obtain compact physics tokens. The resulting tokens capture diverse latent physical states while preserving fine-scale spatial structures, enabling efficient global interaction through factorized attention across space and physical states. The projected representation further supports efficient decoding at arbitrary query locations. Extensive experiments on diverse benchmarks, ranging from two-dimensional PDE dynamics to industrial-scale three-dimensional CFD simulations, demonstrate that Physics Transformer accurately captures fine-grained physical structures and achieves state-of-the-art predictive performance. These results establish function projection as a practical and effective foundation for designing Transformer architectures for PDE solving.
Guoze Sun, Rui Zhang, Jiankai Tang +4
Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China · School of Mechanics and Engineering Science, Peking University, Beijing, China
We introduce ECHO, a transformer-operator framework for generating million-point PDE trajectories. While existing neural operators (NOs) have shown promise for solving partial differential equations, they remain limited in practice due to poor scalability on dense grids, error accumulation during dynamic unrolling, and task-specific design. ECHO addresses these challenges through three key innovations. (i) It employs a hierarchical convolutional encode-decode architecture that achieves a 100 × spatio-temporal compression while preserving fidelity on mesh points. (ii) It incorporates a training and adaptation strategy that enables high-resolution PDE solution generation from sparse input grids. (iii) It adopts a generative modeling paradigm that learns complete trajectory segments, mitigating long-horizon error drift. The training strategy decouples representation learning from downstream task supervision, allowing the model to tackle multiple tasks such as trajectory generation, forward and inverse problems, and interpolation. The generative model further supports both conditional and unconditional generation. We demonstrate state-of-the-art performance on million-point simulations across diverse PDE systems featuring complex geometries, high-frequency dynamics, and long-term horizons.
Armand Kassaï Koupaï, Lise Le Boudec, Patrick Gallinari
Sorbonne Université, CNRS, ISIR, 75005 Paris, France · Criteo AI Lab, Paris, France