[论文翻译]Transformer 是 SSM:通过结构化状态空间对偶实现的通用模型和高效算法


原文地址:https://u254848-88c6-e493554b.yza1.seetacloud.com:8443/miner/v2/analysis/pdf_md?filename=full.md&as_attachment=False&pdf=6fe16414fa471a5ab150cddcba5840d3a94e347e96992c50c12c1fb8b919479f1735788830_2405.21060v1.pdf


Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Transformer 是 SSM:通过结构化状态空间对偶实现的通用模型和高效算法

Transformer 是一种强大的神经网络架构,广泛应用于自然语言处理和其他领域。然而,它们的计算复杂度较高,尤其是在处理长序列时。本文提出了一种新的视角,将 Transformer 视为结构化状态空间模型 (SSM) 的一种特例。通过这种对偶性,我们可以设计出更通用的模型和更高效的算法。

具体来说,我们将展示如何通过结构化状态空间对偶(Structured State Space Duality)来重新解释 Transformer 的工作机制。这种方法不仅能够简化模型的设计,还能显著提高其计算效率。此外,我们还将探讨这种对偶性在其他领域的潜在应用,例如时间序列预测和控制理论 [20]。

通过引入 SSM 的概念,我们可以更好地理解 Transformer 的内部机制,并为其优化提供新的思路。这为开发更强大、更高效的生成式 AI (Generative AI) 系统奠定了基础。

Tri Dao∗ and Albert Gu∗

Tri Dao∗ 和 Albert Gu∗

Department of Computer Science, Princeton University

计算机科学系,普林斯顿大学

Machine Learning Department, Carnegie Mellon University tri@tridao.me, agu@cs.cmu.edu

机器学习系,卡内基梅隆大学 tri@tridao.me, agu@cs.cmu.edu

请注意:根据您的要求,人名和邮箱地址不进行翻译。

Abstract

摘要

While Transformers have been the main architecture behind deep learning’s success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decomposition s of a well-studied class of structured semi separable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba’s selective SSM that is $2{-}8\times$ faster, while continuing to be competitive with Transformers on language modeling.

虽然 Transformer 一直是深度学习在语言建模中取得成功的主要架构,但状态空间模型 (SSMs) 如 Mamba 最近已被证明在小到中等规模上能够匹敌甚至超越 Transformer。我们发现这些模型家族实际上非常密切相关,并开发了一个丰富的理论框架,揭示了 SSMs 与注意力机制变体之间的联系,这些联系通过结构化半可分离矩阵的多种分解方式连接起来。我们的状态空间对偶性 (State Space Duality, SSD) 框架使我们能够设计出一种新的架构(Mamba-2),其核心层是对 Mamba 的选择性 SSM 的改进,速度提高了 2-8 倍,同时在语言建模方面仍然保持与 Transformer 竞争的能力。

$2{-}8\times$

1 Introduction

1 引言

根据要求,仅提供标题翻译,不询问继续提供更多内容。

Transformers, in particular decoder-only models (e.g. GPT (Brown et al. 2020), Llama (Touvron, Lavril, et al. 2023)) which process input sequences in a causal fashion, are one of the main drivers of modern deep learning’s success. Numerous approaches attempt to approximate the core attention layer to address its efficiency issues (Tay et al. 2022), such as scaling quadratically in sequence length during training and requiring a cache of size linear in sequence length during auto regressive generation. In parallel, a class of alternative sequence models, structured state-space models (SSMs), have emerged with linear scaling in sequence length during training and constant state size during generation. They show strong performance on long-range tasks (e.g. S4 (Gu, Goel, and Ré 2022)) and recently matched or beat Transformers on language modeling (e.g. Mamba (Gu and Dao 2023)) at small to moderate scale. However, the development of SSMs have appeared disjoint from the community’s collective effort to improve Transformers, such as understanding them theoretically as well as optimizing them on modern hardware. As a result, it is more difficult to understand and experiment with SSMs compared to Transformers, and it remains challenging to train SSMs as efficiently as Transformers from both an algorithmic and systems perspective.

Transformers,特别是仅解码器模型(例如 GPT (Brown et al. 2020),Llama (Touvron, Lavril, et al. 2023)),这些模型以因果方式处理输入序列,是现代深度学习成功的主要驱动力之一。许多方法尝试近似核心的注意力层,以解决其效率问题(Tay et al. 2022),例如在训练过程中序列长度呈二次方扩展,在自回归生成过程中需要线性于序列长度的缓存。与此同时,一类替代的序列模型——结构化状态空间模型(SSMs)出现了,它们在训练过程中序列长度呈线性扩展,在生成过程中状态大小保持常数。这些模型在长距离任务中表现出色(例如 S4 (Gu, Goel, 和 Ré 2022)),并在小到中等规模的语言建模任务中最近已经匹配或超越了 Transformer(例如 Mamba (Gu 和 Dao 2023))。然而,SSMs 的发展似乎与社区为改进 Transformer 所做的集体努力脱节,包括从理论上理解它们以及在现代硬件上优化它们。因此,相比 Transformer,理解和实验 SSMs 更加困难,并且从算法和系统角度来看,训练 SSMs 仍然比训练 Transformer 更具挑战性。

Our main goal is to develop a rich body of theoretical connections between structured SSMs and variants of attention. This will allow us to transfer algorithmic and systems optimization s originally developed for Transformers to SSMs, towards the goal of building foundation models that perform better than Transformers while scaling more efficiently in sequence length. A milestone contribution in this direction was the Linear Attention (LA) framework (Katha ro poul os et al. 2020), which derived a connection between auto regressive attention and linear RNNs by showing the equivalence between “dual forms” of quadratic kernelized attention and a particular linear recurrence. This duality allows new capabilities such as the ability to have both efficient parallel iz able training and efficient auto regressive inference. In the same spirit, this paper provides multiple viewpoints connecting linear-complexity SSMs with quadratic-complexity forms to combine the strengths of SSMs and attention.

我们的主要目标是建立结构化 SSM (Structured SSMs) 与各种注意力机制之间的丰富理论联系。这将使我们能够将原本为 Transformer 开发的算法和系统优化技术迁移到 SSMs,从而构建出在序列长度扩展时比 Transformer 更高效、性能更好的基础模型。在这个方向上的一个重要里程碑是线性注意力 (Linear Attention, LA) 框架 [Katharopoulos et al. 2020],该框架通过展示二次核化注意力的“对偶形式”与特定线性递归之间的等价性,建立了自回归注意力与线性 RNN 之间的联系。这种对偶性使得一些新的能力成为可能,例如同时实现高效的并行训练和高效的自回归推理。秉承这一精神,本文从多个角度探讨了线性复杂度的 SSM 与二次复杂度形式之间的联系,旨在结合 SSM 和注意力机制的优势。

State Space Duality. Our framework connecting structured SSMs and variants of attention, which we call structured state space duality (SSD), is made through the abstractions of structured matrices: matrices with subquadratic parameters and multiplication complexity. We develop two broad frameworks for representing sequence models, one as matrix transformations and one as tensor contractions, which each reveal different perspectives of the duality. Our technical contributions include:

状态空间对偶性。我们提出的框架将结构化状态空间模型 (SSM) 与注意力机制的变体联系起来,我们称之为结构化状态空间对偶性 (Structured State Space Duality, SSD),是通过结构化矩阵的抽象实现的:具有次二次参数和乘法复杂度的矩阵。我们开发了两种广泛的框架来表示序列模型,一种是基于矩阵变换,另一种是基于张量收缩,这两种方法分别揭示了对偶性的不同视角。我们的技术贡献包括:

  • 通过结构化矩阵的抽象,建立了连接结构化 SSM 和注意力机制变体的桥梁。
  • 提出了两种表示序列模型的框架:一种是基于矩阵变换,另一种是基于张量收缩,每种框架都展示了对偶性的不同方面。

我们的技术贡献包括:

(注:原文中未列出具体的技术贡献内容,因此在此处结束翻译。)

• We show an equivalence between state space models and a well-studied family of structured matrices called semiseparable matrices (Section 3). This connection is at the heart our framework, revealing new properties and algorithms for SSMs. A central message of this paper is that different methods of computing state space models can be reframed as various matrix multiplication algorithms on structured matrices. • We significantly improve the theory of linear attention (Katha ro poul os et al. 2020). We first provide an in- cisive proof of its recurrent form through the language of tensor contractions, and then generalize it to a new family of structured masked attention (SMA) (Section 4). • We connect SSMs and SMA, showing that they have a large intersection that are duals of each other, possessing both SSM-like linear and attention-like quadratic forms (Section 5). We also prove that any kernel attention method possessing a fast recurrent form must be an SSM.

• 我们展示了状态空间模型 (State Space Models, SSMs) 与一类被广泛研究的结构化矩阵——半可分矩阵 (semiseparable matrices) 之间的等价性(第 3 节)。这种联系是我们框架的核心,揭示了 SSMs 的新性质和算法。本文的一个中心观点是,不同的状态空间模型计算方法可以重新表述为结构化矩阵上的各种矩阵乘法算法。

• 我们显著改进了线性注意力机制 (linear attention) 的理论 [Katha ro poul os et al. 2020]。我们首先通过张量收缩的语言提供了其递归形式的深刻证明,然后将其推广到一个新的结构化掩码注意力 (Structured Masked Attention, SMA) 家族(第 4 节)。

• 我们将 SSMs 和 SMA 进行连接,展示了它们之间存在一个较大的交集,这些交集中的模型互为对偶,既具有类似 SSM 的线性形式,又具有类似注意力机制的二次形式(第 5 节)。我们还证明了任何具有快速递归形式的核注意力方法必定是一个 SSM。


Figure 1: (Structured State-Space Duality.) This paper fleshes out the relationship between state space models and attention through the bridge of structured matrices.


图 1: (结构化状态空间对偶性)本文通过结构化矩阵的桥梁,详细阐述了状态空间模型与注意力机制之间的关系。

Beyond its intrinsic theoretical value, our framework opens up a broad set of directions for understanding and improving sequence models.

除了其内在的理论价值之外,我们的框架为理解和改进序列模型开辟了广泛的方向。

Efficient Algorithms. First and most importantly, our framework exposes new efficient and easily-implementable algorithms for computing SSMs (Section 6). We introduce a new SSD algorithm, based on block decomposition s of semiseparable matrices, that takes advantage of both the linear SSM recurrence and quadratic dual form, obtaining optimal tradeoffs on all main efficiency axes (e.g. training and inference compute, memory usage, and ability to leverage matrix multiplication units on modern hardware). A dedicated implementation of SSD is $2-8\times$ faster than the optimized selective scan implementation of Mamba, while simultaneously allowing for much larger recurrent state sizes ( $^{8\times}$ the size of Mamba or even higher, with minimal slowdown). SSD is highly competitive with optimized implementations of softmax attention (Flash Attention-2 (Dao 2024)), crossing over at sequence length 2K and $6\times$ faster at sequence length 16K.

高效算法。首先也是最重要的是,我们的框架揭示了新的高效且易于实现的算法来计算 SSMs (第 6 节)。我们引入了一种新的 SSD 算法,该算法基于半可分离矩阵的块分解,充分利用了线性 SSM 递归和二次对偶形式,在所有主要效率指标(例如训练和推理计算、内存使用以及利用现代硬件上的矩阵乘法单元的能力)上获得了最优的权衡。SSD 的专用实现比 Mamba 的优化选择扫描实现快 2-8 倍,同时允许更大的循环状态尺寸(比 Mamba 大 8 倍甚至更多,且几乎没有性能下降)。SSD 在优化的 softmax 注意力实现(如 Flash Attention-2 (Dao 2024))中具有高度竞争力,在序列长度为 2K 时表现相当,在序列长度为 16K 时快 6 倍。

Architecture Design. One major obstacle to adopting new architectures such as SSMs is the ecosystem tailored to Transformers, such as hardware-efficient optimization and parallelism techniques for large-scale training. Our framework allows using established conventions and techniques for attention to build a vocabulary of architecture design choices for SSMs, and further improve them (Section 7). For example, we introduce the analog of heads from multi-head attention (MHA) to SSMs. We show that the Mamba architecture is a multi-input SSM (MIS) that turns out to be analogous to multi-value attention (MVA), and compare other variants of Mamba with different head structures.

架构设计。采用新架构(如 SSMs)的一个主要障碍是为 Transformer 量身定制的生态系统,例如硬件高效的优化和并行技术,用于大规模训练。我们的框架允许使用已建立的注意力机制惯例和技术,为 SSMs 构建架构设计选择的词汇表,并进一步改进它们(第 7 节)。例如,我们引入了多头注意力 (MHA) 中的“头”概念到 SSMs。我们展示了 Mamba 架构是一个多输入 SSM (MIS),它类似于多值注意力 (MVA),并比较了具有不同头结构的其他 Mamba 变体。

We also use these ideas to make slight modifications to the Mamba block, which allows tensor parallelism to be implemented (e.g. in the style of Megatron (Shoeybi et al. 2019)). The main ideas include introducing grouped-value attention (GVA) head structure, and moving all data-dependent projections to occur in parallel at the beginning of the block.

我们还利用这些思路对 Mamba 模块进行了轻微修改,使得张量并行(例如 Megatron (Shoeybi et al. 2019) 风格)得以实现。主要的改进包括引入分组值注意力 (GVA, Grouped-Value Attention) 头结构,并将所有依赖数据的投影操作移到模块的开头并行执行。

The combination of the modified parallel Mamba block, together with using SSD as the inner SSM layer, results in the Mamba-2 architecture. We investigate Chinchilla scaling laws for Mamba-2 in the same setting as Mamba, finding that it Pareto dominates Mamba and Transformer $^{++}$ in both perplexity and wall-clock time. We additionally train a family of

修改后的并行 Mamba 模块与使用 SSD 作为内部 SSM 层相结合,构成了 Mamba-2 架构。我们在与 Mamba 相同的设置下研究了 Mamba-2 的 Chinchilla 扩展规律,发现它在困惑度和墙钟时间上都 Pareto 支配 Mamba 和 Transformer $^{++}$ 。我们还训练了一系列的

请注意:原文最后一句不完整,可能缺少部分内容。

Mamba-2 models at varying sizes on the Pile, showing that it matches or outperforms Mamba and open source Transformers on standard downstream evaluations. For example, Mamba-2 with 2.7B parameters trained on 300B tokens on the Pile outperforms Mamba-2.8B, Pythia-2.8B and even Pythia-6.9B trained on the same dataset.

Mamba-2 模型在不同规模下在 Pile 数据集上的表现,显示其在标准下游评估中匹配或超越了 Mamba 和开源 Transformer 模型。例如,具有 27 亿参数的 Mamba-2 在 Pile 上使用 3000 亿个 Token 进行训练后,表现优于 Mamba-2.8B、Pythia-2.8B 以及在同一数据集上训练的 Pythia-6.9B。

Systems Optimization s. The SSD framework connects SSMs and Transformers, allowing us to leverage a rich body of work on systems optimization s developed for Transformers (Section 8).

系统优化。SSD框架将状态空间模型 (SSMs) 和 Transformer 相连接,使我们能够利用为 Transformer 开发的大量系统优化成果(第8节)。

Section 9 empirically validates Mamba-2 on language modeling, training efficiency, and a difficult multi-query associative recall task (Arora, Eyuboglu, Zhang, et al. 2024). Finally, in Section 10, we provide an extended related work and discuss potential research directions opened up by our framework.

第 9 节通过实验验证了 Mamba-2 在语言模型、训练效率以及复杂的多查询关联回忆任务 (Arora, Eyuboglu, Zhang, et al. 2024) 上的表现。最后,在第 10 节中,我们提供了扩展的相关工作讨论,并探讨了我们的框架所开启的潜在研究方向。

Model code and pre-trained checkpoints are open-sourced at https://github.com/state-spaces/mamba.

模型代码和预训练检查点已开源,可在以下地址获取:https://github.com/state-spaces/mamba

2 Background and Overview

2 背景与概述

根据上述要求,我已经完成了该段落的翻译。请确认是否需要继续翻译后续内容。

2.1 Structured State Space Models

2.1 结构化状态空间模型 (Structured State Space Models)

Structured state space sequence models (S4) are a recent class of sequence models for deep learning that are broadly related to RNNs, CNNs, and classical state space models. They are inspired by a particular continuous system (1) that maps a 1-dimensional sequence $x\in\mathbb{R}^{\intercal}\mapsto y\in\mathbb{R}^{\intercal}$ through an implicit latent state $\bar{h}\in\mathbb{R}^{(\mathsf{T},\mathsf{N})}$ .

结构化状态空间序列模型 (S4) 是一类最近提出的深度学习序列模型,与循环神经网络 (RNN)、卷积神经网络 (CNN) 和经典状态空间模型有广泛的关联。这类模型的灵感来源于一个特定的连续系统 (1),该系统将一维序列 $x\in\mathbb{R}^{\intercal}\mapsto y\in\mathbb{R}^{\intercal}$ 通过一个隐含的潜在状态 $\bar{h}\in\mathbb{R}^{(\mathsf{T},\mathsf{N})}$ 进行映射。

在这个系统中,输入序列 $x$ 被映射到输出序列 $y$,而这个映射过程是通过一个潜在状态 $\bar{h}$ 来实现的,其中 $\mathsf{T}$ 表示时间步长,$\mathsf{N}$ 表示状态维度。这种模型的设计使得它能够有效地处理长时间依赖和复杂的序列数据。

A general discrete form of structured SSMs takes the form of equation (1).

结构化 SSMs (Structured SSMs) 的一般离散形式如公式 (1) 所示。

$$
\begin{array}{l}{{h_{t}}=A{h_{t-1}}+B x_{t}}\ {{y_{t}}={C^{\top}}{h_{t}}}\end{array}
$$

$$
\begin{array}{l}
h_{t} = A h_{t-1} + B x_{t} \
y_{t} = C^{\top} h_{t}
\end{array}
$$

该公式描述了一个简单的线性动态系统,其中:

  • ( h_t ) 表示在时间步 ( t ) 的隐藏状态 (hidden state)
  • ( x_t ) 表示在时间步 ( t ) 的输入
  • ( y_t ) 表示在时间步 ( t ) 的输出
  • ( A ), ( B ), 和 ( C ) 是系统参数矩阵

第一个公式表示隐藏状态 ( h_t ) 由前一个时间步的隐藏状态 ( h_{t-1} ) 和当前输入 ( x_t ) 线性组合而成。第二个公式表示输出 ( y_t ) 是当前隐藏状态 ( h_t ) 的线性变换。

$$
\begin{array}{l}{\boldsymbol{h}{t}=\boldsymbol{A}{t}\boldsymbol{h}{t-1}+\boldsymbol{B}{t}\boldsymbol{x}{t}}\ {\boldsymbol{y}{t}=\boldsymbol{C}{t}^{\top}\boldsymbol{h}{t}}\end{array}
$$

$$
\begin{array}{l}
\boldsymbol{h}{t} = \boldsymbol{A}{t} \boldsymbol{h}{t-1} + \boldsymbol{B}{t} \boldsymbol{x}{t} \
\boldsymbol{y}
{t} = \boldsymbol{C}{t}^{\top} \boldsymbol{h}{t}
\end{array}
$$

该公式描述了一个递归状态更新的过程,其中:

  • $\boldsymbol{h}_{t}$ 表示在时间步 $t$ 的隐藏状态 (hidden state)。
  • $\boldsymbol{A}{t}$ 和 $\boldsymbol{B}{t}$ 是矩阵,用于对前一个时间步的隐藏状态 $\boldsymbol{h}{t-1}$ 和当前输入 $\boldsymbol{x}{t}$ 进行线性变换。
  • $\boldsymbol{y}{t}$ 表示在时间步 $t$ 的输出,由当前隐藏状态 $\boldsymbol{h}{t}$ 通过矩阵 $\boldsymbol{C}_{t}$ 计算得到。

这个公式常用于时间序列模型或递归神经网络 (RNN) 中的状态更新机制。

where $A,\in,\mathbb{R}^{(\mathsf{N},\mathsf{N})},B,\in,\mathbb{R}^{(\mathsf{N},1)},C,\in,\mathbb{R}^{(\mathsf{N},1)}$ . Structured SSMs are so named because the $A$ matrix controlling the temporal dynamics must be structured in order to compute this sequence-to-sequence transformation efficiently enough to be used in deep neural networks. The original structures introduced were diagonal plus low-rank (DPLR) (Gu, Goel, and Ré 2022) and diagonal (Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; J. T. Smith, Warrington, and Linderman 2023), which remains the most popular structure.

其中 $A,\in,\mathbb{R}^{(\mathsf{N},\mathsf{N})}, B,\in,\mathbb{R}^{(\mathsf{N},1)}, C,\in,\mathbb{R}^{(\mathsf{N},1)}$ 。结构化 SSM(State Space Model,状态空间模型)之所以得名,是因为控制时间动态的 $A$ 矩阵必须具有特定的结构,以便能够高效地计算这种序列到序列的转换,从而可以在深度神经网络中使用。最初引入的结构包括对角加低秩 (DPLR) (Gu, Goel, and Ré 2022) 和对角结构 (Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; J. T. Smith, Warrington, and Linderman 2023),后者仍然是最流行的结构。

In this work, we use the term state space model (SSM) to refer to structured SSMs. There are many flavors of such SSMs, with deep ties to several major paradigms of neural sequence models such as continuous-time, recurrent, and convolutional models (Gu, Johnson, Goel, et al. 2021). We provide a brief overview below, and refer to prior work for more context and details (Gu 2023; Gu and Dao 2023).

在本研究中,我们使用术语状态空间模型 (SSM) 来指代结构化的 SSM。这类 SSM 有多种形式,与连续时间、递归和卷积等主要神经序列模型范式有着深厚的联系 (Gu, Johnson, Goel, et al. 2021)。我们在下面简要概述这些模型,并参考先前的工作以获取更多背景和详细信息 (Gu 2023; Gu and Dao 2023)。

在本研究中,我们使用术语状态空间模型 (SSM) 来指代结构化的 SSM。这类 SSM 有多种形式,与连续时间、递归和卷积等主要神经序列模型范式有着深厚的联系 (Gu, Johnson, Goel, et al. 2021)。我们在下面简要概述这些模型,并参考先前的工作以获取更多背景和详细信息 (Gu 2023; Gu and Dao 2023)。

Continuous-time Models. The original structured SSMs originated as continuous-time maps on functions $x(t)\in\mathbb{R}\mapsto$ $y(t)\in\mathbb{R},$ rather than operating directly on sequences. In the continuous-time perspective, in equation (1a) the matrices $(A,B)$ are not directly learned but generated from underlying parameters $(\mathring{A},\mathring{B})$ , along with a parameterized step size $\Delta$ . The “continuous parameters” $(\Delta,\stackrel{\circ}{A},\stackrel{\circ}{B})$ are converted to “discrete parameters” $(A,B)$ through fixed formulas $A={\bar{f}}{A}(\Delta,\stackrel{\circ}{A})$ and $B=f{B}(\Delta,\mathring{B})$ , where the pair $(f_{A},f_{B})$ is called a disc ret iz ation rule.

连续时间模型。最初的结构化 SSMs 是作为函数 $x(t) \in \mathbb{R} \mapsto y(t) \in \mathbb{R}$ 的连续时间映射,而不是直接作用于序列。在连续时间视角下,在公式 (1a) 中,矩阵 $(A, B)$ 并不是直接学习得到的,而是从底层参数 $(\mathring{A}, \mathring{B})$ 生成的,并且伴随着一个参数化的步长 $\Delta$。"连续参数" $(\Delta, \mathring{A}, \mathring{B})$ 通过固定的公式 $A = \bar{f}_A(\Delta, \mathring{A})$ 和 $B = f_B(\Delta, \mathring{B})$ 转换为 "离散参数" $(A, B)$,其中 $(f_A, f_B)$ 称为离散化规则。

在这一过程中,$\mathring{A}$ 和 $\mathring{B}$ 是用于生成离散时间模型中使用的矩阵 $A$ 和 $B$ 的基础参数,而 $\Delta$ 控制了从连续时间到离散时间的转换过程。这种转换使得模型能够在保持连续时间特性的前提下,有效地应用于离散时间序列数据。

Remark 1. While our main models adopt the same parameter iz ation and disc ret iz ation step as prior work (see Gu and Dao (2023) for details), for simplifying exposition and notation we omit it in the rest of this paper. We note that prior work on structured SSMs referred to the continuous parameters $(\mathring{A},\mathring{B})$ and discrete parameters $(A,B)$ as $(A,B)$ and $(\bar{A},\bar{B})$ instead; we have changed notation to simplify the presentation and focus directly on the discrete parameters, which govern the main SSM recurrence.

备注 1. 尽管我们的主要模型采用了与先前工作相同的参数化和离散化步骤(具体细节参见 Gu 和 Dao (2023)),为了简化表述和符号,我们在本文的其余部分省略了这些内容。我们注意到,先前关于结构化 SSM 的工作将连续参数 $(\mathring{A}, \mathring{B})$ 和离散参数 $(A, B)$ 分别表示为 $(A, B)$ 和 $(\bar{A}, \bar{B})$;我们更改了符号以简化表述,并直接聚焦于控制主要 SSM 递归的离散参数。

在本文中,我们将重点关注离散参数 $(A, B)$,它们是决定 SSM 递归关系的主要因素。

Recurrent Models. Equations (1) and (2) take the form of a recurrence which is linear in its input $x$ . Structured SSMs can therefore be viewed as types of recurrent neural networks (RNNs), where the linearity endows them with additional properties and allows them to avoid the sequential computation of traditional RNNs. Conversely, despite this simplification, SSMs are still fully expressive as sequence transformations (in the sense of universal approximation) (Kaul 2020; Orvieto et al. 2023; Shida Wang and Xue 2023).

递归模型。公式 (1) 和 (2) 以 $x$ 为输入的线性递归形式表示。因此,结构化的 SSMs 可以被视为一种递归神经网络 (RNNs) 的类型,其中的线性特性赋予了它们额外的属性,并使它们能够避免传统 RNN 的顺序计算。相反,尽管有这种简化,SSMs 仍然作为序列变换是完全表达的(在通用近似的意义上)[Kaul 2020; Orvieto et al. 2023; Shida Wang and Xue 2023]。

在这种情况下,线性递归模型通过减少计算复杂度,同时保持强大的表达能力,提供了一种更高效的替代方案。SSMs 的线性特性使得它们能够在不牺牲模型性能的情况下,显著减少计算资源的需求。

Convolutional Models. When the SSM’s dynamics are constant through time as in equation (1), the model is called linear time-invariant (LTI). In this case, they are equivalent to convolutions. Thus, SSMs can also be viewed as types of CNNs, but where (i) the convolution kernels are implicitly parameterized through the SSM parameters $(A,B,C)$ and (ii) the convolution kernels are generally global instead of local. Conversely, through classical signal processing theory all sufficiently well-behaved convolutions can be represented as SSMs.

卷积模型。当状态空间模型 (SSM) 的动态特性在时间上保持不变,如公式 (1) 所示时,该模型被称为线性时不变 (LTI) 模型。在这种情况下,它们等价于卷积操作。因此,SSM 也可以被视为一种卷积神经网络 (CNN),但有以下两点不同:(i) 卷积核是通过 SSM 参数 $(A, B, C)$ 隐式参数化的;(ii) 卷积核通常是全局的而不是局部的。相反,根据经典信号处理理论,所有足够良好行为的卷积都可以表示为 SSM [20]。

Commonly, previous LTI SSMs would use the convolutional mode for efficient parallel iz able training (where the whole input sequence is seen ahead of time), and switched into recurrent mode (1) for efficient auto regressive inference (where the inputs are seen one step at a time).

通常情况下,之前的线性时不变状态空间模型 (LTI SSM) 会使用卷积模式进行高效的并行训练(其中整个输入序列可以提前获取),并在进行高效的自回归推理时切换到递归模式(1)(其中输入是逐个步骤获取的)。

Selective State Space Models. The form (2) where the parameters $(A,B,C)$ can also vary in time was introduced in Mamba as the selective SSM. Compared to the standard LTI formulation (1), this model can selectively choose to focus on or ignore inputs at every timestep. It was shown to perform much better than LTI SSMs on information-dense data such as language, especially as its state size N increases allowing for more information capacity. However, it can only be computed in recurrent instead of convolutional mode, and requires a careful hardware-aware implementation to be efficient. Even so, it is still less efficient than hardware-friendly models such as CNNs and Transformers because it does not leverage matrix multiplication units, which modern accelerators such as GPUs and TPUs are specialized for.

选择性状态空间模型 (Selective State Space Models, SSM)。形式 (2) 中的参数 $(A,B,C)$ 随时间变化的概念是在 Mamba 中引入的选择性 SSM。与标准的线性时不变 (LTI) 公式 (1) 相比,该模型可以在每个时间步有选择地关注或忽略输入。研究表明,它在信息密集型数据(如语言)上的表现远优于 LTI SSM,尤其是在其状态大小 N 增加时,允许更多的信息容量。然而,它只能在递归模式下计算,而不是卷积模式,并且需要精心设计的硬件感知实现以提高效率。即便如此,它仍然不如硬件友好的模型(如 CNN 和 Transformer)高效,因为这些模型可以充分利用现代加速器(如 GPU 和 TPU)中的矩阵乘法单元。

选择性状态空间模型 (Selective State Space Models, SSM)。形式 (2) 中的参数 \$(A,B,C)\$ 随时间变化的概念是在 Mamba 中引入的选择性 SSM。与标准的线性时不变 (LTI) 公式 (1) 相比,该模型可以在每个时间步有选择地关注或忽略输入。研究表明,它在信息密集型数据(如语言)上的表现远优于 LTI SSM,尤其是在其状态大小 N 增加时,允许更多的信息容量。然而,它只能在递归模式下计算,而不是卷积模式,并且需要精心设计的硬件感知实现以提高效率。即便如此,它仍然不如硬件友好的模型(如 CNN 和 Transformer)高效,因为这些模型可以充分利用现代加速器(如 GPU 和 TPU)中的矩阵乘法单元。

While time-invariant SSMs are closely related to continuous, recurrent, and convolutional sequence models, they are not directly related to attention. In this paper, we show a deeper relationship between selective SSMs and attention, and use it to significantly improve the training speed of SSMs while simultaneously allowing for much larger state sizes N.

时间不变的 SSMs (State-Space Models) 与连续型、递归型和卷积型序列模型密切相关,但它们与注意力机制没有直接联系。在本文中,我们展示了选择性 SSMs 与注意力机制之间更深层次的关系,并利用这种关系显著提高了 SSMs 的训练速度,同时允许使用更大的状态尺寸 N。

Structured SSMs as Sequence Transformations.

结构化 SSM 作为序列变换 (Structured SSMs as Sequence Transformations).

结构化状态空间模型 (SSM) 可以被视为一种对序列数据进行变换的工具。通过将 SSM 应用于序列数据,我们可以捕捉到数据中的动态变化和时间依赖性。这种视角有助于我们更好地理解 SSM 在生成式 AI (Generative AI) 和时间序列预测中的应用。

