efficient attention

Effiziente Aufmerksamkeit

Dieses Repository enthält die offizielle Implementierung von Experimenten, die in durchgeführt wurden

• EVA: Effiziente Aufmerksamkeit über Kontrollvariationen (ICLR 2023)
• Lara: Lineare Komplexität Randomisierter Selbstbekämpfungsmechanismus (ICML 2022)

? Repo -Struktur:

• efficient-attention : Eine kleine in sich geschlossene Codebasis, die verschiedene effiziente Aufmerksamkeitsmechanismen implementiert. Weitere Informationen finden Sie in der Verwendung.
• vit : Codebasis für Bildklassifizierungsexperimente , die aus angepasst sind
  • Deit für die Dateistruktur und
  • PVT für PVTV2 -Modellklassen.
• fairseq : Eine modifizierte Gabel von Fairseq für Sprachaufgaben, einschließlich maschineller Übersetzung und autoregressiver Sprachmodellierung .
• main.sh : Ein Bash -Skript zum Starten aller Experimente.
  • Siehe das Skript für die Liste der Argumente.
  • Beachten Sie, dass Argumente nach -e True direkt an den Schulungsbefehl übergeben werden. Sie können benutzerdefinierte Argumente an den Schulungsbefehl übergeben, indem Sie sie nach -e True anhängen.

Abhängigkeiten

Führen Sie die folgenden Befehle aus, um die erforderlichen Abhängigkeiten zu installieren (in einer virtuellen Umgebung empfohlen), um die Umgebung einzurichten:

 # install packages
pip install -r requirements.txt
# install efficient-attention library
pip install -e efficient-attention

# OPTIONAL: install fairseq library for running language tasks
cd fairseq
python3 setup.py build develop
cd ..

Die Umgebung wird mit Python 3.8.10, Pytorch 1.12.0 und CUDA 11.3 getestet. Beachten Sie auch, dass unsere Gabel von Fairseq mehrere Dateien in der ursprünglichen Codebasis verändert. Die Verwendung neuerer Versionen von Fairseq kann zu unerwarteten Abhängigkeitskonflikten führen.

Grundlegende Verwendung der effizienten Aufmerksamkeitsbibliothek

efficient-attention ist eine kleine in sich geschlossene Codebasis, die mehrere effiziente Aufmerksamkeitsmechanismen sammelt.

Übergeben aufmerksamkeitsspezifische Argumente an Argparse

• Für Argumente, die für jeden Aufmerksamkeitsmechanismus spezifisch sind, überprüfen Sie bitte die Klassenmethode add_attn_specific_args() in der entsprechenden Python -Datei.
• Um diese Argumente an den argparse -Parser weiterzugeben, folgen Sie dem folgenden Code -Snippet:

 import argparse
from efficient_attention import AttentionFactory
# ...
parser = argparse . ArgumentParser ()
parser . add_argument ( '--attn-name' , default = 'softmax' , type = str , metavar = 'ATTN' ,
                        help = 'Name of attention model to use' )
# ...
temp_args , _ = parser . parse_known_args ()
# add attention-specific arguments to the parser
# struct_name: name of the inner namespace to store all attention-specific arguments
# prefix: prefix to prepend to all argument names
#         for example, if prefix = encoder-attn, then for the argument --window-size 
#         we need to pass --encoder-attn-window-size
#         this is useful to avoid argument name conflicts.
efficient_attention . AttentionFactory . add_attn_specific_args ( parser , temp_args . attn_name , struct_name = "attn_args" , prefix = "" )
# parse arguments to a namespace that supports nested attributes
args = parser . parse_args ( namespace = efficient_attention . NestedNamespace ())
# now we can access the attention-specific arguments via args.attn_args
print ( args . attn_args . window_size )

Erstellen Sie ein effizientes Aufmerksamkeitsmodul

In einer torch.nn.Module -Klasse können Sie wie folgt ein effizientes Aufmerksamkeitsmodul erstellen:

 # we might want to pass attention-specific arguments to the attention module
# along with other related arguments
attn_args = {
    ** vars ( args . attn_args ),
    ** {
    'dim' : args . embed_dim , 
    'num_heads' : args . num_heads , 
    'qkv_bias' : args . qkv_bias , 
    'attn_drop' : args . attn_drop_rate , 
    'proj_drop' : args . drop_rate ,
    }
}
self . attn = AttentionFactory . build_attention ( attn_name = attn_name , attn_args = attn_args )

# the module can then be used as a normal function as
x = self . attn ( x )

