Skip to content

PIQuIL/RydbergGPT

Repository files navigation

RydbergGPT

A large language model (LLM) for Rydberg atom array physics. Manuscript available on arXiv.

Table of contents

Quick Start

Configuration

Theconfig.yaml is used to define the hyperparameters for :

  • Model architecture
  • Training settings
  • Data loading
  • Others

Training

To train RydbergGPT locally, execute the main.py with :

python main.py --config_name=config_small.yaml

Installation

Clone the repository using the following command :

git clone https://github.com/PIQuIL/RydbergGPT

Create a conda environment

conda create --name rydberg_env python=3.11

and finally install via pip in developer mode:

cd RydbergGPT
pip install -e .

Documentation

Documentation is implemented with MkDocs and available at https://piquil.github.io/RydbergGPT.

Architecture

Rydberg System

Consider the standard Rydberg Hamiltonian of the form :

$$\begin{align} & \hat{H}_{\mathrm{Rydberg}} = \sum_{i < j} V(\lVert \mathbf{R}_i - \mathbf{R}_j \rVert \ ; R_b, \Omega) \hat{n}_i \hat{n}_j - \sum_{i} \Delta_i \hat{n}_i + \sum_{i} \frac{\Omega}{2} \sigma_i^{(x)} \\\ & V(\lVert \mathbf{R}_i - \mathbf{R}_j \rVert \ ; R_b, \Omega) = \frac{R_b^6 \Omega}{\lVert \mathbf{R}_i - \mathbf{R}_j \rVert^6} \end{align}$$

Here, $V_{ij}$ = blockade interaction strength between atoms $i$ and $j$, $R_b$ = blockade radius in units of the lattice spacing, $\hat{n}_i$ = number operator at ion $i$, $\mathbf{R}_i$ = the position of atom $i$ in units of the lattice spacing, $\Delta_i$ = detuning at atom $i$, $\Omega_i$ = Rabi frequency at atom $i$.

Transformer

Vanilla transformer architecture taken from Attention is All You Need.

  • $\mathbf{x} =$ experimental settings
  • $\sigma_i =$ one-hot encoding of measured qubit $i$
  • $p_{\theta}(\sigma_i | \sigma_{&lt; i}) =$ neural network conditional probability distribution of qubit $i$

The transformer encoder represents the Rydberg Hamiltonian with a sequence.
The transformer decoder represents the corresponding ground state wavefunction.

Data

Consider setting $\Omega = 1$ and varying the other Hamiltonian parameters independently :

$$\begin{align} L &= [5, 6, 11, 12, 15, 16] \\\ \delta / \Omega &= [-0.36, -0.13, 0.93, 1.05, 1.17, 1.29, 1.52, 1.76, 2.94, 3.17] \\\ R_b / a &= [1.05, 1.15, 1.3] \\\ \beta \Omega &= [0.5, 1, 2, 4, 8, 16, 32, 48, 64] \end{align}$$

Data available on Pennylane Datasets