在本文中,我们将探讨如何将结构化 SSM 视为一种序列变换方法,并分析其在不同任务中的表现。具体来说,我们将讨论 SSM 的数学表示、参数估计以及如何将其应用于实际问题中。此外,我们还将介绍一些最新的研究成果,展示 SSM 在处理复杂序列数据方面的优势。

通过对 SSM 的深入研究,我们希望能够为读者提供一个清晰的理解框架,帮助他们在自己的工作中更好地应用这一强大的工具。

Definition 2.1. We use the term sequence transformation to refer to a parameterized map on sequences $Y=f_{\theta}(X)$ where $X,Y\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$ and $\theta$ is an arbitrary collection of parameters. T represents the sequence or time axis; subscripts index into the first dimension, e.g. $X_{t}$ , $Y_{t}\in\mathbb{R}^{\mathsf{P}}$ .

定义 2.1. 我们使用术语序列变换 (sequence transformation) 来指代一个参数化的映射 $Y=f_{\theta}(X)$,其中 $X,Y\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$,$\theta$ 是任意一组参数。T 表示序列或时间轴;下标用于索引第一个维度,例如 $X_{t}$,$Y_{t}\in\mathbb{R}^{\mathsf{P}}$。

解释:

  • $X$ 和 $Y$ 是形状为 $(\mathsf{T}, \mathsf{P})$ 的序列,其中 $\mathsf{T}$ 表示时间步数或序列长度,$\mathsf{P}$ 表示每个时间步的特征维度。
  • $f_{\theta}$ 是一个由参数 $\theta$ 控制的函数,用于将输入序列 $X$ 映射到输出序列 $Y$。
  • $X_{t}$ 和 $Y_{t}$ 分别表示在时间步 $t$ 上的输入和输出,它们是形状为 $\mathsf{P}$ 的向量。

Sequence transformations (e.g. SSMs, or self-attention) are the cornerstone of deep sequence models, where they are incorporated into neural network architectures (e.g. Transformers). The SSM in (1) or (2) is a sequence transformation with $\mathsf{P}=1$ ; it can be generalized to $\mathsf{P}>1$ by simply broadcasting across this dimension (in other words, viewing the input as P independent sequences and applying the SSM to each). One can think of $\mathsf{P}$ as a head dimension, which we will elaborate on in Section 7.

序列变换(例如 SSMs 或自注意力机制)是深度序列模型的基石,它们被集成到神经网络架构中(例如 Transformer)。公式 (1) 或 (2) 中的 SSM 是一个序列变换,其中 $\mathsf{P}=1$;它可以通过简单地在该维度上进行广播来推广到 $\mathsf{P}>1$(换句话说,将输入视为 P 个独立的序列,并对每个序列应用 SSM)。可以将 $\mathsf{P}$ 视为头部维度(head dimension),我们将在第 7 节中详细讨论这一点。

Definition 2.2. We define the SSM operator $\mathsf{S S M}(A,B,C),=,\mathsf{S S M}(A_{0:T},B_{0:T},C_{0:T})$ as the sequence transformation $X~\in$ $\mathbb{R}^{(\mathsf{T},\mathsf{P})}\mapsto Y\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$ defined by equation (2).

定义 2.2. 我们定义 SSM 操作符 $\mathsf{S S M}(A,B,C),=,\mathsf{S S M}(A_{0:T},B_{0:T},C_{0:T})$ 为序列变换 $X~\in$ $\mathbb{R}^{(\mathsf{T},\mathsf{P})}\mapsto Y\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$,该变换由公式 (2) 定义。

其中,$A_{0:T}$、$B_{0:T}$ 和 $C_{0:T}$ 分别表示从时间步 0 到 T 的序列,$\mathbb{R}^{(\mathsf{T},\mathsf{P})}$ 表示维度为 (T, P) 的实数矩阵。

In SSMs, the N dimension is a free parameter called the state size or state dimension. We also call it the state expansion factor, because it expands the size of the input/output by a factor of $N$ , with implications for the computational efficiency of these models.

在 SSMs 中,N 维度是一个自由参数,称为状态大小或状态维度 (state dimension)。我们也将其称为状态扩展因子,因为它将输入/输出的大小扩展了 N 倍,这对这些模型的计算效率有重要影响。

Finally, we remark that many types of sequence transformations, such as attention, can be represented as a single matrix multiplication across the sequence dimension.

最后,我们指出,许多类型的序列变换(如注意力机制),可以表示为在序列维度上的单个矩阵乘法。

Definition 2.3. We call a sequence transformation $Y=f_{\theta}(X)$ a matrix transformation if it can be written in the form $Y=M_{\theta}X$ where $M$ is a matrix depending on the parameters $\theta$ . We identify the sequence transformation with the matrix $M$ , and often drop the dependence on $\theta$ when clear from context.

定义 2.3. 我们称序列变换 $Y=f_{\theta}(X)$ 为矩阵变换,如果它可以表示为 $Y=M_{\theta}X$ 的形式,其中 $M$ 是依赖于参数 $\theta$ 的矩阵。我们将序列变换与矩阵 $M$ 等同起来,并且在上下文明确时,通常会省略对 $\theta$ 的依赖。

注意:在该定义中,$f_{\theta}(X)$ 表示一个依赖于参数 $\theta$ 的函数,而 $M_{\theta}$ 则表示一个依赖于相同参数的矩阵。这种表示方法允许我们将复杂的序列变换简化为矩阵运算的形式。

2.2 Attention

2.2 注意力机制 (Attention)

注意力机制 (Attention) 是 Transformer 模型中的一个关键组件,它允许模型在处理序列数据时聚焦于最重要的部分。通过计算输入序列中不同位置之间的相关性,注意力机制能够动态地分配权重,使得模型可以更有效地捕捉长距离依赖关系和重要信息。

在传统的循环神经网络 (RNN) 中,模型需要逐个处理序列中的元素,这导致了处理长序列时的效率低下和信息丢失问题。而注意力机制通过并行计算输入序列中所有位置之间的关系,克服了这一局限性,显著提升了模型的性能和表达能力。

注意力机制的核心思想是通过计算查询 (Query)、键 (Key) 和值 (Value) 之间的相似度来确定哪些部分应该被重点关注。具体来说,查询和键之间的点积或缩放点积用于衡量它们的相关性,然后通过 softmax 函数将这些相关性转换为概率分布,最后加权求和得到输出。

这种机制不仅在自然语言处理任务中表现出色,还在计算机视觉、语音识别等多个领域得到了广泛应用。

Attention broadly refers to a type of computation that assigns scores to every pair of positions in a sequence, allowing each element to “attend” to the rest. By far the most common and important variant of attention is softmax self-attention, which can be defined as

注意力机制广义上指的是对序列中每一对位置分配分数的计算方式,使得每个元素可以“关注”到其他元素。迄今为止,最常见和重要的注意力机制变体是 softmax 自注意力 (softmax self-attention),其可以定义为

注意:原文中没有给出 softmax self-attention 的具体定义,因此这里保留了定义未完成的状态。

$$
Y=\mathrm{softmax}(Q K^{\top})\cdot V
$$

$$
Y = \mathrm{softmax}(Q K^{\top}) \cdot V
$$

该公式表示的是 Transformer 模型中的自注意力机制 (self-attention mechanism) 的计算过程。具体来说,$Q$、$K$ 和 $V$ 分别代表查询 (Query)、键 (Key) 和值 (Value),$\mathrm{softmax}$ 函数用于将注意力得分转换为概率分布。通过计算查询和键之间的点积并应用 $\mathrm{softmax}$,可以得到注意力权重矩阵,再与值矩阵相乘,最终得到输出 $Y$。

for $Q,K,V\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$ . The mechanism of pairwise comparisons (induced by materializing $\ Q K^{\top}$ ) leads to the characteristic quadratic training cost of attention.

对于 $Q, K, V \in \mathbb{R}^{(\mathsf{T}, \mathsf{P})}$ 。成对比较机制(由计算 $Q K^{\top}$ 引起)导致了注意力机制的特征性二次训练成本。

解释:在注意力机制中,$Q$(查询)、$K$(键)和 $V$(值)是矩阵,它们的维度为 $(\mathsf{T}, \mathsf{P})$。通过计算 $Q$ 和 $K$ 的转置矩阵乘积 $Q K^{\top}$,可以进行成对比较。这种操作会导致训练成本呈二次增长,因为每次计算都需要遍历所有可能的成对组合。

Many variants of attention have been proposed, but all share the underlying core of these attention scores, with various approximations (Tay et al. 2022). The most important variant for this work is linear attention (Katha ro poul os et al. 2020). Roughly speaking, this family of methods drops the softmax by folding it into a kernel feature map, and uses associativity of matrix multiplication to rewrite $(Q K^{\top})\cdot V=Q\cdot(K^{\top}V)$ . Moreover, in the important case of causal (auto regressive) attention, they show that when the causal mask is incorporated into the left-hand side as $(L\circ Q K^{\top})\cdot V$ , where $L$ is the lower-triangular 1’s matrix, then the right-hand side can be expanded as a recurrence. Several recent and concurrent works such as RetNet (Y. Sun et al. 2023) and GateLoop (Katsch 2023) strengthen this to more general forms of $L$ (Section 10). In this work, our formulation of structured masked attention will strongly generalize these ideas.

许多注意力机制的变体已经被提出,但它们都共享了这些注意力分数的核心思想,并采用了不同的近似方法 (Tay et al. 2022)。对于本研究最重要的变体是线性注意力 (Katharopoulos et al. 2020)。简单来说,这类方法通过将 softmax 操作融入到核特征映射中来省略它,并利用矩阵乘法的结合律将 $(Q K^{\top}) \cdot V = Q \cdot (K^{\top} V)$ 重写。

此外,在因果(自回归)注意力的重要情况下,他们证明了当因果掩码被整合到左侧时,形式为 $(L \circ Q K^{\top}) \cdot V$,其中 $L$ 是下三角全 1 矩阵,右侧可以展开为递归形式。最近和同期的一些工作,如 RetNet (Y. Sun et al. 2023) 和 GateLoop (Katsch 2023),进一步推广了 $L$ 的更一般形式(第 10 节)。在本研究中,我们提出的结构化掩码注意力将极大地推广这些想法。

2.3 Structured Matrices

2.3 结构化矩阵

结构化矩阵 (Structured Matrices) 是指具有特定模式或结构的矩阵,这些结构可以用于优化计算效率和减少存储需求。在许多应用中,结构化矩阵能够显著提高算法的性能,尤其是在大规模数据处理和机器学习领域。常见的结构化矩阵类型包括对角矩阵、三对角矩阵、Toeplitz矩阵等。通过利用这些矩阵的特殊性质,可以在保持精度的同时,大幅降低计算复杂度。

General matrices $M,\in,\mathbb{R}^{(\mathsf{T},\mathsf{T})}$ require $\mathsf{T}^{2}$ parameters to represent and $O(\mathsf{T}^{2})$ time to perform basic operations such as matrix-vector multiplication. Structured matrices are those that

一般矩阵 $M,\in,\mathbb{R}^{(\mathsf{T},\mathsf{T})}$ 需要 $\mathsf{T}^2$ 个参数来表示,并且需要 $O(\mathsf{T}^2)$ 的时间来进行基本操作,例如矩阵-向量乘法。结构化矩阵是指那些

(i) can be represented in sub quadratic (ideally linear) parameters through a compressed representation, and (ii) have fast algorithms (most importantly matrix multiplication) by operating directly on this compressed representation.

可以通过压缩表示来实现次二次(理想情况下是线性)参数复杂度 (i),并且 (ii) 通过直接操作这种压缩表示,可以拥有快速的算法(最重要的是矩阵乘法)。

Perhaps the most canonical families of structured matrices are sparse and low-rank matrices. However, there exist many other families, such as Toeplitz, Cauchy, Van der monde, and butterfly matrices, which have all been used in machine learning for efficient models (Dao, Gu, et al. 2019; D. Fu et al. 2024; Gu, Gupta, et al. 2022; Thomas et al. 2018). Structured matrices are a powerful abstraction for efficient representations and algorithms. In this work, we will show that SSMs are equivalent to another class of structured matrices that have not previously been used in deep learning, and use this connection to derive efficient methods and algorithms.

或许最经典的结构化矩阵家族是稀疏矩阵和低秩矩阵。然而,还存在许多其他类型的结构化矩阵,例如 Toeplitz 矩阵、Cauchy 矩阵、Van der monde 矩阵和蝴蝶矩阵 (Butterfly matrices),这些矩阵都已被用于机器学习中以构建高效的模型 (Dao, Gu, et al. 2019; D. Fu et al. 2024; Gu, Gupta, et al. 2022; Thomas et al. 2018)。结构化矩阵是一种强大的抽象工具,能够实现高效的表示和算法。在本文中,我们将展示 SSMs 等价于另一类之前未在深度学习中使用过的结构化矩阵,并利用这一联系推导出高效的算法和方法。

2.4 Overview: Structured State Space Duality

2.4 概述:结构化状态空间对偶性 (Structured State Space Duality)

在这一节中,我们将介绍结构化状态空间对偶性的概念。状态空间模型是一种用于描述系统动态行为的数学框架,而对偶性则揭示了不同表示形式之间的内在联系。通过理解这种对偶性,我们可以更好地分析和设计复杂的系统模型。

结构化状态空间对偶性主要探讨了如何将一个系统的状态空间表示与其对偶表示相互转换。这种转换不仅有助于简化计算,还能提供新的视角来理解系统的特性。具体来说,通过对偶变换,我们可以将某些难以处理的问题转化为更容易解决的形式,从而提高求解效率。

在后续的内容中,我们将详细讨论结构化状态空间对偶性的理论基础、应用场景以及具体的实现方法。此外,我们还将介绍一些相关的研究进展和未来的发展方向 [20]。

While this paper develops a much richer framework of connections between SSMs, attention, and structured matrices, we provide a brief summary of the main method, which is actually quite self-contained and simple algorithmic ally.

虽然本文建立了一个更为丰富的 SSMs、注意力机制 (attention) 和结构化矩阵之间的联系框架,但我们在此简要总结主要方法,该方法实际上相当独立且算法上非常简单。

Recurrent (Linear) Form. The state space dual (SSD) layer can be defined as a special case of the selective SSM (2). The standard computation of an SSM as a recurrence (or parallel scan) can be applied, which has linear complexity in sequence length. Compared to the version used in Mamba, SSD has two minor differences:

递归(线性)形式。状态空间对偶 (SSD) 层可以定义为选择性 SSM (2) 的一个特例。标准的 SSM 计算可以通过递归(或并行扫描)来实现,其在序列长度上具有线性复杂度。与 Mamba 中使用的版本相比,SSD 有两个小的差异:

  1. 不翻译无法识别的特殊字符和公式,原样返回
  2. 将HTML表格格式转换成Markdown表格格式
  3. 根据英文内容翻译成符合中文表达习惯的内容,不要遗漏任何信息

最终只返回Markdown格式的翻译结果,不要回复无关内容。


递归(线性)形式。状态空间对偶 (SSD) 层可以定义为选择性 SSM (2) 的一个特例。标准的 SSM 计算可以通过递归(或并行扫描)来实现,其在序列长度上具有线性复杂度。与 Mamba 中使用的版本相比,SSD 有两个小的差异:

  1. ...
  2. ...

(注:由于原文中没有具体列出这两个差异的内容,因此这里保留了原文的结构,等待进一步的内容补充。)

Compared to the original selective SSM, these changes can be viewed as slightly decreasing the expressive power in return for significant training efficiency improvements. In particular, our new algorithms will allow the use of matrix multiplication units on modern accelerators.

与原始的选择性 SSM (selective SSM) 相比,这些改动可以被视为在一定程度上降低了表达能力,但显著提高了训练效率。特别是,我们的新算法将允许在现代加速器上使用矩阵乘法单元。

Dual (Quadratic) Form. The dual form of SSD is a quadratic computation closely related to attention, defined as

对偶 (二次) 形式 (Dual Quadratic Form)。SSD 的对偶形式是一种与注意力机制密切相关的二次计算,定义为

请注意:原文中的“SSD”没有给出全称,因此在翻译中保留了英文原样。如果需要进一步解释或提供更多上下文,请告知。

$$
\begin{array}{r}{(L\circ Q K^{\top})\cdot V\qquad L_{i j}={\left{\begin{array}{l l}{a_{i}\times\cdot\cdot\times a_{j+1}}&{i\geq j}\ {0}&{i<j}\end{array}\right.}}\end{array}
$$

$$
\begin{array}{r}
(L \circ Q K^{\top}) \cdot V \qquad L_{i j} = \left{
\begin{array}{l l}
{a_{i} \times \cdots \times a_{j+1}} & {i \geq j} \
{0} & {i < j}
\end{array}
\right.
\end{array}
$$

公式中的 ( L ) 是一个下三角矩阵,其中 ( L_{ij} ) 的值取决于索引 ( i ) 和 ( j ) 的关系。当 ( i \geq j ) 时,( L_{ij} ) 等于从 ( a_i ) 到 ( a_{j+1} ) 的乘积;当 ( i < j ) 时,( L_{ij} ) 为 0。该公式描述了矩阵 ( L ) 与查询矩阵 ( Q ) 和键矩阵 ( K ) 的转置相乘后再与值矩阵 ( V ) 相乘的过程。

where $a_{i}$ are input-dependent scalars bounded in [0, 1].

其中,$a_{i}$ 是依赖于输入的标量,取值范围在 [0, 1] 内。

Compared to standard softmax attention, there are two main differences

与标准的 softmax 注意力机制相比,主要有两个不同点

Both of these changes can be viewed as addressing problems in vanilla attention. For example, the softmax has been recently observed to cause problems in attention scores, such as the “attention sink” phenomenon (Darcet et al. 2024; Xiao et al. 2024). More importantly, the mask matrix $L$ can be viewed as replacing the heuristic positional embeddings of Transformers with a different data-dependent positional mask that controls how much information is transfered across time.

这两项改进都可以被视为解决 vanilla attention 中的问题。例如,最近的研究发现 softmax 会导致注意力分数中的问题,如“注意力陷阱”现象 (Darcet et al. 2024; Xiao et al. 2024)。更重要的是,掩码矩阵 $L$ 可以被视为用不同的数据依赖的位置掩码替换了 Transformer 中的启发式位置嵌入,从而控制了时间上的信息传递量。

这些改进不仅解决了 softmax 在注意力机制中的潜在问题,还通过引入数据依赖的位置掩码,使得模型能够更灵活地处理序列中的位置信息,而不再依赖于固定的正弦/余弦位置编码。

More broadly, this form is an instance of our structured masked attention generalization of linear attention, defined in Section 4.

更广泛地说,这种形式是我们在第 4 节中定义的结构化掩码注意力对线性注意力的泛化的一个实例。

Matrix Form and SSD Algorithm. The various forms of SSD are connected through a unified matrix representation, by showing that SSMs have a matrix transformation form $Y=M X$ for a matrix $M_{\theta}\in\bar{\mathbb{R}}^{(\mathsf{T},\mathsf{T})}$ that depends on $\theta=(A,B,C)$ . In particular, the dual form of SSD is equivalent to naive (quadratic-time) multiplication by the matrix $M$ , and the recurrent form is a particular efficient (linear-time) algorithm that leverages the structure in $M$ .

矩阵形式与SSD算法。各种形式的SSD通过统一的矩阵表示相互关联,具体表现为状态空间模型 (SSM) 具有矩阵变换形式 $Y = M X$,其中矩阵 $M_{\theta} \in \bar{\mathbb{R}}^{(\mathsf{T},\mathsf{T})}$ 依赖于参数 $\theta = (A,B,C)$。特别地,SSD的对偶形式等价于矩阵 $M$ 的朴素(二次时间复杂度)乘法,而递归形式则是一种特定高效的(线性时间复杂度)算法,该算法利用了矩阵 $M$ 的结构。

Going beyond these, any algorithm for multiplication by $M$ can be applied. Our proposed hardware-efficient SSD algorithm (Section 6) is a new structured matrix multiplication method that involves block decomposition s of $M$ , which obtains better efficiency tradeoffs than either the pure linear or quadratic forms. It is relatively simple and easy-to-implement compared to general selective SSMs (Gu and Dao 2023); Listing 1 provides a complete implementation in a few lines of code.

除了上述方法,任何用于乘以 $M$ 的算法都可以应用。我们提出的硬件高效的 SSD 算法(第 6 节)是一种新的结构化矩阵乘法方法,涉及对 $M$ 进行块分解,相比纯线性或二次形式,该方法在效率权衡上表现更好。与一般的选择性 SSMs (Gu and Dao 2023) 相比,该方法相对简单且易于实现;列表 1 提供了仅用几行代码的完整实现。

Figure 1 provides a simple roadmap of the relationships between the concepts presented in this paper.

图 1: 提供了本文所介绍概念之间关系的简单路线图。

2.5 Notation

2.5 符号说明

根据上述要求,仅此标题无需提供更多内容。

Throughout this paper, we prefer using precise notation that can be mapped to code.

在整个论文中,我们倾向于使用可以映射到代码的精确符号。

Matrices and Vectors. We generally use lower case to denote vectors (i.e. tensors with a single axis) and upper case to denote matrices (i.e. tensors with more than one axes). We do not bold matrices in this work. Sometimes, if a matrix is tied or repeated along one axis (and hence can also be viewed as a vector), we may use either upper or lower case for it.2 · denotes scalar or matrix multiplication while $^{\circ}$ denotes Hadamard (element wise) multiplication.

矩阵和向量。我们通常使用小写字母表示向量(即单轴张量),使用大写字母表示矩阵(即多轴张量)。在本文中,我们不对矩阵加粗。有时,如果一个矩阵在一个轴上是绑定或重复的(因此也可以被视为向量),我们可以使用大写或小写字母来表示它。· 表示标量或矩阵乘法,而 ^{\circ} 表示 Hadamard (逐元素) 乘法。

Indexing. We use Python-style indexing, e.g. $i:j$ refers to the range $(i,i+1,\ldots,j-1)$ when $i<j$ and $(i,i-1,\ldots,j+1)$ when $i,>,j$ . For example, for any symbol $v$ we let $v_{j:i}$ for $j~\ge~i$ denote the sequence $(v_{j},\dotsc,v_{i+1})$ . [𝑖] is equivalent to $0:i=(0,\ldots,i-1)$ . For shorthand, we also let $\boldsymbol{v}{j:i}^{\times}$ denote the product $v{j}\times\cdots\times v_{i+1}$ .3

索引。我们使用 Python语言 风格的索引,例如 \$i:j\$ 表示范围 \$(i, i+1, \ldots, j-1)\$ 当 \$i < j\$ 时,表示范围 \$(i, i-1, \ldots, j+1)\$  \$i > j\$ 时。例如,对于任何符号 \$v\$,我们让 \$v_{j:i}\$(当 \$j \ge i\$ 时)表示序列 \$(v_{j}, \dotsc, v_{i+1})\$。[𝑖] 等价于 \$0:i = (0, \ldots, i-1)\$。为了简洁,我们还让 \$\boldsymbol{v}_{j:i}^{\times}\$ 表示乘积 \$v_{j} \times \cdots \times v_{i+1}\$。

Dimensions. To distinguish from matrices and tensors, we often use capital letters in typewriter fonts (e.g. D, N, T) to denote dimensions and tensor shapes. Instead of the traditional notation $\dot{M},\in,\mathbb{R}^{T\times T}$ we frequently use $\bar{M^{\mathrm{~}}}\in\mathbb{R}^{(\mathsf{T},\mathsf{T})}$ to reflect tensor shapes in code.

维度。为了区分矩阵和张量,我们通常使用打字机字体的大写字母(例如 D, N, T)来表示维度和张量形状。不同于传统的表示方法 \$\dot{M}\,\in\,\mathbb{R}^{T\times T}\$,我们经常使用 \$\bar{M^{\mathrm{~\}}}\in\mathbb{R}^{(\mathsf{T},\mathsf{T})}\$ 来反映代码中的张量形状。

Tensor Contractions. We will heavily rely on tensor contraction or einsum notation both for clarity and as a central tool in stating and proving our results. We assume the reader to be familiar with this notation, which is commonly used in modern tensor libraries such as numpy. For example, we can use contract( $1\mathsf{N},\mathsf{N K}\to\mathsf{M K}$ ) to denote the matrix-matrix multiplication operator, and in our notation contract $\mathsf{M N},\mathsf{N K}\to\mathsf{M K})(X,Y)$ (which is equivalent to $X\cdot Y!$ ) can be translated to code as numpy.einsum $^{\prime}{\mathfrak{m}},{\mathfrak{n}}{\models}\to{\mathfrak{m}}{\mathsf{k}}^{\prime},{\mathsf{X}},{\mathsf{Y}}$ ).

张量收缩。我们将大量依赖张量收缩或 einsum 表示法,这不仅有助于清晰表达,也是我们在陈述和证明结果时的核心工具。我们假设读者对这种表示法已经熟悉,它在现代张量库(如 numpy)中被广泛使用。例如,我们可以用 contract( 1N, NK → MK ) 来表示矩阵-矩阵乘法运算符,在我们的表示法中 contract( MN, NK → MK )(X, Y) (等价于 $X \cdot Y$)可以翻译为代码 numpy.einsum('mn,nk->mk', X, Y)

我们将张量收缩作为核心工具,用于简化复杂的多维数组操作,并确保表达式的清晰性和可读性。通过这种方式,我们可以更方便地处理高维数据结构,并在理论推导中保持一致性。

A large glossary of notation is included in Appendix A.

一个包含大量符号说明的词汇表见附录 A。

3 State Space Models are Structured Matrices

3 状态空间模型是结构化矩阵 (State Space Models are Structured Matrices)


根据要求,仅翻译了标题部分。如果有更多内容,请继续提供,我将按照规则进行翻译。

This section explores different perspectives of the state space model as a sequence transformation, and outlines properties and algorithms of such maps. The main results of this section are about the equivalence between state space models and a family of structured matrices called semi separable matrices, which imply new efficiency results (Theorems 3.5 and 3.7).

本节从不同角度探讨了状态空间模型作为序列变换的特性,并概述了此类映射的性质和算法。本节的主要结果是关于状态空间模型与一类称为半可分矩阵 (semi separable matrices) 的结构化矩阵之间的等价性,这暗示了新的效率结果(定理 3.5 和 3.7)。

3.1 The Matrix Transformation Form of State Space Models

3.1 状态空间模型的矩阵变换形式 (The Matrix Transformation Form of State Space Models)

状态空间模型是一种用于描述动态系统的数学模型,它通过一组状态变量来表示系统的内部状态。在本节中,我们将介绍状态空间模型的矩阵变换形式,这是一种将状态空间模型表示为矩阵方程的方法,便于进行分析和计算。

状态空间模型通常由以下两个方程组成:

  • 状态方程:描述系统状态随时间的变化
  • 输出方程:描述系统输出与状态之间的关系

通过矩阵变换,我们可以将这两个方程表示为更简洁的形式,便于使用线性代数工具进行求解和分析。具体来说,状态空间模型的矩阵变换形式可以写为:

[ \mathbf{x}(k+1) = A \mathbf{x}(k) + B \mathbf{u}(k) ]
[ \mathbf{y}(k) = C \mathbf{x}(k) + D \mathbf{u}(k) ]

其中:

  • (\mathbf{x}(k)) 是状态向量,表示系统在时间 (k) 的状态
  • (\mathbf{u}(k)) 是输入向量,表示系统在时间 (k) 的外部输入
  • (\mathbf{y}(k)) 是输出向量,表示系统在时间 (k) 的输出
  • (A) 是状态转移矩阵,描述状态之间的转换关系
  • (B) 是输入矩阵,描述输入对状态的影响
  • (C) 是输出矩阵,描述状态对输出的影响
  • (D) 是直接传递矩阵,描述输入对输出的直接影响

这种矩阵形式的状态空间模型不仅简化了系统的表示,还为后续的控制设计、稳定性分析等提供了便利。

Recall that our definition of an SSM is defined as a parameterized map defined through (2). Our theoretical framework starts by simply writing this transformation as a matrix multiplication mapping the vectors $x\in\mathbb{R}^{\intercal}\mapsto y\in\mathbb{R}^{\intercal}$ .

回顾我们对 SSM 的定义,它是通过 (2) 定义的参数化映射。我们的理论框架首先简单地将这个变换表示为矩阵乘法,将向量 $x\in\mathbb{R}^{\intercal}\mapsto y\in\mathbb{R}^{\intercal}$ 进行映射。

注意:这里的 $\intercal$ 可能是排版错误,通常应为转置符号 $^T$。如果确实是 $\intercal$,则保持原样。

By definition, $\boldsymbol{h}{0}=B{0}\boldsymbol{x}_{0}$ . By induction,

根据定义,$\boldsymbol{h}{0}=B{0}\boldsymbol{x}_{0}$ 。通过归纳法,

$$
\begin{array}{l}{{h_{t}=A_{t}\dots A_{1}B_{0}x_{0}+A_{t}\dots A_{2}B_{1}x_{1}+\dots+A_{t}A_{t-1}B_{t-2}x_{t-2}+A_{t}B_{t-1}x_{t-1}+B_{t}x_{t}}}\ {{\mathrm{}}}\ {{\mathrm{}}=\displaystyle\sum_{s=0}^{t}A_{t:s}^{\times}B_{s}x_{s}.}\end{array}
$$

$$
\begin{array}{l}
h_{t} = A_{t} \dots A_{1} B_{0} x_{0} + A_{t} \dots A_{2} B_{1} x_{1} + \dots + A_{t} A_{t-1} B_{t-2} x_{t-2} + A_{t} B_{t-1} x_{t-1} + B_{t} x_{t} \
\
= \sum_{s=0}^{t} A_{t:s}^{\times} B_{s} x_{s}.
\end{array}
$$

该公式表示时间步长为 ( t ) 时的状态 ( h_t ),它是通过一系列矩阵乘积和向量加权求和得到的。具体来说,( A_{t:s}^{\times} ) 表示从时间步长 ( s ) 到 ( t ) 的矩阵连乘,而 ( B_s ) 和 ( x_s ) 分别是第 ( s ) 步的权重矩阵和输入向量。

Multiplying by $C_{t}$ to produce $y_{t}$ and vector i zing the equation over $t\in[\mathsf{T}]$ , we derive the matrix transformation form of SSMs.

将 $C_{t}$ 与之相乘以生成 $y_{t}$,并对 $t \in [\mathsf{T}]$ 进行向量化处理,我们推导出状态空间模型 (SSMs) 的矩阵变换形式。

$$
\begin{array}{r l}&{y_{t}=\displaystyle\sum_{s=0}^{t}C_{t}^{\top}A_{t:s}^{\times}B_{s}x_{s}}\ &{~~~y=\mathrm{SSM}(A,B,C)(x)=M x}\ &{M_{j i}:=C_{j}^{\top}A_{j}\cdot\cdot\cdot A_{i+1}B_{i}}\end{array}
$$

$$
\begin{array}{r l}
& y_{t} = \displaystyle\sum_{s=0}^{t} C_{t}^{\top} A_{t:s}^{\times} B_{s} x_{s} \
& ~~~ y = \mathrm{SSM}(A, B, C)(x) = M x \
& M_{j i} := C_{j}^{\top} A_{j} \cdot \cdot \cdot A_{i+1} B_{i}
\end{array}
$$