Bildklassifizierung auf ImageNet1k

Datenvorbereitung

Wir folgen dem ähnlichen Setup, der Deit ähnelt, um den ImageNet-Datensatz vorzubereiten. Laden Sie ImageNet Zug- und VAL -Bilder herunter und platzieren Sie sie in die folgende Verzeichnisstruktur, damit sie mit den Torchvision datasets.ImageFolder kompatibel sein kann.

 /path/to/imagenet/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class2/
      img4.jpeg

Schulung und Bewertung

Die folgenden Befehle werden zum Training und Bewertung verschiedener Vision -Transformatoren mit LARA/EVA verwendet. Es wird angenommen, dass das Training mit 8 GPUs durchgeführt wird.

Imagnet -Klassifizierung auf Deit (Sequenzlänge 784 (Suffix:*_ P8)/196 (Suffix:*_ P16))

LARA/EVA in verschiedenen Deit -Architekturen zu verwenden:

 # LARA: DeiT-tiny-p8
bash main.sh -m evit_tiny_p8 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name lara --mis-type mis-opt --proposal-gen pool-mixed --alpha-coeff 2.0 --num-landmarks 49

# LARA: DeiT-tiny-p16
bash main.sh -m evit_tiny_p16 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name lara --mis-type mis-opt --proposal-gen pool-mixed --alpha-coeff 2.0 --num-landmarks 49

# LARA: DeiT-small-p16
bash main.sh -m evit_small_p16 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name lara --mis-type mis-opt --proposal-gen pool-mixed --alpha-coeff 2.0 --num-landmarks 49

# EVA: DeiT-tiny-p8
bash main.sh -m evit_tiny_p8 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name eva --num-landmarks 49 --adaptive-proj default --window-size 7 --attn-2d --use-rpe

# EVA: DeiT-tiny-p16
bash main.sh -m evit_tiny_p16 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name eva --num-landmarks 49 --adaptive-proj default --window-size 7 --attn-2d --use-rpe

# EVA: DeiT-small-p16
bash main.sh -m evit_small_p16 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name eva --num-landmarks 49 --adaptive-proj default --window-size 7 --attn-2d --use-rpe

Imagnet -Klassifizierung auf PVTV2 -B3 (Sequenzlänge: 3136 -> 784 -> 196 -> 49)

LARA/EVA in PVTV2 -Architekturen anzupassen:

 # LARA Attention
bash main.sh -m pvt_medium2 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 1.0 --drop-path-rate 0.3 --warmup-epochs 10 --seed 1 --attn-name lara --pool-module-type dense --mis-type mis-opt --proposal-gen pool-mixed --num-landmarks 49 --alpha-coeff 2.0 --repeated-aug

# EVA Attention
bash main.sh -m pvt_medium2 -p < dir-of-imagenet-data > -g 8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --drop-path-rate 0.3 --warmup-epochs 10 --seed 1 --attn-name eva --num-landmarks 49 --adaptive-proj default --window-size 7 --attn-2d --use-rpe --repeated-aug

Die Verwendung anderer Aufmerksamkeitsmechanismen:

Alternativ möchten Sie andere Aufmerksamkeitsmechanismen ausprobieren:

 # Softmax Attention
bash main.sh -m evit_tiny_p8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name softmax
# RFA/Performer
bash main.sh -m evit_tiny_p8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name performer --proj-method favorp --approx-attn-dim 64
# Local Attention
bash main.sh -m evit_tiny_p8 -d imagenet -e TRUE --dist-eval --num-workers 16 --clip-grad 5.0 --warmup-epochs 10 --seed 1 --attn-name local --window-size 7 --attn-2d --use-rpe

Sprachaufgaben

Datenvorbereitung

Wir verwenden den Standard-Vorverarbeitungsprozessing Fairseq, um die Daten für Sprachaufgaben zu erstellen.

• Für die maschinelle Übersetzung folgen Sie bitte hier, um sich auf die binarisierten WMT'14 EN-DE Daten vorzubereiten.
• Für autoregressive Sprachmodellierung folgen Sie hier, um den Wikitext-103 Datensatz zu verarbeiten.

Ausbildung

• -r <resume-ckpt-DIR> Gibt das Verzeichnis an, in dem Ihre Kontrollpunkte während des Trainings gespeichert sind und zur Wiederaufnahme des Trainings verwendet werden können.
• Beachten Sie, dass alle aufmerksamkeitsspezifischen Argumente mit dem Präfix --encoder-attn- (für Encoder-Seite) / --decoder-attn- (für Decoder-Side) zugeordnet werden müssen. Siehe die folgenden Beispiele.

