News

Teaching a neural operator to steer a molecule

Announcing “Inverse Design of Quantum Control Sequences with Fourier Neural Operators,” a collaboration across UCLA, Caltech, and NVIDIA.

Today, our group at UCLA’s NarangLab is announcing new work that uses a Fourier Neural Operator to design the control sequences quantum applications need to run. With researchers at Caltech and NVIDIA, the approach learns the quantum dynamics of a molecule and uses that knowledge to identify control protocols in state spaces too large for conventional optimization. The work is being highlighted at IEEE Quantum Week in Toronto as part of NVIDIA’s announcement of an open, programmable platform for designing and testing fault-tolerant quantum computing applications. Here’s what we did and why it matters.

Large language models (LLMs) have a blind spot that the field has been pointing out for a while: they have limited understanding of the physical world. Nowhere is that gap more relevant than in the quantum realm, where the complexity of the dynamics, and the cost of steering them, grows exponentially with system size. We can’t prompt our way to a control pulse sequence. We need a model that has actually learned the underlying physics.

That is the core of our new paper. We show that a Fourier Neural Operator can learn the quantum dynamics of a real molecule well enough to design control protocols for it, in a state space far too large for conventional optimization.

The problem: getting a molecule into one state

Before we can do precision spectroscopy on a molecular ion, we need to get it into a single, well-defined quantum state. That sounds like a simple prerequisite. For a polyatomic molecule like hydronium (H₃O⁺), it is closer to the whole problem.

Hydronium is worth the trouble. Its inversion transitions are sensitive to variations in fundamental constants, which makes it a candidate probe for physics beyond the Standard Model. But at 20 K, a trapped hydronium ion starts out smeared across hundreds of rotational and hyperfine levels. The job of quantum control is to find a sequence of laser pulses that funnels that thermal mess into one state. In the subspace we study, that means navigating a high-dimensional Hilbert space, and every candidate pulse you want to evaluate requires a full numerical propagation of the dynamics. Traditional optimal-control methods work well for small systems, but they lean on exactly that kind of repeated propagation, and the cost grows quickly with dimension.

Replace the simulator with a surrogate

Our approach is to stop calling the simulator inside the optimization loop. Instead, we train a Fourier Neural Operator (FNO) to be the simulator.

The FNO takes the molecule’s current population distribution plus the parameters of a single Raman sideband pulse (its two-photon frequency and polarization) and predicts the full population trajectory over the pulse window in one forward pass. Two design choices made this work at scale:

  • A physics-informed embedding. Rather than feeding the model a raw laser frequency, we encode it as detunings from the molecular transitions it’s actually close to. The model gets the physics that is relevant and skips the rest.
  • Exploiting symmetry. The σ⁺ and σ⁻ polarization channels are mirror images of each other in the molecule’s diagonal basis, so one trained surrogate covers both without retraining.

The result is a surrogate that reproduces reference trajectories from NVIDIA’s CUDA-Q Dynamics with high fidelity while running up to ~10⁷× faster than the GPU-accelerated numerical propagation it learned from. This is GPU compute at two levels: CUDA-Q generates the training data by propagating the full 888-dimensional dynamics, and the FNO then compresses that into a model that answers the same question in a fraction of the time on the same hardware. The speedup is not against a slow CPU baseline; it is against an already-accelerated simulator. Just as important, the surrogate is differentiable, which opens the door to gradient-based refinement of the control parameters.

From forward model to inverse design

With a fast, differentiable forward model in hand, we built a planner on top of it: the FNO stochastic pulse-measurement planner, or FNO-SPMP. At each step the surrogate scores every candidate pulse in a grid of frequencies and durations, the planner samples from the high-scoring pool, applies the pulse, simulates the motional-state measurement and reset, and repeats. Because the experiment measures the shared motional mode rather than the molecule directly, the planner has to track the molecular populations itself, and the FNO is what makes that tractable. The best candidate sequences are then optionally polished by gradient descent on their continuous parameters and validated with Monte Carlo rollouts.

On hydronium, FNO-SPMP produces pulse sequences that reach a target-state population of 0.98 with a success rate of up to 86.2%. Compared against a reinforcement-learning baseline in the same discrete control space, it roughly doubled the success rate, used about half as many pulses, and cut sequence-generation time from around 10 hours to 10–20 minutes.

Why this matters beyond simple molecules

The specific numbers are encouraging, but the more general point is the workflow. Operator learning gives us a surrogate that is fast enough to search over a large control space and differentiable enough to refine what we find. That combination is what lets inverse design reach into Hilbert spaces that were previously too large for direct optimization. We think this pattern will carry over to other molecular ions, and to other quantum platforms where the state space outgrows the simulator.

It is also a small answer to the blind spot we started with. AI models that learn the physics directly, rather than the language about it, can do things in the physical world that language models cannot.

The paper is on arXiv: arXiv:2608.03702. Thanks to Anastasia Pipi, Valentin Duruisseaux, Emily Been, Xuecheng Tao, Taylor Patti, Anima Anandkumar and to the NVIDIA team for CUDA-Q.