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

Mushtaq Ali; Muzamil Tariq; Niaz Ali Khan

Source-linked AI summary

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

Mushtaq Ali, Muzamil Tariq, Niaz Ali Khan

arXiv:2608.15112v1Aug 15, 2026cs.LGmath-phquant-ph

TL;DR

Traditional TDSE solvers are computationally intensive, while Transformer approaches lack rigorous probability conservation; this paper introduces a hard-constrained architecture that preserves probability and offers computationally efficient quantum-state prediction.

Problem
TDSE approximations must maintain wavefunction normalization, but complex potentials make analytical solutions intractable and existing numerical methods impose computational or stability costs.
Method
The Probability-Preserving Transformer embeds hard normalization into Transformer architectures, using multi-token temporal representations to model phase dynamics while conserving probability.
Results
PPT strictly enforces probability conservation while requiring significantly fewer computational resources than the Crank–Nicolson solver for learned quantum-state prediction.
Takeaways & Limitations
PPT is supported as a physically constrained surrogate for efficient quantum-state prediction, rather than a replacement for the high-fidelity solver generating reference data.
Takeaways & Limitations
Scaling to more complex systems requires Sobolev training to address derivative discrepancies relevant to momentum and kinetic-energy calculations.

Abstract

from arXiv · show

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.

I. INTRODUCTION · II. NUMERICAL METHODS

The paper motivates a probability-preserving Transformer for TDSE simulation by treating probability conservation as an architectural requirement. It numerically frames the problem using a one-dimensional infinite potential well benchmark and its sequential wavefunction dynamics.

I. INTRODUCTION: The TDSE governs quantum-state evolution, while |ψ(x, t)|^2 gives the probability density of finding a particle at position x and time t.It underpins analyses ranging from cavity quantum electrodynamics and ultrafast phenomena to tunneling and electron transport.
I. INTRODUCTION: Numerical solvers and machine-learning approximations must enforce wavefunction normalization throughout temporal evolution to remain physically consistent.Numerical discretization is essential when analytical solutions are intractable for complex potentials.
I. INTRODUCTION: The proposed probability-preserving Transformer embeds probability conservation as an intrinsic architectural constraint rather than a soft penalty.It explicitly normalizes the predicted complex wavefunction at every forward pass.
I. INTRODUCTION: A multi-token temporal representation is designed to capture phase transformations while retaining exact probability conservation and self-attention expressivity.The architecture targets quantum dynamics without sacrificing the expressive power of self-attention.
II. NUMERICAL METHODS: The numerical framework begins with a 1D infinite potential well as an ideal benchmark for evaluating the core TDSE physics.The benchmark provides a setting for the proposed approach before later implementation and performance analyses.
II. NUMERICAL METHODS: Because each quantum state determines its subsequent evolution, the TDSE naturally forms a sequential process suited to Transformer self-attention over temporal dependencies.The numerical-methods setup connects TDSE evolution with the modeling structure of Transformer architectures.
II. NUMERICAL METHODS: The benchmark adopts natural units ℏ = 1, m = 1, and L = 1, with Dirichlet boundaries ψ(0, t) = 0 and ψ(1, t) = 0.The initial state is chosen as the ground state, as specified in the numerical setup.

A. Numerical Data Generation and Preprocessing

The dataset uses Crank–Nicolson reference solutions verified for unit total probability, with uniform spatial-temporal discretization and ground-state initialization. Complex wavefunctions are converted to real-valued features, reorganized for full-wavefunction processing, and normalized through a hard probability constraint.

Reference data generation: Crank–Nicolson generates high-fidelity reference solutions whose total probability strictly integrates to unity.This ensures training data are physically consistent.
Discretization and initialization: The spatial domain uses 100 uniformly spaced grid points, while the temporal domain uses 200 steps with constant increment ∆t = 0.001.The system is initialized with the theoretical ground-state solution and homogeneous Dirichlet boundary conditions.
Numerical propagation: At each time step, a tridiagonal matrix algorithm solves the implicit linear system for the next wavefunction.The matrices are determined by the spatial discretization and simulation parameters.
Wavefunction representation: The complex wavefunction is separated into real and imaginary parts, yielding four-feature real-valued entries that are recombined at inference.This representation accommodates real-valued Transformer architectures while retaining the underlying quantum state.
Probability-preserving preprocessing: Full-wavefunction inputs enable global probability evaluation, while a hard normalization layer rescales predictions to enforce unit normalization.Normalization is applied after decomposing the complex wavefunction into real and imaginary parts.

1. Single-token baseline

The single-token baseline uses one time scalar to predict the full wavefunction, testing whether conventional neural architectures preserve probability without explicit constraints. For 100 spatial grid points and 200 time steps, each target contains 200 real and imaginary components.

