An implementation of the transformer quantum state (TQS), a multi-purpose model for quantum many-body problems. Contains everything necessary to reproduce the results in our paper.
arXiv link: arXiv:2208.01758
This repository contains code adapted from the official PyTorch repository and autograd-hacks. Special thanks to the open-source community!
A pre-trained model for the transverse field Ising Hamiltonian is provided. By customizing the code, you can also train your own model on different Hamiltonians.
Requirements: PyTorch, TeNPy (for DMRG simulations), SciPy>=1.7.1 (for predicting parameters)
To train a TQS from scratch:
python3 main.py
In this example, we are training on the Ising model. You can choose different Hamiltonians defined in Hamiltonian.py, or define your own Hamiltonian.
For fine-tuning, load a pre-trained model first, specify the parameters you want to fine-tune on, and set fine_tuning=True
To evaluate the performance of the pre-trained TQS:
python3 test.py
Again, this example is dedicated to the Ising model, and computes its ground state energy and magnetization.
To predict field strengths from experimental measurements:
python3 param_prediction.py
The experimental measurements are provided in data/, which are simulated using DMRG.
Also,
python3 visualization.py
will plot Fig. 2 in our paper.