Can AI solve a foundational problem in quantum mechanics?

At Primateria, our goal is to build simulators that predict the properties of molecules and materials with experimental fidelity before we make them. AI coding agents could accelerate the development of these simulators, so we created Wavefunction Bench to test their ability to improve scientific simulation software. The benchmark asks agents to improve Orbformer, a neural network model that approximates solutions to the electronic Schrödinger equation for small molecules, producing wavefunctions from which molecular properties can be calculated. Solving this equation accurately for complex systems at practical computational cost remains one of the greatest challenges in computational quantum mechanics.

We gave each agent eight hours on a single GPU to modify Orbformer and train a wavefunction for a specified molecule. The tasks span nine molecular systems, including transition metal compounds outside Orbformer's pretraining curriculum. We independently evaluate each submitted wavefunction and score its variational energy against a baseline and a molecule specific target. Targets for LiH and CO₂ come from high accuracy theoretical calculations; the other targets combine calculated atomic energies with corrected experimental atomization energies.

Under our success criterion of a valid measured energy less than 1.594 mHa above the target, successful trials occurred on three of the nine molecular systems: LiH, NaCl, and MgCl₂. All eight models passed on LiH at least once, with 31 of 32 trials passing. GPT-6 Astra, Claude Opus 5.5, and Muse Spark 1.3 passed on NaCl; Astra and Opus also passed on MgCl₂. No selected trial passed on the other six systems.

The task and scoring

Each of the eight models received a task specific instruction. For titanium dioxide, the prompt was:

Use Orbformer to train a neural wavefunction for Titanium dioxide molecule (TiO₂). Use python -m evaluation (source in /workspace/evaluation/) to validate your lowest reliable variational energy, and every time you get a new best, export it to /submission/checkpoint.npz with python -m evaluation export <checkpoint>. The molecule definition is available at /workspace/tio2-ecp.yaml.

You have 8 hours available. Do not stop until you get to -90.3449097116213 Ha or lower.

When a model finishes, we load its submitted checkpoint and measure its variational energy using 512 Monte Carlo walkers, 200 equilibration steps, and 4,000 measurement steps. We test electron-exchange behavior and numerical consistency, and report total spin as a diagnostic when the export includes signed amplitudes. The scope of these checks is described below. We set the score to zero when required checks fail, including when the energy estimates from the two halves of the measurement differ by more than four standard errors.

The measured energy $E$ is then placed on a scale between the baseline and the target, as shown in equation (1) below.

$$\text{score} = \frac{E_{\text{baseline}} - E}{E_{\text{baseline}} - E_{\text{target}}} \tag{1}$$

Some trials reached a continuous score of 0.99, but are counted as failures under our binary metric because their measured energy remained at least 1.594 mHa (1 kcal/mol) above the target. Closing that final gap can be disproportionately difficult: the last few millihartrees may require more effort than most of the improvement that came before.

Success is defined relative to our chosen reference energies. Uncertainty in those references is not included in scoring and can change borderline pass/fail outcomes. Passing the cutoff alone does not establish experimental accuracy. We describe the reference construction and limitations below.

We also explored using bond dissociation energies (BDE) for scoring, but their interpretation was complicated by unequal convergence of molecular and fragment wavefunctions and cancellation of their errors. We use variational total energy as the primary measure of solver improvement. Establishing accuracy for bond breaking and chemical reactions requires additional evaluation of relative energies.

Why wavefunctions are a useful test

A molecular wavefunction describes the joint quantum state of its electrons at a fixed nuclear geometry. It lets us calculate quantities such as energy and electron density while accounting for correlations between electrons. Representing these correlations accurately can be especially difficult during bond breaking and in transition metal compounds, where several electronic arrangements may compete.

Neural wavefunctions represent this joint state with a trainable network. Their accuracy and computational cost depend on choices about representation, optimization, and sampling of the wavefunction during training. These are all implementation choices that coding agents can modify. The resulting wavefunction is therefore an output we can evaluate independently, which is done using the specified physical model, as described below.

The reward function for our AI Agents: the variational energy

One very fortunate feature of quantum mechanics is the variational principle, which gives us a way to compare valid trial wavefunctions by their energy. In equation (2), $|\Psi\rangle$ is the wavefunction, $\hat H_M$ is the Hamiltonian (the energy operator) for a molecular system $M$ at a fixed nuclear geometry, and $E_0$ is its exact ground-state energy.

$$\frac{\langle\Psi\mid\hat{H}_{M}\mid\Psi\rangle}{\langle\Psi\mid\Psi\rangle} \geq E_0(M) \tag{2}$$

Equation (2) says the exact energy expectation of a valid trial wavefunction cannot fall below the ground-state energy for the same Hamiltonian. Lowering this expectation brings the energy closer to that limit, giving agents a physical objective to optimize. Our Monte Carlo estimates of this quantity have sampling uncertainty, so reliable sampling and the validation checks described below are essential when comparing submissions.

