https://github.com/jorgemunozl/psiformer_torch
self-attention neural torch ansatz for electronic structure calculations introduced in the paper “A Self-attention Ansatz for ab initio Quantum Chemistry” from Google Deep Mind.
https://github.com/jorgemunozl/psiformer_torch
attention-mechanism computational-physics computational-quantum-mechanics google-deepmind kfca mlp natural-gradient quantum-mechanics torch training transformer wandb wave-equation-solver
Last synced: about 1 month ago
JSON representation
self-attention neural torch ansatz for electronic structure calculations introduced in the paper “A Self-attention Ansatz for ab initio Quantum Chemistry” from Google Deep Mind.
- Host: GitHub
- URL: https://github.com/jorgemunozl/psiformer_torch
- Owner: jorgemunozl
- Created: 2025-07-19T14:47:33.000Z (11 months ago)
- Default Branch: main
- Last Pushed: 2026-04-13T20:48:24.000Z (about 2 months ago)
- Last Synced: 2026-04-13T22:35:04.288Z (about 2 months ago)
- Topics: attention-mechanism, computational-physics, computational-quantum-mechanics, google-deepmind, kfca, mlp, natural-gradient, quantum-mechanics, torch, training, transformer, wandb, wave-equation-solver
- Language: TeX
- Homepage: https://jorgemunozl.github.io/portfolio/project_transformers
- Size: 29.1 MB
- Stars: 12
- Watchers: 0
- Forks: 1
- Open Issues: 10
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
Solving the Schrödinger equation with Transformers
This repository demonstrates an experimental approach to approximating solutions to the many-electron Schrödinger equation using transformer architectures. It contains reference code and utilities to build and run scale experiments, for more details on how this repository works and was created look [PsiFormer](src/PsiFormer_Torch%20Documentation.md) for the full documentation. Watch the [paper](assets/research/main.pdf) for a detailed theoretical framework result,the [beamer](assets/research/beamer.pdf) for a more friendly explanation or the [poster](assets/research/poster.pdf) for a resume of this work.
> My sincerest thanks to Angel A. Flores and Alejandro D. Paredes for their feedback and helpful technical discussion. I would also like to thank the Faculty of Computer Science at UNI for the GPU used to make this project.
## Table of contents
- [What this does](#what-this-does)
- [Key ideas (short)](#key-ideas-short)
- [Quick install (CPU-only)](#quick-install-cpu-only)
- [Usage](#usage)
- [Files and layout](#files-and-layout)
- [Example workflow](#example-workflow)
- [Tips and caveats](#tips-and-caveats)
- [Reproducibility](#reproducibility)
- [References and further reading](#references-and-further-reading)
- [License](#license)
## What this does
- Uses transformer-style attention to model electronic wavefunction structure.
- Provides utilities and example scripts to run problems and evaluate model performance against simple quantum chemistry references.
This project is intended for research .It is not production-ready for large-scale quantum chemistry simulations but can be used to prototype ideas and compare modeling choices.
## Key ideas (short)
- Represent electronic configurations or basis expansions as sequences and
learn interactions using attention layers.
- Use permutation-equivariant input encodings or ordered basis sequences to
capture antisymmetry constraints through learned modules and loss terms.
- Train with physics-informed losses (energy expectation, cusp conditions,
or density overlaps) and supervised or self-supervised pretraining.
## Quick install (CPU-only)
If you use `uv` (fast installer), you can install a CPU-only PyTorch wheel and
other dependencies like this:
```bash
# install uv if you don't have it
pip install -U uv
# Sync dependencies
uv sync
```
Then, install PyTorch CPU-only version:
```bash
uv pip install --index-url https://download.pytorch.org/whl/cpu torch
```
## Usage
1. Prepare data / basis encodings for the target system.
2. Edit `config.py` to set model size, number of heads, block (sequence)
length, and training hyperparameters.
3. Run the training/experiment script (example):
Replace the example command above with your own runner or flags as needed.
## Files and layout
- `src/` — main source code
- `main.py` — training / experiment entrypoint
- `utils.py` — helper utilities, device selection, etc.
- `config.py` — config and hyperparameters
- `pyproject.toml` — Python project metadata
- `template.tex` — report template (optional)
## Example workflow
1. Choose a small system (2–10 electrons) and a compact basis set.
2. Build sequence encodings for orbitals/electron coordinates.
3. Train the transformer to predict minimal-energy coefficients or approximate the wavefunction amplitude for sampled configurations.
4. Evaluate energy and compare against reference (Hartree–Fock, small CI).
## Tips and caveats
- Enforcing antisymmetry exactly (Slater determinants) is nontrivial when
using sequence models; consider hybrid approaches (learn corrections to a
Slater determinant) or include antisymmetry in the loss.
- Transformers scale quadratically with sequence length; start with small
basis sizes and toy molecules.
- Use physics-informed losses wherever possible to improve sample efficiency.
## References and further reading
- Vaswani et al., "Attention Is All You Need" (transformers).
- A Self Attention Ansatz for Quantum Chemistry
- Ab-Initio Solution of the Many-Electron Schrodinger Equation With Deep Neural Networks
- QMC Torch Molecular Wave functions with Neural Components for energy and force calculations.