[论文解读] Hungry Hungry Hippos: Towards Language Modeling with State Space Models

Daniel Y. Fu, Tri Dao+4 more|arXiv (Cornell University)|Dec 28, 2022

Topic Modeling被引用 116

一句话总结

本文介绍 H3，一种新的状态空间模型层，旨在弥合语言建模中对注意力的表达能力差距，以及 FlashConv，一种硬件感知的基于 FFT 的加速方法，使在长序列上对 SSM 的训练和推理更高效。

ABSTRACT

State space models (SSMs) have demonstrated state-of-the-art sequence modeling performance in some modalities, but underperform attention in language modeling. Moreover, despite scaling nearly linearly in sequence length instead of quadratically, SSMs are still slower than Transformers due to poor hardware utilization. In this paper, we make progress on understanding the expressivity gap between SSMs and attention in language modeling, and on reducing the hardware barrier between SSMs and attention. First, we use synthetic language modeling tasks to understand the gap between SSMs and attention. We find that existing SSMs struggle with two capabilities: recalling earlier tokens in the sequence and comparing tokens across the sequence. To understand the impact on language modeling, we propose a new SSM layer, H3, that is explicitly designed for these abilities. H3 matches attention on the synthetic languages and comes within 0.4 PPL of Transformers on OpenWebText. Furthermore, a hybrid 125M-parameter H3-attention model that retains two attention layers surprisingly outperforms Transformers on OpenWebText by 1.0 PPL. Next, to improve the efficiency of training SSMs on modern hardware, we propose FlashConv. FlashConv uses a fused block FFT algorithm to improve efficiency on sequences up to 8K, and introduces a novel state passing algorithm that exploits the recurrent properties of SSMs to scale to longer sequences. FlashConv yields 2$ imes$ speedup on the long-range arena benchmark and allows hybrid language models to generate text 2.4$ imes$ faster than Transformers. Using FlashConv, we scale hybrid H3-attention language models up to 2.7B parameters on the Pile and find promising initial results, achieving lower perplexity than Transformers and outperforming Transformers in zero- and few-shot learning on a majority of tasks in the SuperGLUE benchmark.

研究动机与目标

• 评估用于语言建模的状态空间模型（SSM）与注意力之间的表达能力差距，采用合成任务。
• 设计一个能够回忆过去的标记并在序列中比较标记的 SSM 层（H3）。
• 提高 SSM 的硬件效率，使其在标准基准测试上达到或超过 Transformer 的性能。
• 展示基于 FlashConv 的大规模 SSM 模型的可扩展性、训练与生成速度提升。

提出的方法

• 引入 H3，这是一个带有移位和对角矩阵的两层 SSM，并通过乘法输入投影实现标记回忆和跨标记比较。
• 在合成语言任务上证明 H3 可达到与注意力相匹配的表现，在 OpenWebText 上与 Transformer 的差距几乎缩小到 0.4 个困惑度。
• 提出一个混合的 H3–注意力模型，包含两层注意力，在 OpenWebText 上的性能比 Transformer 高出 1.0 个困惑度。
• 开发 FlashConv，一种与块 FFT 和状态传递算法相结合的基于 FFT 的卷积融合方法，以将 SSM 扩展到长序列和更长的模型。
• 表明 FlashConv 在长序列上可带来高达 2× 的加速，并使基于 SSM 的文本生成比 Transformer 快 2.4×。
• 将混合 H3–注意力模型规模扩展到 Pile 上的多达 27 亿参数，并报告困惑度和 SuperGLUE 性能。
• 提供推理吞吐量对比，显示混合 H3 在同等规模下比 Transformer 快最多 2.4×。

实验结果

研究问题

• RQ1 SSMs 是否能够在合成任务上展现与注意力相同的表达能力，以用于语言建模？
• RQ2 哪些体系结构变化（如移位/对角 SSM 和乘法交互）能够使 SSM 回忆过去的标记并在序列中进行比较？
• RQ3 像 FlashConv 这样的面向硬件的算法如何在现代加速器上缩小 SSM 与注意力之间的运行时差距？
• RQ4 混合 H3–注意力模型在标准基准上是否能达到与 Transformer 相媲美的困惑度以及零-shot/少量-shot 表现？
• RQ5 FlashConv 是否能够将基于 SSM 的模型扩展到十亿参数量级，同时保持可接受的训练与推理速度提升？

主要发现

• H3 在合成语言任务上可达到与注意力相匹配的表现，并在 OpenWebText 上与 Transformer 的差距几乎缩小到 0.4 个困惑度。
• 一个包含两层注意力的混合 H3–注意力模型在 OpenWebText 上比 Transformer 高出 1.0 个困惑度。
• 在 Pile 上训练的大小为 1.25 亿至 27 亿参数的混合 H3–注意力模型，在困惑度上优于 Transformer 基线，并在 SuperGLUE 零-shot/少量-shot 表现上表现出色。
• FlashConv 在长序列上可实现最高 2× 的加速，并使基于 SSM 的生成比 Transformer 快 2.4×。
• 扩展到 Pile 上 27 亿参数的基于 SSM 的模型在若干设定下获得比 Transformer 更低的困惑度，并显示出与零-shot/少量-shot 结果的竞争力。
• 与 Transformer 基线和其他 SSM 方法进行对比时，FlashConv 在 Long Range Arena 上实现了最先进的速度。