To compute a wavefunction's energy, we use variational Monte Carlo (VMC). The energy is an average over all possible electron positions that we estimate from samples. Using the Metropolis–Hastings algorithm, we draw electron configurations $\mathbf x$ with probability proportional to $|\Psi_\theta(\mathbf x)|^2$, where $\theta$ represents the network's weights. For each sample, we compute the local energy $E_L$. Averaging the local energy over the samples gives the variational energy, as shown in equation (3).

$$E_L(\mathbf x) = \frac{\hat H \Psi_\theta(\mathbf x)}{\Psi_\theta(\mathbf x)}, \qquad E(\theta) \approx \frac{1}{M}\sum_{i=1}^{M} E_L(\mathbf x_i) \tag{3}$$

Backpropagation gives $\nabla_\theta \log|\Psi_\theta|$ for each sample. Weighting these by how far each sample's local energy sits from the average gives a stochastic estimate of the energy gradient, as shown in equation (4).

$$\nabla_\theta E \approx \frac{2}{M}\sum_{i=1}^{M}\big(E_L(\mathbf x_i) - \bar E_L\big)\,\nabla_\theta \log|\Psi_\theta(\mathbf x_i)| \tag{4}$$

Neural wavefunction methods, including Orbformer, follow this gradient with natural-gradient optimizers such as KFAC or SPRING, which account for how a change in the weights changes the wavefunction itself. Below is a visualization of Monte Carlo sampling of the wavefunction during Orbformer training:

How did the models tackle this problem?

• On its Cu₂ runs, GPT-6 Astra explored explicit Hartree-Fock orbital references, learned electron-correlation factors, and backflow corrections to electron coordinates. It tested and refined multiple competing designs.

• On its best CrCl run, Opus 5.5 improved orbital initialization and tested different KFAC optimizer steps by reducing the step size. The final score was 0.988.

• After changing its Hartree-Fock initialization and making its determinant branches identical, one MiMo V2.6 Pro run on CuF reported -45.8 million Ha. The trial's final submission failed independent validation and scored 0. The agent also tried to fix the instability during the run, but when we looked at the trace, we could not find a specific reason for the behavior. We also did not find evidence of deliberate reward hacking.

• Out of all the models, only GPT-6 Astra and Opus improved on the baseline of Cu₂, reaching best scores of 0.915 and 0.756 respectively. While neither model reached the target, none of the other six models produced a validated improvement.

• The baseline energy for Cu₂ came from 1,000 stock training steps. Across the other models' 24 runs, eight produced valid energies worse than baseline while 16 failed validation.

• For LiH, its 0.35 mHa baseline-to-target gap is already smaller than the fixed 1.594 mHa success tolerance. A valid result can therefore pass without improving on the baseline. LiH's small number of electrons, pretrained support for Li and H, and its 0.35 mHa baseline-to-target gap are plausible advantages that each model had. We have not isolated the effect of pretraining.

• Grok often struggled to deliver a valid final submission. 20 of its 36 runs failed independent validation, which was the highest count among the models tested. In one TiO₂ trial, an optimizer rewrite ran into repeated calls to library APIs that did not exist.

Checking physical and numerical consistency

Orbformer guarantees antisymmetry by design. Because agents can modify its architecture, we test electron exchanges to check that their changes preserve this property. Fixing the numbers of spin-up and spin-down electrons does not guarantee a definite total spin, so spin purity requires a separate diagnostic.

Antisymmetry

Electrons are fermions. Exchanging two electrons with the same spin must reverse the sign of the wavefunction while leaving its probability density unchanged:

$$\Psi(\ldots,\mathbf r_i,\ldots,\mathbf r_j,\ldots) = -\Psi(\ldots,\mathbf r_j,\ldots,\mathbf r_i,\ldots)$$

Our test For signed exports, we exchange sampled pairs of same-spin electrons and check that the sign reverses while the magnitude stays unchanged, within numerical tolerances. The LiH and CO₂ exports in these trials contain only the magnitude, so their exchange checks test magnitude invariance alone. That limited check cannot establish antisymmetry.

Spin purity

Fixing the numbers of spin-up and spin-down electrons determines the spin projection, but does not necessarily give the wavefunction a definite total spin. It can still contain contributions from unwanted spin states, known as spin contamination. A state with definite total spin $S$ satisfies:

$$\hat S^2 \Psi = S(S+1)\hbar^2 \Psi$$

Our diagnostic For the seven ECP tasks, we separately estimate $\langle\hat{\mathbf{S}}^2\rangle$ and compare it with the value expected for the specified spin. We report this measurement to identify spin contamination; it does not affect the score. We did not measure it for LiH or CO₂ because their exports lacked the required sign information. The verifier implementation records which checks were performed.

Sampling and numerical consistency

For $N$ electrons, the energy involves an integral over $3N$ spatial coordinates. A grid with $m$ points along each coordinate would require $m^{3N}$ evaluations. Monte Carlo avoids constructing this grid by randomly sampling electron configurations and averaging their local energies. Collecting more samples improves the precision of the energy estimate.

Our tests The verifier generates its own samples from the submitted wavefunction. An initial equilibration period lets the sampling settle before we begin measuring the energy. We check for nonfinite values and statistically significant drift. The reported energy uses unclipped local-energy values.