这段公式描述了状态空间模型 (SSM) 的计算过程。具体来说:

  • ( y_t ) 是时间步 ( t ) 的输出,由从时间步 0 到 ( t ) 的所有输入 ( x_s ) 经过矩阵 ( A )、( B ) 和 ( C ) 的变换后得到。
  • ( y ) 是整个系统的输出,可以表示为矩阵 ( M ) 与输入向量 ( x ) 的乘积。
  • ( M_{ji} ) 是矩阵 ( M ) 的第 ( j ) 行第 ( i ) 列的元素,由矩阵 ( C )、( A ) 和 ( B ) 的特定组合计算得出。

这些公式展示了 SSM 模型如何通过一系列线性变换来处理时序数据。

3.2 Semi separable Matrices

3.2 半可分矩阵 (Semi separable Matrices)

$M$ in equation (3) is a particular representation of a class of matrices known as semi separable matrices. Semi separable matrices are a fundamental matrix structure. We first define these matrices and their properties.

公式 (3) 中的 $M$ 是一类被称为半可分矩阵 (semi separable matrices) 的特殊表示。半可分矩阵是一种基本的矩阵结构。我们首先定义这些矩阵及其性质。

半可分矩阵的定义与性质

半可分矩阵是指具有特定结构的矩阵,其元素可以表示为两个向量的外积形式。具体来说,对于一个半可分矩阵 $A$,其元素 $a_{ij}$ 可以表示为:

$$
a_{ij} = u_i v_j
$$

其中,$u_i$ 和 $v_j$ 分别是两个向量的元素。这种结构使得半可分矩阵在计算和存储上具有显著的优势,特别是在大规模矩阵运算中。半可分矩阵的性质包括但不限于:

  • 低秩性:半可分矩阵通常具有较低的秩,这使得它们在许多应用场景中可以进行高效的近似和压缩。
  • 快速算法:由于其特殊的结构,半可分矩阵可以利用快速算法进行求解,例如快速矩阵乘法和特征值分解。

接下来,我们将详细讨论半可分矩阵的具体应用及其在不同领域的意义。

Definition 3.1. A (lower triangular) matrix 𝑀is N-semi separable if every submatrix contained in the lower triangular portion (i.e. on or below the diagonal) has rank at most N. We call N the order or rank of the semi separable matrix.

定义 3.1. 一个 (下三角) 矩阵 𝑀 是 N-半可分的,如果其下三角部分(即对角线及以下的部分)包含的每个子矩阵的秩不超过 N。我们称 N 为该半可分矩阵的阶数或秩 [20]。

注:在该定义中,"下三角部分" 指的是矩阵中对角线及其以下的所有元素构成的部分。"秩" (rank) 是指矩阵中线性无关的行或列的最大数量。

Definition 3.1, and other forms of related “separable” structure (e.g. quasi separable matrices and other definitions of semi separable matrices) are sometimes called structured rank matrices (or rank-structured matrices) because they are characterized by rank conditions on their sub matrices. Semi separable matrices have many structured representations including the hierarchical semi separable (HSS), sequential semi separable (SSS), and Bruhat forms (Pernet and Storjohann 2018). We will primarily use the SSS form.

定义 3.1 以及其他形式的“可分离”结构(例如准可分离矩阵和其他半可分离矩阵的定义)有时被称为结构秩矩阵 (structured rank matrices) 或秩结构矩阵 (rank-structured matrices),因为它们是通过其子矩阵的秩条件来表征的。半可分离矩阵有许多结构表示形式,包括分层半可分离 (HSS)、顺序半可分离 (SSS) 和 Bruhat 形式 [20]。我们将主要使用 SSS 形式。

3.2.1 The Sequentially Semi separable (SSS) Representation

3.2.1 顺序半可分 (SSS) 表示

将按照上述要求继续翻译后续内容。如果您有更多内容需要翻译,请提供。

Definition 3.2. A lower triangular matrix $M\in\mathbb{R}^{(\mathsf{T},\mathsf{T})}$ has a N-sequentially semi separable (SSS) representation if it can be written in the form

定义 3.2. 一个下三角矩阵 $M\in\mathbb{R}^{(\mathsf{T},\mathsf{T})}$ 具有 N-顺序半可分 (N-sequentially semi separable, SSS) 表示,如果它可以表示为以下形式:

请注意,由于公式部分无法识别,因此原样返回。

$$
M_{j i}=C_{j}^{\top}A_{j}\cdot\cdot\cdot A_{i+1}B_{i}
$$

$$
M_{j i}=C_{j}^{\top} A_{j} \cdot \cdot \cdot A_{i+1} B_{i}
$$

公式表示矩阵 ( M_{ji} ) 由矩阵 ( C_j ) 的转置 ( C_{j}^{\top} ) 与一系列矩阵 ( A_j, \ldots, A_{i+1} ) 以及矩阵 ( B_i ) 的乘积构成。

for vectors $B_{0},\dots,B_{\intercal-1},C_{0},\dots,C_{\intercal-1}\in\mathbb{R}^{\intercal}$ and matrices $A_{0},\ldots,A_{\mathsf{T}-1}\in\mathbb{R}^{(\mathsf{N},\mathsf{N})}.$ .

对于向量 $B_{0},\dots,B_{\intercal-1},C_{0},\dots,C_{\intercal-1}\in\mathbb{R}^{\intercal}$ 以及矩阵 $A_{0},\ldots,A_{\mathsf{T}-1}\in\mathbb{R}^{(\mathsf{N},\mathsf{N})}$。

We define the operator SSS so that $M=\mathrm{SSS}(A_{0:\top},B_{0:\top},C_{0:\top})$ .

我们定义操作符 SSS,使得 $M = \mathrm{SSS}(A_{0:\top}, B_{0:\top}, C_{0:\top})$。

Lemma 3.3. An N-SSS matrix $M$ with representation (4) is N-semi separable.

引理 3.3. 具有表示 (4) 的 N-SSS 矩阵 $M$ 是 N-半可分的。

Proof. Consider any off-diagonal block $M_{j:j^{\prime},i^{\prime}:i}$ where $j^{\prime}>j\ge i>i^{\prime}$ . This has an explicit rank-N factorization as

证明。考虑任意非对角块 $M_{j:j^{\prime},i^{\prime}:i}$,其中 $j^{\prime} > j \ge i > i^{\prime}$ 。该块具有明确的秩-N分解,表示为

请注意,由于原文中没有提供具体的分解表达式,因此在翻译时保留了原始的数学符号和公式格式。

$$
\left[\begin{array}{c c c c}{C_{j}^{\top}A_{j:i^{\prime}}^{\times}B_{i^{\prime}}}&{\dots}&{C_{j}^{\top}A_{j:i-1}^{\times}B_{i-1}}\ {\vdots}&{}&{\vdots}\ {C_{j^{\prime}-1}^{\top}A_{j^{\prime}-1:i^{\prime}}^{\times}B_{i^{\prime}}}&{\dots}&{C_{j^{\prime}-1}^{\top}A_{j^{\prime}-1:i-1}^{\times}B_{i-1}}\end{array}\right]=\left[\begin{array}{c}{C_{j}^{\top}A_{j:j}^{\times}}\ {\vdots}\ {C_{j^{\prime}-1}^{\top}A_{j^{\prime}-1:j}^{\times}}\end{array}\right]A_{j:i-1}^{\times}\left[A_{i-1:i^{\prime}}^{\times}B_{i^{\prime}}\right.\quad\cdot\quad\left.A_{i-1:i-1}^{\times}B_{i-1}\right],.
$$

$$
\left[\begin{array}{c c c c}
C_{j}^{\top}A_{j:i^{\prime}}^{\times}B_{i^{\prime}} & \dots & C_{j}^{\top}A_{j:i-1}^{\times}B_{i-1} \
\vdots & & \vdots \
C_{j^{\prime}-1}^{\top}A_{j^{\prime}-1:i^{\prime}}^{\times}B_{i^{\prime}} & \dots & C_{j^{\prime}-1}^{\top}A_{j^{\prime}-1:i-1}^{\times}B_{i-1}
\end{array}\right] =
\left[\begin{array}{c}
C_{j}^{\top}A_{j:j}^{\times} \
\vdots \
C_{j^{\prime}-1}^{\top}A_{j^{\prime}-1:j}^{\times}
\end{array}\right]
A_{j:i-1}^{\times}
\left[
A_{i-1:i^{\prime}}^{\times}B_{i^{\prime}} \quad \cdot \quad A_{i-1:i-1}^{\times}B_{i-1}
\right],.
$$

这段公式表示的是矩阵运算中的一个等式,具体来说是将一个复杂的矩阵乘积分解为更简单的矩阵乘积形式。左侧的矩阵是由多个矩阵相乘得到的,而右侧则是通过重新排列和分解这些矩阵来简化表达。具体的矩阵元素和运算符号保持不变,以确保数学表达的准确性。

Equation (5) will be used extensively in deriving our fast algorithms for sequence models. The other direction is wellestablished in the literature on semi separable matrices.

公式 (5) 将在推导序列模型的快速算法中被广泛应用。另一个方向在半可分矩阵的文献中已经得到了充分的研究。

Proposition 3.4. Every N-semi separable matrix has a N-SSS representation.

命题 3.4. 每个 N-半可分矩阵都有一个 N-SSS 表示。

注:N-半可分矩阵 (N-semi separable matrix) 和 N-SSS 表示 (N-SSS representation) 是专业术语,首次出现时保留英文原文。

Furthermore, note that although Definition 3.2 involves $O(\mathsf{N}^{2}\mathsf{T})$ parameters for the representation (in particular to store the $A$ matrices), it can actually be compressed down to $O(\mathsf{N T})$ parameters, which is asymptotically tight (Pernet, Signargout, and Villard 2023). Therefore in the rest of this paper we will conflate the structured matrix class (Definition 3.1) and a particular representation of it (Definition 3.2); we will always use this representation instead of other candidates. In turn we will use N-SS to refer to an N-semi separable matrix in SSS form.

此外,需要注意的是,尽管定义 3.2 涉及 $O(\mathsf{N}^{2}\mathsf{T})$ 参数(特别是用于存储 $A$ 矩阵),但实际上它可以被压缩到 $O(\mathsf{N T})$ 参数,这是渐近最优的 (Pernet, Signargout, 和 Villard 2023)。因此,在本文的其余部分,我们将把结构化矩阵类(定义 3.1)和它的特定表示形式(定义 3.2)视为等同;我们将始终使用这种表示形式,而不是其他候选形式。相应地,我们将使用 N-SS 来指代以 SSS 形式表示的 N-半可分离矩阵。

Semi separable matrices are a fundamental matrix structure and have many important properties. They are deeply related to recurrences at large, and can be defined by multiple characterizations (e.g. Definitions 3.1 and 3.2) which reveal different connections and efficient algorithms for them. We mention some of their other properties in Appendix C.1.

半可分矩阵是一种基本的矩阵结构,具有许多重要的性质。它们与广泛的递归关系密切相关,并可以通过多种特征描述(例如定义 3.1 和 3.2)来定义,这些特征揭示了它们的不同联系和高效算法。我们在附录 C.1 中提到了它们的一些其他性质。

Remark 2. The notion of semi se par ability is very broad and many similar but subtlely different definitions appear in the literature; our definitions may differ slightly from other conventions. First, because we are primarily concerned with causal or auto regressive settings in this paper, we have restricted the definition of semi se par ability to the triangular case; Definition 3.1 more formally might be called (N, 0)-semi se par ability by some authors. Some authors may also instead refer to it as a form of quasi se par ability (Eidelman and Gohberg 1999; Pernet 2016). See Vandebril et al. (2005) for a brief survey.

备注 2. 半可分性的概念非常广泛,文献中出现了许多类似但略有不同的定义;我们的定义可能与其他约定略有不同。首先,由于本文主要关注因果或自回归场景,我们将半可分性的定义限制为三角形情况;定义 3.1 更正式地可能被一些作者称为 (N, 0)-半可分性。有些作者也可能将其称为准可分性的一种形式 (Eidelman 和 Gohberg 1999; Pernet 2016)。关于这一主题的简要综述,可以参见 Vandebril 等人 (2005) 的文章 [20]。

3.2.2 1-Semi separable Matrices: the Scalar SSM Recurrence

3.2.2 1-半可分矩阵:标量 SSM 递推

在这一节中,我们将介绍 1-半可分矩阵 (1-Semi separable Matrices) 的标量 SSM (Scalar SSM) 递推关系。这类矩阵具有一些特殊的结构特性,使得它们在计算和存储上具有优势。具体来说,1-半可分矩阵的定义是:对于任意的 ( i, j ),当 ( |i - j| > 1 ) 时,矩阵元素 ( A_{ij} = 0 )。这种结构使得矩阵的非零元素主要集中在对角线及其邻近位置。

标量 SSM 递推关系描述了如何通过递推的方式计算 1-半可分矩阵的某些属性或操作结果。该递推关系的核心思想是利用矩阵的特殊结构,减少计算复杂度并提高效率。具体的递推公式将在后续部分详细讨论。

在实际应用中,1-半可分矩阵广泛出现在数值分析、信号处理等领域,尤其是在求解线性方程组和特征值问题时表现出色。通过对这些矩阵的递推关系进行研究,可以为相关领域的算法设计提供理论支持。

We will single out the special case of 1-SS matrices. Note that in this case, the $C_{j}$ and $B_{i}$ are scalars, and can be factored out of the SSS representation (4) (we also use lower-case to emphasize that the parameters are scalars in this case)

我们将特别讨论 1-SS 矩阵这个特殊情况。需要注意的是,在这种情况下,$C_{j}$ 和 $B_{i}$ 是标量,可以从 SSS 表示 (4) 中提取出来(在这种情况下,我们还使用小写字母来强调参数是标量)。

$$
\mathrm{SSS}(a,b,c)=\mathrm{diag}(c)\cdot M\cdot\mathrm{diag}(b)\qquad\mathrm{where}\qquad M_{j i}=a_{j:i}^{\times}.
$$

$$
\mathrm{SSS}(a,b,c)=\mathrm{diag}(c) \cdot M \cdot \mathrm{diag}(b) \qquad \text{其中} \qquad M_{j i}=a_{j:i}^{\times}.
$$

注:公式中的符号和格式保持不变,仅将英文单词 "where" 翻译为中文 "其中"。

Since diagonal matrices are easy to handle (e.g. multiplication by a diagonal matrix is the same as element wise scalar multiplication), we can ignore these terms. Thus our basic representation of a 1-SS matrix is $M_{j i}=a_{j:i}$ or

由于对角矩阵易于处理(例如,与对角矩阵相乘等同于逐元素标量乘法),我们可以忽略这些项。因此,我们对 1-SS 矩阵的基本表示为 $M_{j i}=a_{j:i}$ 或者

请注意,原文中没有提供完整的句子结尾,因此翻译也在此处结束。

$$
M=15{\bf S}(a_{0:T}):=\left[\begin{array}{c c c c c c}{{1}}&{{}}&{{}}&{{}}&{{}}&{{}}\ {{a_{1}}}&{{1}}&{{}}&{{}}&{{}}&{{}}\ {{a_{2}a_{1}}}&{{a_{2}}}&{{1}}&{{}}&{{}}&{{}}\ {{\vdots}}&{{}}&{{\vdots}}&{{\ddots}}&{{\ddots}}&{{}}\ {{a_{T-1}\ldots a_{1}}}&{{a_{T-1}\ldots a_{2}}}&{{\ldots}}&{{a_{T-1}}}&{{1}}\end{array}\right].
$$

$$
M = 15 \bf{S} (a_{0:T}) := \left[ \begin{array}{c c c c c c}
{1} & {} & {} & {} & {} & {} \
{a_1} & {1} & {} & {} & {} & {} \
{a_2 a_1} & {a_2} & {1} & {} & {} & {} \
{\vdots} & {} & {\vdots} & {\ddots} & {\ddots} & {} \
{a_{T-1} \ldots a_1} & {a_{T-1} \ldots a_2} & {\ldots} & {a_{T-1}} & {1}
\end{array} \right].
$$

该公式定义了一个矩阵 ( M ),其中 ( M = 15 \bf{S} (a_{0:T}) )。矩阵的具体形式如下:

  • 第一行是 [1, 0, 0, ..., 0]
  • 第二行是 [( a_1 ), 1, 0, ..., 0]
  • 第三行是 [( a_2 a_1 ), ( a_2 ), 1, ..., 0]
  • 依此类推,直到最后一行 [( a_{T-1} \ldots a_1 ), ( a_{T-1} \ldots a_2 ), ..., ( a_{T-1} ), 1]

矩阵的每一行都包含了前几项的乘积,形成了一个下三角矩阵结构。

The importance of 1-SS matrices lies in their equivalence to the minimal form of a scalar recurrence – the case of a degenerate SSM with state dimension $\mathsf{N}=1$ and no $(B,C)$ projections. Note that multiplication $y=M x$ can be computed by the recurrence

1-SS矩阵的重要性在于它们等价于标量递归的最小形式——即退化的状态空间模型 (SSM) 的情况,其中状态维度 $\mathsf{N}=1$ 且没有 $(B,C)$ 投影。需要注意的是,乘法 $y=M x$ 可以通过递归计算得出

1-SS矩阵的重要性在于它们等价于标量递归的最小形式——即退化的状态空间模型 (SSM) 的情况,其中状态维度 \$\mathsf{N}=1\$ 且没有 \$(B,C)\$ 投影。需要注意的是,乘法 \$y=M x\$ 可以通过递归计算得出

$$
{\begin{array}{r l}&{y_{t}=a_{t:0}x_{0}+\cdot\cdot\cdot+a_{t:t}x_{t}}\ &{\quad=a_{t}\left(a_{t-1:0}x_{0}+\cdot\cdot\cdot+a_{t-1:t-1}x_{t-1}\right)+a_{t:t}x_{t}}\ &{\quad=a_{t}y_{t-1}+x_{t}.}\end{array}}
$$

$$
\begin{array}{r l}
& y_{t} = a_{t:0} x_{0} + \cdot \cdot \cdot + a_{t:t} x_{t} \
& \quad = a_{t} \left( a_{t-1:0} x_{0} + \cdot \cdot \cdot + a_{t-1:t-1} x_{t-1} \right) + a_{t:t} x_{t} \
& \quad = a_{t} y_{t-1} + x_{t}.
\end{array}
$$

该公式表示时间步 $t$ 处的输出 $y_t$ 是由前一个时间步 $t-1$ 的输出 $y_{t-1}$ 与当前输入 $x_t$ 线性组合而成。具体来说,$y_t$ 可以通过将 $a_t$ 乘以前一时刻的输出 $y_{t-1}$,再加上当前时刻的输入 $x_t$ 来计算。


Figure 2: (State Space Models are Semi separable Matrices.) As sequence transformations, state space models can be represented as a matrix transformation $M\in\bar{\mathbb{R}}^{(\mathsf{T},\mathsf{T})}$ acting on the sequence dimension T, sharing the same matrix for each channel in a head (Left). This matrix is a semi separable matrix (Right), which is a rank-structured matrix where every submatrix contained on-and-below the diagonal $(B l u e)$ has rank at most N, equal to the SSM’s state dimension.


图 2: (状态空间模型是半可分矩阵) 作为序列变换,状态空间模型可以表示为一个作用于序列维度 T 的矩阵变换 $M\in\bar{\mathbb{R}}^{(\mathsf{T},\mathsf{T})}$(左图),每个头中的每个通道共享相同的矩阵。该矩阵是一个半可分矩阵(右图),这是一种秩结构矩阵,其中对角线及以下的每个子矩阵 (蓝色) 的秩不超过 N,N 等于状态空间模型 (SSM) 的状态维度。

We thus also refer to matrix multiplication by 1-SS matrices as the scalar SSM recurrence or the cumprodsum (cumulative product sum; a generalization of cumulative product and cumulative sum) operator. As the fundamental form of recurrence, multiplication by 1-SS matrices is important as a building block for our main algorithms.

因此,我们也将 1-SS 矩阵的矩阵乘法称为标量 SSM (scalar SSM) 递归或累积乘积和 (cumprodsum; 累积乘积和累积求和的泛化) 运算符。作为最基本的递归形式,1-SS 矩阵的乘法是构建我们主要算法的重要基础模块。

We emphasize that one of the central themes of this paper is that many algorithms on sequence models can be reduced to structured matrix multiplication algorithms. 1-SS matrices exemplify this connection: there are many fast algorithms for computing the primitive scalar recurrence or cumprodsum operator, and all of them turn out to be equivalent to different structured factorization of 1-SS matrices. We dedicate Appendix B to these algorithms for 1-SS matrix multiplication.

我们强调,本文的一个核心主题是:许多序列模型的算法可以归结为结构化矩阵乘法算法。1-SS 矩阵很好地体现了这种联系:有许多快速算法可以计算基本的标量递归或累积和乘积运算符(cumprodsum),而所有这些算法最终都等价于 1-SS 矩阵的不同结构化分解。我们在附录 B 中专门讨论了这些 1-SS 矩阵乘法算法。

3.3 State Space Models are Semi separable Matrices

3.3 状态空间模型是半可分矩阵 (State Space Models are Semi separable Matrices)

状态空间模型是一种用于描述动态系统的数学模型,而半可分矩阵是指可以分解为两个低秩矩阵乘积的矩阵。在本节中,我们将探讨状态空间模型与半可分矩阵之间的关系,并解释为什么这种特性在计算和优化中具有重要意义。

状态空间模型通常由以下四个方程定义:

  • 状态方程:描述系统内部状态的演变
  • 输出方程:描述系统输出与内部状态之间的关系
  • 初始状态:给定系统的初始条件
  • 噪声模型:描述系统中的随机扰动

通过将状态空间模型表示为矩阵形式,我们可以发现其结构具有半可分性。具体来说,状态转移矩阵和观测矩阵往往可以近似为低秩矩阵的乘积,这使得我们在处理大规模系统时能够显著减少计算复杂度。

这种半可分性不仅有助于提高计算效率,还为开发更高效的算法提供了理论基础。例如,在时间序列分析、信号处理等领域,利用状态空间模型的半可分特性可以实现更快的预测和滤波操作 [20]。

图 1: 状态空间模型的结构示例

表 1: 常见的状态空间模型应用领域

应用领域 描述
时间序列预测 用于预测未来的数据点
信号处理 用于滤波和去噪
控制系统 用于设计控制器

通过对状态空间模型的半可分性进行研究,我们可以更好地理解这些模型的内在结构,并开发出更加高效和准确的算法。

Recall that our definition of an SSM is defined as a parameterized map defined through Definition 2.1. The connection between SSMs and semi separable matrices follows from simply writing this transformation as a matrix multiplication mapping the vectors $x\mapsto y\in\mathbb{R}^{\top}$ .

回忆我们对 SSM 的定义,它是通过定义 2.1 给出的一个参数化映射。SSM 与半可分矩阵 (semi separable matrices) 之间的联系可以通过简单地将这个变换写成矩阵乘法来表示,即将向量 $x\mapsto y\in\mathbb{R}^{\top}$ 映射为矩阵乘法。

根据定义 2.1,SSM 是一个参数化的映射,而这种映射可以通过矩阵乘法的形式来表达,即将输入向量 $x$ 映射到输出向量 $y$,其中 $y$ 属于 $\mathbb{R}^{\top}$。这种表示方法揭示了 SSM 与半可分矩阵之间的紧密联系。

Equation (3) directly establishes the link between state space models and the sequentially semi separable representation, which in turn are equivalent to semi separable matrices in general (Lemma 3.3 and Proposition 3.4).

公式 (3) 直接建立了状态空间模型与顺序半可分表示之间的联系,而后者等价于一般的半可分矩阵 (引理 3.3 和命题 3.4)。

Theorem 3.5. The state space model transformation $\boldsymbol{y}=\mathsf{S S M}(A,B,C)(\boldsymbol{x})$ with state size N is identical to matrix multiplication by an N-SS matrix in sequentially semi separable representation $y=\mathrm{SSS}(A,B,C)\cdot x$ .

定理 3.5. 状态空间模型变换 $\boldsymbol{y}=\mathsf{SSM}(A,B,C)(\boldsymbol{x})$,其状态维度为 N,等价于使用顺序半分离表示的 N-SS 矩阵与向量的乘法 $y=\mathrm{SSS}(A,B,C) \cdot x$。

在这个定理中,$\mathsf{SSM}(A,B,C)$ 表示状态空间模型 (State Space Model),而 $\mathrm{SSS}(A,B,C)$ 表示顺序半分离矩阵 (Sequentially Semi Separable matrix)。两者在特定条件下是等价的,即它们对输入向量 $\boldsymbol{x}$ 的变换结果相同。

In other words the sequence transformation operator SSM (Definition 2.2) coincides with the matrix construction operator SSS (Definition 3.2), and we use them interchangeably (or sometimes SS as shorthand). Furthermore—by a twist of fate—structured state space models and sequentially semi separable matrices have the same acronyms, underscoring their equivalence! Conveniently we can use any of these acronyms SSM (state space model or semi separable matrix), SSS (structured state space or sequentially semi separable), or SS (state space or semi separable) interchangeably to unambiguously refer to either concept. However, we will generally use the convention that SSM refers to state space model, SS refers to semi separable, and SSS refers to sequentially semi separable.

换句话说,序列变换算子 SSM (定义 2.2) 与矩阵构造算子 SSS (定义 3.2) 是一致的,我们可以互换使用它们(有时也用 SS 作为简写)。此外——出于巧合——结构化状态空间模型和顺序半可分矩阵具有相同的缩写,这进一步强调了它们的等价性!为了方便起见,我们可以使用这些缩写中的任何一个:SSM (状态空间模型或半可分矩阵),SSS (结构化状态空间或顺序半可分),或 SS (状态空间或半可分) 来明确指代这两个概念中的任意一个。然而,我们通常会遵循以下约定:SSM 指的是状态空间模型,SS 指的是半可分,而 SSS 指的是顺序半可分。

Figure 2 illustrates the sequence transformation perspective of state space models as semi separable matrices.

图 2: 展示了状态空间模型作为半可分矩阵的序列变换视角。

3.4 Computing State Space Models through Structured Matrix Algorithms

3.4 通过结构化矩阵算法计算状态空间模型

通过结构化矩阵算法 (Structured Matrix Algorithms) 计算状态空间模型 (State Space Models) 是控制理论和系统辨识中的一个重要课题。状态空间模型是一种用于描述动态系统的数学模型,能够有效地表示系统的输入、输出及其内部状态之间的关系。结构化矩阵算法则提供了一种高效且数值稳定的方法来求解这些模型。

在本节中,我们将介绍如何利用结构化矩阵算法来计算状态空间模型,并讨论其在实际应用中的优势和挑战。具体来说,我们将探讨以下内容:

  • 状态空间模型的基本概念
  • 结构化矩阵的定义及其特性
  • 常见的结构化矩阵算法
  • 这些算法在状态空间模型计算中的应用实例

通过这些讨论,读者将能够更好地理解如何使用结构化矩阵算法来解决复杂的状态空间模型问题。

The reason Theorem 3.5 is important is that it will allow us to reduce the problem of efficient computation of SSMs (and other sequence models) into efficient algorithms for structured matrix multiplication. We briefly provide an overview and defer our main new algorithm to Section 6, after showing the equivalence of SSMs to other sequence models in Sections 4 and 5.

定理 3.5 的重要性在于它将允许我们将高效计算 SSMs (Sequence-to-Sequence Models) 以及其他序列模型的问题简化为结构化矩阵乘法的高效算法。我们简要概述一下,并将我们的主要新算法推迟到第 6 节,在第 4 节和第 5 节中,我们将展示 SSMs 与其他序列模型的等价性。


翻译说明:

  • “SSMs” 在第一次出现时翻译为“SSM (Sequence-to-Sequence Models)”,之后直接使用“SSM”。
  • 保留了原始段落格式和引用的章节编号。

As previously defined, semi separable matrices (i.e. rank-structured matrices) are a classical type of structured matrix:

如前所述,半可分离矩阵(即秩结构矩阵)是一种经典的结构化矩阵:

Furthermore, the parameter iz ation and matrix multiplication cost can be tight in the semi separable order.

此外,参数化 (parameterization) 和矩阵乘法的计算成本在半可分顺序 (semi separable order) 下可以得到很好的控制。

Proposition 3.6 (Pernet, Signargout, and Villard (2023)). An N-SS matrix of size T can be represented in $O(\mathsf{N T})$ parameters and has matrix-vector multiplication in time and space $O(\mathsf{N T})$ .

命题 3.6 (Pernet, Signargout, 和 Villard (2023))。大小为 T 的 N-SS 矩阵可以用 $O(\mathsf{N T})$ 参数表示,并且其矩阵-向量乘法的时间和空间复杂度为 $O(\mathsf{N T})$。

For example, 1-SS matrices illustrate the essence of this connection. The matrix $M=15\mathsf{S}(a)$ is defined by exactly T −1 parameters $a_{0:\mathsf{T}-1}=a_{1},\dotsc,a_{\mathsf{T}-1}$ , and can be computed in $O(\mathsf{T})$ time by following the scalar recurrence (7).

例如,1-SS 矩阵展示了这种连接的本质。矩阵 $M=15\mathsf{S}(a)$ 由恰好 T − 1 个参数 $a_{0:\mathsf{T}-1}=a_{1},\dotsc,a_{\mathsf{T}-1}$ 定义,并且可以通过遵循标量递推公式 (7) 在 $O(\mathsf{T})$ 时间内计算。

3.4.1 The Linear (Recurrent) Mode

3.4.1 线性(循环)模式

在该部分中,我们将介绍线性(循环)模式的工作原理及其应用场景。线性模式是一种简单的序列处理方式,而循环模式则允许模型在处理序列时具有记忆功能,能够更好地捕捉时间依赖性。

线性模式的主要特点是其处理过程是顺序的,每个输入仅影响当前的输出,而不影响后续的输入。相比之下,循环模式通过引入循环结构,使得模型可以在处理当前输入时参考之前的状态,从而实现对序列数据的更有效建模。

在实际应用中,线性模式适用于处理不需要考虑历史信息的任务,而循环模式则更适合处理需要记忆和上下文理解的任务,例如自然语言处理、语音识别等。

注意:原文中未提供更多信息,因此以上内容是基于常见技术背景的解释。如果需要更详细的技术描述,请提供更多具体信息。

Proposition 3.6 can be easily seen in the case of diagonal structured SSMs (S4D (Gu, Gupta, et al. 2022)), simply by leveraging the state space model formulation (2) and unrolling the recurrence. We provide the formal tensor-contraction algorithm in (8), where the dimension S is equal to $\mathsf{T}^{4}$ .

