RecurFormer：并非所有Transformer注意力头都需要自注意力
计算与语言
2024-10-18 v1 人工智能
机器学习
摘要
Transformer 架构的大型语言模型 (LLM) 在建模复杂语言模式方面表现优异，但在推理阶段尤其是面对长文本输入时，由于注意力机制的内存开销，面临显著的计算成本问题。我们观察到，某些注意力头呈现注意力权重集中于查询 token 附近 tokens 的分布，称为 recency aware (复近感知)，其关注局部和短程依赖。基于此洞察，我们提出 RecurFormer，一种新型架构，将这些注意力头替换为线性循环神经网络 (RNN)，具体为 Mamba 架构。该替换在不驱逐 tokens 的情况下显著降低 cache 大小，从而维持生成质量。RecurFormer 通过剩余的注意力头保留建模长程依赖的能力，并允许使用持续训练方式复用预训练 Transformer-based LLM 的权重。实验表明，RecurFormer 在保持原模型性能的同时，显著提升了推理效率。我们的方案为 Transformer-based LLM 推理的计算挑战提供了实用解决方案，对于处理长输入的任务尤为吸引人。
关键词
引用
@article{arxiv.2410.12850,
title = {RecurFormer: Not All Transformer Heads Need Self-Attention},
author = {Ruiqing Yan and Linghan Zheng and Xingbo Du and Han Zou and Yufeng Guo and Jianfei Yang},
journal= {arXiv preprint arXiv:2410.12850},
year = {2024}
}