For TiO₂, each submitted wavefunction is equilibrated for 200 iterations, after which its energy is measured for 4,000 iterations using 512 Monte Carlo walkers. Each walker stores a configuration of the 24 electrons treated explicitly by the model and is updated between energy measurements.

Explore Astra's best TiO₂ wavefunction

Since a many-electron wavefunction gives the joint probability of measuring all electron positions at once, we cannot plot a wavefunction directly. Instead, we can freeze all but one electron at a given sampled configuration and plot the probability of finding the last electron. Using this method, we plot a wavefunction of TiO₂ optimized by GPT-6 Astra for 12 hours.

The cloud shows where the unfrozen electron is likely to be, given the positions of the other frozen electrons. Dragging the position bar moves one of those frozen electrons around a small circle and shows how the cloud of the unfrozen electron responds. The gray surfaces are where $\Psi$ passes through zero and flips sign.

Where our benchmark target energies come from

Seven tasks use effective core potentials (ECPs), which replace core electrons with an effective interaction. For these tasks, we construct total-energy targets by subtracting corrected experimental atomization energies from calculated energies of the separated atoms. TiO₂ uses our measured atomic VMC energies; the other six use published ccECP atomic references. LiH and CO₂ use the theoretical total-energy references listed in the leaderboard.

$$E_{\mathrm{target}}(\mathrm{molecule}) = \sum_{A} E_{\mathrm{reference}}(A) - D_{\mathrm{electronic}}$$

The corrections use published thermochemical and spectroscopic inputs. The molecule configurations record the values, sources, and approximations for each target.

Temperature

Thermochemical quantities reported at finite temperature include thermal contributions from molecular and atomic motion and excitation.

Our correction TiO₂ starts from a 298.15 K atomization enthalpy. We account for molecular translation, rotation, and vibration, and atomic translation and electronic excitation. The recorded enthalpy increments are 11.18 kJ/mol for TiO₂ and 20.99 kJ/mol for Ti + 2 O; their difference gives a −0.1017 eV correction to the atomization energy. The other six thermochemical targets start from 0 K dissociation energies or electronic well depths and need no thermal correction.

Zero-point vibration

Even at 0 K, nuclei retain vibrational energy, while our electronic calculations hold the nuclei fixed.

Our correction If the source reports a dissociation energy $D_0$ from the lowest vibrational level, we add the zero-point energy to obtain the electronic well depth:

$$D_e = D_0 + E_{\mathrm{ZPE}}$$

For TiO₂, we approximate the zero-point energy as half the sum of the NIST gas-phase vibrational frequencies, 959, 323, and 944 cm⁻¹, giving 0.138 eV. MgCl₂ uses the same approximation with its reported fundamentals, including the doubly degenerate bend. NaCl uses a published spectroscopic zero-point energy; CrCl and CuF use diatomic spectroscopic constants, with an estimated anharmonic term for CrCl. The Cu₂ and CuH sources already report $D_e$, so no extra zero-point term is added. The CuH reference combines spectroscopy with calculations for unobserved levels.

Spin–orbit coupling

Interactions between electronic spin and orbital motion split atomic terms into levels with different energies.

Our correction Our ECP calculations omit explicit spin–orbit coupling. We use published atomic fine-structure levels to shift from each atom's lowest level to the average over its ground term, weighted by the degeneracy of each level. For TiO₂, the Ti and two O contributions add 0.0469 eV to the atomization energy.

For TiO₂, these steps convert 13.2395 eV into an electronic atomization energy of $13.2395 - 0.1017 + 0.1380 + 0.0469 = 13.3227$ eV. These constructed targets inherit uncertainty from thermochemical data, atomic reference energies, correction approximations, and the ECP treatment. That uncertainty is separate from the molecular Monte Carlo standard error. The 1.594 mHa success cutoff is applied to the central target values and does not include reference uncertainty.

As a sensitivity check, we compared the cached energies with the existing ccECP CCSD(T)/CBS calculations at the configured geometries. Using those calculations as alternative targets changes the NaCl pass count from 6/32 to 0/32 and MgCl₂ from 3/32 to 0/32. These alternative references are themselves approximate and nonvariational. The leaderboard retains the original targets; this comparison shows how the reference choice affects its binary outcomes.

Appendix

Pareto Frontier

Although open-source and non-frontier models are highly cost-efficient, their capabilities for scientific discovery remain far behind those of frontier models.

Scaling test-time compute

We compared GPT-6 Astra and Claude Opus 5.5 on TiO₂ with time budgets of 4, 8, 12, and 16 hours. We observed no improvement in variational energy beyond 12 hours.

Harness Differences

With GPT-6 Astra, Mini SWE Agent slightly outperformed Codex while using 80% more tokens. Other models showed noticeable regressions.

Acknowledgements

Thanks to Adam Foster for providing valuable feedback on this blog post.

Citation

@misc{wavefunctionbench_2026,
  title        = {{Can AI solve a foundational problem in quantum mechanics?}},
  author       = {Primateria},
  year         = {2026},
  howpublished = {\url{https://www.primateria.com/blog/wavefunctionbench/}},
  note         = {Blog post by Primateria}
}