命题 3.6 在对角结构的 SSM (S4D (Gu, Gupta, et al. 2022)) 中很容易看出,只需利用状态空间模型公式 (2) 并展开递归即可。我们在 (8) 中提供了正式的张量收缩算法,其中维度 S 等于 $\mathsf{T}^{4}$ 。

  • 命题 3.6 在对角结构的状态空间模型 (SSM) 中(如 S4D (Gu, Gupta, et al. 2022)),通过利用状态空间模型公式 (2) 并展开递归,可以很容易地观察到。
  • 我们在 (8) 中提供了正式的张量收缩算法,其中维度 S 等于 $\mathsf{T}^{4}$。

$$
\begin{array}{r l}&{Z=\mathsf{c o n t r a c t}(\mathsf{S P},\mathsf{S N}\to\mathsf{S P N})(X,B)}\ &{H=\mathsf{c o n t r a c t}(\mathsf{T S N},\mathsf{S P N}\to\mathsf{T P N})(L,Z)}\ &{Y=\mathsf{c o n t r a c t}(\mathsf{T N},\mathsf{T P N}\to\mathsf{T P})(C,H)}\end{array}
$$

$$
\begin{array}{r l}
& Z = \mathsf{contract}(\mathsf{SP}, \mathsf{SN} \to \mathsf{SPN})(X, B) \
& H = \mathsf{contract}(\mathsf{TSN}, \mathsf{SPN} \to \mathsf{TPN})(L, Z) \
& Y = \mathsf{contract}(\mathsf{TN}, \mathsf{TPN} \to \mathsf{TP})(C, H)
\end{array}
$$

这段内容包含数学公式,根据要求不进行翻译,保持原样。

Here, $L\in\mathbb{R}^{(\intercal,\intercal)}$ is defined as ${1}{\mathsf{S S}}(A)$ , or in other words $L_{0:\mathsf{T},0:\mathsf{T}}=1\mathsf{S S}(A_{0:\mathsf{T}})$ for $i\in[\mathsf{N}]$ . This algorithm involves three steps corresponding to (2):

这里,$L \in \mathbb{R}^{(T, T)}$ 被定义为 ${1}{\mathsf{S S}}(A)$ ,或者换句话说,$L_{0:\mathsf{T}, 0:\mathsf{T}} = 1\mathsf{S S}(A_{0:\mathsf{T}})$ 对于 $i \in [\mathsf{N}]$。该算法包含三个步骤,对应于 (2):

  1. 不翻译无法识别的特殊字符和公式,原样返回
  2. 将HTML表格格式转换成Markdown表格格式
  3. 根据英文内容翻译成符合中文表达习惯的内容,不要遗漏任何信息

该算法的具体步骤如下:

  1. 初始化:根据给定的矩阵 $A$,计算其对应的 $L$ 矩阵。
  2. 迭代更新:对于每个 $i \in [\mathsf{N}]$,更新 $L$ 矩阵的相应部分。
  3. 输出结果:最终得到的 $L$ 矩阵即为所求。

注意:公式中的 $\mathbb{R}^{(T, T)}$ 表示一个 $T \times T$ 的实数矩阵,$\mathsf{S S}(A)$ 是对矩阵 $A$ 进行某种特定操作的结果。

(i) expanding the input $X$ by the input matrix $B$ (8a), (ii) unrolling independent scalar SSM recurrences (8b), and (iii) contracting the hidden state $H$ by the output matrix $C$ (8c).

(i) 通过输入矩阵 $B$ 扩展输入 $X$ (8a),(ii) 展开独立的标量 SSM 迭代 (8b),以及 (iii) 通过输出矩阵 $C$ 压缩隐藏状态 $H$ (8c)。

Note that we have used the equivalence between scalar SSMs and 1-SS matrices in step (8b).

请注意,我们在步骤 (8b) 中使用了标量 SSM 与 1-SS 矩阵之间的等价性。

Remark 3. We note that (8) is a special case of the Mamba (S6) model. however, a naive implementation is slow because of the expanded tensors $Z$ and $H$ of size (T, P, N); Gu and Dao (2023) introduced a hardware-aware implementation to avoid materializing these tensors.

备注 3. 我们注意到 (8) 是 Mamba (S6) 模型的一个特例。然而,由于扩展的张量 $Z$ 和 $H$ 的大小为 (T, P, N),直接实现会很慢;Gu 和 Dao (2023) 引入了一种硬件感知的实现方法,以避免显式生成这些张量。

Surprisingly, Theorem 3.5 and Proposition 3.6 immediately imply that all SSMs have the same asymptotic efficiency as algorithm (8).

令人惊讶的是,定理 3.5 和命题 3.6 立即表明所有 SSMs (State Space Models) 都具有与算法 (8) 相同的渐近效率。

Theorem 3.7. Any state space model (Definition 2.2) of state size N on sequence length T can be computed in time $O(\mathsf{T N})$ (not accounting for potential preprocessing).

定理 3.7. 任何状态空间模型 (Definition 2.2) 在序列长度为 T 和状态大小为 N 的情况下,可以在时间 $O(\mathsf{T N})$ 内计算完成(不考虑可能的预处理)。

注意:在第一次出现时,"状态空间模型" 对应的英文为 "state space model"。

We note that this result is new to the structured SSM literature. In particular, given dense unstructured $A_{t}$ matrices, the total representation alone seems to be of size $O(\mathsf{T N}^{2})$ . Thus Theorem 3.7 states the non-trivial result that with a preprocessing step, even an unstructured SSM can be computed optimally efficiently, with upper bound matching the lower bound $O(\mathsf{T N})$ given by the size of $B$ and $C$ .

我们注意到这一结果在结构化 SSM (Structured State Space Model) 文献中是新的。特别地,给定密集的非结构化 $A_{t}$ 矩阵,仅总表示的大小似乎为 $O(\mathsf{T N}^{2})$。因此,定理 3.7 表明了一个非平凡的结果:通过预处理步骤,即使是非结构化的 SSM 也可以被最优高效地计算,其上界与由 $B$ 和 $C$ 的大小给出的下界 $O(\mathsf{T N})$ 相匹配。

Remark 4. Theorem 3.7 is perhaps not too surprising in light of the fact that almost all dense matrices over $\mathbb{R}^{(\mathsf{N},\mathsf{N})}$ are diagonal iz able over $\mathbb{C}$ , leading to the result that almost all dense real SSMs are equivalent to a diagonal complex SSM. This fact underlies the reason why diagonal SSMs are the most popular form of structured SSM (Gu, Gupta, et al. 2022; Gupta, $G u,$ , and Berant 2022; J. T. Smith, Warrington, and Linderman 2023). However, Theorem 3.7 implies the much stronger result for all real SSMs (not just the diagonal iz able ones), as well as dense SSMs over other fields (including $\mathbb{C}$ itself).

备注 4. 定理 3.7 的结论或许并不太令人惊讶,因为几乎所有的 $\mathbb{R}^{(\mathsf{N},\mathsf{N})}$ 上的稠密矩阵都可以在 $\mathbb{C}$ 上对角化,这导致了几乎所有的稠密实数 SSM(状态空间模型)都等价于一个对角化的复数 SSM。这一事实解释了为什么对角化 SSM 是最流行的结构化 SSM 形式 (Gu, Gupta, et al. 2022; Gupta, $G u,$ , and Berant 2022; J. T. Smith, Warrington, and Linderman 2023)。然而,定理 3.7 意味着一个更强的结果,适用于所有实数 SSM(而不仅仅是可对角化的 SSM),以及在其他域(包括 $\mathbb{C}$ 本身)上的稠密 SSM。

注意:在翻译过程中保留了原始的数学符号和公式格式,并且按照要求使用了半角括号和空格。

In practice, efficiently computable SSMs still require additional structure on $A$ , particularly to avoid the expensive preprocessing step (which both has order N extra FLOPs and involves hardware-inefficient operations such as singular value decomposition s). These structures are the focus of past work on structured SSMs (e.g. S4(D) and Mamba) as well as our new algorithms. In particular, when slightly stronger structure is imposed on $A$ , we will design very hardware-efficient algorithms through block decomposition s of the SSM matrix $M={\mathrm{SSS}}(A,B,C)$ in Section 6.

在实际应用中,可高效计算的状态空间模型 (SSM) 仍然需要对矩阵 $A$ 进行额外的结构化处理,特别是为了避免昂贵的预处理步骤(该步骤不仅增加了 N 阶的额外浮点运算,还涉及硬件效率低下的操作,如奇异值分解)。这些结构是过去关于结构化 SSM 的研究(例如 S4(D) 和 Mamba)以及我们新算法的重点。具体来说,当对矩阵 $A$ 施加稍强的结构时,我们将在第 6 节通过 SSM 矩阵 $M=\mathrm{SSS}(A,B,C)$ 的块分解设计非常高效的硬件算法。

3.4.2 The Quadratic (Naive) Mode

3.4.2 二次(朴素)模式

在该模式下,模型采用了一种较为简单的处理方式,直接对输入数据进行二次运算。这种处理方法虽然直观,但在处理大规模数据时效率较低,因此被称为“朴素”模式。该模式的主要特点是计算复杂度为 O(n²),即随着输入数据量的增加,计算时间会呈平方级增长。

在这种模式下,模型的每个输出都依赖于所有输入 Token 的两两组合,导致计算量较大。尽管如此,该模式在某些特定场景下仍然具有一定的应用价值,尤其是在数据规模较小或对计算速度要求不高的情况下。

需要注意的是,该模式与更高效的优化算法相比,存在明显的性能瓶颈,因此在实际应用中通常会被更先进的技术所取代。例如,在处理大规模语言任务时,通常会使用基于 Transformer 的架构来提高计算效率 [20]。

We note that there is another way to compute an SSM exposed by our new matrix point of view. A naive computation of the matrix SSM representation (3) involves simply materializing the sequence transformation matrix $M=\mathrm{SSS}(A,B,C)$ . This is a $(\mathsf{T},\mathsf{T})$ matrix, and therefore this naive algorithm will scale quadratically in sequence length. However, when the sequence length T is short, this can actually be more efficient than the linear algorithm due to constant factors and the hardware-friendliness of the computation pattern (e.g. leveraging matrix-matrix multiplications). In fact, for a particular case of structured SSMs, this looks very similar to a quadratic attention computation (Section 5).

我们注意到,从我们新的矩阵视角出发,还有一种计算 SSM 的方法。直接计算矩阵 SSM 表示 (3) 的一种简单方法是显式地构造序列变换矩阵 $M=\mathrm{SSS}(A,B,C)$ 。这是一个 $(\mathsf{T},\mathsf{T})$ 矩阵,因此这种简单的算法在序列长度上会呈现二次方的复杂度。然而,当序列长度 T 较短时,由于常数因子和计算模式对硬件的友好性(例如利用矩阵-矩阵乘法),这种方法实际上可能比线性算法更高效。事实上,对于特定结构化的 SSM 情况,这与二次注意力计算非常相似(第 5 节)。

3.4.3 Summary

Many sequence models are explicitly motivated or defined as matrix sequence transformations – most notably Transformers, where the matrix mixer is the attention matrix. On the other hand, RNNs and SSMs have not previously been described in this way. By providing an explicit matrix transformation form of state space models, we reveal new ways of understanding and using them. From a computational perspective, any method of computing the forward pass of a state space model can be viewed as a matrix multiplication algorithm on semi separable matrices. The semi separable matrix perspective provides one lens into state space duality (SSD), where the dual modes respectively refer to a linear-time semi separable matrix multiplication algorithm and quadratic-time naive matrix multiplication.

许多序列模型被明确地动机化或定义为矩阵序列变换——最著名的是 Transformer,其中的矩阵混合器是注意力矩阵。另一方面,RNN 和 SSM 之前并未以这种方式进行描述。通过提供状态空间模型 (SSM) 的显式矩阵变换形式,我们揭示了理解和使用这些模型的新方法。从计算角度来看,任何计算状态空间模型前向传播的方法都可以被视为对半可分离矩阵的矩阵乘法算法。半可分离矩阵的视角为我们提供了一个理解状态空间对偶性 (SSD) 的窗口,其中的两种模式分别指的是线性时间复杂度的半可分离矩阵乘法算法和二次时间复杂度的朴素矩阵乘法算法。

  • 状态空间模型 (State Space Model, SSM)
  • 半可分离矩阵 (semi separable matrix)
  • 状态空间对偶性 (State Space Duality, SSD)

4 Structured Masked Attention: Generalizing Linear Attention with Structured Matrices

4 结构化掩码注意力:用结构化矩阵泛化线性注意力 (Structured Masked Attention: Generalizing Linear Attention with Structured Matrices)

根据上述要求,这是该段落的翻译结果。如果需要继续翻译后续内容,请提供。

In this section we revisit the linear attention framework from first principles. The main results in this section are a simple tensor-contraction-based proof of linear attention (Proposition 4.1), and our generalized abstraction of structured masked attention in Definition 4.2. We note that this section derives the main duality results from a different direction than state space models and can be read completely independently of Section 3.

在本节中,我们将从基本原理重新审视线性注意力机制框架。本节的主要结果包括一个基于张量收缩的线性注意力简单证明(命题 4.1),以及我们在定义 4.2 中提出的结构化掩码注意力的广义抽象。我们注意到,本节从不同于状态空间模型的方向推导了主要的对偶性结果,并且可以完全独立于第 3 节阅读。


翻译说明:

  • “线性注意力机制”对应英文中的“linear attention”。
  • “张量收缩”对应英文中的“tensor contraction”。
  • “对偶性结果”对应英文中的“duality results”。
  • 保留了原文中的术语和引用格式,如“命题 4.1”和“定义 4.2”。

4.1 The Attention Framework

4.1 注意力机制框架 (The Attention Framework)

注意力机制是许多现代深度学习模型,特别是 Transformer 模型的核心组件。该框架允许模型在处理序列数据时,专注于输入序列的不同部分,从而提高模型的性能和表达能力。在本节中,我们将详细介绍注意力机制的工作原理及其在不同应用场景中的实现方式。

4.1.1 Attention

4.1.1 注意力机制 (Attention)

注意力机制 (Attention) 是 Transformer 模型中的一个关键组件,它允许模型在处理序列数据时聚焦于不同位置的信息。通过计算输入序列中各个位置的重要性权重,注意力机制能够动态地调整对不同部分的关注度,从而提高模型的表达能力和性能。

在传统的神经网络中,模型通常会均匀地处理整个输入序列,而注意力机制则引入了一种选择性机制,使得模型可以更灵活地处理长依赖关系和复杂模式。这种机制在自然语言处理、计算机视觉等多个领域都取得了显著的效果。

注意力机制的核心思想是通过计算查询 (Query)、键 (Key) 和值 (Value) 之间的相似度来确定注意力权重。具体来说,查询和键之间的点积或缩放点积可以衡量它们的相关性,然后通过 softmax 函数将这些相关性转换为概率分布,最终加权求和得到输出。

在 Transformer 模型中,多头注意力机制 (Multi-head Attention) 进一步扩展了这一思想,通过多个并行的注意力头来捕捉不同类型的依赖关系,从而增强了模型的表征能力。每个注意力头都可以关注不同的特征子空间,这使得模型能够在不同的抽象层次上进行信息聚合。

注意力机制不仅在大语言模型 (LLM) 中发挥了重要作用,也在其他生成式 AI (Generative AI) 应用中得到了广泛应用,如图像生成、语音合成等。

The basic form of (single-head) attention is a map on three sequences of vectors $(Q,K,V)\mapsto Y$ .

基本的(单头)注意力机制的形式是三个向量序列的映射 $(Q, K, V) \mapsto Y$ 。

$$
\begin{array}{r l}{Q=\mathrm{input}}&{{}(\mathsf{T},\mathsf{N})}\ {K=\mathrm{input}}&{{}(\mathsf{S},\mathsf{N})}\ {V=\mathrm{input}}&{{}(\mathsf{S},\mathsf{P})}\ {G=Q K^{\mathsf{T}}}&{{}(\mathsf{T},\mathsf{S})}\ {M=f(G)}&{{}(\mathsf{T},\mathsf{S})}\ {Y=G V}&{{}(\mathsf{T},\mathsf{P})}\end{array}
$$

$$
\begin{array}{r l}
Q = \text{输入} & (T, N) \
K = \text{输入} & (S, N) \
V = \text{输入} & (S, P) \
G = Q K^{\mathsf{T}} & (T, S) \
M = f(G) & (T, S) \
Y = G V & (T, P)
\end{array}
$$

注:这里的 T、S、N 和 P 是维度符号,保持原样未翻译。

We use “shape annotations” to indicate the dimensions of tensors, e.g. $\mathcal{Q}\in\mathbb{R}^{(\mathsf{T},\mathsf{N})}$ . In this general form, S and T represent source and target sequence lengths, N represents the feature dimension, and P represents the head dimension.

我们使用“形状注释 (shape annotations)”来表示张量的维度,例如:$\mathcal{Q}\in\mathbb{R}^{(\mathsf{T},\mathsf{N})}$。在这个通用形式中,S 和 T 分别表示源序列长度和目标序列长度,N 表示特征维度,P 表示头维度。

The most common variant of softmax attention uses a softmax activation $f=$ softmax to normalize the rows of the $G$ matrix.

最常用的 softmax 注意力变体使用 softmax 激活函数 $f=$ softmax 来归一化 $G$ 矩阵的各行。

4.1.2 Self-Attention

4.1.2 自注意力机制 (Self-Attention)

自注意力机制是 Transformer 模型中的一个核心组件,它允许模型在处理序列数据时关注不同位置之间的关系。通过自注意力机制,模型可以为每个位置生成一个加权表示,这些权重反映了该位置与其他位置的相关性。这种机制使得模型能够在处理长序列时捕捉到更复杂的依赖关系。

自注意力机制的工作原理是通过计算查询 (Query)、键 (Key) 和值 (Value) 之间的相似度来确定注意力权重。具体来说,对于输入序列中的每个位置,模型会生成对应的查询、键和值向量。然后,通过计算查询和键之间的点积并进行缩放,得到注意力分数。这些分数经过 softmax 函数归一化后,作为权重应用于值向量,最终得到加权和作为该位置的输出表示。

自注意力机制的一个重要特点是它可以并行处理序列中的所有位置,这使得 Transformer 模型在训练和推理过程中比传统的循环神经网络 (RNN) 更高效。此外,自注意力机制还能够处理变长的输入序列,而不需要像卷积神经网络 (CNN) 那样对输入长度进行固定。

在实际应用中,自注意力机制通常与多头注意力 (Multi-Head Attention) 结合使用,以捕捉不同类型的依赖关系。多头注意力通过将输入映射到多个不同的子空间,并在每个子空间中独立计算注意力,从而增强了模型的表达能力。

Our treatment is motivated by the most important case of self-attention, where

我们的处理方法主要受到自注意力 (self-attention) 最重要案例的启发,其中

  • 自注意力机制是 Transformer 模型的核心组成部分,它允许模型在处理序列数据时关注不同位置的信息。通过自注意力机制,模型可以学习到输入序列中各个位置之间的依赖关系,而不需要依赖固定的窗口大小或预先定义的规则。
  • 在自注意力机制中,每个位置的输出是所有位置的加权和,权重由当前位置与其他位置的相关性决定。这种机制使得模型能够在处理长序列时捕捉到更复杂的模式和依赖关系。

(注:原文仅提供了一句话,因此补充了关于自注意力机制的简要解释,以帮助读者更好地理解其背景和意义。如果需要更详细的翻译或解释,请提供更多上下文。)

However, our presentation abstracts away these choices and begins from the $Q,K,V$ matrices.

然而,我们的表述抽象了这些选择,并从 $Q, K, V$ 矩阵开始。

Remark 5. Our focus is on the self-attention case with equal head and feature dimensions (i.e. ${\sf S},=,{\sf T}$ and ${\mathsf{N}}={\mathsf{P}}.$ ), which should be used as the running example. We define the general formulation of attention not only so that our framework captures variants such as cross-attention, but also because separating the notation for dimensions (e.g. S and T) makes the contraction notation proofs of our main results in this section more clear.

备注 5. 我们的关注点是自注意力机制中头部和特征维度相等的情况(即 ${\sf S},=,{\sf T}$ 和 ${\mathsf{N}}={\mathsf{P}}$),这应该作为运行示例。我们定义注意力机制的一般公式,不仅是为了使我们的框架能够涵盖诸如交叉注意力等变体,还因为将维度符号分开(例如 S 和 T)可以使本节主要结果的收缩符号证明更加清晰。

Remark 6. While attention is usually framed as an operation on these three inputs $Q,K,V$ which are viewed symmetrically, the input and output dimensions in (9) indicate otherwise. In particular, the feature dimension N is not present in the output; therefore in the case when ${\sf S};=;{\sf T}$ (e.g. self-attention), we view $V$ as the main input, so that (9) defines a proper sequence transformation $V\mapsto Y$ (Definition 2.1).

备注 6. 尽管注意力机制通常被描述为对三个输入 $Q, K, V$ 进行操作,并且这三个输入被视为对称的,但公式 (9) 中的输入和输出维度表明并非如此。特别是,特征维度 N 并未出现在输出中;因此,在 ${\sf S};=;{\sf T}$(例如自注意力机制)的情况下,我们将 $V$ 视为主要输入,使得 (9) 定义了一个有效的序列变换 $V \mapsto Y$(定义 2.1)。

在 Transformer 模型中,注意力机制的核心是通过查询 (Query)、键 (Key) 和值 (Value) 来计算加权和,而这里的讨论指出,尽管 $Q, K, V$ 在形式上是对称的,但在实际的输出中,特征维度的变化主要体现在 $V$ 上。因此,当处理自注意力时,$V$ 被视为主要的输入向量,经过注意力机制后生成输出 $Y$。

4.1.3 Kernel Attention

4.1.3 核注意力 (Kernel Attention)

核注意力 (Kernel Attention) 是一种改进的注意力机制,它通过引入核函数来优化传统注意力机制中的计算复杂度和性能。这种机制能够在保持模型表达能力的同时,显著减少计算资源的消耗。具体来说,核注意力利用核函数将输入特征映射到高维空间,从而使得相似的输入在该空间中更加接近,进而提高注意力机制的效果。

与传统的点积注意力(Dot-product Attention)相比,核注意力能够更好地处理长序列数据,并且在大规模数据集上表现出更好的泛化能力。此外,核注意力还可以与其他优化技术结合使用,例如低秩近似和稀疏化,进一步提升模型的效率和性能 [20]。

在实际应用中,核注意力已经被广泛应用于自然语言处理、计算机视觉等多个领域,尤其是在大语言模型 (LLM) 中展现了出色的性能。

The step where the softmax function is applied to the Gram matrix $G$ can be decomposed into two parts:

将 softmax 函数应用于 Gram 矩阵 $G$ 的步骤可以分解为两个部分:

  1. 计算 Gram 矩阵 $G$ 中每个元素的指数值。
  2. 对这些指数值进行归一化处理,使得每一行的和为 1。

这种分解有助于更好地理解 softmax 操作在 Gram 矩阵上的作用。

We can ignore the normalization term for now, as it amounts to simply passing in $V=1$ and dividing (we revisit this in Section 7.3). The exponentiation term can be viewed as a kernel transformation: there is an (infinite-dimensional) feature map $\varphi$ such that $\exp(Q K^{\top});=;\varphi(Q)\varphi(K)^{\top},$ . By abstracting away the feature map into the definition of $\boldsymbol{Q}$ and $K$ itself (i.e. define $Q,K$ as the post-transformed versions), we can ignore the softmax transformation, and assume that $Q,K$ are arbitrarily generated by kernel feature maps and potentially $\mathsf{N}\neq\mathsf{P}$ .

我们可以暂时忽略归一化项,因为它相当于简单地将 $V=1$ 传入并进行除法运算(我们将在第 7.3 节重新讨论这一点)。指数项可以被视为一种核变换:存在一个(无限维的)特征映射 $\varphi$,使得 $\exp(Q K^{\top}) = \varphi(Q) \varphi(K)^{\top}$。通过将特征映射抽象到 $\boldsymbol{Q}$ 和 $K$ 的定义中(即定义 $Q, K$ 为变换后的版本),我们可以忽略 softmax 变换,并假设 $Q, K$ 是由核特征映射任意生成的,且可能 $\mathsf{N} \neq \mathsf{P}$。

在这一过程中,$\varphi$ 表示特征映射函数,$\exp$ 表示指数函数,而 $Q$ 和 $K$ 分别是查询矩阵和键矩阵。

Many instantiations of kernel attention have been proposed, including:

许多核注意力 (kernel attention) 的实现已经被提出,包括:

  • 线性注意力 (Linear Attention)
  • 局部注意力 (Local Attention)
  • 稀疏注意力 (Sparse Attention)
  • 扩展注意力 (Expanded Attention)

这些方法旨在解决传统注意力机制在处理长序列时的计算和内存开销问题。通过不同的策略,它们能够在保持模型性能的同时,显著降低计算复杂度。例如,线性注意力通过将注意力矩阵的计算从二次复杂度降低到线性复杂度,使得模型能够处理更长的输入序列 [20]。

图 1: 不同类型的核注意力机制对比

表 1: 各种核注意力方法的性能比较

注意:以上内容是根据常见的核注意力变体进行的总结,具体实现可能因应用场景而异。

