Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers

📄 arXiv: 2407.09941v1 📥 PDF

作者: Sukjun Hwang, Aakash Lahoti, Tri Dao, Albert Gu

分类: cs.LG, cs.AI

发布日期: 2024-07-13

💡 一句话要点

Hydra：通过广义矩阵混合器实现双向状态空间模型，显著提升非因果任务性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 双向模型 状态空间模型 矩阵混合器 序列建模 非因果任务

📋 核心要点

  1. 现有序列模型（如Transformer）计算复杂度高，难以有效处理长序列，且缺乏对双向信息的有效建模。
  2. 论文提出一种基于广义矩阵混合器的双向状态空间模型Hydra，通过准可分离矩阵实现高效的双向信息融合。
  3. 实验表明，Hydra在GLUE和ImageNet等基准测试中优于Transformer和BERT等模型，验证了其有效性。

📝 摘要（中文）

本文研究了一种统一的矩阵混合器视角，用于观察序列混合器，该混合器可以概念化为输入序列上的线性映射。该框架涵盖了广泛的序列模型，包括Transformer的自注意力以及最近的强大替代方案，如结构化状态空间模型（SSM）。通过结构化矩阵类的属性，可以理解效率和表达性等下游特性。本文确定了一个关键的矩阵参数化轴，称为序列对齐，它增加了矩阵混合器的灵活性和性能，从而深入了解了Transformer和Mamba等最新SSM的强大性能。此外，矩阵混合器框架提供了一种系统的方法来开发具有所需属性的序列混合器，从而能够开发几种新的亚二次序列模型。特别地，本文提出了一种Mamba模型的自然双向扩展（Hydra），参数化为准可分离矩阵混合器，该模型在非因果任务上表现出优于包括Transformer在内的其他序列模型的性能。作为注意力层的直接替代品，Hydra在GLUE基准测试中优于BERT 0.8个点，在ImageNet上优于ViT 2%的Top-1准确率。

🔬 方法详解

问题定义：现有序列模型，特别是Transformer，在处理长序列时面临计算复杂度高的挑战。此外，许多任务需要利用双向信息，而传统的自回归模型难以有效捕捉这种信息。因此，需要一种既能高效处理长序列，又能有效利用双向信息的序列建模方法。

核心思路：论文的核心思路是将序列混合器视为输入序列上的线性映射，并通过矩阵混合器来统一不同的序列模型。通过分析矩阵混合器的性质，特别是序列对齐，可以设计出更高效和更具表达能力的序列模型。针对双向建模的需求，论文提出了Mamba模型的双向扩展Hydra，利用准可分离矩阵来实现高效的双向信息融合。

技术框架：Hydra模型基于矩阵混合器框架，包含序列混合器和通道混合器。序列混合器采用准可分离矩阵进行参数化，实现双向状态空间建模。整体流程包括：输入序列经过线性变换后，进入序列混合器进行双向信息融合；然后，经过通道混合器进行非线性变换；最后，输出序列。

关键创新：Hydra的关键创新在于提出了基于准可分离矩阵的双向状态空间模型，实现了高效的双向信息融合。与传统的自回归模型相比，Hydra能够同时利用过去和未来的信息，从而更有效地处理非因果任务。此外，Hydra的计算复杂度低于Transformer，使其能够处理更长的序列。

关键设计：Hydra的关键设计包括：1) 使用准可分离矩阵来参数化序列混合器，实现高效的双向状态空间建模；2) 采用与Mamba模型类似的硬件感知算法，优化计算效率；3) 通过实验选择合适的矩阵维度和激活函数，以获得最佳性能。

📊 实验亮点

Hydra在GLUE基准测试中优于BERT 0.8个点，在ImageNet图像分类任务中，Hydra的Top-1准确率比ViT高2%。这些结果表明，Hydra在各种任务中都具有优越的性能。特别是在非因果任务上，Hydra的表现明显优于其他序列模型，验证了其双向建模能力的有效性。

🎯 应用场景

Hydra模型具有广泛的应用前景，包括自然语言处理、计算机视觉和语音识别等领域。它可以应用于机器翻译、文本摘要、图像分类、目标检测和语音识别等任务。由于其高效性和双向建模能力，Hydra特别适用于需要处理长序列和利用双向信息的任务，例如视频理解和时间序列分析。

📄 摘要（原文）

A wide array of sequence models are built on a framework modeled after Transformers, comprising alternating sequence mixer and channel mixer layers. This paper studies a unifying matrix mixer view of sequence mixers that can be conceptualized as a linear map on the input sequence. This framework encompasses a broad range of well-known sequence models, including the self-attention of Transformers as well as recent strong alternatives such as structured state space models (SSMs), and allows understanding downstream characteristics such as efficiency and expressivity through properties of their structured matrix class. We identify a key axis of matrix parameterizations termed sequence alignment, which increases the flexibility and performance of matrix mixers, providing insights into the strong performance of Transformers and recent SSMs such as Mamba. Furthermore, the matrix mixer framework offers a systematic approach to developing sequence mixers with desired properties, allowing us to develop several new sub-quadratic sequence models. In particular, we propose a natural bidirectional extension of the Mamba model (Hydra), parameterized as a quasiseparable matrix mixer, which demonstrates superior performance over other sequence models including Transformers on non-causal tasks. As a drop-in replacement for attention layers, Hydra outperforms BERT by 0.8 points on the GLUE benchmark and ViT by 2% Top-1 accuracy on ImageNet.