Learning Faster Physics with AI

Imagine trying to optimize a quantum memory.

You adjust the power of a laser, slightly change the timing of a control pulse, and ask:

What will happen?

Traditionally, the answer comes from solving a large system of coupled differential equations describing how light propagates through a cloud of atoms.

These equations are well understood and extremely accurate.

They're also computationally expensive.

Running a single simulation can take seconds or minutes depending on the complexity of the system. That may not sound like much until an optimization algorithm needs to evaluate tens of thousands of possible parameter combinations.

The physics itself becomes the limiting factor.

As AI begins playing a larger role in scientific research, this creates an interesting paradox:

We want AI to accelerate science, but AI must wait for physics simulations before making every decision.

Many scientific workflows follow essentially the same loop:

Choose parameters
  -> Predict system behavior
  -> Evaluate performance
  -> Update parameters
  -> Repeat

If prediction is slow, optimization is slow.

If prediction becomes nearly instantaneous, entirely new possibilities emerge:

• AI-assisted experiment optimization
• Real-time feedback during experiments
• Digital twins of laboratory systems
• Reinforcement learning for scientific discovery
• Autonomous laboratories

Instead of spending most of the time waiting for simulations, researchers can spend more time exploring ideas.

At first glance, this sounds like a standard machine learning problem.

Collect many simulations.

Train a neural network.

Predict future simulations.

Unfortunately, physical systems are very different from images or text.

In the quantum systems studied here:

• Light propagates continuously through space.
• Atoms interact with that light.
• External control fields modify those interactions.
• Everything evolves together over time.

Every point in space influences nearby regions, and those interactions accumulate as light travels through the medium.

Large transformer models are excellent at capturing long-range dependencies, but applying standard self-attention everywhere becomes computationally expensive.

The architecture that works well for language isn't necessarily the best architecture for physics.

One observation motivated this work:

Most physical interactions are local.

A point inside a physical system is usually influenced most strongly by its nearby surroundings.

Instead of allowing every spatial location to attend equally to every other location, I designed a transformer architecture that performs attention primarily within local spatial regions while still allowing information to flow across the entire system.

This architecture, called Regional Attention, dramatically reduces computational cost while preserving the important physical interactions needed for accurate prediction.

Rather than adapting an existing language model directly to physics, the architecture is designed around the structure of physical systems themselves.

To evaluate this idea, I trained the model on two very different quantum systems.

A single quantum bit

The first task involved learning the dynamics of a driven two-level quantum system.

Although relatively simple, it provides an excellent benchmark for evaluating whether the model can accurately learn quantum evolution.

A realistic quantum memory

The second task is considerably more challenging.

A pulse of light enters a cloud of atoms, slows down dramatically through a phenomenon called Electromagnetically Induced Transparency (EIT), and can even be temporarily stored before being retrieved later.

Modeling this process requires simultaneously capturing:

• light propagation
• atomic dynamics
• external control pulses
• spatial structure

Instead of simply memorizing examples, the model learns to approximate the underlying physical dynamics governing the system.

After training, the model successfully predicted the evolution of quantum systems under experimental conditions it had never seen before.

It generalized beyond its training data and accurately reproduced system dynamics for new control sequences and parameter settings.

Most importantly, it accomplished this while requiring only a single forward pass through a neural network instead of repeatedly solving complicated differential equations.

Accuracy was important.

Speed was transformative.

Compared with conventional numerical solvers, the learned surrogate achieved inference speeds up to three orders of magnitude faster while maintaining high prediction fidelity.

This fundamentally changes how these simulations can be used.

Instead of waiting for a simulation every time an optimization algorithm asks a question, the AI model can provide answers almost immediately.

That opens the door to optimization strategies that were previously computationally impractical.

Although this work focuses on quantum optics, I think the broader idea extends much further.

Many scientific disciplines rely on repeatedly solving expensive simulations:

• quantum technologies
• photonics
• materials science
• chemistry
• fluid dynamics
• biological systems

As machine learning models become better at approximating physical processes, learned surrogate models may become a standard component of scientific computing.

Rather than replacing physics, they make physics interactive.

This project changed how I think about scientific computation.

For decades, optimization has typically looked like this:

Optimization algorithm
  -> Physics simulator
  -> Optimization algorithm
  -> Physics simulator

I increasingly believe a different workflow will become common:

Physics
  -> Train a surrogate model
  -> Optimization happens inside the learned model

The expensive computation happens once during training.

After that, prediction becomes almost instantaneous.

For me, this represents one small step toward a much larger vision:

Building autonomous laboratories that can reason about experiments, predict outcomes, and continuously improve scientific workflows.

• Designed a transformer architecture tailored for spatially structured physical systems.
• Introduced regional attention to efficiently model local physical interactions.
• Learned quantum dynamics directly from simulation data.
• Demonstrated strong generalization beyond the training distribution.
• Achieved up to 1,000x faster inference compared with traditional numerical simulation.

Learning Spatially Structured Open Quantum Dynamics with Regional-Attention Transformers

Authors: Dounan Du, Eden Figueroa

Paper on arXiv