Hydra：通过广义矩阵混合器实现双向状态空间模型
机器学习
2024-07-16 v1 人工智能
摘要
广泛的序列模型构建于以 Transformer 为模型的框架，包含交替的序列混合器和通道混合器层。本文研究了序列混合器的统一矩阵混合器视角，可概念化为对输入序列的线性映射。该框架涵盖了广泛知名序列模型，包括 Transformer 的自注意力机制以及最近的强大替代方案，如结构状态空间模型（SSMs），并允许通过其结构矩阵类的属性来理解下游特征，如效率和表达性。我们识别出一种关键的参数化轴，称为序列对齐，增加了矩阵混合器的灵活性和性能，提供了对 Transformer 以及最近的 SSMs 如 Mamba 强大性能的洞见。此外，矩阵混合器框架提供了开发具有所需属性的序列混合器的系统方法，使我们能够开发出几个新的亚二次序列模型。特别是，我们提出了 Mamba 模型的一种自然的双向扩展（Hydra），参数化为准可分离矩阵混合器，在非因果任务上 outperform 了包括 Transformer 在内的其他序列模型。作为注意力层的即插即用替换，Hydra 在 GLUE 基准测试中超过 BERT 0.8 分，在 ImageNet 上超过 ViT 2% 的 Top-1 准确率。
关键词
引用
@article{arxiv.2407.09941,
title = {Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers},
author = {Sukjun Hwang and Aakash Lahoti and Tri Dao and Albert Gu},
journal= {arXiv preprint arXiv:2407.09941},
year = {2024}
}