quantum-friendly-transformer

This repo serves as numerical experiments of quantum-friendly-transformer where we spectral-normalize / frobenius-normalize layers of transformers to understand how the block encoding affects model performance.

Setup

To set up the environment, we recommend using a virtual environment or conda environment. The following instructions assume you are using a conda environment.

bash setup.sh

All dependencies should be installed in approximately 5 minutes. This may vary depending on your internet speed and the performance of your machine.

To access Hugging Face datasets, you need to set up your Hugging Face token. You can do this by running:

huggingface-cli login

or by setting the HF_TOKEN environment variable with your token.

Matrix Norms in Transformers

Before experiments showcasing the capability of normalized transformers, we provide scripts to check the norms of pretrained models, check alpha_* files:

1. To check the norms of inputs S, in alpha_s:

   To see the plots, run plot_data.ipynb

   The data is generated by running:

   python alpha_s.py {model_name} {max_tokens} {batch_size} {random}
   

   where:

   model_name: ['bert', 'bart', 'roberta', 'distilgpt', 'gpt2', 'gpt', 'llama2-7b', 'tinyllama', 'mistral7b'] max_tokens: integer for maximum number of tokens, set to 512.

   batch_size: varies among models, larger models will need a smaller batch size to deal with OOM issues. If batch_size == -1, it means there is no batch_size and we do not truncate the tokens.

   random: 1 or 0 as a boolean variable to run the random dataset (1) or the MMLU dataset (0)

2. To check the norms of parameters W, in alpha_w: The statistics of the norms of parameters W are generated by running:

   python wqkv_norm_per_layer.py
   

   for a specific layer, you can check the norms of parameters W in that layer by running:

   python wqkv_norm.py
   

Runtime:

All the above scripts utilize no GPU and on a commercial laptop. Once you have relevant models downloaded from huggingface, the runtime should be less than 10 minutes for each script. The runtime may vary depending on the model and the number of tokens.

Training Normalized Transformers

All experiments run on a single NVIDIA A100 SXM4 GPU (40 GiB HBM2, driver 535.154.05, CUDA 12.2) paired with an AMD EPYC 7713 processor—of which 16 threads are allocated from the system’s 64-core/128-thread CPU (1.5–2.0 GHz base, boosting to ≈3.1 GHz; 32 KiB L1d/L1i, 512 KiB L2, 32 MiB L3 cache)—and 110 GiB of DDR4 RAM, all under Ubuntu 20.04 LTS.

After setting up the environment, make sure you have changed output directory in files under src/quantum_friendly_transformer/trainer to your desired location,

To train single-layer normalized transformers on the Genomic Benchmarks non-tata promoter dataset(0-1 classification task), use:

python src/quantum_friendly_transformer/trainer/train_genomic_bench.py

This result is reported in our paper and key hyperparameters are:

• --learning-rate: 8e-4
• --num-epochs: 100
• --early-stopping: 10
• --weight-decay: 0.01
• --lr-scheduler-type: cosine

We noted that more intricate hyperparameter tuning may be needed for the normalized models, as the optimization is slightly off the sweet point of current optimizers. For more details, please refer to the codebase.

Runtime:

The training of 3 single-layer transformers (vanilla, spectral normalized, frobenius normalized) on the Genomic Benchmarks dataset with given hyperparameters takes approximately 1 hour the specified machine.

To fine-tune the transformer on the same dataset with specific layer normalized, use:

python src/quantum_friendly_transformer/trainer/fine_tune_sn_model_genomic_bench.py

To train a frobenius normalized DNABert on the same dataset, use:

python src/quantum_friendly_transformer/trainer/fine_tune_multilayer_genomic_bench.py

This result is reported in our paper and key hyperparameters are:

• --learning-rate: 2e-4
• --num-epochs: 200
• --early-stopping: 25
• --weight-decay: 0.01
• --lr-scheduler-type: cosine

For more details, please refer to the codebase.

Runtime:

The training of a multi-layer normalized Bert model on the Genomic Benchmarks dataset with given hyperparameters takes approximately 2 hour on the specified machine.

To train an one-layer transformer on the Genome Understanding Evaluation (GUE) notata promoter dataset(0-1 classification task), use:

python src/quantum_friendly_transformer/trainer/train_gue.py

To fine-tune the transformer on the same dataset with specific layer normalized, use:

python src/quantum_friendly_transformer/trainer/fine_tune_sn_model_gue.py

To train a frobenius normalized DNABert on the same dataset, use:

python src/quantum_friendly_transformer/trainer/fine_tune_multilayer_gue.py

To train an one-layer transformer on the conll 2003 dataset(POS classification task), use:

python src/quantum_friendly_transformer/trainer/train_conll2003.py

To fine-tune the transformer on the same dataset with specific layer normalized, use:

python src/quantum_friendly_transformer/trainer/fine_tune_sn_model_conll2003.py