LLM2D
基于快速张量卷积和分块对角化的高效状态空间模型
EfficientState Space Model viaFast Tensor Convolutionand Block Diagonalization
作者: Tongyi Liang, Han-Xiong Li
发布日期: 10/8/2024
arXiv ID: oai:arXiv.org:2402.15290v3

摘要

现有的模型在对长序列建模时,在平衡性能和计算效率方面遇到了瓶颈。虽然状态空间模型(SSM)在处理长序列任务方面取得了显著成功,但它仍然面临着参数数量庞大的问题。为了进一步提高SSM的效率,我们提出了一种基于多输入多输出SSM的新型状态空间层,称为高效SSM(eSSM)。我们的eSSM建立在多输入多输出(MIMO)SSM的卷积表示之上。我们提出了一系列有效的策略来提高计算效率。系统矩阵的对角化首先解耦了原始系统。然后,基于快速傅里叶变换提出了一种快速张量卷积。此外,SSM的块对角化进一步减少了模型参数,提高了模型灵活性。大量的实验结果表明,所提出的模型在多个数据库上的性能与最先进的模型(如S4)相当,并且明显优于Transformers和LSTM。在模型效率基准测试中,eSSM的参数仅为LSTM的12.89%和Mamba的13.24%。eSSM的训练速度比LSTM快3.94倍,比Mamba快1.35倍。代码可在以下地址获取:\href{https://github.com/leonty1/essm}{https://github.com/leonty1/essm}。