Single-token formulation: Each input is a single time scalar, providing a baseline for testing whether conventional neural architectures inherently preserve probability without explicit constraints.The raw point-wise data are transformed to reconstruct the full wavefunction after sorting spatial coordinates.
Data representation: Each time step maps one scalar input to a 200-value target containing 100 real and 100 imaginary wavefunction components.The dataset covers 200 time steps, with outputs corresponding to 100 spatial grid points.
Data representation: The dataset is represented as Xs ∈R200×1 and Ys ∈R200×200, with rows corresponding to 200 discrete Crank–Nicolson time steps.Each row of Xs contains one temporal token, while each row of Ys contains the complete wavefunction at that time.

2. Multi-token baseline

The multi-token baseline groups consecutive time steps into sliding windows so Transformer self-attention can learn temporal dependencies and phase dynamics. It predicts the next wavefunction state in one forward pass using 195-step input and target matrices.

Windowed temporal representation: Five consecutive time steps are grouped into sliding windows, enabling self-attention to learn temporal dependencies and phase dynamics.Dataset engineering is required because initial inputs are unit-length sequences that cannot inherently model temporal relationships.
Windowed temporal representation: The multi-token format lets self-attention capture dependencies across previous states, while the Transformer predicts the future state in a single forward pass.This contrasts with numerical solvers that integrate the TDSE step-by-step.
Data dimensions: The representation uses Xm as a 195 steps ×5 scalar input matrix and Ym as a 195 steps ×200 wavefunction-values target matrix.These dimensions define the multi-token data representation.

III. NUMERICAL FRAMEWORK AND IMPLEMENTATION

The numerical framework evaluates the proposed PPT through six model configurations, beginning with an MLP baseline that learns quantum dynamics but violates probability conservation. A hard-normalization layer then enforces the quantum normalization condition automatically, preserving probability to numerical precision without penalty-based hyperparameters.

Baseline evaluation: The evaluation progresses through six model configurations, beginning with an MLP baseline to test neural-network learnability and inherent probability conservation.The MLP assesses whether standard architectures can learn the data while preserving the physical constraint.
Baseline evaluation: The MLP reproduces general quantum dynamics but predicts total probability values from 1.039418 to 1.045530 instead of the required P = 1.The increasing values across evaluated samples show that minimizing prediction loss alone does not guarantee probability conservation.
Training dynamics: The MLP loss decreases from approximately 0.5 to below 0.05 within the initial 50 epochs, demonstrating rapid capture of dominant wavefunction features.Training is reported over 1,000 epochs, after which the network reaches a robust converged state for inference.
Hard normalization: The hard-normalization layer computes the predicted wavefunction norm over the spatial grid and divides its real and imaginary components by a normalization factor.The implementation adds ε = 10^-12 for numerical stability before rescaling and reconstructing the wavefunction.
Hard normalization: Because α^2 ≈ P, the normalization condition is automatically satisfied for every prediction, preserving probability to numerical precision without penalty-based hyperparameters.This hard constraint makes conservation independent of raw network outputs and avoids additional loss penalties for enforcing probability.

A. Multi-token PPT Architecture

The Multi-token PPT maps sequences of temporal inputs to complete quantum states through a four-stage Transformer architecture adapted for temporal evolution in physical systems.

The Multi-token PPT processes a sequence of temporal inputs and maps them to a complete quantum state.
The architecture comprises temporal embedding, positional encoding, Transformer-based temporal feature extraction, and constrained state reconstruction.
Its sequence-processing framework follows Transformer models while adapting them to temporal evolution in physical systems.

1. Input processing and temporal embedding · 2. Transformer encoder

The model encodes five consecutive time values into 64-dimensional temporal representations, augments them with learnable positional embeddings, and processes them through two stacked Transformer encoder layers. Multi-head self-attention models relationships among temporal tokens to extract information from preceding observations.

1. Input processing and temporal embedding: The input consists of five consecutive time values, [ti−4, ti−3, ti−2, ti−1, ti].
1. Input processing and temporal embedding: These five scalar values form an input tensor with shape (batch size, 5).
1. Input processing and temporal embedding: Each time value is independently projected into a dmodel = 64-dimensional representation using a linear embedding layer.
1. Input processing and temporal embedding: Learnable positional embeddings are added to preserve temporal-token ordering and distinguish relative observation positions.
2. Transformer encoder: The embedded sequence passes through two stacked Transformer encoder layers.
2. Transformer encoder: Each encoder layer uses multi-head self-attention to model relationships among the five temporal tokens.
2. Transformer encoder: Self-attention enables the network to extract temporal information from the preceding sequence.
2. Transformer encoder: The temporal-dependency modeling approach is consistent with Transformer-based methods for physical systems and partial differential equations.

3. Quantum state reconstruction · 4. Hard probability-preserving constraint

The architecture reconstructs a 200-dimensional quantum state from temporally encoded features, with separate real and imaginary wavefunction components. A hard normalization layer then enforces probability conservation directly within the network while preserving the state representation and combining this constraint with self-attention.