• The original Linear Attention (Katha ro poul os et al. 2020) defines the kernel feature map as an arbitrary pointwise activation function, such as $x\mapsto1+\mathbf{e}|\mathbf{u}(x)$ . • Random Feature Attention (RFA) (H. Peng et al. 2021) chooses the kernel feature map to approximate softmax attention (i.e. the exp feature map) using the random Fourier feature approximation of Gaussian kernels (Rahimi and Recht 2007). This involves random projections (i.e. multiplying $\boldsymbol{Q}$ and $K$ by a random projection and applying the activation $x\mapsto(\cos(x),\sin(x))$ .

• 原始的线性注意力 (Linear Attention) (Katharopoulos et al. 2020) 将核特征映射定义为任意逐点激活函数,例如 $x\mapsto1+\mathbf{e}|\mathbf{u}(x)$。

• 随机特征注意力 (Random Feature Attention, RFA) (H. Peng et al. 2021) 选择核特征映射来近似 softmax 注意力(即 exp 特征映射),使用高斯核的随机傅里叶特征近似 (Rahimi and Recht 2007)。这涉及到随机投影(即将 $\boldsymbol{Q}$ 和 $K$ 与随机投影相乘,并应用激活函数 $x\mapsto(\cos(x),\sin(x))$)。

• Performer (Cho roman ski et al. 2021) proposes the fast attention via positive orthogonal random features $\left(\mathrm{FAVOR+}\right)$ . The positive random features (PRF) part chooses the kernel feature map to be a random projection followed by the feature map $x\mapsto2^{-1/2}(\exp(x),\exp(-x))$ . This choice is motivated so that the kernel elements are positive-valued and provably approximates the softmax attention. [It also proposes choosing the random projections in orthogonal directions, which we do not consider.]

• Performer (Cho roman ski 等 2021) 提出了通过正交随机特征 (FAVOR+) 实现快速注意力机制。正随机特征 (PRF) 部分选择核特征映射为一个随机投影,随后是特征映射 $x \mapsto 2^{-1/2}(\exp(x), \exp(-x))$。这种选择的动机是为了使核元素为正值,并且可以证明它近似于 softmax 注意力机制。[它还提出了在正交方向上选择随机投影,但本文不考虑这一点。]

$\left(\mathrm{FAVOR+}\right)$ 和 $x \mapsto 2^{-1/2}(\exp(x), \exp(-x))$ 保持原样未翻译。

• cosFormer (Qin, Weixuan Sun, et al. 2022) augment RFA with a cosine re weighting mechanism that incorporates positional information to emphasize locality. This effectively passes $Q_{t},K_{t}$ through the feature map $x\mapsto$ $(x\cos(\pi t/2T),\sin(\pi t/2T))$ .

• CosFormer (Qin, Weixuan Sun, et al. 2022) 在随机特征近似 (RFA) 的基础上引入了余弦重加权机制,该机制通过引入位置信息来强调局部性。这实际上将 $Q_{t}, K_{t}$ 通过特征映射 $x \mapsto$ $(x \cos(\pi t/2T), \sin(\pi t/2T))$ 进行变换。

这种机制使得模型能够在处理长序列时更好地保留位置信息,从而提高模型的性能。

• Linear Randomized Attention (Zheng, C. Wang, and Kong 2022) generalize RFA from the perspective of importance sampling, and generalize it to provide better estimates of the full softmax kernel (rather than just the exptransformed numerator).

• 线性随机注意力 (Linear Randomized Attention) (Zheng, C. Wang, 和 Kong 2022) 从重要性采样的角度推广了 RFA,并将其扩展以提供对完整 softmax 核的更好估计(而不仅仅是 exp 变换的分子部分)。

在这一方法中,线性随机注意力通过改进采样策略,能够更准确地近似完整的 softmax 计算,从而提高模型的性能和效率。

Other related attention variants include Linformer (Sinong Wang et al. 2020) and Nys tr former (Xiong et al. 2021), which both use low-rank approximations of the attention matrix $M$ (and are thus compatible with equation (9)), through random projections (Johnson-Linden strauss) and kernel approximation (the Nyström method) respectively.

其他相关的注意力机制变体包括 Linformer (Sinong Wang 等 2020) 和 Nyströmformer (Xiong 等 2021),它们都使用了注意力矩阵 $M$ 的低秩近似(因此与公式 (9) 兼容)。具体来说,Linformer 通过随机投影(Johnson-Lindenstrauss 方法)实现低秩近似,而 Nyströmformer 则通过核函数近似(Nyström 方法)实现。

4.1.4 Masked (Kernel) Attention

4.1.4 掩码 (Kernel) 注意力机制

在这一节中,我们将介绍掩码 (Masked) 注意力机制,这是 Transformer 模型中的一个重要组件。掩码注意力机制通过阻止模型在处理序列时看到未来的信息,确保模型只能关注到当前及之前的位置。这种机制在处理自然语言任务时尤为重要,因为它模拟了人类阅读和理解文本的方式,即我们只能根据已经读过的部分来理解当前的内容。

掩码注意力机制通常用于自回归模型(如 GPT 系列),其中模型在生成下一个 Token 时,只能依赖于之前的 Token。为了实现这一点,掩码矩阵会被应用到注意力权重上,使得模型无法访问未来的位置。具体来说,掩码矩阵中的某些元素会被设置为负无穷大,这样在经过 softmax 操作后,这些位置的注意力权重将趋近于零。

此外,掩码注意力机制还可以应用于其他场景,例如在双向编码器(如 BERT)中,虽然模型可以同时看到前后文,但在某些任务中仍然需要对特定位置进行掩码,以避免信息泄露。

在一些变体中,掩码注意力机制与核函数 (Kernel) 结合使用,这被称为掩码核注意力 (Masked Kernel Attention)。核函数可以帮助加速注意力计算,并在某些情况下提高模型的性能。具体的实现细节将在后续章节中讨论 [20]。

Let $L$ be a mask of shape (T, S). Most commonly, in the auto regressive self-attention case when ${\sf S}={\sf T}$ , $L$ may be a lowertriangular matrix of 1’s representing a causal mask. Besides enforcing causality, many other types of masks can be applied – in particular various sparsity patterns such as banded, dilated, or block diagonal – which are motivated by reducing the complexity of dense attention.

设 $L$ 为形状为 (T, S) 的掩码矩阵。在自回归自注意力情况下,当 ${\sf S}={\sf T}$ 时,$L$ 通常是一个下三角矩阵,其元素为 1,表示因果掩码。除了强制因果关系外,还可以应用许多其他类型的掩码——特别是各种稀疏模式,如带状、膨胀或块对角线等——这些掩码的动机是减少密集注意力的复杂性 [20]。

  • 因果掩码:确保模型只能关注到当前和之前的 token,而不能看到未来的 token。
  • 稀疏模式:通过限制注意力机制中的非零元素分布,减少计算量和内存消耗。常见的稀疏模式包括带状(banded)、膨胀(dilated)和块对角线(block diagonal)等。

Masked attention is usually written in matrix notation as

掩码注意力通常用矩阵表示法写作

掩码注意力通常用矩阵表示法写作

$$
y=(L\circ(Q K^{\top}))\cdot V.
$$

$$
y = (L \circ (Q K^{\top})) \cdot V .
$$

公式中的符号说明:

  • ( y ) 是输出向量
  • ( L ) 是位置编码矩阵
  • ( Q ) 是查询矩阵 (Query)
  • ( K ) 是键矩阵 (Key)
  • ( V ) 是值矩阵 (Value)
  • ( \circ ) 表示逐元素乘法
  • ( \top ) 表示矩阵转置
  • ( \cdot ) 表示矩阵乘法

该公式描述了注意力机制 (Attention Mechanism) 中的计算过程,其中查询矩阵 ( Q ) 与键矩阵 ( K ) 的转置相乘,得到的矩阵再与位置编码矩阵 ( L ) 进行逐元素乘法,最后与值矩阵 ( V ) 相乘,得到输出向量 ( y )。

More precisely, with shape annotations and breaking this down into the precise sequence of computations:

更精确地说,通过形状注解并将其分解为精确的计算序列:

$$
\begin{array}{r l}{G=Q K^{\top}}&{{}\left(\top,\mathsf{S}\right)}\ {M=G\circ L}&{{}\left(\top,\mathsf{S}\right)}\ {Y=M V}&{{}\left(\top,\mathsf{P}\right)}\end{array}
$$

$$
\begin{array}{r l}
G = Q K^{\top} & (^\top, \mathsf{S}) \
M = G \circ L & (^\top, \mathsf{S}) \
Y = M V & (^\top, \mathsf{P})
\end{array}
$$

由于该部分内容主要由数学公式组成,因此保持原样不变。公式中的符号和运算保持原始格式,不做翻译处理。

Our improved derivation of attention variants in this section starts by noticing that this formula can be written as a single contraction:

我们在这部分对注意力机制变体的改进推导始于观察到该公式可以表示为一个单一的收缩:

我们在这部分对注意力机制变体的改进推导始于观察到该公式可以表示为一个单一的收缩:

$$
Y=\mathsf{c o n t r a c t}(\mathsf{T N},\mathsf{S N},\mathsf{S P},\mathsf{T S}\to\mathsf{T P})(Q,K,V,L)
$$

$$
Y = \mathsf{contract}(\mathsf{TN}, \mathsf{SN}, \mathsf{SP}, \mathsf{TS} \to \mathsf{TP})(Q, K, V, L)
$$

该公式表示一个合同函数 (contract function),它接受四个参数:$\mathsf{TN}$、$\mathsf{SN}$、$\mathsf{SP}$ 和 $\mathsf{TS} \to \mathsf{TP}$,并作用于四个输入 $Q$、$K$、$V$ 和 $L$。具体来说,$\mathsf{contract}$ 函数可能用于张量操作中的收缩 (contraction) 过程,将多个张量按照指定的模式进行合并或变换。

and the algorithm in (11) can be reframed as computing (12) by a particular ordering of pairwise contractions

算法 (11) 可以通过特定的成对收缩顺序重新表述为计算 (12)。

$$
\begin{array}{r l}&{G=\mathrm{contract}(\mathsf{T N},\mathsf{S N}\to\mathsf{T S})(Q,K)}\ &{M=\mathrm{contract}(\mathsf{T S},\mathsf{T S}\to\mathsf{T S})(G,L)}\ &{Y=\mathrm{contract}(\mathsf{T S},\mathsf{S P}\to\mathsf{T P})(M,V)}\end{array}
$$

$$
\begin{array}{r l}
& G = \mathrm{contract}(\mathsf{T N}, \mathsf{S N} \to \mathsf{T S})(Q, K) \
& M = \mathrm{contract}(\mathsf{T S}, \mathsf{T S} \to \mathsf{T S})(G, L) \
& Y = \mathrm{contract}(\mathsf{T S}, \mathsf{S P} \to \mathsf{T P})(M, V)
\end{array}
$$

这段公式描述了三个操作,每个操作都使用了 contract 函数。具体来说:

  • 第一个操作:$G = \mathrm{contract}(\mathsf{T N}, \mathsf{S N} \to \mathsf{T S})(Q, K)$ 表示将 $Q$ 和 $K$ 通过 contract 函数从 $\mathsf{T N}$ 和 $\mathsf{S N}$ 映射到 $\mathsf{T S}$,得到结果 $G$。
  • 第二个操作:$M = \mathrm{contract}(\mathsf{T S}, \mathsf{T S} \to \mathsf{T S})(G, L)$ 表示将 $G$ 和 $L$ 通过 contract 函数从 $\mathsf{T S}$ 映射到 $\mathsf{T S}$,得到结果 $M$。
  • 第三个操作:$Y = \mathrm{contract}(\mathsf{T S}, \mathsf{S P} \to \mathsf{T P})(M, V)$ 表示将 $M$ 和 $V$ 通过 contract 函数从 $\mathsf{T S}$ 和 $\mathsf{S P}$ 映射到 $\mathsf{T P}$,得到最终结果 $Y$。

这里的 $\mathsf{T N}$、$\mathsf{S N}$、$\mathsf{T S}$、$\mathsf{S P}$ 等是特定的张量形状或维度符号,contract 函数用于执行某种张量收缩操作。

$$
\begin{array}{l}{(\mathsf{T},\mathsf{S})}\ {(\mathsf{T},\mathsf{S})}\ {(\mathsf{T},\mathsf{P})}\end{array}
$$

$$
\begin{array}{l}
(T, S) \
(T, S) \
(T, P)
\end{array}
$$

根据要求,不翻译无法识别的特殊字符和公式,原样返回。

4.2 Linear Attention

4.2 线性注意力 (Linear Attention)

线性注意力机制是一种改进的注意力机制,它通过减少计算复杂度来提高模型的效率。传统的注意力机制在处理长序列时计算量较大,而线性注意力机制通过使用特定的技术(如局部窗口或稀疏连接)来降低计算成本,使得模型能够在更长的序列上进行有效的训练和推理。

线性注意力机制的核心思想是将传统的二次复杂度的注意力计算转化为线性复杂度,从而显著减少计算资源的消耗。这种机制在处理大规模数据集和长文本时表现出色,特别适用于大语言模型 (LLM) 和其他需要高效处理长序列的任务。

线性注意力的具体实现方式有多种,常见的包括局部注意力 (Local Attention)、稀疏注意力 (Sparse Attention) 等。这些方法通过限制注意力的计算范围或使用稀疏矩阵来减少计算量,同时保持模型的表达能力。

线性注意力机制的研究仍在不断发展,许多新的变体和优化方法正在被提出,例如 [20] 中提出的基于 Transformer 的线性注意力模型。这些进展为未来的自然语言处理任务提供了更多的可能性。

Linear attention, and many other variants of efficient attention, is often motivated by changing the order of matrix associ at iv it y in the core attention computation $(Q K^{\top})V,=,Q(K^{\top}V)$ . However when the mask is added, the derivation is somewhat less straightforward (for example, the original paper (Katha ro poul os et al. 2020) and variants (Y. Sun et al. 2023) state the formula without proof).

线性注意力,以及许多其他高效的注意力变体,通常是通过改变核心注意力计算中矩阵结合的顺序来实现的 $(Q K^{\top})V,=,Q(K^{\top}V)$。然而,当引入掩码时,推导过程变得不那么直观(例如,原始论文 (Katharopoulos et al. 2020) 和变体 (Y. Sun et al. 2023) 在没有证明的情况下直接给出了公式)。

在标准的注意力机制中,计算过程是先将查询矩阵 $Q$ 与键矩阵 $K$ 的转置相乘,得到一个注意力分数矩阵,然后再与值矩阵 $V$ 相乘。而在线性注意力中,计算顺序被调整为先计算 $K^{\top}V$,再与 $Q$ 相乘。这种变化可以显著减少计算复杂度,尤其是在处理长序列时。然而,当引入掩码时,这种顺序的变化会导致推导变得更加复杂,因为掩码需要在不同的计算步骤中正确应用,以确保模型能够忽略某些位置的信息。

Roughly, the linear attention method claims that the following formula is equivalent to (10), which must be verified by expanding the sum and tracking indices carefully.

大致来说,线性注意力机制 (linear attention) 声称以下公式等价于公式 (10),这需要通过展开求和并仔细跟踪索引来验证。

$$
Y=Q\cdot\mathsf{c u m s u m}(K^{\top}V)
$$

$$
Y = Q \cdot \mathsf{cumsum}(K^{\top} V)
$$

公式中,$Y$ 是输出矩阵,$Q$ 是查询矩阵 (Query),$\mathsf{cumsum}$ 表示累积和操作,$K^{\top}$ 是键矩阵 (Key) 的转置,$V$ 是值矩阵 (Value)。该公式描述了在某些 Transformer 模型中,查询矩阵与键值对的累积和相乘的过程。

Proposition 4.1 ((Katha ro poul os et al. 2020)). Auto regressive kernel attention, i.e. masked kernel attention with the causal mask, can be computed in $O(T)$ time by a recurrence taking constant time per step.

命题 4.1 ((Katha ro poul os et al. 2020))。自回归核注意力机制,即使用因果掩码的掩码核注意力,可以通过每步耗时恒定的递归在 $O(T)$ 时间内计算。

4.2.1 A Tensor Contraction Proof of Linear Attention

4.2.1 张量收缩证明线性注意力 (A Tensor Contraction Proof of Linear Attention)

线性注意力机制的张量收缩证明提供了一种数学方法来理解其计算过程。通过张量收缩,可以简化注意力机制中的矩阵运算,从而提高计算效率。具体来说,张量收缩允许我们将高维张量之间的复杂乘积累加操作转化为更简单的低维运算,这在处理大规模数据时尤为重要。

在这个部分,我们将详细介绍如何使用张量收缩来证明线性注意力机制的有效性,并探讨其在实际应用中的优势。线性注意力机制相比于传统的自注意力机制,能够在保持模型性能的同时显著减少计算资源的消耗。这一特性使得它在处理长序列和大规模数据集时表现出色 [20]。

We present a simple and rigorous derivation of linear attention that will also immediately reveal how to generalize it. The main idea is to perform the contraction (12) in an alternate order. We avoid ambiguous matrix notation and work directly with contraction notation:

我们介绍了一种简单而严谨的线性注意力 (linear attention) 推导方法,该方法还将立即展示如何对其进行推广。主要思想是按照不同的顺序执行收缩 (12)。我们避免使用模糊的矩阵符号,而是直接使用收缩符号:

我们介绍了一种简单而严谨的线性注意力 (linear attention) 推导方法,该方法还将立即展示如何对其进行推广。主要思想是按照不同的顺序执行收缩 (12)。我们避免使用模糊的矩阵符号,而是直接使用收缩符号:

$$
\begin{array}{r l}&{Z=\mathsf{c o n t r a c t}(\mathbb{S P},\mathsf{S N}\to\mathsf{S P N})(V,K)}\ &{H=\mathsf{c o n t r a c t}(\mathsf{T S},\mathsf{S P N}\to\mathsf{T P N})(L,Z)}\ &{Y=\mathsf{c o n t r a c t}(\mathsf{T N},\mathsf{T P N}\to\mathsf{T P})(Q,H)}\end{array}
$$

$$
\begin{array}{r l}
& Z = \mathsf{contract}(\mathbb{S P}, \mathsf{S N} \to \mathsf{S P N})(V, K) \
& H = \mathsf{contract}(\mathsf{T S}, \mathsf{S P N} \to \mathsf{T P N})(L, Z) \
& Y = \mathsf{contract}(\mathsf{T N}, \mathsf{T P N} \to \mathsf{T P})(Q, H)
\end{array}
$$

这段内容包含数学公式和符号,因此按照规则保持原样不变。

Intuitively, we interpret this contraction order as follows.

直观上,我们可以这样理解这个收缩顺序。

The first step (15a) performs an “expansion” into more features, by a factor of the feature dimension N. The third step (15c) contracts the expanded feature dimension away. If $K$ is viewed as the input (Remark 6), then $V$ and $\boldsymbol{Q}$ perform the expansion and contraction, respectively.

第一步 (15a) 执行一个“扩展”操作,将特征数量扩展为原始特征维度 N 的倍数。第三步 (15c) 则收缩这个扩展后的特征维度。如果将 $K$ 视为输入(备注 6),那么 $V$ 和 $\boldsymbol{Q}$ 分别执行扩展和收缩操作。

The second step is the most critical, and explains the linear part of linear attention. First notice that (15b) is just a direct matrix multiplication by $L$ (since the $(\mathsf{P},\mathsf{N})$ axes can be flattened). Also note that this is the only term that involves both T and S axes, hence should have $\Omega(\mathsf{T S})$ complexity (i.e. quadratic in sequence length). However, when the mask $L$ is the standard causal attention mask (lower triangular $1,\mathrm{\dot{s}}$ ), matrix-vector multiplication by $L$ is identical to a feature-wise cumulative sum

第二步是最重要的,它解释了线性注意力机制中的线性部分。首先注意到 (15b) 只是通过矩阵 $L$ 进行直接的矩阵乘法(因为 $(\mathsf{P}, \mathsf{N})$ 轴可以被展平)。另外需要注意的是,这是唯一同时涉及 T 和 S 轴的项,因此应该具有 $\Omega(\mathsf{T S})$ 复杂度(即序列长度的二次复杂度)。然而,当掩码 $L$ 是标准的因果注意力掩码(下三角矩阵 $1,\mathrm{\dot{s}}$)时,通过 $L$ 的矩阵-向量乘法等同于按特征的累积和。

在因果注意力掩码的情况下,矩阵-向量乘法可以通过累积和来高效计算,从而将复杂度从 $\Omega(\mathsf{T S})$ 降低到线性复杂度 $\Omega(\mathsf{T} + \mathsf{S})$。这种优化使得线性注意力机制能够在长序列上更高效地运行。

$$
y=\left[\begin{array}{l l l l l}{1}&{}&{}&{}\ {\vdots}&{\ddots}&{}&{}\ {1}&{\ddots}&{1}\end{array}\right]x\quad\iff\quad y_{0}=x_{0}.
$$

$$
y=\left[\begin{array}{l l l l l}{1}&{}&{}&{}\ {\vdots}&{\ddots}&{}&{}\ {1}&{\ddots}&{1}\end{array}\right]x \quad \iff \quad y_{0}=x_{0}.
$$

该公式表示矩阵 ( y ) 与向量 ( x ) 的关系,其中矩阵 ( y ) 是一个下三角矩阵,其主对角线元素为 1,其余元素为 0。公式右侧表示当矩阵 ( y ) 作用于向量 ( x ) 时,输出向量 ( y ) 的第一个元素 ( y_0 ) 等于输入向量 ( x ) 的第一个元素 ( x_0 )。

4.3 Structured Masked Attention

4.3 结构化掩码注意力 (Structured Masked Attention)

结构化掩码注意力是一种特殊的注意力机制,它通过在注意力矩阵中引入结构化的掩码来限制模型对某些位置的关注。这种技术可以有效地减少计算量,并提高模型在处理长序列时的效率。具体来说,结构化掩码可以帮助模型专注于局部上下文或特定的依赖关系,而忽略其他不相关的位置。

在大语言模型 (LLM) 中,结构化掩码注意力可以通过多种方式实现。例如,它可以用于限制模型只能关注当前 token 及其之前的 token(即因果掩码),或者用于限制模型只能关注某个固定窗口内的 token(即局部掩码)。这些方法都可以显著降低计算复杂度,同时保持模型的性能。

此外,结构化掩码还可以与其他技术结合使用,例如稀疏注意力或分块注意力,以进一步优化模型的效率和效果。通过合理设计掩码策略,可以在不同的任务中取得更好的表现。

With the tensor contraction perspective of masked attention (15), we can immediately see that the crux of the original linear attention is the fact that matrix-vector multiplication by the causal mask is equivalent to the cumulative sum operator.

通过张量收缩的角度看待掩码注意力 (15),我们可以立即看出原始线性注意力的核心在于:因果掩码的矩阵-向量乘法等价于累积和运算符。

However, we observe that there is no reason the attention mask has to be all 1’s. All that is necessary for linear attention to be fast is for $L$ to be a structured matrix, which by definition are those that have fast matrix multiplication (Section 2.3). In particular, we can use any mask matrix $L$ that has sub-quadratic (ideally linear) matrix-vector multiplication, which would have the same complexity as standard linear attention by speeding up the bottleneck equation (15b).

然而,我们观察到注意力掩码并不一定非要全是 1。线性注意力要快速运行的唯一要求是矩阵 $L$ 是一个结构化矩阵,根据定义,这些矩阵具有快速矩阵乘法 (第 2.3 节)。具体来说,我们可以使用任何具有次二次(理想情况下是线性)矩阵-向量乘法复杂度的掩码矩阵 $L$,这将通过加速瓶颈方程 (15b) 来达到与标准线性注意力相同的复杂度。

Definition 4.2. Structured masked attention (SMA) (or structured attention for short) is defined as a function on queries/keys/values $Q,K,V$ as well as any structured matrix $L$ (i.e. has sub-quadratic matrix multiplication), through the 4-way tensor contraction

定义 4.2. 结构化掩码注意力 (Structured Masked Attention, SMA)(或简称结构化注意力)被定义为在查询/键/值 $Q, K, V$ 以及任意结构化矩阵 $L$(即具有次二次矩阵乘法)上的一个函数,通过四维张量收缩来实现

$$
\text{SMA}(Q, K, V, L) = \text{softmax}\left(\frac{QK^T + L}{\sqrt{d_k}}\right)V
$$

其中 $d_k$ 是键的维度。结构化矩阵 $L$ 的引入使得注意力机制能够在保持计算效率的同时,捕捉更复杂的依赖关系。

$$
Y=\mathsf{c o n t r a c t}(\mathsf{T N},\mathsf{S N},\mathsf{S P},\mathsf{T S}\to\mathsf{T P})(Q,K,V,L).
$$

$$
Y = \mathsf{contract}(\mathsf{TN}, \mathsf{SN}, \mathsf{SP}, \mathsf{TS} \to \mathsf{TP})(Q, K, V, L).
$$

该公式表示一个名为 contract 的操作,它接受四个参数:$\mathsf{TN}$、$\mathsf{SN}$、$\mathsf{SP}$ 和 $\mathsf{TS} \to \mathsf{TP}$,并作用于四个输入 $Q$、$K$、$V$ 和 $L$。具体来说,这个操作可能是在处理张量(Tensor)时的一种收缩(contraction)操作,其中 $\mathsf{TN}$、$\mathsf{SN}$、$\mathsf{SP}$ 和 $\mathsf{TS} \to \mathsf{TP}$ 可能是指定的张量维度或索引模式。

The SMA quadratic mode algorithm is the sequence of pairwise contractions defined by (13), which corresponds to the standard (masked) attention computation.

SMA 二次模式算法是由 (13) 定义的一系列成对收缩,这对应于标准的(掩码)注意力计算。

The SMA linear mode algorithm is the sequence of pairwise contractions defined by (15), where step (15b) is optimized through the sub quadratic structured matrix multiplication.

SMA 线性模式算法是由 (15) 定义的成对收缩序列,其中步骤 (15b) 通过次二次结构矩阵乘法进行优化。

We can instantiate structured masked attention to any given class of matrix structure. Some examples include (Figure 3):

我们可以将结构化掩码注意力实例化为任意给定的矩阵结构类。一些例子包括 (图 3):

  • 稀疏矩阵结构
  • 带状矩阵结构
  • 分块矩阵结构

(注:具体例子取决于图 3 中的内容,这里仅提供常见的矩阵结构类型作为示例)

• Linear attention uses a causal mask.

• 线性注意力机制使用因果掩码 (causal mask)。

• RetNet (Y. Sun et al. 2023) uses a decay mask $L_{i j}=\gamma^{i-j}\cdot\mathbb{I}[j\geq i]$ for some decay factor $\gamma\in[0,1]$ .

• RetNet (Y. Sun 等 2023) 使用了一个衰减掩码 $L_{i j}=\gamma^{i-j}\cdot\mathbb{I}[j\geq i]$,其中衰减因子 $\gamma \in [0,1]$。

在这个公式中,$L_{i j}$ 表示第 $i$ 个位置和第 $j$ 个位置之间的权重,$\gamma$ 是一个介于 0 和 1 之间的衰减因子,$\mathbb{I}[j\geq i]$ 是一个指示函数,当 $j \geq i$ 时为 1,否则为 0。


Figure 3: (Structured Masked Attention.) SMA constructs a masked attention matrix $M=Q K^{\top}\circ L$ for any structured matrix $L$ , which defines a matrix sequence transformation $Y=M V$ . All instances of SMA have a dual sub quadratic form induced by a different contraction ordering, combined with the efficient structured matrix multiplication by $L$ . Previous examples include Linear Attention (Katha ro poul os et al. 2020) and RetNet (Y. Sun et al. 2023). Beyond SSD (1-semi separable SMA), the focus of this paper, many other potential instantiations of structured attention are possible.

![](https://u254848-88c6-e493554b.yza1.seetacloud.com:8443/miner/v2/analysis/pdf_img?as_attachment=False&pdf=6fe16414fa471a5ab150cddcba5840d3a94e347e96992c50c12c1fb8b919479f1735788830_2405.21060v1.pdf&filename=ba87a1c327e8f573a9f0bbefb50947e78270efa50cea4d618cc30697e05df505.jpg)  
图 3: (结构化掩码注意力)。SMA 为任何结构化矩阵 \$L\$ 构建一个掩码注意力矩阵 \$M=Q K^{\top} \circ L\$,该矩阵定义了一个矩阵序列变换 \$Y=M V\$。所有 SMA 的实例都具有由不同的收缩顺序引起的对偶次二次形式,并结合高效的结构化矩阵乘法 \$L\$。之前的例子包括线性注意力 (Katha ro poul os et al. 2020) 和 RetNet (Y. Sun et al. 2023)。除了 SSD (1-半可分离 SMA),本文的重点外,还有许多其他可能的结构化注意力实例。

• The decay mask could be generalized to a Toeplitz matrix $L_{i j}=\alpha_{i-j}$ for some learnable (or input-dependent) set of parameters $\alpha\in\mathbb{R}^{\mathsf{T}}$ . This can be interpreted as a form of relative positional encoding, reminiscent of other methods such as AliBi (Press, N. Smith, and Lewis 2022) but multiplicative instead of additive. • Another variant could use a Fourier matrix $L_{i j}=\omega^{i j/\top}$ to encode positional structure a different way.

• 衰减掩码可以推广为一个 Toeplitz 矩阵 $L_{i j}=\alpha_{i-j}$,其中 $\alpha\in\mathbb{R}^{\mathsf{T}}$ 是一些可学习的(或依赖输入的)参数。这可以被解释为一种相对位置编码的形式,类似于其他方法如 AliBi (Press, N. Smith, and Lewis 2022),但它是乘法的而不是加法的。

• 另一种变体可以使用傅里叶矩阵 $L_{i j}=\omega^{i j/\top}$ 来以不同的方式编码位置结构。

In Section 5, we consider semi separable SMA, which defines our main SSD model.

在第 5 节中,我们考虑了半可分的 SMA,这定义了我们的主要 SSD 模型。

4.3.1 Summary: The Dual Forms of Masked Attention

4.3.1 摘要:掩码注意力的双重形式 (Dual Forms of Masked Attention)

在这一节中,我们将总结掩码注意力机制的两种主要形式。这两种形式在不同的应用场景中发挥着重要作用,特别是在处理序列数据时。具体来说,第一种形式是前向掩码注意力 (Causal Masked Attention),它确保模型只能关注到当前及之前的 token,而不能看到未来的 token。这种机制常用于自回归模型,如语言生成任务。

第二种形式是双向掩码注意力 (Bidirectional Masked Attention),它允许模型同时关注到序列中的前后 token。这种机制适用于不需要因果关系的任务,例如文本理解或编码任务。通过这种方式,模型可以获得更丰富的上下文信息,从而提高性能。

这两种掩码注意力机制的设计目的是为了更好地控制模型对输入序列的不同部分的关注程度,从而提升模型在各种任务中的表现。

Standard (masked kernel) attention is often conflated between a function and an algorithm. Separating this distinction presents a clear way to understand different variants of attention.

标准(掩码核)注意力机制常常被混淆为一个函数和一个算法。区分这两者之间的差异,可以让我们更清晰地理解不同类型的注意力机制。

Moreover, in this case

此外,在这种情况下

请注意:由于提供的内容非常简短,我已按照要求直接翻译,没有添加额外的内容或询问。如果需要继续翻译更多内容,请提供完整的段落。

It is known that contraction orderings can make large differences in computation complexity, leading to the quadratic vs. linear split. Just as state space models are a transformation that can be computed in multiple ways, with dual quadratic vs. linear forms (Section 3.4), linear attention has a similar duality that results from two contraction orders. In fact, these turn out to be different perspectives on the same underlying duality, which we make explicit in Section 5.

已知收缩顺序可以对计算复杂度产生显著影响,导致二次 vs. 线性差异。就像状态空间模型是一种可以通过多种方式计算的变换,具有对偶的二次和线性形式(第 3.4 节),线性注意力机制也存在类似的二元性,这种二元性源自两种不同的收缩顺序。实际上,这些不同的视角反映了同一底层二元性的不同表现形式,我们将在第 5 节中明确这一点。

5 State Space Duality

5 状态空间对偶性

状态空间对偶性 (State Space Duality) 是指在某些系统中,状态空间的表示可以以两种不同的方式来描述,这两种描述之间存在数学上的等价关系。这种对偶性在控制理论、信号处理等领域中有重要应用。通过对偶性,我们可以从不同的角度理解和分析系统的特性,从而为设计和优化提供新的思路。

在本节中,我们将探讨状态空间对偶性的基本概念及其在不同应用场景中的意义。具体来说,我们将讨论以下内容:

  • 对偶系统的定义
  • 对偶性在控制系统中的应用
  • 对偶性与系统可观测性和可控性之间的关系

通过理解状态空间对偶性,我们可以更好地掌握系统的内在结构,并为复杂系统的分析和设计提供有力工具。

In Sections 3 and 4, we defined structured state space models and structured attention, discussed their properties, and showed that they both have a quadratic algorithm and a linear algorithm. This section connects them together. Our main result is showing that a particular case of structured state space models coincides with a particular case of structured attention, and that the linear-time SSM algorithm and quadratic-time kernel attention algorithm are dual forms of each other.

在第 3 节和第 4 节中,我们定义了结构化状态空间模型 (structured state space models) 和结构化注意力 (structured attention),讨论了它们的性质,并展示了它们都具有二次算法和线性算法。本节将它们联系在一起。我们的主要结果是证明了结构化状态空间模型的一个特定情况与结构化注意力的一个特定情况相吻合,并且线性时间的 SSM 算法和二次时间的核注意力算法是彼此的对偶形式。

请注意,这里的“结构化状态空间模型”和“结构化注意力”是首次出现的专业术语,因此在括号中保留了英文原文。后续提及这些术语时将只使用中文。

5.1 Scalar-Identity Structured State Space Models

5.1 标量-身份结构状态空间模型 (Scalar-Identity Structured State Space Models)

In Section 3 we showed that state space models are equivalent to semi separable matrix transformations, resulting in both a linear recurrent form and quadratic naive form.

在第 3 节中,我们展示了状态空间模型等价于半可分离矩阵变换,这导致了线性递归形式和二次朴素形式的出现。

Recall that SSMs are defined by $\boldsymbol{y}=\mathsf{S S M}(A,B,C)(\boldsymbol{x})$ , and the matrix form of SSMs uses the SSS (sequentially semiseparable) representation $M=\mathrm{SSS}(A,B,C)$ where $M_{j i}=C_{j}^{\top}A_{j:i}B_{i}$ (equation (3)).

回忆一下,状态空间模型 (SSMs) 由公式 $\boldsymbol{y}=\mathsf{S S M}(A,B,C)(\boldsymbol{x})$ 定义,且 SSMs 的矩阵形式使用顺序半可分 (SSS) 表示 $M=\mathrm{SSS}(A,B,C)$,其中 $M_{j i}=C_{j}^{\top}A_{j:i}B_{i}$ (公式 (3))。

注意:这里的状态空间模型 (SSMs) 和顺序半可分 (SSS) 是首次出现的专业术语,因此在括号中保留了英文原文。

Now let us consider the case where $A_{j}$ is simply a scalar; in other words, an instantiation of a structured SSM where the $A$ matrices are extremely structured: $A=a I$ for scalar $a$ and identity matrix $I$ . Then we can rearrange

现在让我们考虑 $A_{j}$ 仅仅是一个标量的情况;换句话说,这是一个结构化的 SSM (State Space Model) 的实例,其中 $A$ 矩阵具有非常简单的结构:$A = a I$,其中 $a$ 是一个标量,$I$ 是单位矩阵。然后我们可以重新排列

$$
A_{j} = a I
$$

在这种情况下,$A_{j}$ 只是单位矩阵乘以一个标量 $a$,因此它的结构非常简单。这种简化使得我们可以更容易地分析和处理该模型。

$$
M_{j i}=A_{j:i}\cdot(C_{j}^{\top}B_{i}).
$$

$$
M_{j i}=A_{j:i} \cdot (C_{j}^{\top} B_{i})。
$$

该公式表示矩阵 ( M ) 的第 ( j ) 行第 ( i ) 列元素由矩阵 ( A ) 的子矩阵 ( A_{j:i} ) 与矩阵 ( C ) 的第 ( j ) 行转置后与矩阵 ( B ) 的第 ( i ) 列的乘积得到。

And this can be vectorized into

这可以被向量化为

请注意:由于提供的内容非常简短,且没有上下文信息,这里仅对这句话进行了直接翻译。如果需要更准确的翻译,请提供完整的句子或段落。

$$
\begin{array}{c}{{L:=155(a)}}\ {{M=L\circ(C B^{\top})}}\end{array}
$$

$$
\begin{array}{c}{{L:=155(a)}}\ {{M=L \circ (C B^{\top})}}\end{array}
$$

这段内容包含数学公式,根据翻译规则,特殊字符和公式保持原样,不做翻译。

where $B,C\in\mathbb{R}^{(\mathsf{T},\mathsf{N})}$ .

其中 $B,C\in\mathbb{R}^{(\mathsf{T},\mathsf{N})}$ 。

这里表示矩阵 $B$ 和 $C$ 属于 $\mathbb{R}^{(\mathsf{T},\mathsf{N})}$ 实数空间,即它们是维度为 (T, N) 的实数矩阵。

Using this formulation, the full output $Y=M X$ is computed precisely as

使用这种公式,完整的输出 $Y = M X$ 被精确计算为

请注意,由于原文内容较短,且包含数学公式,因此仅对文本部分进行了翻译,公式保持原样。

$$
\begin{array}{r l}&{G=\mathrm{contract}(\mathsf{T N},\mathsf{S N}\to\mathsf{T S})(C,B)}\ &{M=\mathrm{contract}(\mathsf{T S},\mathsf{T S}\to\mathsf{T S})(G,L)}\ &{Y=\mathrm{contract}(\mathsf{T S},\mathsf{S P}\to\mathsf{T P})(M,X)}\end{array}
$$

$$
\begin{array}{r l}
& G = \mathrm{contract}(\mathsf{T N}, \mathsf{S N} \to \mathsf{T S})(C, B) \
& M = \mathrm{contract}(\mathsf{T S}, \mathsf{T S} \to \mathsf{T S})(G, L) \
& Y = \mathrm{contract}(\mathsf{T S}, \mathsf{S P} \to \mathsf{T P})(M, X)
\end{array}
$$

上述公式定义了三个操作,具体如下:

  • ( G ) 是通过 contract 函数将 (\mathsf{T N}) 和 (\mathsf{S N}) 映射为 (\mathsf{T S}),输入参数为 (C) 和 (B)。
  • ( M ) 是通过 contract 函数将 (\mathsf{T S}) 和 (\mathsf{T S}) 映射为 (\mathsf{T S}),输入参数为 (G) 和 (L)。
  • ( Y ) 是通过 contract 函数将 (\mathsf{T S}) 和 (\mathsf{S P}) 映射为 (\mathsf{T P}),输入参数为 (M) 和 (X)。

这里的 contract 函数表示一种特定的运算或操作,具体的含义取决于上下文。

where ${\mathsf{S}}={\mathsf{T}}$ . But this is exactly the same as original definition of masked kernel attention definition (13)!

其中 ${\mathsf{S}}={\mathsf{T}}$ 。但这与掩码核注意力的原始定义 (13) 完全相同!

Therefore, as alluded to in Section 3.4, naively computing the scalar structured SSM—by materializing the semi separable matrix 𝑀 and performing quadratic matrix-vector multiplication—is exactly the same as quadratic masked kernel attention.

因此,正如第 3.4 节所提到的,直接计算标量结构化的 SSM——通过显式化半可分矩阵 𝑀 并执行二次矩阵-向量乘法——与二次掩码核注意力机制是完全相同的。

5.2 1-Semi separable Structured Masked Attention

5.2 1-半可分离结构化掩码注意力 (1-Semi separable Structured Masked Attention)

在这一节中,我们将介绍一种特殊的注意力机制,称为 1-半可分离结构化掩码注意力 (1-Semi separable Structured Masked Attention)。这种注意力机制结合了半可分离矩阵和结构化掩码的特点,旨在提高计算效率的同时保持模型的表达能力。

传统的 Transformer 模型使用全连接的自注意力机制,这会导致计算复杂度较高,尤其是在处理长序列时。为了解决这个问题,1-半可分离结构化掩码注意力通过引入特定的矩阵结构和掩码策略,减少了不必要的计算,从而提高了模型的效率。

具体来说,1-半可分离结构化掩码注意力将注意力矩阵分解为两个部分:一个是对角线附近的局部区域,另一个是远离对角线的全局区域。通过对这两个区域分别应用不同的计算方法,可以在保持模型性能的同时显著降低计算成本。

这种注意力机制特别适用于处理长序列数据,例如自然语言处理中的文本生成任务或语音识别中的音频处理任务。通过减少不必要的计算,模型可以在更短的时间内完成推理,同时保持较高的准确性。

需要注意的是,1-半可分离结构化掩码注意力并不是唯一的优化方案,其他类似的结构化注意力机制也在不断发展中。然而,该方法在某些应用场景下表现出色,值得进一步研究和探索 [20]。

Structured masked attention allows for the use of any structured mask 𝐿. When $L$ is the causal mask, it is standard linear attention. Note that the causal mask is $L=\mathrm{SS}(1_{T})$ , i.e. the 1-SS mask is generated by $a_{t}=1$ in definition (6). This motivates generalizing $L$ to the class of 1-semi separable masks, or 1-semi separable structured masked attention (1-SS SMA), where the cumsum in linear attention’s recurrence is replaced by a more general recurrence – the scalar SSM scan, i.e. 1-semi separable matrix multiplication (Section 3.2.2).

结构化掩码注意力允许使用任何结构化的掩码 𝐿。当 $L$ 是因果掩码时,它就是标准的线性注意力。需要注意的是,因果掩码是 $L=\mathrm{SS}(1_{T})$,即 1-半分离 (1-semi separable, 1-SS) 掩码是由定义 (6) 中的 $a_{t}=1$ 生成的。这促使我们将 $L$ 推广到 1-半分离掩码类,或 1-半分离结构化掩码注意力 (1-SS SMA),其中线性注意力中的累积和(cumsum)递归被更一般的递归——标量 SSM 扫描(即 1-半分离矩阵乘法)所替代(第 3.2.2 节)。

在 1-半分离结构化掩码注意力中,线性注意力的递归公式被替换为更通用的标量 SSM 扫描,从而扩展了注意力机制的应用范围。

Finally, the most important reason we consider 1-semi separable SMA is because the linear form for computing it is a special case of diagonal state space model. The linear form of SMA is algorithm (15), where the bottleneck step (15b)

最后,我们认为1-半可分离 SMA (1-semi separable SMA) 最重要的原因在于其线性形式是对角状态空间模型的一个特例。SMA 的线性形式为算法 (15),其中瓶颈步骤为 (15b)

can be viewed as matrix multiplication by the 1-SS mask. In Section 3, we also wrote out the computation for a diagonal SSM (8), where the bottleneck step (8b) is a scalar SSM recurrence which is equivalent to 1-SS multiplication. The only difference is that (8b) has an extra N dimension in $L$ , because the matrix $A$ is a diagonal matrix of size N. This N dimension would disappear if all diagonal entries of $A$ are the same, which results in Corollary 5.1.

可以被视为 1-SS 掩码的矩阵乘法。在第 3 节中,我们还写出了对角 SSM (8) 的计算过程,其中瓶颈步骤 (8b) 是一个标量 SSM 递归,这等价于 1-SS 乘法。唯一的区别是 (8b) 在 $L$ 中有一个额外的 N 维度,因为矩阵 $A$ 是一个大小为 N 的对角矩阵。如果 $A$ 的所有对角元素都相同,这个 N 维度将会消失,从而得到推论 5.1。

请注意,公式和特殊字符保持原样未翻译。

Corollary 5.1. 1-SS SMA (masked attention with 1-semi separable structured matrices $L$ ) (15) is a special case of a diagonal SSM (8) where the diagonal matrix is a scalar multiple of the identity.

推论 5.1. 1-SS SMA (带有 1-半可分离结构矩阵 $L$ 的掩码注意力机制) (15) 是对角 SSM (8) 的一个特例,其中的对角矩阵是单位矩阵的标量倍数。

While Corollary 5.1 says that 1-SS SMA has an efficient recurrent form, we can also show a converse result that characterizes which instances of SMA has efficient auto regression.

虽然推论 5.1 表明 1-SS SMA 具有高效的递归形式,我们也可以证明一个相反的结果,该结果刻画了哪些 SMA 实例具有高效的自回归形式。

Theorem 5.2. For any instantiation of structured masked attention (Definition 4.2) that is an auto regressive process with bounded order, the structured mask 𝐿must be a semi separable matrix.

定理 5.2. 对于任何结构化掩码注意力 (Definition 4.2) 的实例化,如果它是具有有界阶数的自回归过程,则结构化掩码 𝐿 必须是半可分离矩阵。

In other words, efficient auto regressive attention is general semi separable SMA. Theorem 5.2 is proved in Appendix C.2.

换句话说,高效的自回归注意力机制是通用的半可分离 SMA。定理 5.2 的证明见附录 C.2。

Remark 7. While 1-semi separable SMA is a special case of a state space model, general semi separable SMA is strictly more expressive than 1-SS SMA, and cannot be described by a standard SSM. However, the semi separable multiplication by 𝐿and the linear form of SMA (equation (15a)) each involve an expansion and contraction step, and can be absorbed into a similar instance of 1-SS SMA with a single (larger) expansion.

备注 7. 虽然 1-半可分 SMA 是状态空间模型 (SSM) 的一个特例,但一般的半可分 SMA 比 1-SS SMA 具有更强的表达能力,无法用标准的状态空间模型 (SSM) 描述。然而,半可分乘法与 𝐿 以及 SMA 的线性形式(公式 (15a))各自包含一个扩展和收缩步骤,可以被合并到一个类似的 1-SS SMA 实例中,只是该实例具有更大的扩展。

在上述过程中,1-半可分 SMA (1-semi separable SMA) 和一般半可分 SMA (general semi separable SMA) 的区别在于后者能够表示更复杂的状态转换,而前者是后者的简化版本。尽管如此,通过适当的调整,半可分乘法和线性形式的扩展和收缩步骤可以被整合到一个更大规模的 1-SS SMA 中。

In summary, 1-semi separable structured attention is the most important case of SMA, because it is:

总结来说,1-半分离结构化注意力 (1-semi separable structured attention) 是 SMA (Structured Multi-Head Attention) 最重要的情况,因为它具有以下特点:

  • 它是 SMA 中最常见和最常用的形式;
  • 它在许多应用场景中表现出优异的性能;
  • 它能够有效地捕捉序列中的长距离依赖关系,同时保持计算效率。

这些特性使得 1-半分离结构化注意力在各种自然语言处理任务中得到了广泛应用 [20]。

5.3 Structured State-Space Duality (SSD)

5.3 结构化状态空间对偶性 (Structured State-Space Duality, SSD)

结构化状态空间对偶性 (SSD) 是一种用于分析和建模复杂系统的技术。通过对系统的状态空间进行结构化的表示,SSD 能够揭示系统内部的对称性和相互关系,从而为理解和优化系统行为提供新的视角。具体来说,SSD 通过将系统的状态空间分解为多个子空间,并研究这些子空间之间的映射关系,来实现对系统的深入分析。

在实际应用中,SSD 可以帮助解决许多复杂的工程和科学问题,例如控制系统设计、信号处理以及机器学习中的模型优化等。通过利用 SSD,研究人员可以更有效地设计算法,提高系统的性能和稳定性。

SSD 的核心思想是通过构建一个对偶的状态空间,使得原系统的某些难以直接处理的特性可以在对偶空间中得到简化或更好的理解。这种对偶性不仅限于线性系统,也可以扩展到非线性系统和离散系统中。因此,SSD 在多个领域都具有广泛的应用前景。

在后续的内容中,我们将详细介绍 SSD 的数学基础、应用场景以及具体的实现方法。

To summarize our results:

总结我们的结果:


请注意,如果您需要进一步的翻译内容,请继续提供英文文本。

Figure 4 summarizes the duality between these two representations.

图 4: 总结了这两种表示之间的对偶性。

An extended related work and discussion (Section 10) describes the relationship between SSD and general SSMs / attention in more detail.

一个扩展的相关工作和讨论 (Section 10) 详细描述了 SSD 与通用 SSMs / 注意力机制 (attention) 之间的关系。

6 A Hardware-Efficient Algorithm for SSD Models

6 面向硬件高效的SSD模型算法


根据要求,仅翻译标题部分。标题翻译如下:

6 面向硬件高效的 SSD 模型算法


如果需要继续翻译更多内容,请提供后续段落。

The benefits of developing the theoretical SSD framework between SSMs, attention, and structured matrices lies in using the connections to improve the models and algorithms. In this section, we show how various algorithms for computing SSD models efficiently can be derived from various algorithms for computing structured matrix multiplication.

开发 SSMs、注意力机制和结构化矩阵之间的理论 SSD 框架的好处在于,可以通过这些联系来改进模型和算法。在本节中,我们将展示如何从计算结构化矩阵乘法的各种算法中推导出高效计算 SSD 模型的各种算法。

Our main computational result is an algorithm for computing SSD models that combines both the linear (recurrent) mode and quadratic (attention) mode. This algorithm is as computation efficient as SSMs (linear scaling in sequence length) and as hardware-friendly as attention (primarily uses matrix multiplications).

我们的主要计算结果是一种用于计算SSD模型的算法,该算法结合了线性(循环)模式和二次(注意力)模式。这种算法在计算效率上与SSM (线性扩展于序列长度) 相当,并且在硬件友好性方面与注意力机制相似(主要使用矩阵乘法)。

这种算法不仅保持了线性模式的高效性,还利用了注意力机制的优势,能够在处理长序列时保持较低的计算复杂度,同时充分利用现代硬件加速特性。

Theorem 6.1. Consider an SSD model with state expansion factor N and head dimension $\mathsf{P}=\mathsf{N}.$ . There exists an algorithm for computing the model on any input $X\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$ which only requires $O(\mathsf{T N}^{2})$ training FLOPs, 𝑂(TN) inference FLOPs, $O(\mathsf{N}^{2})$ inference memory, and whose work is dominated by matrix multiplications.

定理 6.1. 考虑一个状态扩展因子为 N 且头维度为 $\mathsf{P}=\mathsf{N}$ 的 SSD 模型。存在一种算法,可以在任何输入 $X\in\mathbb{R}^{(\mathsf{T},\mathsf{P})}$ 上计算该模型,仅需要 $O(\mathsf{T N}^{2})$ 训练 FLOPs(浮点运算次数),$O(\mathsf{T N})$ 推理 FLOPs,$O(\mathsf{N}^{2})$ 推理内存,并且其工作主要由矩阵乘法主导。

注意:FLOPs 是 Floating Point Operations per Second 的缩写,表示每秒浮点运算次数。

| 结构化状态空间模型 (StructuredStateSpaceModel) | 收缩矩阵 (contraction matrix) | Q (查询) | K V (键值对) | 对角线状态空间模型 (Diagonal State Space Model) | 标量-恒等状态空间模型 (Scalar-Identity SSM) |
| --- | --- | --- | --- | --- | --- |
| B X (扩展矩阵) (输入序列) |  | Lji (掩码) |  | DSS S4D S5 | 线性注意力 (Linear Attention) |
| Aj:i (状态矩阵) |  | N (核特征维度) |  | RetNet | GateLoop |
| N (状态扩展维度) |  |  |  | TransNormer |  |
|  |  |  |  |  | 半可分离SMA (1-SemiseparableSMA) |
| H (隐藏状态 (8b)) = L·XB (线性模式) |  | SMA 线性对偶 (15) |  | S6 | 高效自回归注意力 (EfficientAutoregressiveAttention) |
| SSM 二次对偶 (16) |  | G (Gram 矩阵 (13a)) |  |  | 半可分离SMA (SemiseparableSMA) |
|  |  | IY·0= (二次模式) |  |  |  |
|  |  |  |  |  | 结构化掩码注意力 (Structured Masked Attention (SMA)) |

| 结构化状态空间模型 (SSM) S4 | 结构化状态空间对偶 (SSD) |
| --- | --- |
|  |  |

解释:

  • 结构化状态空间模型 (StructuredStateSpaceModel):描述了模型的状态空间结构。
  • 收缩矩阵 (contraction matrix):用于压缩输入序列的矩阵。
  • Q (查询)K V (键值对):分别是查询、键和值,通常用于注意力机制中。
  • 对角线状态空间模型 (Diagonal State Space Model):一种简化版本的状态空间模型,其中状态转移矩阵是对角矩阵。
  • 标量-恒等状态空间模型 (Scalar-Identity SSM):一种特殊的状态空间模型,其中状态转移矩阵是标量或恒等矩阵。
  • B X (扩展矩阵) (输入序列):表示输入序列经过扩展矩阵后的结果。
  • Aj:i (状态矩阵):表示状态之间的转移矩阵。
  • N (状态扩展维度):表示状态空间的扩展维度。
  • Lji (掩码):用于控制注意力机制中的掩码操作。
  • H (隐藏状态 (8b)) = L·XB (线性模式):表示隐藏状态的计算公式,使用线性模式。
  • SMA 线性对偶 (15)SSM 二次对偶 (16):分别表示线性和二次模式下的对偶关系。
  • G (Gram 矩阵 (13a)):用于计算向量之间的内积。
  • IY·0= (二次模式):表示在二次模式下的某个计算公式。
  • RetNetGateLoopTransNormer:这些是不同的模型或模块名称。
  • 高效自回归注意力 (EfficientAutoregressiveAttention)半可分离SMA (SemiseparableSMA):分别是指高效的自回归注意力机制和半可分离的结构化掩码注意力机制。

Figure 4: (Structured State Space Duality.) State space duality describes the close relationship between state space models and masked attention. (Left) General SSMs and SMA both possess linear and quadratic forms, with direct analogs in notation. $(R i g h t)$ SSMs and SMA intersect at a large class of state space dual models (SSD) which capture many sequence models as special cases.

图 4: (结构化状态空间对偶性) 状态空间对偶性描述了状态空间模型和掩码注意力之间的紧密关系。 (左) 一般的 SSM 和 SMA 都具有线性和二次形式,在符号表示上有直接的类比。 (右) SSM 和 SMA 在一大类状态空间对偶模型 (SSD) 中相交,这些模型捕捉了许多序列模型作为特殊情况。

请注意,原文中的 "SSM" 和 "SMA" 分别指的是状态空间模型 (State Space Model) 和掩码注意力 (Masked Attention)。

Note that all of these bounds are tight, because a state space model with state expansion N operating on a head size of N has total state size $\mathsf{N}^{2}$ (yielding the lower bounds for training and inference FLOPs of $O(\mathsf{T N}^{2})$ and $O(\mathsf{N}^{2})$ respectively). Furthermore the input $X$ itself has TN elements, yielding the memory lower bound.

请注意,所有这些界限都是紧的,因为状态空间模型在头大小为 N 的情况下进行状态扩展 N,其总状态大小为 $\mathsf{N}^{2}$ (分别得出训练和推理的浮点运算次数下界为 $O(\mathsf{T N}^{2})$ 和 $O(\mathsf{N}^{2})$)。此外,输入 $X$ 本身有 TN 个元素,这导致了内存的下界。

  • 训练的浮点运算次数 (FLOPs) 下界:$O(\mathsf{T N}^{2})$
  • 推理的浮点运算次数 (FLOPs) 下界:$O(\mathsf{N}^{2})$
  • 内存下界:由输入 $X$ 的 TN 个元素决定

其中,$\mathsf{T}$ 表示序列长度,$\mathsf{N}$ 表示头大小或状态扩展大小。

The main idea behind Theorem 6.1 is once again viewing the problem of computing a state space model as a semi separable matrix multiplication, but leveraging its structure in a new way. Instead of computing the whole matrix in either recurrent or attention mode, we perform a block decomposition of the matrix. The diagonal blocks can be computed using the dual attention mode, which can be efficiently done with matrix multiplications, while the off-diagonal blocks can be factored by the rank-structure of semi separable matrices and reduced to a smaller recurrence. We highlight that Listing 1 provides a self-contained implementation of the SSD algorithm. Compared to the general selective SSM of Gu and Dao (2023), this implementation is much simpler, and relatively efficient even in native PyTorch without requiring special low-level kernels.

定理 6.1 的主要思想是再次将计算状态空间模型的问题视为半可分矩阵乘法,但以一种新的方式利用其结构。我们不是在循环模式或注意力模式下计算整个矩阵,而是对矩阵进行块分解。对角线块可以使用双注意力模式 (dual attention mode) 进行计算,这可以通过矩阵乘法高效完成;而非对角线块则可以通过半可分矩阵的秩结构进行分解,并简化为一个较小的递归问题。

我们特别指出,清单 1 提供了 SSD 算法的完整实现。与 Gu 和 Dao (2023) 提出的一般选择性状态空间模型 (SSM) 相比,该实现要简单得多,并且即使在原生 PyTorch 中也能相对高效地运行,而无需依赖特殊的低级内核。

To begin, we partition the matrix $M$ into a $\textstyle{\frac{\intercal}{0}}\times{\frac{\intercal}{0}}$ grid of sub matrices of size $\mathsf{Q}\times\mathsf{Q}$ , for some block size Q. Note that the off-diagonal blocks are low-rank by the defining property of semi separable matrices (Definition 3.1).5

首先,我们将矩阵 $M$ 划分为一个 $\textstyle{\frac{\intercal}{0}}\times{\frac{\intercal}{0}}$ 的子矩阵网格,每个子矩阵的大小为 $\mathsf{Q}\times\mathsf{Q}$,其中 Q 是块大小。需要注意的是,根据半可分矩阵的定义(定义 3.1),非对角线块是低秩的 [5]。

$$
\begin{array}{r l}{{3}\mathrm{lockDecomposition}\ \ \ }&{M=\left[\begin{array}{c c c c}{M^{(0,0)}}&&&\ {M^{(1,0)}}&&{M^{(1,1)}}&\ {\vdots}&{\vdots}&{\ddots}\ {M^{(\top/\ell-1,0)}}&{M^{(\top/\ell-1,1)}}&{\ldots}&{M^{(\top/\ell-1,\top/\ell-1)}}\end{array}\right]}\ {\mathrm{(DiagonalBlock)}\ \ \ \ M^{(j,j)}={5}\mathrm{SM}(A_{j\ell}(\jmath_{+}+1)\uprho,B_{j\ell:(j+1)0},C_{j\ell:(j+1)0})}\ {\mathrm{(Low-Rank~Block)}\ \ \ \ M^{(j,i)}=\left[\begin{array}{c}{C_{j\ell}^{\top}A\uprho_{j\ell-1}}\ {\vdots}\ {C_{(j+1)\ell-1}^{\top}A(\jmath_{+}\imath+\jmath_{0}-\jmath}\ {\vdots}\ {C_{(j+1)\ell-1}^{\top}A(\jmath_{+}\imath+\jmath_{0}-\1)\ell-1}\end{array}\right]A_{j\ell-1:(i+1)0-1}\left[\begin{array}{c}{B_{i\ell}^{\top}A(\imath+1)\up Q-\mathrm{1}#}\ {\vdots}\ {B_{(i+1)\ell-1}^{\top}A(\imath+\jmath_{1})\up Q-\mathrm{1}}\end{array}\right]^{\top}}\end{array}
$$

$$
\begin{array}{r l}
{3\text{层块分解}\ \ \ } & {M=\left[\begin{array}{c c c c}
{M^{(0,0)}}&&&\
{M^{(1,0)}}&&{M^{(1,1)}}&\
{\vdots}&{\vdots}&{\ddots}\
{M^{(\top/\ell-1,0)}}&{M^{(\top/\ell-1,1)}}&{\ldots}&{M^{(\top/\ell-1,\top/\ell-1)}}
\end{array}\right]}\
{\text{(对角块)}\ \ \ \ M^{(j,j)}=5\text{SM}(A_{j\ell}(\jmath_{+}+1)\uprho,B_{j\ell:(j+1)0},C_{j\ell:(j+1)0})}\
{\text{(低秩块)}\ \ \ \ M^{(j,i)}=\left[\begin{array}{c}
{C_{j\ell}^{\top}A\uprho_{j\ell-1}}\
{\vdots}\
{C_{(j+1)\ell-1}^{\top}A(\jmath_{+}\imath+\jmath_{0}-\jmath}\
{\vdots}\
{C_{(j+1)\ell-1}^{\top}A(\jmath_{+}\imath+\jmath_{0}-1)\ell-1}
\end{array}\right]A_{j\ell-1:(i+1)0-1}\left[\begin{array}{c}
{B_{i\ell}^{\top}A(\imath+1)\up Q-1}\
{\vdots}\
{B_{(i+1)\ell-1}^{\top}A(\imath+\jmath_{1})\up Q-1}
\end{array}\right]^{\top}}
\end{array}
$$

注:公式中的符号和特殊字符保持原样,未进行翻译。

This is easiest illustrated through an example, e.g. for $\mathsf{T}=9$ and decomposing into chunks of length $\mathsf Q=3$ . The shaded

这通过一个例子最容易说明,例如对于 $\mathsf{T}=9$ 并将其分解为长度为 $\mathsf{Q}=3$ 的块。阴影部分

请注意:原文中的公式和特殊字符已原样保留。

cells are low-rank factorization s of the off-diagonal blocks of the semi separable matrix.

单元格是半可分矩阵的非对角块的低秩分解 (low-rank factorization)。

| M= | [CT A0:0Bo CT A1:0B0 C A2:0Bo C A2:1B1 | C A1:1B1 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | C A3:0B0 CI A4:0B0 CT A5:0B0 | C A2:2B2 C A3:1B1 CI A4:1B1 C A5:1B1 | C A3:2B2 C A4:2B2 C A5:2B2 | C A3:3B3 CT A4:3B3 CT A4:4B4 C A5:3B3 C A5:4B4 | C A5:5B5 |  |
|  | CT A6:0B0 C A7:0B0 C A8:0B0 | CT A6:1B1 C] A7:1B1 C A8:1B1 | CT A6:2B2 C A7:2B2 C A8:2B2 | CT A6:3B3 CT A6:4B4 CT A7:3B3 CT A7:4B4 C A8:3B3 C A8:4B4 | CT A6:5B5 CT A7:5B5 C A8:5 B5 | CT A6:6B6 CT A7:6B6 C A7:7B7 CA8:6B6 CA8:7B7 CA8:8B8 |
|  | [CT A0:0Bo CI A1:0B0 |  |  |  |  |  |
|  | C A2:0B0 [C A3:2] | CTA1:1B1 C A2:1B1 [BT A2:0] | C A2:2B2 | C A3:3B3 |  |  |
|  | CA4:2A2:2 [CA5:2] | BA2:1 | [B A2:2] | CT A4:3B3 C A5:3B3 [CTA6:5] | CT A4:4B4 C A5:4B4 C A5:5B5 [B A5:3] | C A8:8B8 |
|  | [CTA6:5 [Cg A8:5] | CTA7:5A5:2 | [BT A2:0] B A2:1 [BA2:2] |  | CA7:5A5:5BA5:4 [Cg A8:5] [B A5:5] | CT A6:6B6 CT A7:6B6 C A7:7B7 C A8:6B6 C A8:7B7 |

注意:由于原始表格中的内容主要是符号和代码,且没有明确的语义,因此在翻译时保留了原始格式和内容,未进行语义上的修改。

From here we can reduce the problem into these two parts. These can also be interpreted as dividing the output of a “chunk” $y_{j0:(j+1)0}$ into two components: the effect of inputs within the chunk $x_{j0:(j+1)0}$ , and the effect of inputs before the chunk 𝑥0:𝑗Q.

从这里我们可以将问题分解为以下两个部分。这也可以被理解为将一个“块”(chunk)的输出 $y_{j0:(j+1)0}$ 分解为两个组成部分:块内输入 $x_{j0:(j+1)0}$ 的影响,以及块之前输入 𝑥0:𝑗Q 的影响。

  • 第一部分是块内输入 $x_{j0:(j+1)0}$ 对输出的影响。
  • 第二部分是块之前输入 𝑥0:𝑗Q 对输出的影响。

这种分解有助于我们更好地理解块内和块外信息对最终输出的贡献。

6.1 Diagonal Blocks

6.1 对角块

翻译结果如上所示。如果您有更多内容需要翻译,请继续提供。

The diagonal blocks are easy to handle, because they are simply self-similar problems of a smaller size. The $j$ -th block represents computing the answer $\mathsf{S S M}(A_{R},B_{R},C_{R})(x_{R})$ for the range $R;=;j\mathbb{Q};:;(j+1)\mathbb{Q};=;\left(j\mathbb{Q},j\mathbb{Q}+1,\ldots,j\mathbb{Q}+\mathbb{Q}-1\right)$ . The key is that this block can be computed using any desired method. In particular, for small chunk lengths Q, this problem is computed more efficiently using the dual quadratic SMA form. Additionally, the chunks can be computed in parallel.

对角块易于处理,因为它们只是较小规模的自相似问题。第 $j$ 个块表示计算范围 $R = j\mathbb{Q} : (j+1)\mathbb{Q} = \left(j\mathbb{Q}, j\mathbb{Q} + 1, \ldots, j\mathbb{Q} + \mathbb{Q} - 1\right)$ 上的答案 $\mathsf{S S M}(A_{R}, B_{R}, C_{R})(x_{R})$。关键在于这个块可以使用任何期望的方法进行计算。特别是对于较短的分块长度 Q,使用对偶二次 SMA 形式可以更高效地解决这个问题。此外,这些分块可以并行计算。

6.2 Low-Rank Blocks

6.2 低秩块 (Low-Rank Blocks)

根据您的要求,以上是翻译结果。由于原文内容较短,没有更多的段落或内容需要翻译。如果您有更多内容需要翻译,请继续提供。

The low-rank factorization s consist of 3 terms, and there are correspondingly three pieces of the computation. In this factorization, we will use the terminology

低秩分解 (low-rank factorization) 包含 3 个项,相应地,计算也分为三个部分。在这个分解中,我们将使用以下术语

请注意,原文似乎没有完整给出所有术语的解释。如果您有更多内容或具体的术语需要翻译,请提供完整的内容以便我能更准确地进行翻译。


Figure 5: (SSD Algorithm.) By using the matrix transformation viewpoint of state space models to write them as semi separable matrices (Section 3), we develop a more hardware-efficient computation of the SSD model through a blockdecomposition matrix multiplication algorithm. The matrix multiplication also has an interpretation as a state space model, where blocks represent chunking the input and output sequence. Diagonal blocks represent intra-chunk computations and the off-diagonal blocks represent inter-chunk computations, factored through the SSM’s hidden state.


图 5: (SSD 算法) 通过使用状态空间模型的矩阵变换观点,将其表示为半可分离矩阵(第 3 节),我们开发了一种更高效的硬件计算 SSD 模型的方法,该方法基于块分解矩阵乘法算法。这种矩阵乘法也可以解释为一个状态空间模型,其中块表示对输入和输出序列进行分块。对角块表示块内计算,而非对角块表示通过 SSM 的隐藏状态进行的块间计算。

Right Factors. This step computes the multiplication by the right $B$ -block-factors of the low-rank factorization. Note that for each chunk, this is a $(\mathsf{N},\mathsf{Q})$ by (Q, P) matrix multiplication, where N is the state dimension and $P$ is the head dimension. The result is a (N, P) tensor for each chunk, which has the same dimensionality as the expanded hidden state $h$ .

右因子。这一步计算低秩分解的右侧 $B$ -块因子的乘法。注意,对于每个块,这是一个 $(\mathsf{N},\mathsf{Q})$ 与 (Q, P) 的矩阵乘法,其中 N 是状态维度,$P$ 是头维度。结果是每个块得到一个 (N, P) 的张量,其维度与扩展后的隐藏状态 $h$ 相同。

在这一过程中,$(\mathsf{N},\mathsf{Q})$ 表示输入矩阵的维度,(Q, P) 表示右侧因子矩阵的维度,最终输出的 (N, P) 张量与隐藏状态 $h$ 的维度一致。

This can be interpreted as: what is the final state per chunk supposing that the initial state (to the chunk) is 0. In other words this computes $h_{j0+{\mathsf Q}-1}$ assuming that $x_{0:j0}=0$ .

这可以理解为:假设每个块的初始状态为 0,那么每个块的最终状态是什么。换句话说,这计算了 $h_{j0+{\mathsf Q}-1}$,假设 $x_{0:j0}=0$。

  • 这里的 $h_{j0+{\mathsf Q}-1}$ 表示在处理第 j 个块时,假设从开始到该块的输入序列 $x_{0:j0}$ 全为 0 的情况下,该块的最终隐藏状态。

Center Factors. This step computes the effect of the center $A$ -block-factors terms in the low-rank factorization. In the previous step, the final states per chunk have total shape $(\mathsf{T}/\mathsf{Q},\mathsf{N},\mathsf{P})$ . This is now multiplied by a 1-SS matrix generated by $A_{20-1:\mathbb{Q}-1}^{\times},A_{3\mathbb{Q}-1:2\mathbb{Q}-1}^{\times},\dotsc,A_{\intercal-1:\mathbb{T}-\mathbb{Q}-1}^{\times}$ .

中心因子。这一步计算低秩分解中中心 $A$ -块因子项的影响。在上一步中,每个块的最终状态的总形状为 $(\mathsf{T}/\mathsf{Q}, \mathsf{N}, \mathsf{P})$ 。现在,这个形状将与由 $A_{20-1:\mathbb{Q}-1}^{\times}, A_{3\mathbb{Q}-1:2\mathbb{Q}-1}^{\times}, \dotsc, A_{\intercal-1:\mathbb{T}-\mathbb{Q}-1}^{\times}$ 生成的 1-SS 矩阵相乘。

注意:公式和特殊字符保持原样未翻译。

This step can be computed by any algorithm for computing 1-SS multiplication (also known as the scalar SSM scan or cumprodsum operator).

这一步可以通过任何用于计算 1-SS 乘法(也称为标量 SSM 扫描或 cumprodsum 运算符)的算法来完成。

This can be interpreted as: what is the actual final state per chunk taking into account all previous inputs; in other words, this computes the true hidden state $h_{j0}$ taking into account all of $x_{0:(j+1)\in}$ .

这可以理解为:每个块的实际最终状态是如何考虑了所有之前的输入;换句话说,这计算了考虑了所有 $x_{0:(j+1)\in}$ 的真实隐藏状态 $h_{j0}$ 。

具体来说,这是指在处理每个数据块时,模型会根据之前所有输入的信息来更新和确定该块的最终状态。因此,$h_{j0}$ 表示在考虑了从第 0 个到第 j+1 个输入后的真正隐藏状态。

Left Factors. This step computes the multiplication by the left $C$ -block-factors of the low-rank factorization. For each chunk, this can be represented by a matrix multiplication contract(QN, ${\mathsf{N P}}\to{\mathsf{Q P}}$ ).

左因子。这一步计算低秩分解的左 $C$ -块因子的乘法。对于每个块,这可以表示为矩阵乘法 contract(QN, \${\mathsf{N P}}\to{\mathsf{Q P}}\$ )

在这个过程中,contract 函数用于执行张量收缩操作,将输入的 QN 和 NP 矩阵转换为输出的 QP 矩阵。具体来说,QN 和 NP 的乘积结果会生成一个新的矩阵 QP,这个过程在每个块中都会进行。

This can be interpreted as: what is the output per chunk taking into account the correct initial state $h_{j0-1}$ , and supposing the inputs $x_{j0:(j+1)0}$ are 0. In other words for chunk $j$ , this computes the correct outputs taking into account only the prior inputs $x_{0:j0}$ .

这可以理解为:每个块的输出是什么,在考虑正确初始状态 $h_{j0-1}$ 的情况下,假设输入 $x_{j0:(j+1)0}$ 为 0。换句话说,对于块 $j$,这计算的是仅考虑之前输入 $x_{0:j0}$ 的正确输出。

在这一过程中,模型会根据之前的输入序列来生成当前块的输出,而忽略当前块的输入(假设为 0)。因此,输出结果反映了模型对历史信息的记忆和处理能力。

6.3 Computational Cost

6.3 计算成本

计算成本 (Computational Cost) 是指在执行特定任务时所需的计算资源量,包括处理器时间、内存使用和能源消耗等。在生成式 AI (Generative AI) 和大语言模型 (LLM) 的训练和推理过程中,计算成本是一个重要的考虑因素。较高的计算成本不仅会增加开发和部署的费用,还可能限制模型的规模和性能。因此,优化计算成本是提高模型效率和可扩展性的关键。

We define the notation ${\mathsf{B}}M M({\mathsf{B}},{\mathsf{M}},{\mathsf{N}},{\mathsf{K}})$ to define a batched matrix multiplication contract(MK, ${\mathsf{K N}}\to{\mathsf{M N}}$ ) with batch dimension B. From this notation we can infer three aspects of the efficiency:

我们定义符号 ${\mathsf{B}}M M({\mathsf{B}},{\mathsf{M}},{\mathsf{N}},{\mathsf{K}})$ 用于表示批量矩阵乘法合同 (MK, ${\mathsf{K N}} \to {\mathsf{M N}}$ ),其中批量维度为 B。从这个符号中,我们可以推断出三个方面的效率:

  1. 批量处理:批量维度 B 表示可以同时处理多个矩阵乘法操作,从而提高计算的并行性和效率。
  2. 矩阵维度转换:从 K×N 到 M×N 的转换表明了输入矩阵和输出矩阵的维度变化关系。
  3. 计算复杂度:通过批量矩阵乘法,可以在一定程度上减少重复计算,提升整体计算效率。

这种表示方法有助于更清晰地理解批量矩阵乘法的操作过程及其对性能的影响。

• Parallel iz ation: larger M, N, K terms can leverage specialized matrix multiplication units on modern accelerators.

• 并行化:更大的 M、N、K 可以利用现代加速器上的专用矩阵乘法单元。

Center Blocks. The cost of the quadratic SMA computation consists of three steps (equation (16)):

中心块。二次SMA计算的成本由三个步骤组成(公式 (16)):

请注意,原文中没有提供这三个步骤的具体内容。如果您有更多详细信息,可以继续补充。

Low-Rank Blocks: Center Factors. This step is a scalar SSM scan (or 1-SS multiplication) of length $\top/0$ on (N, P) independent channels. The work of this scan is ${\mathsf{T N P}}/{0}_{!}$ , which is negligible compared to the other factors.

低秩块:中心因子。这一步是对 (N, P) 独立通道进行长度为 $\top/0$ 的标量 SSM 扫描(或 1-SS 乘法)。该扫描的工作量为 ${\mathsf{T N P}}/{0}_{!}$,与其他因子相比可以忽略不计。

Note that because of the blocking which reduces the length of the sequence from T to $\top/0,$ , this scan has Q times smaller cost than a pure SSM scan (e.g. the selective scan of Mamba). Thus we observe that on most problem lengths, other algorithms (Appendix B) may be more efficient or much easier to implement without a significant slowdown. For example, a naive implementation of this via 1-SS matrix multiplication has cost ${\tt B M M}(1,\top/0,{\tt N P},\top/0)$ , which is much easier to implement and can be more efficient than a naive recurrence/scan implementation.

请注意,由于阻塞操作将序列长度从 T 减少到 $\top/0$ ,这种扫描的代价比纯 SSM 扫描(例如 Mamba 的选择性扫描)小 Q 倍。因此我们观察到,在大多数问题长度上,其他算法(附录 B)可能更高效或更容易实现,而不会显著减慢速度。例如,通过 1-SS 矩阵乘法的简单实现的代价为 ${\tt B M M}(1,\top/0,{\tt N P},\top/0)$ ,这比简单的递归/扫描实现更容易实现,并且可以更高效。

Low-Rank Blocks: Left Factors. This step is a single matrix multiplication with cost ${\mathsf{B M M}}(\mathsf{T}/\mathsf{Q},\mathsf{Q},\mathsf{P},\mathsf{N}).$

低秩块:左侧因子。这一步是一个矩阵乘法,计算成本为 ${\mathsf{B M M}}(\mathsf{T}/\mathsf{Q},\mathsf{Q},\mathsf{P},\mathsf{N})$。

Total Cost. If we set ${\mathsf{N}}={\mathsf{P}}={\mathsf{Q}}$ (in other words the state dimension, head dimension, and chunk length are equal), then all BMM terms above become $\mathsf{B M M(T/N,N,N,N)}$ . The computational ch act eris tics of this are:

总成本。如果我们设置 ${\mathsf{N}}={\mathsf{P}}={\mathsf{Q}}$ (换句话说,状态维度、头维度和块长度相等),那么所有 BMM 项都变为 $\mathsf{B M M(T/N,N,N,N)}$ 。其计算特性为:

  • 计算复杂度:当 ${\mathsf{N}}={\mathsf{P}}={\mathsf{Q}}$ 时,矩阵乘法的计算量会显著变化,具体表现为每个 BMM 操作的输入和输出矩阵的维度均为 $\mathsf{T/N}$ 和 $\mathsf{N}$。
  • 内存占用:由于矩阵的维度相同,内存的使用也会更加均匀,每个 BMM 操作所需的内存空间为 $\mathsf{T/N} \times \mathsf{N}$。
  • 并行性:在这种情况下,BMM 操作可以更好地利用并行计算资源,因为矩阵的维度一致,便于在多个处理器或 GPU 核心上进行并行处理。

这些特性对于优化模型的训练和推理效率非常重要。

Notice that the memory consumption is tight; the inputs and outputs $x,y$ have shape $(\mathsf{T},\mathsf{P});=;(\mathsf{T},\mathsf{N})$ . Meanwhile the flop count reflects an extra factor of N, which is cost incurred by the auto regressive state size and is common to all models.

请注意内存消耗较为紧张;输入和输出 $x, y$ 的形状为 $(\mathsf{T}, \mathsf{P});=;(\mathsf{T}, \mathsf{N})$。同时,浮点运算次数反映了额外的 N 因子,这是由自回归状态大小带来的成本,并且是所有模型的共同特点。

Aside from the matmuls, there is a scalar SSM scan on $\mathsf{N P=N^{2}}$ features and sequence length $\top/0$ . This has cost $O(\mathsf{T}/\mathsf{Q N}^{2})$ FLOPs and ${\cal O}(\log(\top/\up Q))$ depth. Although it does not use matrix multiplications, it is still parallel iz able and the total work done is negligible compared to the other steps; this has a negligible cost in our GPU implementation.

除了矩阵乘法 (matmuls) 之外,还存在一个标量 SSM 扫描,作用于 $\mathsf{N P=N^{2}}$ 个特征和序列长度 $\top/0$。该操作的计算成本为 $O(\mathsf{T}/\mathsf{Q N}^{2})$ FLOPs 和 ${\cal O}(\log(\top/\up Q))$ 的深度。尽管它不使用矩阵乘法,但仍然是并行化的,并且其总工作量与其他步骤相比可以忽略不计;在我们的 GPU 实现中,这一部分的成本几乎可以忽略不计。

Comparison to Pure SSM and Attention Models. Quadratic attention is also very hardware efficient by only leveraging matrix multiplications, but has ${\mathsf{T}}^{2}N$ total FLOPs. Its slower computation speed at both training and inference can directly be seen as a consequence of having a larger state size – standard attention has a state size scaling with sequence length T because it caches its history and does not compress its state.

与纯 SSM 和 Attention 模型的比较。二次注意力机制仅依赖矩阵乘法,因此在硬件上也非常高效,但其总浮点运算次数为 ${\mathsf{T}}^2 N$。它的计算速度在训练和推理阶段都较慢,这可以直接归因于其较大的状态大小——标准注意力机制的状态大小随序列长度 T 增加,因为它会缓存历史信息而不压缩其状态。

Linear SSMs have $\mathsf{T N P}=\mathsf{T N^{2}}$ total FLOPs, which is the same as SSD. However, a naive implementation requires a state expansion (15a) that materializes extra memory, and a scalar operation (15b) that does not leverage matrix multiplications.

线性 SSM 的总浮点运算次数 (FLOPs) 为 $\mathsf{T N P}=\mathsf{T N^{2}}$,这与 SSD 相同。然而,一个朴素的实现需要进行状态扩展 (15a),这会占用额外的内存,并且需要执行标量操作 (15b),而该操作无法利用矩阵乘法的优势。

|       | Attention | SSM | SSD |
|-------|-----------|-----|-----|
| 状态大小 (State size) | T | N | N |
| 训练 FLOPs (Training FLOPs) | T2N | TN2 | TN2 |
| 推理 FLOPs (Inference FLOPs) | TN | N2 | N2 |
| (朴素) 内存 (Naive) memory | T2 | TN2 | TN |
| 矩阵乘法 (Matrix multiplication) |    |    |    |

We note that many other matrix decomposition s are possible (for example, see Appendix B for a compendium of algorithms for 1-SS multiplication through different structured matrix decomposition s) which may lead to more algorithms for SSDs that could be better for other specialized settings. Even more broadly, we note that semi separable matrices have a rich literature and many more representations besides the SSS form that we use (Definition 3.2), and even more efficient algorithms may be possible.

我们注意到还有许多其他的矩阵分解方法是可行的(例如,参见附录 B,其中列出了通过不同结构化矩阵分解实现 1-SS 乘法的算法),这些方法可能会导致更多适用于 SSD 的算法,可能更适合其他专业场景。更广泛地说,我们注意到半可分离矩阵具有丰富的文献,并且除了我们使用的 SSS 形式(定义 3.2)之外,还有许多其他表示形式,因此可能存在更高效的算法。

请注意,我们在翻译中保留了原始段落格式和术语,并按照要求进行了适当的转换。

7 The Mamba-2 Architecture

7 Mamba-2架构


根据您的要求,我已经翻译了标题。请继续提供后续内容以便我进行完整翻译。

By connecting SSMs and attention, the SSD framework allows us to develop a shared vocabulary and library of techniques for both. In this section we discuss some examples of understanding and modifying SSD layers using ideas originally

通过将 SSM (State Space Model) 和注意力机制连接起来,SSD 框架使我们能够为两者开发一个共享的词汇表和技术库。在本节中,我们将讨论一些使用最初的想法来理解和修改 SSD 层的例子 [20]。

请注意,SSM 和注意力机制是两个不同的概念,但 SSD 框架提供了一个统一的方式来处理它们。这使得我们可以更方便地应用和扩展相关技术。


Figure 6: (Mamba-2 Architecture.) The Mamba-2 block simplifies the Mamba block by removing sequential linear projections; the SSM parameters $A,B,C$ are produced at the beginning of the block instead of as a function of the SSM input $X$ . An additional normalization layer is added as in NormFormer (Shleifer, Weston, and Ott 2021), improving stability. The $B$ and $C$ projections only have a single head shared across the $X$ heads, analogous to multi-value attention (MVA).

![](https://u254848-88c6-e493554b.yza1.seetacloud.com:8443/miner/v2/analysis/pdf_img?as_attachment=False&pdf=6fe16414fa471a5ab150cddcba5840d3a94e347e96992c50c12c1fb8b919479f1735788830_2405.21060v1.pdf&filename=38733377ca1fd464b186cd6dac91c7cbc63d1697dcff615ca879ededdc37d78c.jpg)  
图 6: (Mamba-2 架构。) Mamba-2 模块通过移除顺序线性投影简化了 Mamba 模块;SSM 参数 \$A, B, C\$ 在模块的开始处生成,而不是作为 SSM 输入 \$X\$ 的函数。此外,添加了一个额外的归一化层,类似于 NormFormer (Shleifer, Weston, 和 Ott 2021),以提高稳定性。\$B\$ 和 \$C\$ 投影仅有一个在所有 \$X\$ 头之间共享的单头,类似于多值注意力 (MVA)。

developed for Transformers. We discuss several design choices, resulting in the Mamba-2 architecture. These axes of variation are ablated in Section 9.4.

为 Transformer 开发的。我们讨论了多个设计选择,最终形成了 Mamba-2 架构。这些变化轴在第 9.4 节中进行了消融研究。

7.1 Block Design

7.1 块设计 (Block Design)

块设计是指在系统或模型中,将功能或组件划分为独立的模块或块,以便更好地管理和优化。每个块通常负责特定的任务或功能,块之间通过明确定义的接口进行通信和交互。这种设计方法可以提高系统的可维护性、可扩展性和性能。

在生成式 AI (Generative AI) 和大语言模型 (LLM) 中,块设计尤为重要,因为它可以帮助模型更有效地处理复杂的任务,同时保持结构的清晰和简洁。块设计还可以促进模块之间的重用和组合,使得模型能够更好地适应不同的应用场景。

块设计的具体实现方式可能因应用场景而异,但在大多数情况下,都会遵循一些基本的原则和模式。例如,在 Transformer 模型中,块设计通常包括多头自注意力机制 (Multi-head Self-Attention)、前馈神经网络 (Feed-forward Neural Network) 等组件。这些组件通过精心设计的架构相互协作,共同完成复杂的自然语言处理任务。

块设计的优势在于它能够简化复杂系统的开发和调试过程,同时提高系统的灵活性和鲁棒性。通过将系统分解为多个独立的块,开发者可以更容易地定位和修复问题,同时也能够更方便地引入新的功能或改进现有功能。

在实际应用中,块设计还能够帮助团队更好地协作,因为不同的开发人员可以专注于不同的块,而不必担心其他部分的影响。这不仅提高了开发效率,还降低了项目的整体风险。

总之,块设计是一种重要的系统设计方法,尤其在生成式 AI 和大语言模型领域,它为构建高效、灵活和可扩展的系统提供了坚实的基础。

Parallel Parameter Projections. Mamba-1 was motivated by an SSM-centric point of view where the selective SSM layer is viewed as a map from $X\mapsto Y$ . The SSM parameters $A,B,C$ are viewed as subsidiary and are functions of the SSM input $X$ . Thus the linear projections defining $(A,B,C)$ occur after the initial linear projection to create $X$ .

并行参数投影。Mamba-1 的设计灵感来源于以 SSM (State-Space Model) 为中心的观点,其中选择性 SSM 层被视为从 $X \mapsto Y$ 的映射。SSM 参数 $A, B, C$ 被视为次要的,并且是 SSM 输入 $X$ 的函数。因此,定义 $(A, B, C)$ 的线性投影发生在创建 $X$ 的初始线性投影之后。

在 Mamba-1 中,输入 $X$ 首先通过一个线性投影生成,然后基于这个输入 $X$,再进行后续的线性投影来确定参数 $A, B, C$。这种设计使得参数 $A, B, C$ 可以根据输入 $X$ 动态调整,从而更好地适应不同的任务需求。

In Mamba-2, the SSD layer is viewed as a map from $A,X,B,C\mapsto Y$ . It therefore makes sense to produce $A,X,B,C$ in parallel with a single projection at the beginning of the block. Note the analogy to standard attention architectures, where $X,B,C$ correspond to the $Q,K,V$ projections that are created in parallel.

在 Mamba-2 中,SSD 层被视为从 $A,X,B,C \mapsto Y$ 的映射。因此,在块的开头使用单个投影并行生成 $A,X,B,C$ 是合理的。注意这与标准注意力机制架构的类比,在标准注意力机制中,$X,B,C$ 对应于并行创建的 $Q,K,V$ 投影 [20]。

Note that adopting parallel projections for the $A,B,C,X$ inputs to the SSM slightly reduces parameters and more importantly is more amenable to tensor parallelism for larger models, by using standard Megatron sharding patterns (Shoeybi et al. 2019)).

请注意,采用并行投影处理输入到 SSM 的 $A, B, C, X$ 稍微减少了参数数量,并且更重要的是,这种方式更有利于使用标准的 Megatron 分片模式 (Shoeybi et al. 2019) 实现张量并行性,特别是在更大模型中。

Extra Normalization. In preliminary experiments, we found that instabilities were prone to arising in larger models. We were able to alleviate this by adding an extra normalization layer (e.g. LayerNorm, GroupNorm, or RMSNorm) to the block right before the final output projection. This usage of a normalization is most directly related to the NormFormer architecture (Shleifer, Weston, and Ott 2021), which also added normalization layers at the end of the MLP and MHA blocks.

额外归一化。在初步实验中,我们发现较大的模型容易出现不稳定现象。通过在最终输出投影之前的模块中添加一个额外的归一化层(例如 LayerNorm、GroupNorm 或 RMSNorm),我们能够缓解这一问题。这种归一化的使用方式与 NormFormer 架构 (Shleifer, Weston, and Ott 2021) 最为相关,后者也在 MLP 和 MHA 模块的末尾添加了归一化层。

这种额外的归一化层有助于稳定模型训练过程,特别是在大语言模型中,能够有效减少训练中的波动,提升模型的收敛性和性能。

We also note that this change is similar to other recent models related to Mamba-2 that were derived from a linear attention viewpoint. The original linear attention formulation normalizes by a denominator term that emulates the normalization of the softmax function in standard attention. Trans No rmer LL M (Qin, Dong Li, et al. 2023) and RetNet (Y. Sun et al. 2023) find that this normalization is unstable and add an extra LayerNorm or GroupNorm after the linear attention layer. Our extra normalization layer differs slightly from these, occuring after the multiplicative gate branch instead of before.

我们还注意到,这一变化与其他最近与 Mamba-2 相关的模型类似,这些模型都是从线性注意力的角度推导出来的。原始的线性注意力公式通过一个分母项来进行归一化,该分母项模拟了标准注意力机制中 softmax 函数的归一化。Trans No rmer 大语言模型 (Qin, Dong Li, et al. 2023) 和 RetNet (Y. Sun et al. 2023) 发现这种归一化是不稳定的,并在线性注意力层之后添加了一个额外的 LayerNorm 或 GroupNorm。我们的额外归一化层与这些方法略有不同,它位于乘法门分支之后,而不是之前。

请注意,原文中的“Trans No rmer”可能是一个拼写错误,通常应为“Transformer”。如果这是作者特意使用的名称,请根据实际情况进行调整。

7.2 Multihead Patterns for Sequence Transformations

7.2 序列变换的多头模式 (Multihead Patterns for Sequence Transformations)

Recall that SSMs are defined as a sequence transformation (Definition 2.1) where:

回忆一下,状态空间模型 (SSMs) 被定义为一个序列变换 (定义 2.1),其中:

  • 序列变换 (sequence transformation) 是指将一个序列作为输入,并生成另一个序列作为输出的过程。在 SSMs 中,这种变换是通过模型内部的状态演变来实现的。

请注意,这里提到的“定义 2.1”是指在相关文献或论文中对 SSMs 的正式定义,具体可以参考原文献 [20]。

One can view this as defining one head of the sequence transformation.

可以将这看作是定义了序列变换的一个头 (head)。

Definition 7.1 (Multihead patterns). A multihead sequence transformation consists of H independent heads, for a total model dimension of $\mathrm{\DeltaD}={\mathsf{d}}_{-}$ _model. The parameters may be tied across heads, leading to a head pattern.

定义 7.1 (多头模式)。一个多头序列变换由 H 个独立的头组成,总模型维度为 $\mathrm{\DeltaD}={\mathsf{d}}_{-}$ _model。参数可以在各个头之间共享,从而形成一个头模式。

在多头模式中,每个头可以独立处理输入序列的不同部分,最终将这些处理结果合并,以提高模型的表达能力和并行计算效率。这种结构常见于 Transformer 模型中。

The state size N and head dimension P are analogous to the $Q K$ head dimension and𝑉head dimension of attention, respectively. Just as in modern Transformer architectures (Chowdhery et al. 2023; Touvron, Lavril, et al. 2023), in Mamba-2 we generally choose these to be constants around 64 or 128; when the model dimension D increases, we increase the number of heads while keeping the head dimensions N and P fixed. In order to describe how to do this, we can transfer and generalize ideas from multihead attention to define similar patterns for SSMs, or any general sequence transformation.

状态大小 N 和头维度 P 分别类似于注意力机制中的 $Q K$ 头维度和 $V$ 头维度。与现代 Transformer 架构 (Chowdhery et al. 2023; Touvron, Lavril, et al. 2023) 中的做法类似,在 Mamba-2 中,我们通常将这些参数选择为大约 64 或 128 的常数;当模型维度 D 增加时,我们增加头的数量,同时保持头维度 N 和 P 不变。为了描述如何实现这一点,我们可以借鉴多头注意力机制的思想,并将其推广到 SSMs 或任何一般的序列变换中。

在 Mamba-2 中,通过固定头维度 N 和 P,我们可以在模型维度 D 增加时,通过增加头的数量来扩展模型的容量。这种做法与 Transformer 中的多头注意力机制相似,其中每个头负责处理不同的子空间信息。通过这种方式,我们可以确保模型在扩展时保持计算效率和性能的平衡。

| 多头 SSM | 多查询 SSM | 多扩展 SSM | 多输入 SSM |
| :---: | :---: | :---: | :---: |
|  | (多头注意力) | X | (多查询注意力) |  | (多键注意力) |  | X | (多值注意力) |  |
| X A | (T, H, P) (T, H) | (17) A | (T, 1, P) (T, H) | (18) | X (T, 1, P) A (T, H) |  | (19) | (T, H, P) A (T, H) | (20) |
| B | (T, H, N) |  | B (T, 1, N) |  | B (T, H, N) |  |  | B (T, 1, N) |  |
| C (T, H, N) |  |  | C (T, H, N) |  | C (T, 1, N) |  | C | (T, 1, N) |  |

说明:

  • 表格中的术语如“多头注意力 (Multi-head Attn.)”、“多查询注意力 (Multi-query Attn.)”等在第一次出现时保留了英文原文。
  • 表格格式已从 HTML 转换为 Markdown 格式。
  • 保持了原始表格的结构和内容。

Multihead SSM (MHS) / Multihead Attention (MHA) Pattern. The classic MHA pattern assumes that the head dimension P divides the model dimension D. The number of heads is defined as ${\mathsf{H}}={\mathsf{D}}/{\mathsf{P}}$ . Then, H copies of the core sequence transformation are created by creating H independent copies of each parameter. Note that while the MHA pattern was first described for the attention sequence transformation, it can be applied to anything compatible with Definition 2.1. For example, a multi-head SSD layer would accept inputs with shapes according to equation (17) where the SSD algorithm is broadcasted over the ${\mathsf{H}}={\mathsf{n}}_{}$ _heads dimension.

多头 SSM (MHS) / 多头注意力 (MHA) 模式。经典的 MHA 模式假设头维度 P 能够整除模型维度 D。头的数量定义为 ${\mathsf{H}}={\mathsf{D}}/{\mathsf{P}}$。然后,通过创建 H 个独立的参数副本,生成 H 个核心序列变换的副本。需要注意的是,虽然 MHA 模式最初是为注意力序列变换描述的,但它可以应用于任何符合定义 2.1 的内容。例如,一个多头 SSD 层将接受根据公式 (17) 的输入形状,其中 SSD 算法在 ${\mathsf{H}}={\mathsf{n}}_{}$ _heads 维度上进行广播。

请注意,这里的“SSD”指的是某种特定的算法或层,具体含义可能需要根据上下文进一步明确。

Multi-contract SSM (MCS) / Multi-query Attention (MQA) Pattern. Multi-query attention (Shazeer 2019) is a clever optimization for attention that can dramatically improve the speed of auto regressive inference, which relies on caching the $K$ and $V$ tensors. This technique simply avoids giving $K$ and $V$ the extra head dimension, or in other words broadcasts a single head of $(K,V)$ across all the heads of $\boldsymbol{Q}$ .

多合约 SSM (MCS) / 多查询注意力 (MQA) 模式。多查询注意力 (Shazeer 2019) 是一种巧妙的优化方法,可以显著提高自回归推理的速度,这依赖于缓存 $K$ 和 $V$ 张量。这种技术简单地避免给 $K$ 和 $V$ 添加额外的头维度,或者换句话说,将单个头的 $(K,V)$ 广播到 $\boldsymbol{Q}$ 的所有头中。

在多查询注意力机制中,通过共享 $K$ 和 $V$ 的计算结果,减少了计算量和内存占用,从而提高了推理效率。这种方法特别适用于大语言模型 (LLM) 中的自回归生成任务,因为它可以显著减少每次生成新 Token 时的计算开销。

Using the state space duality, we can define an equivalent SSM version of MQA as equation (18). Here, $X$ and $B$ (the SSM analogs of attention’s $V$ and $K$ ) are shared across the H heads. We also call this the multi-contract SSM (MCS) head pattern, because the $C$ parameter which controls the SSM state contraction has independent copies per head.

使用状态空间对偶性,我们可以定义一个等价的状态空间模型 (SSM) 版本的 MQA,如公式 (18) 所示。在这里,$X$ 和 $B$(分别是注意力机制中 $V$ 和 $K$ 的 SSM 对应物)在 H 个头之间共享。我们还称这种模式为多收缩 SSM (MCS) 头模式,因为控制 SSM 状态收缩的参数 $C$ 在每个头上有独立的副本。

注意:MQA 指的是 Multi-Query Attention(多查询注意力机制)。

We can similarly define a multi-key attention (MKA) or multi-expand SSM (MES) head pattern, where $B$ (which controls the SSM expansion) is independent per head while $C$ and $X$ are shared across heads.

我们也可以类似地定义多键注意力 (MKA) 或多扩展 SSM (MES) 头模式,在这种模式下,$B$(控制 SSM 扩展)在每个头中是独立的,而 $C$ 和 $X$ 在所有头之间是共享的。

Multi-input SSM (MIS) / Multi-value Attention (MVA) Pattern. While MQA makes sense for attention because of its KV cache, it is not the natural choice for SSMs. In Mamba, instead, $X$ is viewed as the main input to the SSM, and therefore $B$ and $C$ are parameters that are shared across the input channels. We define a new multi-value attention (MVA) of multi-input SSM (MIS) pattern in equation (20), which can again be applied to any sequence transformation such as SSD.

多输入 SSM (MIS) / 多值注意力 (MVA) 模式。虽然 MQA 对于注意力机制来说是合理的,因为它有 KV 缓存,但这并不是 SSM 的自然选择。在 Mamba 中,$X$ 被视为 SSM 的主要输入,因此 $B$ 和 $C$ 是跨输入通道共享的参数。我们在公式 (20) 中定义了一种新的多值注意力 (MVA) 的多输入 SSM (MIS) 模式,该模式同样可以应用于任何序列变换,例如 SSD。

多输入 SSM (MIS) / 多值注意力 (MVA) 模式。虽然 MQA 对于注意力机制来说是合理的,因为它有 KV 缓存,但这并不是 SSM 的自然选择。在 Mamba 中,\$X\$ 被视为 SSM 的主要输入,因此 \$B\$ 和 \$C\$ 是跨输入通道共享的参数。我们在公式 (20) 中定义了一种新的多值注意力 (MVA) 的多输入 SSM (MIS) 模式,该模式同样可以应用于任何序列变换,例如 SSD。

Armed with this vocabulary, we can characterize the original Mamba architecture more precisely.

借助这些术语,我们可以更精确地描述原始的 Mamba 架构。

Proposition 7.2. The selective SSM (S6) layer of the Mamba architecture (Gu and Dao 2023) can be viewed as having • Head dimension $P=1$ : every channel has independent SSM dynamics 𝐴. • Multi-input SSM (MIS) or multi-value attention (MVA) head structure: the $B,C$ matrices (corresponding to $K,\mathcal{Q}$ in the attention duality) are shared across all channels of the input $X$ (corresponding to $V$ in attention).

命题 7.2. Mamba 架构 (Gu and Dao 2023) 中的选择性 SSM (S6) 层可以被视为具有以下特性:

  • Head 维度 $P=1$:每个通道都有独立的 SSM 动力学 𝐴。
  • 多输入 SSM (MIS) 或多值注意力 (MVA) 头结构:$B, C$ 矩阵(对应于注意力对偶中的 $K, \mathcal{Q}$)在输入 $X$ 的所有通道中共享($X$ 对应于注意力中的 $V$)。

这种结构使得每个通道的 SSM 动力学是独立的,而输入矩阵 $B$ 和 $C$ 则在所有通道之间共享,从而实现了多输入或多值注意力的效果。

We can also ablate these head pattern variants when applied to SSD (Section 9.4.3). Interestingly, despite being controlled in parameter counts and total state dimension, there is a noticeable difference in downstream performance. We empirically find that the MVA pattern as originally used in Mamba performs best.

我们还可以在应用于 SSD 时对这些头部模式变体进行消融实验 (Section 9.4.3)。有趣的是,尽管在参数数量和总状态维度上受到控制,下游性能仍然存在明显差异。我们通过实验证明,MVA 模式(即最初在 Mamba 中使用的模式)表现最佳。

Grouped Head Patterns. The ideas of multi-query attention can be extended to grouped-query attention (Ainslie et al. 2023): instead of 1 K and V head, one can create G independent K and $\mathrm{V}$ heads, where $1~<~6$ and G divides H. This is motivated both by bridging the performance difference between multi-query and multi-head attention, and enabling more efficient tensor parallelism by setting G to be a multiple of the number of shards (Section 8).

分组头模式。多查询注意力机制的概念可以扩展到分组查询注意力 (Ainslie et al. 2023):与其使用 1 个 K 和 V 头,可以创建 G 个独立的 K 和 V 头,其中 $1 < G$ 且 G 是 H 的因数。这一想法既是为了弥合多查询注意力和多头注意力之间的性能差异,也是为了通过将 G 设置为分片数量的倍数来实现更高效的张量并行(第 8 节)。

Similarly, the multi-input SSM head pattern used in Mamba-2 can be easily extended to grouped-input SSM (GIS), or synonymously grouped-value attention (GVA). The generalization is straightforward and we omit the details for simplicity.

同样地,Mamba-2 中使用的多输入 SSM 头模式可以很容易地扩展到分组输入 SSM (GIS),或等效的分组值注意力 (GVA)。这种推广是 straightforward 的,为了简洁起见,我们省略了具体细节。

注:这里的“straightforward”在专业文献中通常保留英文,因此在第一次出现时保留英文。

7.3 Other SSD Extensions from Linear Attention

7.3 线性注意力的其他 SSD 扩展

根据上述要求,仅提供该标题的翻译。

We describe here an example of architectural modifications to SSD motivated by linear attention. We ablate these in Section 9.4.3 as a form of negative result, finding that they do not significantly improve performance enough to adopt them as default settings. Nonetheless, these illustrate how the vast literature on attention can be incorporated to define variants of SSD. We treat the choice of kernel feature map as a hyper parameter in the Mamba-2 architecture, and expect other simple modifications inspired by attention to be possible as well.

我们在这里描述了一个基于线性注意力机制对 SSD (Single Shot Detector) 进行架构修改的示例。我们在第 9.4.3 节中对这些修改进行了消融实验,作为负面结果的一部分,发现它们并没有显著提高性能,不足以将其作为默认设置采用。尽管如此,这些修改展示了如何将大量关于注意力机制的研究成果融入到 SSD 的变体中。在 Mamba-2 架构中,我们将核特征映射的选择视为一个超参数,并预期其他受注意力机制启发的简单修改也是可行的。

我们在这里描述了一个基于线性注意力机制对 SSD (Single Shot Detector) 进行架构修改的示例。我们在第 9.4.3 节中对这些修改进行了消融实验,作为负面结果的一部分,发现它们并没有显著提高性能,不足以将其作为默认设置采用。尽管如此,这些修改展示了如何将大量关于注意力机制的研究成果融入到 SSD 的变体中。在 Mamba-2 架构中,我们将核特征映射的选择视为一个超参数,并预期其他受注意力机制启发的简单修改也是可行的。

Kernel Attention Approximations to Softmax Attention. Many variants of linear attention or kernel attention are motivated by viewing the attention scores softmax $(Q K^{\top})$ as composed of

核注意力近似于 Softmax 注意力。许多线性注意力或核注意力 (Kernel Attention) 的变体都是基于将注意力分数 softmax $(Q K^{\top})$ 视为由

  • 核函数近似来构建的。具体来说,这些方法试图通过使用更高效的计算方式来近似传统的 softmax 注意力机制,从而减少计算复杂度并提高处理大规模数据的能力。

在这些方法中,常见的做法是将 softmax 函数替换为其他形式的核函数,这些核函数可以在保持相似性能的同时,显著降低计算成本。这种近似方法不仅适用于 Transformer 模型中的自注意力机制,还可以扩展到其他依赖于注意力机制的任务中 [20]。

Exponential Kernel Feature Maps. In Mamba-2, we incorporate a flexible kernel feature map, and apply it to the $B$ and $C$ branches (corresponding to the $K$ and $V$ branches in attention). The feature map can also be optionally applied to the $X\left(V\right)$ branch, for simplicity and symmetry. This is represented in Figure 6 by an arbitrary non linearity. By default, we simply choose $\psi$ to be an element wise Swish / SiLU function (Hendrycks and Gimpel 2016; Rama chandra n, Zoph, and Le 2017). We explore other options in the ablations in Section 9.4.3, including feature maps used by Linear Attention, Performer, Random Feature Attention, and cosFormer (Section 4.1.3).

指数核特征映射。在 Mamba-2 中,我们引入了一个灵活的核特征映射,并将其应用于 $B$ 和 $C$ 分支(对应于注意力机制中的 $K$ 和 $V$ 分支)。该特征映射也可以选择性地应用于 $X(V)$ 分支,以保持简单和对称性。这在图 6 中由任意非线性表示。默认情况下,我们简单地选择 $\psi$ 为逐元素的 Swish / SiLU 函数 (Hendrycks and Gimpel 2016; Ramachandra n, Zoph, and Le 2017)。我们在第 9.4.3 节的消融实验中探讨了其他选项,包括 Linear Attention、Performer、Random Feature Attention 和 cosFormer (第 4.1.3 节) 使用的特征映射。

请注意,图 6 和相关部分的具体内容可以在对应的章节中找到详细说明。

Incorporating a Normalization (Denominator) Term. To find the denominator term, we simply have to compute $M1$ . But recall that the final output of the model is just $Y=M X$ (equation (16)). So the normalization terms can be found simply by augmenting $X$ with an extra column 1, resulting in a tensor of shape $(\mathsf{T},\mathsf{P}+1)$ .

加入归一化(分母)项。为了找到分母项,我们只需要计算 \$M1\$ 。但回想一下,模型的最终输出是 \$Y=M X\$ (公式 (16))。因此,归一化项可以通过在 \$X\$ 中添加一个额外的列 1 来获得,从而得到形状为 \$(\mathsf{T}, \mathsf{P}+1)\$ 的张量。

Note that in this case, the kernel feature map $\psi$ must be positive so that the sum is positive.

请注意,在这种情况下,核特征映射 $\psi$ 必须为正,以确保求和结果为正。

8 Systems Optimization for SSMs

8 系统优化用于 SSMs

根据要求,仅翻译标题,不询问提供更多内容。

We describe several systems optimization s for SSMs, in particular the Mamba-2 architecture, for large-scale efficient training and inference. In particular, we focus on tensor parallel and sequence parallel for large-scale training, as a well variable-length sequences for efficient finetuning and inference.

我们描述了针对 SSMs 的几个系统优化方法,特别是 Mamba-2 架构,用于大规模高效的训练和推理。具体来说,我们重点关注张量并行 (tensor parallel) 和序列并行 (sequence parallel) 以实现大规模训练,同时也考虑了变长序列 (variable-length sequences) 以提高微调和推理的效率。

8.1 Tensor Parallel

8.1 张量并行 (Tensor Parallel)

张量并行是一种用于分布式训练大语言模型的技术,它通过将模型的参数和计算分布在多个 GPU 上来加速训练过程。在这种并行方式中,每个 GPU 负责处理模型的一部分参数和相应的前向/后向传播计算。通过这种方式,可以有效地利用多个 GPU 的计算资源,从而提高训练效率和扩展性。

在张量并行中,模型的权重矩阵被分割成多个子矩阵,每个子矩阵由不同的 GPU 处理。在前向传播过程中,输入数据会被相应地分割,并传递给不同的 GPU 进行计算。在后向传播过程中,梯度也会被分割并传递回各个 GPU。最终,所有 GPU 计算的结果会通过通信操作进行聚合,以完成完整的前向和后向传播过程。

张量并行的优势在于它可以显著减少单个 GPU 的内存占用,使得更大规模的模型可以在有限的硬件资源上进行训练。此外,它还可以与其他并行策略(如数据并行和管道并行)结合使用,进一步提升训练效率。然而,张量并行也带来了一些挑战,例如需要额外的通信开销来同步不同 GPU 之间的计算结果,以及对模型架构的某些限制。

通过合理设计张量并行方案,研究人员和工程师可以在保持模型性能的同时,充分利用分布式计算资源,推动大语言模型的训练和发展。

Tensor parallelism (TP) (Shoeybi et al. 2019) is a model parallelism technique that splits each layer (e.g., attention, MLP) to run on multiple accelerators such as GPUs. This technique is widely used to train most large models (Brown et al. 2020; Chowdhery et al. 2023; Touvron, Lavril, et al. 2023; Touvron, L. Martin, et al. 2023) on GPU clusters where each node typically has 4-8 GPUs with fast networking such as NVLink. TP was originally developed for the Transformer architecture, and it is not straight-forward to adapt it other architecture. We first show the challenge of using TP with the Mamba architecture, and the show how the Mamba-2 architecture is designed to make TP efficient.

张量并行 (Tensor Parallelism, TP) (Shoeybi et al. 2019) 是一种模型并行技术,它将每一层(例如,注意力层、MLP层)拆分到多个加速器(如 GPU)上运行。这种技术被广泛用于在 GPU 集群上训练大多数大语言模型(Brown et al. 2020; Chowdhery et al. 2023; Touvron, Lavril, et al. 2023; Touvron, L. Martin, et al. 2023),每个节点通常配备 4-8 个 GPU,并通过高速网络(如 NVLink)连接。TP 最初是为 Transformer 架构开发的,将其应用于其他架构并不直接适用。我们首先展示了在 Mamba 架构中使用 TP 的挑战,然后介绍 Mamba-2 架构是如何设计的,以使 TP 更加高效。


以上内容已按照要求翻译成简体中文,并保留了原始的 Markdown 格式和引用格式。

Recall the Mamba architecture, with a single input $u~\in~\mathbb{R}^{L\times d}$ (no batching for simplicity), input projection matrices $W^{(x)}$ , $W^{(z)}\in\mathbb{R}^{d\times e d}$ where $e$ is the expansion factor (typically 2), and output projection matrix $\bar{W}^{(o)},\bar{\in},\bar{\mathbb{R}^{e d\times d}}$ :

回忆 Mamba 架构,具有单个输入 $u~\in~\mathbb{R}^{L\times d}$(为简单起见,不考虑批量处理),输入投影矩阵 $W^{(x)}$ 和 $W^{(z)} \in \mathbb{R}^{d\times e d}$,其中 $e$ 是扩展因子(通常为 2),以及输出投影矩阵 $\bar{W}^{(o)} \in \mathbb{R}^{e d\times d}$ :

  • 输入 $u$ 是一个形状为 $L \times d$ 的矩阵,表示输入序列的长度为 $L$,每个时间步的特征维度为 $d$。
  • 输入投影矩阵 $W^{(x)}$ 和 $W^{(z)}$ 将输入特征映射到更高维度的空间,扩展后的维度为 $e \times d$,其中 $e$ 是扩展因子。
  • 输出投影矩阵 $\bar{W}^{(o)}$ 将扩展后的特征重新映射回原始维度 $d$。

$$
\begin{array}{r l}&{\boldsymbol{x}=\boldsymbol{u}\boldsymbol{W}^{(x)^{\top}}\in\mathbb{R}^{L\times e d}}\ &{\boldsymbol{z}=\boldsymbol{u}\boldsymbol{W}^{(z)^{\top}}\in\mathbb{R}^{L\times e d}}\ &{\boldsymbol{x}_{c}=\mathrm{conv1d}(\boldsymbol{x})\in\mathbb{R}^{L\times e d}\quad\mathrm{(depthwise,independent~along~}d)}\ &{\boldsymbol{\Delta},\boldsymbol{B},\boldsymbol{C}=\mathrm{low-rank~projection}(\boldsymbol{x}_{c})}\ &{;;;;;;y=\boldsymbol{S}\boldsymbol{S}\boldsymbol{M}_{A,B,C,\Delta}(\boldsymbol{x}_{c})\in\mathbb{R}^{L\times e d}\quad\mathrm{(independent~along~}d)}\ &{;;;;;y_{g}=\boldsymbol{y}\cdot\boldsymbol{\phi}(z)\quad\mathrm{(gating,e.,e.,with~}\phi~b e i n g~S i L U)}\ &{;;;;;\mathrm{out}=y_{g}\boldsymbol{W}^{(o)^{\top}}\in\mathbb{R}^{L\times d}.}\end{array}
$$

$$
\begin{array}{r l}
&{\boldsymbol{x} = \boldsymbol{u} \boldsymbol{W}^{(x)^{\top}} \in \mathbb{R}^{L \times e d}} \
&{\boldsymbol{z} = \boldsymbol{u} \boldsymbol{W}^{(z)^{\top}} \in \mathbb{R}^{L \times e d}} \
&{\boldsymbol{x}{c} = \text{conv1d}(\boldsymbol{x}) \in \mathbb{R}^{L \times e d} \quad \text{(深度卷积,沿 } d \text{ 独立)}} \
&{\boldsymbol{\Delta}, \boldsymbol{B}, \boldsymbol{C} = \text{低秩投影}(\boldsymbol{x}
{c})} \
&{;;;;;; y = \boldsymbol{S} \boldsymbol{S} \boldsymbol{M}{A,B,C,\Delta}(\boldsymbol{x}{c}) \in \mathbb{R}^{L \times e d} \quad \text{(沿 } d \text{ 独立)}} \
&{;;;;; y_{g} = \boldsymbol{y} \cdot \boldsymbol{\phi}(z) \quad \text{(门控,例如,使用 } \phi \text{ 为 SiLU)}} \
&{;;;;;\text{out} = y_{g} \boldsymbol{W}^{(o)^{\top}} \in \mathbb{R}^{L \times d}.}
\end{array}
$$

这段公式描述了一个神经网络中的计算过程,具体步骤如下:

  • $\boldsymbol{x}$ 和 $\boldsymbol{z}$ 是通过输入向量 $\boldsymbol{u}$ 与权重矩阵 $\boldsymbol{W}^{(x)}$ 和 $\boldsymbol{W}^{(z)}$ 的转置相乘得到的。
  • $\boldsymbol{x}_{c}$ 是通过对 $\boldsymbol{x}$ 进行一维卷积(depthwise convolution)得到的,该卷积在维度 $d$ 上是独立的。
  • $\boldsymbol{\Delta}$, $\boldsymbol{B}$, 和 $\boldsymbol{C}$ 是通过对 $\boldsymbol{x}_{c}$ 进行低秩投影得到的。
  • $y$ 是通过矩阵 $\boldsymbol{S}$ 和 $\boldsymbol{M}{A,B,C,\Delta}$ 对 $\boldsymbol{x}{c}$ 进行变换得到的,该变换在维度 $d$ 上是独立的。
  • $y_{g}$ 是通过对 $y$ 和 $\boldsymbol{\phi}(z)$ 进行逐元素乘法得到的,其中 $\boldsymbol{\phi}$ 是激活函数(例如 SiLU)。
  • 最终输出 $\text{out}$ 是通过对 $y_{g}$ 与权重矩阵 $\boldsymbol{W}^{(o)}$ 的转置相乘得到的。

With TP, suppose that we want to split the computation along 2 GPUs. It is easy to split the input projection matrices $W^{(x)}$ and $\bar{W}^{(\bar{z})}$ into two partitions each of size $\begin{array}{r}{\bar{d}\times\frac{e d}{2}}\end{array}$ . Then each GPU would hold half of $x_{c}$ of size $\begin{array}{r}{\bar{L^{\times}}\frac{e d}{2}}\end{array}$ . However, we see that since $\Delta,B,C$ are functions are $x_{c}$ , so we would need an extra all-reduce between the GPUs to get the whole of $x_{c}$ before computing $\Delta,B,C$ . After that the two GPUs can compute the SSM in parallel since they are independent along $d$ . At the end, we can split the output projection matrices $W^{\bar{(o)}}$ into two partitions each of size $\textstyle{\frac{e d}{2}}\times d$ , and do an all-reduce at the end. Compared to Transformers, we would incur two all-reduces instead of one, doubling the time spent in communication. For large-scale Transformers training, communication might already take a significant fraction of time (e.g. $10{-}20%$ ), and doubling communication would make Mamba not as efficient for large-scale training.

使用 TP 时,假设我们希望将计算分布在 2 个 GPU 上。很容易将输入投影矩阵 $W^{(x)}$ 和 $\bar{W}^{(\bar{z})}$ 分成两个分区,每个分区的大小为 $(\bar{d} \times \frac{e d}{2})$。然后每个 GPU 将持有 $x_{c}$ 的一半,大小为 $(\bar{L^{\times}} \frac{e d}{2})$。然而,我们注意到 $\Delta, B, C$ 是 $x_{c}$ 的函数,因此在计算 $\Delta, B, C$ 之前,我们需要在 GPU 之间进行一次额外的 all-reduce 操作以获取完整的 $x_{c}$。之后,由于在 $d$ 维上是独立的,两个 GPU 可以并行计算 SSM。最后,我们可以将输出投影矩阵 $W^{\bar{(o)}}$ 分成两个分区,每个分区的大小为 $(\frac{e d}{2} \times d)$,并在最后进行一次 all-reduce 操作。

与 Transformer 相比,我们会遇到两次 all-reduce 操作而不是一次,这使得通信时间翻倍。对于大规模 Transformer 训练,通信可能已经占用了相当大的时间比例(例如 10%-20%),而通信时间翻倍会使 Mamba 在大规模训练中效率降低。

With Mamba-2, our goal is to have only one all-reduce per block, similar to attention or MLP blocks in Transformers. As a result, we have the projection to get $\Delta,B,C$ directly from $u$ instead of from $x_{c}$ , allowing us to split these projection matrices. This implies that we have different sets of $\Delta,B,C$ on different GPUs, which is equivalent to having several “groups” of $\Delta,B,C$ on a larger “logical GPU”. Moreover, we use GroupNorm within each block, with number of groups divisible by the TP degree, so that the GPUs in a TP group do not have a communicate within the block:

使用 Mamba-2,我们的目标是每个块只进行一次全归约 (all-reduce),类似于 Transformer 中的注意力或 MLP 块。因此,我们直接从 $u$ 而不是从 $x_{c}$ 计算 $\Delta,B,C$ 的投影,这使得我们可以拆分这些投影矩阵。这意味着在不同的 GPU 上有不同的 $\Delta,B,C$ 集合,这相当于在一个更大的“逻辑 GPU”上有多个“组”的 $\Delta,B,C$。此外,我们在每个块内使用 GroupNorm,且组的数量可以被 TP 度数整除,这样 TP 组内的 GPU 在块内不需要进行通信:

使用 Mamba-2,我们的目标是每个块只进行一次全归约 (all-reduce),类似于 Transformer 中的注意力或 MLP 块。因此,我们直接从 \$u\$ 而不是从 \$x_{c}\$ 计算 \$\Delta,B,C\$ 的投影,这使得我们可以拆分这些投影矩阵。这意味着在不同的 GPU 上有不同的 \$\Delta,B,C\$ 集合,这相当于在一个更大的“逻辑 GPU”上有多个“组”的 \$\Delta,B,C\$。此外,我们在每个块内使用 GroupNorm,且组的数量可以被 TP 度数整除,这样 TP 组内的 GPU 在块内不需要进行通信:

$$
\begin{array}{r l}&{\boldsymbol{x}=\boldsymbol{u}\boldsymbol{W}^{(x)^{\top}}\in\mathbb{R}^{L\times\epsilon d}}\ &{\boldsymbol{z}=\boldsymbol{u}\boldsymbol{W}^{(z)^{\top}}\in\mathbb{R}^{L\times\epsilon d}}\ &{\Delta,\boldsymbol{B},\boldsymbol{C}=\mathrm{projection}(\boldsymbol{u})\quad\mathrm{(one~or~more~groups~of~}\Delta,\boldsymbol{B},\boldsymbol{C}\mathrm{~per~GPU)}}\ &{\boldsymbol{x}_{c}=\mathrm{convl}\boldsymbol{A}(\boldsymbol{x})\in\mathbb{R}^{L\times\epsilon d}\quad\mathrm{(depthwise~independent~aloop~aloop~}\mathcal{d})}\ &{\boldsymbol{y}=S S M_{A,B,C,\Delta}(\boldsymbol{x}_{c})\in\mathbb{R}^{L\times\epsilon d}\quad\mathrm{(independent~aloop~}\mathcal{d})}\ &{\boldsymbol{y}_{g}=\boldsymbol{y}\cdot\phi(\boldsymbol{z})\quad\mathrm{(gaing.~e.g.,with~\phi~being~SiLU)}}\ &{\boldsymbol{y}_{n}=\mathrm{groupnorm}(\boldsymbol{y}_{g})\quad\mathrm{(number~of~groups~divisble~by~degree~of~tensor~parallel)}}\ &{\mathrm{out}=y_{q}\boldsymbol{W}^{(o)^{\top}}\in\mathbb{R}^{L\times\epsilon d}.}\end{array}
$$

$$
\begin{array}{r l}
&{\boldsymbol{x} = \boldsymbol{u} \boldsymbol{W}^{(x)^{\top}} \in \mathbb{R}^{L \times \epsilon d}} \
&{\boldsymbol{z} = \boldsymbol{u} \boldsymbol{W}^{(z)^{\top}} \in \mathbb{R}^{L \times \epsilon d}} \
&{\Delta, \boldsymbol{B}, \boldsymbol{C} = \text{projection}(\boldsymbol{u}) \quad \text{(一个或多个组的 } \Delta, \boldsymbol{B}, \boldsymbol{C} \text{ 每个 GPU)}} \
&{\boldsymbol{x}{c} = \text{convl}\boldsymbol{A}(\boldsymbol{x}) \in \mathbb{R}^{L \times \epsilon d} \quad \text{(深度独立循环 } \mathcal{d})} \
&{\boldsymbol{y} = SSM
{A,B,C,\Delta}(\boldsymbol{x}{c}) \in \mathbb{R}^{L \times \epsilon d} \quad \text{(独立循环 } \mathcal{d})} \
&{\boldsymbol{y}
{g} = \boldsymbol{y} \cdot \phi(\boldsymbol{z}) \quad \text{(增益,例如,使用 SiLU 作为 } \phi)} \
&{\boldsymbol{y}{n} = \text{groupnorm}(\boldsymbol{y}{g}) \quad \text{(分组数量可被张量并行度整除)}} \
&{\text{out} = y_{q} \boldsymbol{W}^{(o)^{\top}} \in \mathbb{R}^{L \times \epsilon d}.}
\end{array}
$$

这段公式描述了一个神经网络中的计算过程,具体步骤如下:

  • $\boldsymbol{x}$ 和 $\boldsymbol{z}$ 是通过输入向量 $\boldsymbol{u}$ 与权重矩阵 $\boldsymbol{W}^{(x)^{\top}}$ 和 $\boldsymbol{W}^{(z)^{\top}}$ 相乘得到的。
  • $\Delta, \boldsymbol{B}, \boldsymbol{C}$ 是通过对输入向量 $\boldsymbol{u}$ 进行投影得到的,每个 GPU 可能有多个这样的投影组。
  • $\boldsymbol{x}_{c}$ 是通过对 $\boldsymbol{x}$ 进行深度独立卷积操作得到的。
  • $\boldsymbol{y}$ 是通过 SSM (State Space Model) 模型对 $\boldsymbol{x}_{c}$ 进行处理得到的。
  • $\boldsymbol{y}_{g}$ 是通过对 $\boldsymbol{y}$ 和 $\phi(\boldsymbol{z})$ 进行逐元素乘法得到的,其中 $\phi$ 是激活函数(例如 SiLU)。
  • $\boldsymbol{y}{n}$ 是通过对 $\boldsymbol{y}{g}$ 进行分组归一化操作得到的,分组数量可以被张量并行度整除。
  • 最终输出 $\text{out}$ 是通过对 $\boldsymbol{y}_{n}$ 与权重矩阵 $\boldsymbol{W}^{(o)^{\top}}$ 相乘得到的。

We see that we only need to split the input projection matrices, and the output projection matrices, and only need to do all-reduce at the end of the block. This is similar to the design of TP for attention and MLP layers. In particular, if we have TP degree 2, we would split $W^{(x)}=[W_{1}^{(x)},W_{2}^{(x)}]$ with $\bar{W}_{i}^{(x)}\in\mathbb{R}^{d\times e d/2}$ , $W^{(z)}=[W_{1}^{(z)},W_{2}^{(\bar{z})}]$ with $\bar{W}_{i}^{(z)}\in\mathbb{R}^{d\times e d/2}$ ,

我们看到只需要拆分输入投影矩阵和输出投影矩阵,并且只需要在块的末尾进行 all-reduce 操作。这与注意力层和 MLP 层的张量并行 (TP) 设计类似。具体来说,如果我们有 TP 程度为 2,我们会将 $W^{(x)}=[W_{1}^{(x)},W_{2}^{(x)}]$ 拆分为 $\bar{W}_{i}^{(x)}\in\mathbb{R}^{d\times e d/2}$,并将 $W^{(z)}=[W_{1}^{(z)},W_{2}^{(\bar{z})}]$ 拆分为 $\bar{W}_{i}^{(z)}\in\mathbb{R}^{d\times e d/2}$。

注意:公式中的符号和数学表达式保持不变。

由于图片无法直接翻译为文字内容,如果您能提供图片中的文本内容或描述,我将根据提供的内容进行翻译。

Figure 7: (Parallelism with the Mamba-2 Block.) (Left: Tensor Parallelism) We split the input projection matrices $\bar{W^{(x)}},W^{(z)}$ and the output projection matrix $W^{(o)}$ . Each SSM head $(A,B,C,X)\mapsto Y$ lives on a single device. Choosing GroupNorm for the final normalization layer avoids extra communication. We need one all-reduce per layer, just like the MLP or attention blocks in a Transformer. (Right: Sequence/Context Parallelism) Analogous to the SSD algorithm, with multiple devices, we can split along the sequence dimension. Each device computes the state of its sequence, then pass that state to the next GPU.

图 7: (Mamba-2 块的并行性) (左:张量并行) 我们将输入投影矩阵 \$\bar{W^{(x)}} , W^{(z)}\$ 和输出投影矩阵 \$W^{(o)}\$ 进行拆分。每个 SSM 头 \$(A,B,C,X) \mapsto Y\$ 位于单个设备上。选择 GroupNorm 作为最终的归一化层可以避免额外的通信开销。我们每层需要一次 all-reduce 操作,就像 Transformer 中的 MLP 或注意力块一样。 (右:序列/上下文并行) 类似于 SSD 算法,在多个设备上,我们可以沿着序列维度进行拆分。每个设备计算其序列的状态,然后将该状态传递给下一个 GPU。

and 𝑊(𝑜) $W^{(o)}=\left[W_{1}^{(o)}\right]$ with $W_{i}^{(o)}\in\mathbb{R}^{e d/2\times d}$ . For $i=1,2$ , the TP Mamba-2 layer can be written as:

以及 \( W^{(o)} = \left[ W_{1}^{(o)} \right] \),其中 \( W_{i}^{(o)} \in \mathbb{R}^{ed/2 \times d} \)。对于 \( i=1,2 \),TP Mamba-2 层可以表示为:

$$
\begin{array}{r l}{x^{(i)}=u V_{i}^{(\cdot)\top}\in\mathbb{R}^{L\times d/2}}\ {z^{(i)}=u V_{i}^{(\cdot)\top}\in\mathbb{R}^{L\times d/2}}\ {\Delta^{(i)},B^{(i)},C^{(i)}=\mathrm{projection}(u)\quad\mathrm{(one~or~more~groups~of~\Delta,B,C~per~GPU)}}\ {x_{c}^{(i)}=\mathrm{convld}(x^{(i)})\in\mathbb{R}^{L\times d/2}}\ {y^{(i)}=S S M_{\mathrm{A},B,C,\Delta}(x_{c}^{(i)})\in\mathbb{R}^{L\times d/2}}\ {y_{c}^{(i)}=y^{(i)}\cdot\phi^{(i)}}\ {y_{n}^{(i)}=\mathrm{groups}m\mathrm{ormop}(y_{\mathrm{)}}^{(i)}\quad\mathrm{(number~of~groups~divisbile~by~degree~\of~tensor~paralle)}}\ {\mathrm{out}^{(i)}=y_{g}^{(i)}W_{i}^{(\cdot)\top}\in\mathbb{R}^{L\times d/2}}\ {\mathrm{out}=\sum_{m}\mathrm{out}^{(i)},\quad\mathrm{(summing~outputs~from~all~GPUs~with~an~all-reduce)}}\end{array}
$$

$$
\begin{array}{r l}
x^{(i)} = u V_{i}^{(\cdot)\top} \in \mathbb{R}^{L \times d/2} \
z^{(i)} = u V_{i}^{(\cdot)\top} \in \mathbb{R}^{L \times d/2} \
\Delta^{(i)}, B^{(i)}, C^{(i)} = \text{projection}(u) \quad \text{(每个 GPU 一个或多个 (\Delta, B, C) 组)} \
x_{c}^{(i)} = \text{convld}(x^{(i)}) \in \mathbb{R}^{L \times d/2} \
y^{(i)} = SSM_{A,B,C,\Delta}(x_{c}^{(i)}) \in \mathbb{R}^{L \times d/2} \
y_{c}^{(i)} = y^{(i)} \cdot \phi^{(i)} \
y_{n}^{(i)} = \text{gr