First Attentions Last：更高效地利用首次注意力进行 Transformer 训练
机器学习
2025-12-09 v2
摘要
随着以十亿规模 Transformer 的训练变得越来越常见，采用多个分布式 GPU 以及并行训练方法已成为标准做法。然而，现有的 Transformer 设计在通信开销方面存在显著问题，尤其是在张量并行（TP）中，每个块的 MHA-MLP 连接需要进行 all-reduce 通信。通过我们的调查，我们表明 MHA-MLP 连接可以被省略以提高效率，而第一层的注意力输出可作为被省略连接的替代信号。基于我们的观察，我们提出了 FAL（First Attentions Last），一种高效的 Transformer 架构，将第一个 MHA 输出重定向到后续层的 MLP 输入，从而消除每个块的 MHA-MLP 连接，消除 all-reduce 通信，并使 MHA 和 MLP 能在单个 GPU 上并行执行。我们还引入了 FAL+，通过将归一化后的首次注意力输出添加到后续层的 MHA 输出中，以增强模型质量。我们的评估表明，FAL 在多 GPU 训练时间上可缩短最高 44%，提高单 GPU 吞吐量最高可达 1.18 倍，并在 perplexity 上优于基线 GPT。FAL+ 在不增加训练时间的情况下实现更低的 perplexity。代码已公开：https://github.com/CASL-KU/FAL
引用
@article{arxiv.2510.14614,
title = {First Attentions Last: Better Exploiting First Attentions for Efficient Transformer Training},
author = {Gyudong Kim and Hyukju Na and Jin Hyeon Kim and Hyunsung Jang and Jaemin Park and Jaegi Hwang and Namkoo Ha and Seungryong Kim and Young Geun Kim},
journal= {arXiv preprint arXiv:2510.14614},
year = {2025}
}