使用正确的嵌入方式，Transformer 可以进行算术运算！arXiv 论文链接

联合项目成员：Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild 和 Tom Goldstein

本仓库包含复现我们研究的代码。它是语言模型训练框架 cramming 的一个分支，经过修改以适用于下一个标记预测任务。

我们在 abacus.py 中提供了算盘嵌入的独立实现。

引用我们的工作

如需引用我们的工作，请使用以下 bibtex：

@article{mcleish2024transformers,
    title={Transformers Can Do Arithmetic with the Right Embeddings}, 
    author={Sean McLeish and Arpit Bansal and Alex Stein and Neel Jain and John Kirchenbauer and Brian R. Bartoldson and Bhavya Kailkhura and Abhinav Bhatele and Jonas Geiping and Avi Schwarzschild and Tom Goldstein},
    journal={arXiv preprint arXiv:2405.17399},
    year={2024}
}

入门指南

我们使用 Python 3.10.4 进行开发，安装步骤如下：

git clone git@github.com:mcleish7/arithmetic.git
cd arithmetic
pip install .

在某些机器上，您可能需要运行：

pip install multiprocess -U
pip install dill -U
pip install apache-beam -U

算术运算

数据集

我们在 Google Drive 上发布了压缩格式的数据集。我们建议您使用压缩版本，直到正确放置在文件系统中。

或者，您可以使用 create_data_split.py 和 shells/generate_and_tokenize_data.sh 中的命令创建自己的数据集。

文件结构

我们建议在 arithmetic 目录内创建另一个名为 cramming-data 的目录。这里将存储模型、日志和数据。

您可以将 cramming 基础目录路径导出到 .bashrc，或手动替换提供的 shell 脚本中的 $cramming_base_dir。

cd arithmetic
mkdir cramming-data
echo 'export cramming_base_dir=MY_BASE_DIR' >> ~/.bashrc
source ~/.bashrc

例如，这可能看起来像：echo 'export cramming_base_dir=~/arithmetic/cramming-data' >> ~/.bashrc

例如，我们的文件系统结构如下：

cramming-generative
└── cramming-data
    ├── addition-train-one
    │    ├── pretrain/<DATE>/<TIME>
    │    │    ├── .hydra
    │    │    │   ├── config.yaml
    │    │    │   ├── hydra.yaml
    │    │    │   └── overrides.yaml
    │    │    └── addition-train-one_pretrain.log
    │    ├── checkpoints/FINAL_<LOSS_VAL>
    │    │    ├── model_config.json
    │    │    ├── model.safetensors
    │    │    └── state_dict.pth
    │    └── downstream
    └── data
        └── arithmetic_data
            ├── +_grid_eval_dataset_reverse_all_tokenized
            └── ... other datasets ...

训练

示例命令在 shells 目录中，按任务组织。

一些命令的解释

在损失中给予样本而非标记相等的重要性：arch.loss_reduction=none
将循环块中的梯度除以循环次数：arch.throttle=True
在等号之前进行掩码：arch.mask_before_equals=True
循环块内的跳跃连接：arch.forward_only_model_with_skip=True
多 GPU：将 python 替换为 torchrun --nproc_per_node=<GPU 数量> --standalone ，并添加 impl.fullgraph=false

位置嵌入：

绝对位置

学习型：arch.embedding.pos_embedding=learned
算盘型：arch.embedding.pos_embedding=abacus

如果您想要算盘中的最大 k 值更大：arch.embedding.max_abacus_len=100，默认值为 100。算盘也以独立方式在 abacus.py 中实现。

相对位置

NoPE：arch.embedding.pos_embedding=None
FIRE：arch.embedding.pos_embedding=None arch.attention.type="self-attention" arch.attention.rotary_embedding="fire"
随机化 FIRE：例如 arch.embedding.pos_embedding=None arch.attention.type="self-attention" arch.attention.rotary_embedding="fire" arch.attention.max_length=128，默认情况下 arch.attention.max_length=0，因此将其设置为长于最大序列长度会在嵌入中引入一些随机性。
RoPE：arch.attention.type="self-attention" arch.attention.rotary_embedding=true

检查点

我们实现了单 GPU 训练检查点，使用方法如下： impl.save_every_n_minutes=60 impl.save_intermediate_model_name='last' 这将每 60 分钟保存一个名为 'last' 的检查点

注意：此功能尚未完全测试多 GPU 情况。我们目前也无法继续训练已用完全部预算的模型。

WandB

您可以将运行记录到您的 weights&biases 账户。只需在命令行或 cramming/config/wandb/default.yaml 中修改 wandb.entity 和 wandb.project。

测试

我们在 shells/evaluation.sh 中展示了示例。

我们在 gen_eval_script.py 中提供了一个非常基础的自动化脚本，它打印出您可能需要进一步编辑的基本命令。

加法

对于加法，我们有一个非常大的可能评估集，我们对 100x100 的网格进行网格搜索，将其分成 20 份，目的是平衡所有 20 份中的前向调用次数。然后我们对操作数长度 100->160 进行进一步评估。

乘法

我们只评估到 25x25，这在一个作业中完成。

排序

排序使用单独的评估文件 sort_eval.py，这是因为评估调用无法并行化，使评估时间更长。评估无法并行化，因为等号的位置对于一个批次来说不是固定的。我们目前在 30 个作业中评估 30x30 的网格，但可以使用这些标志减少作业数量：max_size_given, start_ind_1_given, start_ind_2_given