面向遗忘 Transformer 的自适应计算剪枝
机器学习
2025-08-13 v2 人工智能
计算与语言
摘要
最近提出的遗忘 Transformer（FoX）在 softmax 注意力中引入了遗忘门，显示出在标准基于 RoPE 的 Transformer 上表现始终更好或持平。值得注意的是，FoX 中的许多注意力头倾向于快速遗忘，导致其在每个时间步的输出主要依赖于局部上下文。基于此观察，我们提出了面向 FoX 的自适应计算剪枝（ACP），该方法动态剪枝由遗忘门导致输入输出依赖性强衰减的计算。具体而言，我们的方法通过动态设置的剪枝阈值实现可证明的安全剪枝，保证被剪枝后的注意力权重可忽略。我们将 ACP 应用于 FoX 的语言模型预训练中，证明其在不同模型规模和上下文长度下，能将 softmax 注意力中的 FLOPs 和内存访问次数约降低 70%，从而实现注意力运行时间约降低 50% 至 70%（或加速 2 倍至 3 倍），且端到端训练吞吐量提升约 10% 至 40%。此外，更长的上下文长度会带来更大的计算节省。所有这些速度提升均无需任何性能退化。我们的代码已公开于 https://github.com/zhixuan-lin/forgetting-transformer。
引用
@article{arxiv.2504.06949,
title = {Adaptive Computation Pruning for the Forgetting Transformer},
author = {Zhixuan Lin and Johan Obando-Ceron and Xu Owen He and Aaron Courville},
journal= {arXiv preprint arXiv:2504.06949},
year = {2025}
}
备注
Published as a conference paper at COLM 2025