How Simulacra AI is transforming quantum chemistry

Currently, Simulacra AI is focused on pre-training an ab initio foundation wave-function model, which they call large wavefunction models (LWM). Taking from the scalability lessons learnt in the well-developed LLMs research field, LWMs are ML models that incorporate first-principles quantum mechanics to predict electron behaviour on a large scale, without the need for any labelled data. This approach provides the best of both worlds: the accuracy of quantum mechanics with the speed and efficiency of AI-driven predictions.

Instead of running full quantum chemistry calculations for every new molecule, Simulacra AI’s pre-trained models generalize across molecular structures, reducing computational costs while maintaining high fidelity.

Unlike traditional ab initio methods, where calculations scale poorly with system size, neural networks trained with the goal of optimizing the underlying wave function can bypass these limitations. This breakthrough allows users to achieve “gold standard” quantum mechanical accuracy for large, complex molecules without the exponential computational cost of traditional post-HF methods such as CC.

The eventual goal is to build billion parameters LWMs and infrastructure for computing highly accurate quantum chemistry properties, including time dynamics into the underlying modelled equations.

Simulacra AI needed to scale its model beyond the 10M parameter limit, which required a distributed training framework capable of handling large datasets and model architectures. Traditional deep learning approaches, such as data parallelism, are not sufficient for LWMs. Instead, model partitioning (sharding) is required.

The key challenges included:

• Handling large-scale training on JAX: Frequent updates break existing implementations, requiring constant maintenance.

• Implementing single program multiple data (SPMD) partitioning: This ensures that model weights, data, and molecular geometries are distributed across processing units.

• Computationally expensive differential operators: These are required for variational energy calculations in quantum chemistry.

• Managing PySCF orbital computations: PySCF needs recompilation for every new molecular geometry, slowing down training.

• Reducing compilation times: JIT compilation could take 2+ hours, making rapid iteration difficult.

To address these challenges, Simulacra AI built its own distributed AI framework with support for:

1. SPMD model partitioning: Simulacra AI splits both data and model weights across multiple GPUs. This allows them to scale beyond 100M parameters, where increasing GPU count results in larger models, not just faster training.

2. Custom equivariant neural networks: Built on Equinox with Clifford group representations, ensuring that models respect SE (3) symmetry which is important in the context of molecular properties computation.

3. Optimization with natural conjugate gradient descent: Given the complexity of variational energy calculations, standard optimizers are not effective. Simulacra AI developed a natural conjugate gradient optimizer, fully integrated with SPMD partitioning.

4. Ahead-of-time (AOT) compilation and caching: Instead of recompiling models dynamically, Simulacra AI caches precomputed steps, reducing compilation time from 2+ hours to just minutes.

5. Precomputed orbitals datasets: Instead of recalculating reference electron orbitals every time, results are stored and retrieved when necessary, improving efficiency.

This compute-first approach allows Simulacra AI to train models (20M parameters) significantly larger than publicly available SOTA benchmarks while maintaining efficiency.

Rather than commercializing the AI model itself, Simulacra AI will focus on selling the high-fidelity molecular properties data it generates. These datasets can serve various purposes such as input for the parametrisation of both classical and machine-learning molecular dynamics force fields, DFT machine learning functionals and direct quantum mechanical property prediction, to name a few.

This enables researchers and industries — including pharmaceuticals, materials science, and chemical engineering — to:

• Accelerate discovery by leveraging high-quality quantum chemistry data.

• Optimize proprietary molecular models with greater accuracy.

• Reduce reliance on costly, compute-heavy quantum simulations by using AI-driven precomputed datasets.

This strategy aligns with the needs of industries that require precise quantum chemistry calculations but may not have the infrastructure to train large-scale models from scratch.

By leveraging distributed computing and AI-first molecular simulation, Simulacra AI is transforming how research in materials science, pharmaceuticals, and chemical engineering is conducted. Their data-driven approach enables:

• More accurate property computation and molecular-dynamics simulations, accelerating and improving the quality of material design and drug discovery.

• Faster training cycles, thanks to optimized model compilation and caching.

• Scalability beyond traditional quantum mechanical methods, unlocking solutions previously computationally infeasible.

Their compute-first approach, leveraging distributed GPUs, model parallelism, and ahead-of-time optimization, represents a significant shift in computational chemistry and AI-driven material and molecular chemistry.