[论文解读] Efficient Long Sequence Modeling via State Space Augmented Transformer

Simiao Zuo, Xiaodong Liu+5 more|arXiv (Cornell University)|Dec 15, 2022

Topic Modeling被引用 7

一句话总结

本文提出SPADE，一种混合Transformer架构，通过在第一层集成状态空间模型（SSM）来捕捉全局依赖关系，同时在后续层使用高效的局部注意力机制进行细粒度的局部建模，从而提升长序列建模能力。该方法实现线性复杂度，在Long Range Arena和语言建模任务上均优于基线模型，并能有效扩展至大规模预训练模型，适用于自然语言理解（NLU）和生成任务。

ABSTRACT

Transformer models have achieved superior performance in various natural language processing tasks. However, the quadratic computational cost of the attention mechanism limits its practicality for long sequences. There are existing attention variants that improve the computational efficiency, but they have limited ability to effectively compute global information. In parallel to Transformer models, state space models (SSMs) are tailored for long sequences, but they are not flexible enough to capture complicated local information. We propose SPADE, short for $\underline{ extbf{S}}$tate s$\underline{ extbf{P}}$ace $\underline{ extbf{A}}$ugmente$\underline{ extbf{D}}$ Transform$\underline{ extbf{E}}$r. Specifically, we augment a SSM into the bottom layer of SPADE, and we employ efficient local attention methods for the other layers. The SSM augments global information, which complements the lack of long-range dependency issue in local attention methods. Experimental results on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method. To further demonstrate the scalability of SPADE, we pre-train large encoder-decoder models and present fine-tuning results on natural language understanding and natural language generation tasks.

研究动机与目标

• 解决标准Transformer在长序列建模中计算成本呈二次方增长的问题。
• 克服局部注意力机制在捕捉长距离依赖关系方面的局限性。
• 将状态空间模型（SSMs）的全局建模优势与高效注意力的局部归纳偏置相结合。
• 设计一种可扩展、高效且有效的架构，用于长上下文NLP任务。
• 通过在下游NLU和生成任务上的预训练与微调，验证该方法的有效性。

提出的方法

• SPADE采用多层Transformer结构，其中第一层为状态空间模型（SSM），用于向序列注入粗粒度的全局信息。
• SSM层后接多层高效局部注意力（如窗口注意力）以细化局部依赖关系。
• SSM以线性时间与空间复杂度运行，支持长序列的高效处理。
• 该方法结合了SSM的全局归纳偏置与注意力的局部建模灵活性，互补其各自缺陷。
• 该架构支持自回归与非自回归序列建模，包括编码器-解码器设置。
• 模型采用标准训练目标进行训练，并在GLUE与摘要生成基准上进行微调。

实验结果

研究问题

• RQ1 结合SSM与局部注意力的混合架构是否能提升Transformer在长距离依赖建模方面的能力？
• RQ2 在Long Range Arena等长序列基准上，将SSM置于第一层是否能提升性能？
• RQ3 SPADE能否在保持线性复杂度的同时，超越二次注意力与稀疏注意力的基线模型？
• RQ4 当在大规模语言建模上进行预训练并在下游任务上微调时，SPADE的可扩展性如何？
• RQ5 SSM层与局部注意力层在整体性能提升中的贡献分别是什么？

主要发现

• SPADE在Long Range Arena基准上达到最先进性能，优于采用全连接、稀疏或低秩注意力的模型。
• 在自回归语言建模任务中，SPADE显著快于原始Transformer，且困惑度得分更优。
• 模型展现出强大的可扩展性，其预训练的编码器-解码器版本在GLUE与摘要生成任务上均取得优异结果。
• 消融实验表明，SSM层对捕捉长距离依赖关系至关重要，而局部注意力层则用于细化局部模式。
• SPADE将8k长度序列的GPU显存消耗降低至80GB以下，使长上下文建模成为现实。
• 即使在结构归纳偏置有限的情况下，该方法仍保持高性能，优于完全依赖位置编码或稀疏模式的模型。