3. Quantum state reconstruction: The final input-token representation summarizes the preceding temporal sequence after temporal feature extraction.This reduces the encoded representation to (batch size, 64).
3. Quantum state reconstruction: A fully connected output layer projects the latent representation into a 200-dimensional state vector.The resulting vector contains 100 real-part and 100 imaginary-part components of the predicted wavefunction.
3. Quantum state reconstruction: The 200-dimensional output represents the real and imaginary components of the predicted wavefunction across the spatial grid.These components are organized as [ψreal, ψimag].
4. Hard probability-preserving constraint: The unconstrained output is passed through a hard probability-preserving normalization layer.This layer is applied after state reconstruction.
4. Hard probability-preserving constraint: Probability conservation is incorporated directly into the network architecture rather than imposed through an additional loss-function penalty.The approach explicitly embeds the physical constraint in the model formulation.
4. Hard probability-preserving constraint: The normalization layer transforms the complete predicted quantum state into a normalized wavefunction satisfying the required probability condition.The final output retains the same 200-dimensional real-and-imaginary representation at all spatial grid points.
4. Hard probability-preserving constraint: The proposed architecture combines temporal sequence modeling through self-attention with explicit enforcement of the probability constraint.The constraint is applied to the complete predicted state while preserving its output dimensionality.

B. Numerical Validation

The multi-token PPT strictly preserves probability normalization while accurately approximating wavefunction evolution relative to Crank–Nicolson reference solutions. Its residual errors reflect surrogate approximation, while the architecture offers computationally efficient, physically constrained prediction.

The multi-token PPT strictly maintains quantum probability normalization while providing an accurate learned approximation of wavefunction evolution.This result supports incorporating physical constraints directly into the architecture rather than relying exclusively on optimization objectives.
5.8011 × 10−4 MSE and 2.4086 × 10−2 RMSE quantify the PPT’s average deviation from Crank–Nicolson reference wavefunctions.The maximum absolute error is 7.4499 × 10−2 across evaluated wavefunction components.
0.999999821186 minimum probability and 1.000000238419 maximum probability surround a mean probability of 1.The standard deviation is 1.0324×10−7, indicating tight concentration around unity.
7.1230 × 10−9 CNM MSE and 5.48 × 10−14 probability deviation establish CNM as an accurate reference, while PPT reduces computational demands and enforces conservation.CNM and PPT are therefore complementary in numerical accuracy and computational role.
The PPT reproduces complex phase rotation, infinite-well boundary vanishing, and the π/2 phase relationship between wavefunction components.A residual 2% to 3% error motivates future Sobolev training to regularize derivatives for energy and momentum calculations.

IV. COMPUTATIONAL PERFORMANCE

The proposed PPT achieves computational efficiency by directly predicting future wavefunctions in one forward pass, avoiding the sequential time stepping required by the CNM. Its single-state inference cost remains nearly constant as prediction horizons increase, while CNM runtime grows approximately linearly.

Inference strategy: The CNM computes every intermediate time step sequentially, whereas the trained PPT maps historical inputs directly to future states in one forward pass.This distinction eliminates iterative time stepping for PPT inference.
Benchmark protocol: The evaluation measured Sequential CNM Runtime, PPT Batched throughput, and PPT Constant-Time Inference over R = 1000 independent runs.Both methods used identical spatial resolution and computational conditions to reduce hardware-related timing variation.
Computational scaling: The PPT’s relative computational advantage becomes more significant for longer temporal predictions because its inference cost remains fixed while CNM cost scales with time.The speedup factor satisfies S > 1 when Transformer efficiency is superior.
Computational scaling: As prediction horizons extend, CNM execution time increases approximately linearly because the solver advances through required intermediate steps.The CNM propagates iteratively from the base condition, producing horizon-dependent computational cost.
Computational scaling: 0.0006 s: PPT single-state inference remains nearly constant across the tested horizons.The trained model predicts the future wavefunction directly from the temporal input sequence, bypassing sequential propagation.

V. CONCLUSION · DATA AVAILABILITY

The study demonstrates a Probability-Preserving Transformer for the 1D TDSE that enforces probability conservation through hard normalization. It concludes that more complex systems require Sobolev training, while the source code, trained models, and datasets are publicly available.

V. CONCLUSION: The Probability-Preserving Transformer solves the 1D time-dependent Schrödinger equation while enforcing probability conservation.The approach uses a hard-normalization layer.
V. CONCLUSION: A hard-normalization layer embeds probability conservation directly into the Transformer architecture.This design is presented as a mechanism for enforcing the physical constraint.
V. CONCLUSION: Embedding physical constraints into Transformer architectures yields physically consistent solutions.This conclusion follows from the study’s reported results.
V. CONCLUSION: Scaling to more complex systems requires moving beyond simple normalization to Sobolev training.The passage identifies Sobolev training as the needed next step for greater system complexity.
V. CONCLUSION: Sobolev training penalizes derivative discrepancies that are essential for extending the approach.The supplied passage truncates immediately after this statement.
DATA AVAILABILITY: The source code, trained models, and datasets used in the study are publicly available in the Repository.The passage identifies these materials as available for access.