Attention Sink下載 - Attention Sink器代碼下載

Attention Sink

Ai源碼

1.0.0

下載

當注意力下沉在語言模型中時：一種經驗觀點

[arxiv]

TL，DR：

在這項工作中，我們研究了LM預訓練中的優化，數據分佈，損失函數和模型結構如何影響注意力下降的出現。

設定

我們在具有40GB內存的A100 GPU上運行所有實驗。我們遵循Tinyllama和Regmix準備環境：

安裝

我們希望您安裝了CUDA 11.8。

每晚安裝Pytorch。

pip install --index-url https://download.pytorch.org/whl/nightly/cu118 --pre ' torch>=2.1.0dev '

從源構建Xformer

注意：截至2023/09/02，Xformers不能為火炬2.1提供預製的二進製文件。您必須從源頭構建它。

pip uninstall ninja -y && pip install ninja -U
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

安裝閃存注意力2和其他融合操作員：

git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention
python setup.py install
cd csrc/rotary && pip install .
cd ../layer_norm && pip install .
cd ../xentropy && pip install .
cd ../.. && rm -rf flash-attention

安裝剩餘的依賴項

 pip install -r requirements.txt tokenizers sentencepiece

安裝其他依賴關係。構建Xformers/flash注意力可能需要> = 5分鐘。不用擔心該過程似乎停滯不前或終端打印出許多警告。

預處理

在訓練模型之前，您需要預處理數據。我們提供易於使用的腳本，用於預處理數據。您可以使用以下命令來預處理數據：

 cd preprocess
bash run_preprocess.sh

默認情況下，您將首先從擁抱面下載regmix-data-sample ，然後預處理數據。 JSONL數據將保存在preprocess/sail/regmix-data-sample目錄中，並且預處理數據將保存在datasets/lit_dataset_regmix目錄中。這包括大約5B令牌，並佔用約20GB的磁盤空間。

WANDB集成

默認情況下，我們使用WANDB收集數據，以避免在本地機器上節省大量的小型型號和日誌。如果要使用wandb，則需要在wandb上創建一個帳戶並獲取API密鑰。然後，您應該在scripts/*.sh ：中設置以下環境變量：

 # wandb project name, entity, and API key
export WANDB_PROJECT=YOUR_PROJECT_NAME
export WANDB_ENTITY=YOUR_WANDB_ENTITY
export WANDB_API_KEY=YOUR_WANDB_API_KEY

注意開源LMS

注意 - 鏈接

我們首先對開源LMS進行全面研究，包括Llama2，Llama3，Mismtral，GPT2，Pythia和Opt。我們提供腳本來計算上述LMS的注意下沉指標：

python eval_sink_open_source.py

我們還提供了腳本來計算隱藏狀態，密鑰和值的L2-norm：

python eval_activations_open_source.py

我們注意到，我們在一個腳本中包括所有30個開源LMS的評估，這可能需要很長時間才能下載模型。請在eval_sink_open_source.py中修改57-63行或eval_activations_open_source.py中的第104-110行，以評估特定的LMS。

訓練前駱駝模型

我們在5B令牌數據上預先培訓一系列的Llama模型，以研究LM訓練前訓練中的優化，數據，損耗函數和模型成績如何影響注意力下水道。

不同的設置可能會導致不同的內存使用情況，並且服務設置與閃光注意力不符。我們可能具有不同的代碼庫，用於在文件夾lit_gpt中進行預訓練和推斷。我們歡迎社區的貢獻使代碼更加簡潔，效率。

除非額外的說明，否則遵循的實驗需要至少40 GB內存的4個GPU。請修改configs/*.yaml中的微批量尺寸。

運行默認設置：

bash scripts/run_defaults_60m.sh

最終檢查點保存在checkpoints/tinyllama_60M/iter-020000-ckpt.pth 。

我們提供所有腳本以在主要論文中重現我們的實驗結果。

優化：請參閱scripts/optimization.md 。外賣：1。有效訓練LMS之後，注意下沉出現。 2。注意下沉率在接受較小學習率的LMS中顯而易見。
數據分佈：請參閱scripts/data_distribution.md 。外賣：1。在LMS接受足夠的培訓數據訓練後，注意下水道出現。 2。如果修改數據分佈，則可以將注意點轉移到其他位置，而不是第一個令牌。
損失功能：請參閱scripts/loss_function.md 。外賣：1。重量衰減鼓勵注意力下降。 2。使用前綴語言建模，注意下沉出現在前綴令牌中，而不是僅第一個令牌。 3。隨著窗戶關注的轉移，注意下沉出現在絕對上，而不是相對的第一個令牌上。較小的窗戶大小可防止注意力下降。

損失功能

模型體系結構：請參閱scripts/model_architecture.md 。外賣：1。位置嵌入，FFN設計，LN位置和多頭設計不會影響注意流的出現。 2。注意水槽的作用更像是關鍵偏見，存儲額外的關注，同時不促進價值計算。 3。放鬆令牌對注意力評分的內在依賴性時，注意力下降不會在LMS中出現。

Model2

Model1

最後，我們將模型大小擴展到1B參數，發現仍然沒有註意力下沉和大量激活。默認情況下，我們使用8個GPU運行以下腳本

bash scripts/run_defaults_1b.sh
bash scripts/run_sigmoid_1b.sh

Bibtex

如果您發現此項目在您的研究中有用，請考慮引用我們的論文：

 @article{gu2024attention,
  title={When Attention Sink Emerges in Language Models: An Empirical View},
  author={Gu, Xiangming and Pang, Tianyu and Du, Chao and Liu, Qian and Zhang, Fengzhuo and Du, Cunxiao and Wang, Ye and Lin, Min},
  journal={arXiv preprint arXiv:2410.10781},
  year={2024}
}