This repository is the official implementation of Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystr"om Method (NeurIPS 2021).

To install requirements in a conda environment:

conda create -n skyformer python=3.6
conda activate skyformer
pip install -r requirements.txt

Note: Specific requirements for data preprocessing are not included here.

Processed files can be downloaded here, or processed with the following steps:

1. Requirements

tensorboard>=2.3.0
tensorflow>=2.3.1
tensorflow-datasets>=4.0.1

2. Download the TFDS files for pathfinder and then set _PATHFINER_TFDS_PATH to the unzipped directory (following google-research/long-range-arena#11)
3. Download lra_release.gz (7.7 GB).
4. Unzip lra-release and put under ./data/.

cd data
wget https://storage.googleapis.com/long-range-arena/lra_release.gz
tar zxvf lra-release.gz 

5. Create a directory lra_processed under ./data/.

mkdir lra_processed
cd ..

6.The directory structure would be (assuming the root dir is code)

./data/lra-processed
./data/long-range-arena-main
./data/lra_release

7. Create train, dev, and test dataset pickle files for each task.

cd preprocess
python create_pathfinder.py
python create_listops.py
python create_retrieval.py
python create_text.py
python create_cifar10.py

Note: most source code comes from LRA repo.

Modify the configuration in config.py and run

python main.py --mode train --attn skyformer --task lra-text

• mode: train, eval
• attn: softmax, nystrom, linformer, reformer, perfromer, informer, bigbird, kernelized, skyformer
• task: lra-listops, lra-pathfinder, lra-retrieval, lra-text, lra-image

@inproceedings{Skyformer,
    title={Skyformer: Remodel Self-Attention with Gaussian Kernel and Nystr\"om Method}, 
    author={Yifan Chen and 
            Qi Zeng and 
            Heng Ji and 
            Yun Yang},
    booktitle={Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December
               6-14, 2021, virtual},
    year={2021}
}