Maschinelle Übersetzung

 # # LARA
CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -d wmt -s lara_8 -g 4 -e TRUE --attn-name-encoder lara --encoder-attn-num-landmarks 8 --encoder-attn-proposal-gen adaptive-1d --encoder-attn-mis-type mis-opt

CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -d wmt -s lara_16 -g 4 -e TRUE --attn-name-encoder lara --encoder-attn-num-landmarks 16 --encoder-attn-proposal-gen adaptive-1d --encoder-attn-mis-type mis-opt

CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -d wmt -s lara_32 -g 4 -e TRUE --attn-name-encoder lara --encoder-attn-num-landmarks 32 --encoder-attn-proposal-gen adaptive-1d --encoder-attn-mis-type mis-opt

# # EVA
CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -d wmt -s eva_8_8 -g 4 -e TRUE --attn-name-encoder eva --encoder-attn-window-size 8 --encoder-attn-num-landmarks 8 --encoder-attn-adaptive-proj no-ln --encoder-attn-use-t5-rpe --encoder-attn-overlap-window

CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -d wmt -s eva_16_8 -g 4 -e TRUE --attn-name-encoder eva --encoder-attn-window-size 16 --encoder-attn-num-landmarks 8 --encoder-attn-adaptive-proj no-ln --encoder-attn-use-t5-rpe --encoder-attn-overlap-window

CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -d wmt -s eva_32_8 -g 4 -e TRUE --attn-name-encoder eva --encoder-attn-window-size 32 --encoder-attn-num-landmarks 8 --encoder-attn-adaptive-proj no-ln --encoder-attn-use-t5-rpe --encoder-attn-overlap-window

Autoregressive Sprachmodellierung

 # Currently, LARA does not support causal masking yet.

# EVA on a 16-layer Transformer LM
CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -m 16layers -d wikitext103 -s eva_128_8_16layers -g 4 -e TRUE --attn-name-decoder causal_eva --decoder-attn-window-size 128 --decoder-attn-causal --decoder-attn-adaptive-proj qk --decoder-attn-chunk-size 8 --decoder-attn-use-t5-rpe

# EVA on a 32-layer Transformer LM
CUDA_VISIBLE_DEVICES=0,1,2,3 bash main.sh -p < dir-of-your-bin-data > -m 32layers -d wikitext103 -s eva_128_8_32layers -g 4 -e TRUE --attn-name-decoder causal_eva --decoder-attn-window-size 128 --decoder-attn-causal --decoder-attn-adaptive-proj qk --decoder-attn-chunk-size 8 --decoder-attn-use-t5-rpe

Generation und Bewertung

Für Generation und Bewertung passieren einfach Argument -i true wenn Sie main.sh nur aufrufen, um das Inferenzverfahren nur durchzuführen. Der Checkpoint-Pfad kann als -c <your-ckpt-path> angegeben werden. Zum Beispiel,

 # Machine Translation
CUDA_VISIBLE_DEVICES=0 bash main.sh -i true -c < your-possibly-avg-checkpoint.pt > -p < dir-of-your-bin-data > -d wmt -g 1

# Autoregressive Language Modeling
CUDA_VISIBLE_DEVICES=0 bash main.sh -i true -c < your-checkpoint_last.pt > -p < dir-of-your-bin-data > -d wikitext103 -g 1

Vorausgebildete Modelle

Wir bieten auch geschulte EVA -Modellkontrollpunkte in OneDrive für maschinelle Übersetzungs- und Sprachmodellierungsaufgaben:

• Wikitext103-EVA-16Layers-LM
• Wikitext103-EVA-32Layers-LM
• WMT14ende-eva-e32_c8-mt
• WMT14ende-eva-e8_c8-mt

Zitat

 @inproceedings { zheng2023efficient ,
  title = { Efficient Attention via Control Variates } ,
  author = { Lin Zheng and Jianbo Yuan and Chong Wang and Lingpeng Kong } ,
  booktitle = { International Conference on Learning Representations } ,
  year = { 2023 } ,
  url = { https://openreview.net/forum?id=G-uNfHKrj46 }
}

 @inproceedings { zheng2022linear ,
  title = { Linear complexity randomized self-attention mechanism } ,
  author = { Lin Zheng and Chong Wang and Lingpeng Kong } ,
  booktitle = { International Conference on Machine Learning } ,
  pages = { 27011--27041 } ,
  year = { 2022 } ,
  organization = { PMLR }
}