TRIME
1.0.0
这是Zexuan Zhong，Tao Lei和Danqi Chen的EMNLP2022纸培训语言模型的存储库。
我们为语言建模提出了一个新的培训目标，该目标是将模型输出与令牌嵌入和内部内存的记忆保持一致。我们还为数据批处理和构建培训记忆设计了新颖的方法，以便我们的模型可以有效地利用远程环境和外部数据存储。
请在我们的论文中找到更多有关此作品的详细信息。
该代码基于以下要求/依赖项（我们在括号中指定了我们在实验中使用的版本）：
您可以安装此项目（基于Fairseq），如下：
pip install --editable .我们在Wikitext-103和ENWIK8数据集上进行实验。请使用get_data.sh下载并预处理数据集。
bash get_data.sh {wikitext-103 | enwik8}处理后的数据集将存储在data-bin/wikitext-103和data-bin/enwik8中。
我们显示了在Wikitext-103上运行预训练的模型的示例，模型尺寸= 247m，段长度= 3072。对于其他实验（例如，具有不同的数据集或模型），我们参考所有实验设置上的脚本run_pretrained_models.md。
Trimelm仅使用本地内存（在输入中使用令牌构建）。它可以看作是轻量化的香草langauge型号。
# download the pre-trained TrimeLM
mkdir pretrained_models ; cd pretrained_models
wget https://nlp.cs.princeton.edu/projects/trime/pretrained_models/wiki103-247M-trime.zip ;
unzip wiki103-247M-trime.zip ; rm -f wiki103-247M-trime.zip
cd ..
# run evaluation
python eval_lm-trime.py data-bin/wikitext-103
--path pretrained_models/wiki103-247M-trime/checkpoint_best.pt
--sample-break-mode complete --max-tokens 3072 --context-window 2560
--softmax-batch 1024 --gen-subset valid --fp16
--max-sentences 1 --knn-keytype last_ffn_input
--use-local --softmax-temp 1.17
# the following output is expected:
# Loss (base 2): 4.0962, Perplexity: 17.10参数：
--use-local使用本地内存指定。--softmax-temp指定计算损耗时使用的温度项。Trimelm_long在推理过程中使用本地内存和长期内存。该模型能够利用较长的上下文，尽管它经过较短的培训。
# download the pre-trained TRIME_long
mkdir pretrained_models ; cd pretrained_models
wget https://nlp.cs.princeton.edu/projects/trime/pretrained_models/wiki103-247M-trime_long.zip ;
unzip wiki103-247M-trime_long.zip ; rm -f wiki103-247M-trime_long.zip
cd ..
# run evaluation
python eval_lm-trime.py data-bin/wikitext-103
--path pretrained_models/wiki103-247M-trime_long/checkpoint_best.pt
--sample-break-mode complete --max-tokens 3072 --context-window 2560
--softmax-batch 1024 --gen-subset valid --fp16
--max-sentences 1 --knn-keytype last_ffn_input
--use-local --use-long --mem-size 12288 --softmax-temp 1.22
# the following output is expected:
# Loss (base 2): 4.0879, Perplexity: 17.01参数：
--use-long使用长期记忆。--mem-size指定本地 +长期内存的大小。Trimelm_ext使用本地内存，长期内存和外部内存。在推断期间，我们在训练集上运行模型以构建外部内存，并使用FAISS库来构建索引，以检索最接近的近距离邻居外部内存。我们还校准了在内存上的分离分布，并与KNN-LM类似（请参阅本文中的详细信息）。
我们首先下载预训练的trimelm_ext：
mkdir pretrained_models ; cd pretrained_models
wget https://nlp.cs.princeton.edu/projects/trime/pretrained_models/wiki103-247M-trime_ext.zip ;
unzip wiki103-247M-trime_ext.zip ; rm -f wiki103-247M-trime_ext.zip
cd ..然后，我们使用训练集生成外部内存（键和值），然后构建Faiss索引：
MODEL_PATH=pretrained_models/wiki103-247M-trime_ext
# generate the external memory (keys and values) using the training set
python eval_lm.py data-bin/wikitext-103
--path ${MODEL_PATH} /checkpoint_best.pt
--sample-break-mode none --max-tokens 3072
--softmax-batch 1024 --gen-subset train
--context-window 2560 --tokens-per-sample 512
--dstore-mmap ${MODEL_PATH} /dstore --knn-keytype last_ffn_input
--dstore-size 103224461
--save-knnlm-dstore --fp16 --dstore-fp16
# build Faiss index
python build_dstore.py
--dstore_mmap ${MODEL_PATH} /dstore
--dstore_size 103224461 --dimension 1024
--faiss_index ${MODEL_PATH} /knn.index
--num_keys_to_add_at_a_time 500000
--starting_point 0 --dstore_fp16 --dist ip现在，我们准备评估该模型：
MODEL_PATH=pretrained_models/wiki103-247M-trime_ext
python eval_lm-trime.py data-bin/wikitext-103
--path ${MODEL_PATH} /checkpoint_best.pt
--sample-break-mode complete --max-tokens 3072 --context-window 2560
--softmax-batch 1024 --gen-subset valid --fp16
--max-sentences 1 --knn-keytype last_ffn_input
--use-local --use-long --mem-size 12288 --softmax-temp 1.25
--use-external --dstore-filename ${MODEL_PATH} /dstore --indexfile ${MODEL_PATH} /knn.index.ip
--probe 32 --dstore-fp16 --faiss-metric-type ip --no-load-keys --k 1024
--use-interp --interp-temp 10.5 --lmbda 0.3
# the following output is expected:
# Loss (base 2): 3.9580, Perplexity: 15.54参数：
--use-external使用外部内存指定。--dstore-filename和indexfile指定数据存储和FAISS索引路径。--use-interp指定使用两个分布之间的线性插值来校准最终概率。--lmbda和--interp-temp使用线性插值时指定温度项和权重。我们列出了Wikitext-103和Enwik8上已发布的预训练模型的性能，以及他们的下载链接。
| 数据集 | 模型 | 开发 | 测试 | 超参数 |
|---|---|---|---|---|
| Wikitext-103 | 三角形 （247m，L = 3072） | 17.10 | 17.76 | --softmax-temp 1.17 |
| Wikitext-103 | trimelm_long （247m，L = 3072） | 17.01 | 17.64 | --softmax-temp 1.22 --mem-size 12288 |
| Wikitext-103 | trimelm_ext （247m，L = 3072） | 15.54 | 15.46 | --softmax-temp 1.25 --mem-size 12288 --interp-temp 10.5 --lmbda 0.3 |
| Wikitext-103 | 三角形 （150m，L = 150） | 24.45 | 25.61 | --softmax-temp 1.03 |
| Wikitext-103 | trimelm_long （150m，L = 150） | 21.76 | 22.62 | --softmax-temp 1.07 --mem-size 15000 |
| Enwik8 | 三角形 （38m，L = 512） | 1.14 | 1.12 | --softmax-temp 1.05 |
| Enwik8 | trimelm_long （38m，L = 512） | 1.08 | 1.05 | --softmax-temp 1.10 --mem-size 24576 |
我们遵循Fairseq的培训配方（例如，优化器，学习率，批次尺寸）进行训练。以不同的方式，我们使用自己的损失功能（由--criterion指定）和数据批处理方法。
我们通过使用不同的数据批处理和内存构建方法训练了三种trimelm。
--criterion trime_loss训练--criterion trime_long_loss或--criterion trime_long_loss_same_device训练--keep-order以批量连续段。trime_long_loss时，我们需要通过--train-mem-size指定内存大小（连续段的数字将为args.train_mem_size/args.tokens_per_sample ）。trime_long_loss_same_device时，我们假设所有连续的段都加载在同一GPU设备中（等效于args.mem_size == args.max_tokens ）。使用trime_long_loss_same_device比使用trime_long_loss更有效，因为它需要更少的跨GPU通信。--criterion trime_ext_loss训练--predefined-batches指定。p ，我们禁用本地内存（即，仅使用其他段的令牌来构建内存）。概率p由--cross-sent-ratio指定这是训练trimelm_ext模型的示例。您可以在Train_script中找到我们在实验中使用的所有培训脚本。
我们在4个NVIDIA RTX3090 GPU上训练模型。
# download the results of bm25 batching
wget https://nlp.cs.princeton.edu/projects/trime/bm25_batch/wiki103-l3072-batches.json -P data-bin/wikitext-103/
python train.py --task language_modeling data-bin/wikitext-103
--save-dir output/wiki103-247M-trime_ext
--arch transformer_lm_wiki103
--max-update 286000 --max-lr 1.0 --t-mult 2 --lr-period-updates 270000 --lr-scheduler cosine --lr-shrink 0.75
--warmup-updates 16000 --warmup-init-lr 1e-07 --min-lr 1e-09 --optimizer nag --lr 0.0001 --clip-norm 0.1
--criterion trime_ext_loss --max-tokens 3072 --update-freq 6 --tokens-per-sample 3072 --seed 1
--sample-break-mode none --skip-invalid-size-inputs-valid-test --ddp-backend=no_c10d --knn-keytype last_ffn_input --fp16
--ce-warmup-epoch 9 --cross-sent-ratio 0.9
--predefined-batches data-bin/wikitext-103/wiki103-l3072-batches.json重要论点：
--arch指定模型体系结构。在我们的实验中，我们一直在使用以下架构。transformer_lm_wiki103 （Wikitext-103的247m型号）transformer_lm_wiki103_150M （Wikitext-103的150m型号）transformer_lm_enwik8 （ENWIK8的38m型号）--criterion指定了计算损失值的函数。请参阅上面有关我们支持哪些功能的描述。--tokens-per-sample指定段长度。--max-tokens指定每个GPU中要加载的令牌数量。--update-freq指定梯度蓄积步骤。--ce-warmup-epoch指定在训练中使用的原始CE损失的数量。--cross-sent-ratio指定概率p以禁用本地内存。--predefined-batches指定了预定范围批处理的文件路径（我们使用BM25进行批处理段）。在使用--criterion trime_ext_loss训练Trimelm_ext模型时，我们使用BM25分数来批处理培训数据。
我们使用Pyserini库来构建BM25索引。可以通过PIP安装库。
pip install pyserini我们首先将培训集中的所有细分市场保存到.json文件中。
mkdir -p bm25/wiki103-l3072/segments
CUDA_VISIBLE_DEVICES=0 python train.py --task language_modeling
data-bin/wikitext-103
--max-tokens 6144 --tokens-per-sample 3072
--arch transformer_lm_wiki103
--output-segments-to-file bm25/wiki103-l3072/segments/segments.json
# Modify --tokens-per-sample for different segment lengths然后，我们使用Pyserini构建BM25索引。
python -m pyserini.index.lucene
--collection JsonCollection
--input bm25/wiki103-l3072/segments
--index bm25/wiki103-l3072/bm25_index
--generator DefaultLuceneDocumentGenerator --threads 1
--storePositions --storeDocvectors --storeRaw接下来，对于每个培训部分，我们使用上面构建的BM25索引搜索类似的段。
python bm25_search.py
--index_path bm25/wiki103-l3072/bm25_index/
--segments_path bm25/wiki103-l3072/segments/segments.json
--results_path bm25/wiki103-l3072/bm25_results
# Use --num_shards and --shard_id; you can parallel the computation of NN search (e.g., --num_shards 20).最后，基于检索结果，我们按组相似段创建批处理。
python bm25_make_batches.py
--results_path bm25/wiki103-l3072/bm25_results
--batch_file data-bin/wikitext-103/wiki103-l3072-batches.json输出文件wiki103-l3072-batches.json包含培训段的索引列表，相邻段可能相似。
批处理文件wiki103-l3072-batches.json可在trimelm_ext培训期间使用参数--predefined-batches 。在培训期间，我们只需从文件中以序列序列进行子列表来获得培训批次。
对于机器翻译代码和实验，请查看子目录。
如果您有与代码或论文有关的任何疑问，或者在使用代码时遇到任何问题，请随时发送电子邮件至Zexuan Zhong ([email protected])或打开问题。请尝试使用详细信息指定问题，以便我们可以更好，更快地帮助您！
如果您在研究中使用我们的代码，请引用我们的工作：
@inproceedings { zhong2022training ,
title = { Training Language Models with Memory Augmentation } ,
author = { Zhong, Zexuan and Lei, Tao and Chen, Danqi } ,
booktitle = { Empirical Methods in Natural Language Processing (EMNLP) } ,
year = { 2022 }
}我们的存储库基于Fairseq，KNNLM和自适应KNN-MT项目。我们感谢作者开源的出色代码！
