JoMA：通过 MLP 与注意力联合动力学揭示多层 Transformer 机理
机器学习
2024-03-18 v3 人工智能
计算与语言
摘要
我们提出联合 MLP/注意力（JoMA）动力学，一种用于理解多层 Transformer 架构训练过程的新型数学框架。该方法通过对 Transformer 中的自注意力层进行积分，得到仅含 MLP 层的修正动力学。JoMA 去除了以往分析中不切实际的假设（例如缺乏残差连接），并预测在存在非线性激活时，注意力先变得稀疏（以学习显著 token）再变得稠密（以学习较不显著的 token），而在线性情况下，其与现有表明注意力随时间变得稀疏的工作一致。我们利用 JoMA 定性解释了当输入 token 由潜在层次生成模型生成时，多层 Transformer 中 token 如何组合形成层次结构。在真实世界数据集（Wikitext2/Wikitext103）上训练的模型以及多种预训练模型（OPT、Pythia）上的实验验证了我们的理论发现。代码见 https://github.com/facebookresearch/luckmatters/tree/yandong3。
引用
@article{arxiv.2310.00535,
title = {JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention},
author = {Yuandong Tian and Yiping Wang and Zhenyu Zhang and Beidi Chen and Simon Du},
journal= {arXiv preprint arXiv:2310.00535},
year = {2024}
}
备注
ICLR'24 camera ready. Improve theorem 3 and theorem 4. Polish writing and add code link