CAST:利用代理令牌聚类自注意力以实现高效 Transformer
机器学习
2024-02-07 v1
摘要
Transformer 架构已被证明是多种任务的强大工具。它基于自注意力机制,这是一种本质上计算昂贵的操作,具有二次计算复杂度:内存使用和计算时间随输入序列长度的增加而呈二次方增长,从而限制了 Transformer 的应用。在本工作中,我们提出了一种新颖的利用代理令牌聚类自注意力机制(CAST),以优化注意力计算并实现高效 Transformer。CAST 利用可学习的代理令牌构建聚类亲和矩阵,用于对输入序列进行聚类并生成新的聚类摘要。然后将每个聚类内部的自注意力与其他聚类的聚类摘要相结合,从而实现整个输入序列的信息流动。CAST 通过将复杂度从降低到(其中是序列长度,是根据聚类数量和每个聚类的样本数确定的常数)提高了效率。我们表明,CAST 在长程序列建模任务上的表现优于或相当于基线 Transformer,同时在时间和内存效率方面比其他高效 Transformer 取得了更高的结果。
引用
@article{arxiv.2402.04239,
title = {CAST: Clustering Self-Attention using Surrogate Tokens for Efficient Transformers},
author = {Adjorn van Engelenhoven and Nicola Strisciuglio and Estefanía Talavera},
journal= {arXiv preprint arXiv:2402.04239},
year = {2024}
}