学习Transformer程序

本仓库包含了我们论文《学习Transformer程序》的代码。该代码可用于训练经过修改的Transformer来解决任务，然后将其转换为人类可读的Python程序。仓库还包含了一些我们为论文中描述的任务学习到的示例程序。更多详细信息请参阅我们的论文。

快速链接

设置
学习程序
- 训练
- 转换为代码
示例程序
问题？
引用

设置

安装PyTorch，然后安装其余依赖：pip install -r requirements.txt。本代码使用Python 3.8和PyTorch 1.13.1版本进行测试。

在我们的自然语言处理任务实验中，我们使用300维的预训练GloVe词嵌入来初始化词嵌入，可以在这里下载（Common Crawl，区分大小写）：

mkdir data
wget https://huggingface.co/stanfordnlp/glove/resolve/main/glove.840B.300d.zip -P data/
unzip data/glove.840B.300d.zip

学习程序

训练

学习Transformer程序的代码可以在src/run.py中找到。例如，以下命令将为sort任务训练一个Transformer程序，使用两层，每层四个分类注意力头，以及独热输入嵌入：

python src/run.py \
     --dataset "sort" \
     --vocab_size 8 \
     --dataset_size 10000 \
     --min_length 1 \
     --max_length 8 \
     --n_epochs 250 \
     --batch_size 512 \
     --lr "5e-2" \
     --n_layers 2 \
     --n_heads_cat 4 \
     --n_heads_num 0 \
     --n_cat_mlps 1 \
     --n_num_mlps 0 \
     --one_hot_embed \
     --count_only \
     --seed 0 \
     --save \
     --save_code \
     --output_dir "output/sort";

这个命令将为CoNLL 2003命名实体识别任务训练一个Transformer程序，学习由四个32维分类变量组成的输入嵌入：

python src/run.py \
     --dataset "conll_ner" \
     --vocab_size 10000 \
     --min_length 1 \
     --max_length 32 \
     --n_epochs 50 \
     --batch_size 32 \
     --lr "5e-2" \
     --n_vars_cat 4 \
     --d_var 32 \
     --n_layers 2 \
     --n_heads_cat 4 \
     --n_heads_num 0 \
     --n_cat_mlps 1 \
     --n_num_mlps 0 \
     --mlp_vars_in 2 \
     --count_only \
     --seed 0 \
     --replace_numbers 1 \
     --glove_embeddings "data/glove.840B.300d.txt" \
     --do_glove 1 \
     --save \
     --save_code \
     --output_dir "output/conll";

请查看src/run.py了解所有可能的参数。训练数据将被生成（对于RASP任务）或从Hugging Face Datasets下载；支持的数据集请参见src/utils/data_utils.py。 scripts目录包含了使用论文中实验设置训练Transformer程序和标准Transformer的脚本。

转换为代码

在训练脚本中使用--save_code标志，以在训练结束时将模型转换为Python程序。要转换已经训练好的模型，请使用src/decompile.py。例如，

python src/decompile.py --path output/sort/ --output_dir programs/sort/

output/sort/应该是训练运行的输出目录。

示例程序

programs目录包含了所有RASP任务的小规模版本以及命名实体识别的示例程序。每个程序定义了一个名为run的函数，该函数接受一个标记序列作为输入，并返回预测标签列表。例如：

>>> from programs.rasp.sort import sort
>>> sort.run(["<s>", "3", "1", "4", "2", "4", "0", "</s>"])
['<s>', '0', '1', '2', '3', '4', '4', '</s>']

programs/rasp包含了每个任务表现最佳的程序，同时使用分类和数值注意力头。 programs/rasp_categorical_only包含了仅使用分类变量的表现最佳程序。 programs/conll_ner包含了一个命名实体识别程序。

问题？

如果您对代码或论文有任何问题，请发邮件给Dan（dfriedman@cs.princeton.edu）或开一个issue。

引用

@inproceedings{
    friedman2023learning,
    title={Learning Transformer Programs},
    author={Dan Friedman and Alexander Wettig and Danqi Chen},
    booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
    year={2023},
    url={https://openreview.net/forum?id=Pe9WxkN8Ff}
}