Hydra模型:超越Transformer的高效双向序列建模新框架

Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers

摘要

本文探讨了一种基于Transformer框架的序列模型统一矩阵混合器视图,该框架包括交替的序列混合器和通道混合器层。论文提出了一种序列混合器的线性映射概念,涵盖了包括Transformer的自注意力机制和最近的结构化状态空间模型(SSM)在内的多种序列模型。通过分析这些模型的结构化矩阵类属性,论文揭示了它们在效率和表达能力上的下游特性。此外,论文引入了一个名为序列对齐的关键矩阵参数化轴,增强了矩阵混合器的灵活性和性能,并对Transformer和最近的SSM如Mamba的强大性能提供了见解。论文还提出了一种新的次二次序列模型,特别是Mamba模型的自然双向扩展(Hydra),其在非因果任务上表现优于包括Transformer在内的其他序列模型。Hydra作为注意力层的直接替代,在GLUE基准测试中超过BERT 0.8分,在ImageNet上超过ViT 2%的Top-1准确率。

原理

论文的核心在于提出了一种新的矩阵混合器框架,该框架将序列混合器视为输入序列上的线性映射。这种框架不仅包括了Transformer的自注意力机制,还涵盖了如MLP-Mixer和FNet等模型,以及结构化状态空间模型(SSM)。通过分析这些模型的矩阵结构特性,论文揭示了它们在计算效率和表达能力上的关键差异。特别是,论文引入了一个名为序列对齐的新属性,该属性通过数据依赖性和可扩展性增强了序列模型的性能。此外,论文通过系统地探索和分类现有的序列模型,提出了新的序列模型,如基于Vandermonde和Cauchy矩阵的混合器,这些模型展示了框架在设计新序列混合器方面的灵活性。

流程

Hydra模型的工作流程基于其矩阵混合器框架,特别是采用了准分离矩阵混合器。准分离矩阵是一种具有重要特性的基本矩阵结构,适用于序列建模。它们概括了线性注意力模型中发现的低秩矩阵混合器和状态空间模型中使用的半分离矩阵。Hydra通过系统的方法,将半分离矩阵自然地扩展为双向模型,解决了状态空间模型在非因果设置中的主要限制。Hydra保持了SSM的强大性能和线性时间计算效率,得益于结构化矩阵乘法算法。与之前通过元素级加法、Hadamard积或串联等临时方法使SSM双向化的尝试不同,Hydra提供了一种更连贯和理论基础更强的进步。

应用

Hydra模型的应用前景广泛,特别是在需要高效处理长序列和复杂数据依赖性的任务中。由于其在非因果任务上的优越性能,Hydra可以作为现有注意力机制的直接替代,广泛应用于自然语言处理和计算机视觉领域。此外,Hydra的矩阵混合器框架为未来序列模型的设计和优化提供了新的视角和工具,特别是在需要平衡计算效率和模型表达能力的场景中。