[论文解读] Simplifying and Understanding State Space Models with Diagonal Linear RNNs

Ankit Gupta, Harsh Mehta+1 more|arXiv (Cornell University)|Dec 1, 2022

Topic Modeling被引用 4

一句话总结

本文提出了对角线线性 RNN（DLR），这是一种简化版的连续时间状态空间模型公式，消除了离散化的需求，从而实现更简洁、更可解释的架构。尽管结构简单，DLR 在长序列任务（包括 Long Range Arena 和原始语音分类）上的性能仍与先前的 SSM 和注意力模型相当或更优，同时保持 O(L) 的推理复杂度。

ABSTRACT

Sequence models based on linear state spaces (SSMs) have recently emerged as a promising choice of architecture for modeling long range dependencies across various modalities. However, they invariably rely on discretization of a continuous state space, which complicates their presentation and understanding. In this work, we dispose of the discretization step, and propose a model based on vanilla Diagonal Linear RNNs ($\mathrm{DLR}$). We empirically show that, despite being conceptually much simpler, $\mathrm{DLR}$ is as performant as previously-proposed SSMs on a variety of tasks and benchmarks including Long Range Arena and raw speech classification. Moreover, we characterize the expressivity of SSMs (including $\mathrm{DLR}$) and attention-based models via a suite of $13$ synthetic sequence-to-sequence tasks involving interactions over tens of thousands of tokens, ranging from simple operations, such as shifting an input sequence, to detecting co-dependent visual features over long spatial ranges in flattened images. We find that while SSMs report near-perfect performance on tasks that can be modeled via $ extit{few}$ convolutional kernels, they struggle on tasks requiring $ extit{many}$ such kernels and especially when the desired sequence manipulation is $ extit{context-dependent}$. Despite these limitations, $\mathrm{DLR}$ reaches high performance on two higher-order reasoning tasks $\mathrm{ListOpsSubTrees}$ and $\mathrm{PathfinderSegmentation} ext{-}\mathrm{256}$ with input lengths $8K$ and $65K$ respectively, and gives encouraging performance on $\mathrm{PathfinderSegmentation} ext{-}\mathrm{512}$ with input length $262K$ for which attention is not a viable choice.

研究动机与目标

• 通过消除使理解和实现复杂化的离散化步骤，简化状态空间模型。
• 提出一种基于原始对角线线性 RNN（DLR）的直接、可解释的公式，作为长序列建模的起点。
• 在合成和真实世界长序列任务上，实证评估 DLR 的表达能力和性能，与 SSM 和基于注意力的模型进行比较。
• 研究当前 SSM 在处理上下文相关和高内核任务时的局限性，并将其表达能力与注意力机制进行对比。
• 提供一个清晰、易访问的基准测试框架，使用 13 个合成序列到序列任务来评估长序列建模能力。

提出的方法

• 提出 DLR 作为对角线线性 RNN 的直接参数化，绕过连续时间状态空间的离散化过程。
• 使用复数域的转移矩阵，并将特征值初始化在单位圆附近，以实现百万步以上的长距离信号传播。
• 引入一种周期性初始化方案，使大小为 N 的 DLR 能够表达任意长度为 N 的卷积核。
• 在一组涵盖从简单序列移位到展平图像中复杂空间推理的 13 个合成任务上评估 DLR。
• 在 Long Range Arena、语音分类和因果语言建模任务上，将 DLR 与 S4、S4D、Transformer 和 LocalAttention 进行基准对比。
• 使用长达 262K 个 token 的输入长度，测试其在极长序列上的可扩展性和性能。

实验结果

研究问题

• RQ1 是否可以通过直接的对角线线性 RNN（DLR）公式，在不引入连续时间离散化复杂性的情况下，实现与离散化 SSM 相当或更优的性能？
• RQ2 在需要长距离依赖和复杂交互的任务上，DLR 的表达能力与 SSM 和基于注意力的模型相比如何？
• RQ3 当前 SSM 在处理上下文相关序列操作或需要大量卷积核的任务时存在哪些局限性？
• RQ4 DLR 是否能在真实世界的长序列任务（如原始语音分类和因果语言建模）中实现具有竞争力的性能和效率？
• RQ5 DLR 参数化的周期性是否能支持一种有原则的初始化方式，从而表达任意卷积核？

主要发现

• 在 Long Range Arena 的所有任务上，包括 Path-X 和 ListOps，DLR 的性能与 S4 和 S4D 相当或更优，且公式更简洁、更可解释。
• 在输入长度为 262K 的 PathfinderSegmentation-512 任务中，尽管在此规模下注意力机制在计算上不可行，DLR 仍表现出色。
• 在 ListOps-SubTrees 和 PathfinderSegmentation-256 任务上，DLR 胜过 LocalAttention，尤其在长序列长度下优势更明显，且吞吐量更高。
• 在 10 类语音命令数据集上，DLR 与 S4 和 S4D 表现相当或更优，以 O(L) 的解码复杂度实现了 97.1% 的准确率。
• 在 PG-19 的因果语言建模任务中，DLR 的测试交叉熵损失与 Transformer 相当（2.44–2.88），参数量为 1.2B，吞吐量为 1.0×。
• 与已知对层归一化位置高度敏感的 Transformer 相比，DLR 对层归一化位置的鲁棒性更强。