[论文翻译]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{groups}m\text{ormop}(y_{\mathrm{)}}^{(i)} \quad \text{(组数可被张量并行度整除)} \
\text{out}^{(i)} = y_{g}^{(i)} W_{i}^{(\cdot)\top} \in \mathbb{R}^{L \times d/2} \
\text{out} = \sum_{m} \text{out}^{(i)}, \quad \text{(通过 all-reduce 汇总所有 GPU 的输出)}
\end{array}
$$

这段公式描述了一个分布式计算过程,其中涉及矩阵运算、投影、卷积、状态空间模型 (SSM) 和归一化等操作。具体步骤如下:

  1. ( x^{(i)} ) 和 ( z^{(i)} ) 是通过将输入 ( u ) 与矩阵 ( V_{i} ) 相乘得到的。
  2. (\Delta^{(i)})、(B^{(i)}) 和 (C^{(i)}) 是通过对 ( u ) 进行投影得到的,每个 GPU 可能有多个这样的组。
  3. ( x_{c}^{(i)} ) 是对 ( x^{(i)} ) 进行卷积操作后的结果。
  4. ( y^{(i)} ) 是通过状态空间模型 (SSM) 对 ( x_{c}^{(i)} ) 进行处理得到的。
  5. ( y_{c}^{(i)} ) 是 ( y^{(i)} ) 与 (\phi^{(i)}) 的逐元素乘积。
  6. ( y_{n}^{(i)} ) 是对 ( y_{c}^{(i)} ) 进行归一化操作的结果,组数需要能够被张量并行度整除。
  7. (\text{out}^{(i)}) 是将 ( y_{n}^{(i)} ) 与权重矩阵 ( W_{i} ) 相乘得到的输出。
  8. 最终的输出 (\text{out}) 是通过 all-reduce 操作汇总所有 GPU 的输出。

We illustrate tensor parallel with Mamba-2 in Figure 7 $(L e f t)$ .

我们使用 Mamba-2 在图 7 (左) 中展示了张量并行。

8.2 Sequence Parallelism

8.2 序列并行性 (Sequence Parallelism)

序列并行性是一种用于加速模型训练的技术,它通过将序列数据分割成多个部分并在不同的计算设备上并行处理这些部分来提高训练效率。这种方法特别适用于处理长序列数据的任务,例如自然语言处理中的文本生成和机器翻译。

在传统的并行化方法中,模型的参数更新是按顺序进行的,这会导致计算资源的浪费和训练时间的增加。而序列并行性通过将序列数据划分为多个片段,并在每个片段上独立执行前向和后向传播,从而减少了等待时间并提高了计算资源的利用率。

序列并行性可以与其它并行化技术(如数据并行性和模型并行性)结合使用,以进一步提升训练效率。然而,序列并行性也带来了一些挑战,例如如何有效地管理不同片段之间的依赖关系,以及如何确保梯度信息在各个片段之间正确传递。

在实际应用中,序列并行性已经被证明可以在大规模训练任务中显著减少训练时间,尤其是在处理长序列数据时表现尤为突出。

For very long sequences, we might need to split the input and activation to different GPUs along the sequence length dimension. There are two main techniques:

对于非常长的序列,我们可能需要沿着序列长度维度将输入和激活拆分到不同的 GPU 上。主要有两种技术:

(此处原文未提供具体技术细节,因此保持空白)

  1. Sequence parallelism (SP) for the residual and normalization operations: first proposed by Kort hik anti et al. (2023), this technique decomposes the all-reduce in TP as reduce-scatter and all-gather. Noticing that the residual and normalization operations are repeated on the same input for all GPUs in the same TP group, SP splits the activation s along the sequence length dimension by performing: reduce-scatter, residual and normalization, then all-gather. Since the Mamba-2 architecture uses the same residual and normalization structure, SP applies without modification. 2. Sequence parallelism for the token-mixing operations (attention or SSM), also known as “context parallelism” (CP). Several techniques have been developed for attention layer (e.g., Ring attention (Liu, Yan, et al. 2024; Liu, Zaharia, and Abbeel 2023)), with sophisticated load-balancing technique (Brandon et al. 2023). The difficulty with sequence parallelism in attention is that we can split queries and keys into block, but each query block needs to interact with key blocks, leading to communication bandwidth quadratic in the number of workers.
1. 序列并行 (Sequence parallelism, SP) 用于残差和归一化操作:该技术最早由 Kort hik anti 等人 (2023) 提出,它将 TP 中的 all-reduce 操作分解为 reduce-scatter 和 all-gather。注意到在同一个 TP 组中的所有 GPU 上,残差和归一化操作对相同的输入重复执行,SP 通过沿序列长度维度分割激活值来实现:reduce-scatter、残差和归一化,然后 all-gather。由于 Mamba-2 架构使用相同的残差和归一化结构,因此 SP 可以直接应用。

2. 序列并行用于 Token 混合操作(如注意力机制或 SSM),也称为“上下文并行” (context parallelism, CP)。针对注意力层已经开发了多种技术(例如,环形注意力 (Ring attention) [Liu, Yan, et al. 2024; Liu, Zaharia, and Abbeel 2023]),并且采用了复杂的负载均衡技术 (Brandon et al. 2023)。在注意力机制中实现序列并行的难点在于,我们可以将查询和键分割成块,但每个查询块需要与键块进行交互,这导致通信带宽随工作节点数量的平方增长。


Figure 8: (Multi-Query Associative Recall (MQAR)). Associative recall tasks are challenging for SSMs, which must memorize all relevant information into their recurrent state. The SSD layer combined with improved architecture allows for much larger state sizes in Mamba-2, which performs significantly better than Mamba-1 and even vanilla attention.


图 8: (多查询关联回忆 (MQAR))。关联回忆任务对状态空间模型 (SSM) 来说具有挑战性,因为它们必须将所有相关信息记忆到其循环状态中。Mamba-2 中的 SSD 层与改进的架构相结合,允许拥有更大的状态尺寸,因此 Mamba-2 的表现显著优于 Mamba-1 甚至传统的注意力机制。

With SSMs, we can split the sequence in a simple manner: each worker takes an initial state, compute the SSM with respect to their inputs, return the final state, and pass that final state to the next worker. The communication bandwidth is linear in the number of workers. This decomposition is exactly the same as the block-decomposition in the SSD algorithm (Figure 5) to split into blocks / chunks. We illustrate this context parallelism in Figure 7 (Right).

使用 SSM,我们可以以一种简单的方式拆分序列:每个工作线程获取一个初始状态,根据其输入计算 SSM,返回最终状态,并将该最终状态传递给下一个工作线程。通信带宽与工作线程的数量成线性关系。这种分解方式与 SSD 算法中的块分解(图 5)完全相同,用于将序列拆分为块/片段。我们在图 7(右)中展示了这种上下文并行性。

8.3 Variable Length

8.3 可变长度

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

While pre training often uses the same sequence lengths for the batch, during finetuning or inference, the model might need to process different input sequences of different lengths. One naive way to handle this case is to right-pad all sequences in the batch to the maximum length, but this can be inefficient if sequences are wildly different lengths. For transformers, sophisticated techniques have been develop to avoid padding and do load-balancing between GPUs (Zeng et al. 2022; Y. Zhai et al. 2023), or packing multiple sequences in the same batch and adjust the attention mask (Ding et al. 2024; Pouransari et al. 2024). With SSMs and Mamba in particular, we can handle variable sequence lengths by simply treating the whole batch as one long sequence, and avoid passing the states between individual sequences. This is equivalent to simply setting $A_{t}=0$ for tokens $t$ at the end of one sequence to prevent it from passing information to the token $t+1$ , which belongs to a different sequence.

在预训练阶段,通常会使用相同的序列长度来处理批次中的所有样本。但在微调或推理过程中,模型可能需要处理不同长度的输入序列。一种简单的方法是将批次中的所有序列右填充到最大长度,但如果序列长度差异很大,这种方法可能会导致效率低下。对于 Transformer 模型,研究人员已经开发了多种复杂的技术来避免填充,并在 GPU 之间进行负载均衡 (Zeng et al. 2022; Y. Zhai et al. 2023),或者通过在同一批次中打包多个序列并调整注意力掩码 (Ding et al. 2024; Pouransari et al. 2024)。

特别是对于 SSMs 和 Mamba,我们可以通过将整个批次视为一个长序列来处理可变长度的序列,并且不需要在各个序列之间传递状态。这相当于在每个序列的末尾,将矩阵 $A_{t}$ 设置为 0,以防止信息从属于该序列的 Token $t$ 传递到属于下一个序列的 Token $t+1$。

9 Empirical Validation

9 实证验证

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

We empirically evaluate Mamba-2 on synthetic recall tasks that have been challenging for recurrent models (Section 9.1), and standard language modeling pre-training and downstream evaluations (Section 9.2). We validate that our SSD algorithm is much more efficient than Mamba-1 (Section 9.3) and comparable to optimized attention for moderate sequence lengths. Finally, we ablate various design choices in the Mamba-2 architecture (Section 9.4).

我们对 Mamba-2 进行了实证评估,测试其在合成回忆任务上的表现,这些任务对循环模型具有挑战性(第 9.1 节),以及标准的语言建模预训练和下游评估(第 9.2 节)。我们验证了我们的 SSD 算法比 Mamba-1 更高效(第 9.3 节),并且在中等序列长度下与优化的注意力机制相当。最后,我们对 Mamba-2 架构中的各种设计选择进行了消融实验(第 9.4 节)。

9.1 Synthetics: Associative Recall

9.1 合成:关联回忆 (Synthetics: Associative Recall)

请注意,由于原文内容较短,我已按照要求完成了翻译,并保留了原始的 Markdown 格式。如果您有更多内容需要翻译,请继续提供。

Synthetic associative recall tasks have been popular for testing the ability of language models to look up information in their context. Broadly, they involve feeding auto regressive models pairs of key-value associations, and then prompting the model to produce the correct completion upon being shown a previously-seen key. The multi-query associative recall (MQAR) task is a particular formulation of this task that requires the model to memorize multiple associations (Arora, Eyuboglu, Timalsina, et al. 2024). The original Mamba paper reported results on related synthetic tasks, in particular Selective Copying (Gu and Dao 2023) and Induction Heads (Olsson et al. 2022), which can be seen as easier associative recall tasks. The MQAR task is also closely related to “phonebook look-up” tasks which has been shown to be challenging for recurrent models such as SSMs, due to their finite state capacity (De et al. 2024; Jelassi et al. 2024).

合成关联回忆任务在测试语言模型从其上下文中查找信息的能力方面非常流行。大致来说,这些任务涉及向自回归模型输入键值对 (key-value pairs),然后提示模型在看到之前见过的键时生成正确的补全内容。多查询关联回忆 (multi-query associative recall, MQAR) 任务是这种任务的一种特定形式,要求模型记住多个关联关系 (Arora, Eyuboglu, Timalsina, et al. 2024)。最初的 Mamba 论文报告了相关合成任务的结果,特别是选择性复制 (Selective Copying) (Gu and Dao 2023) 和归纳头 (Induction Heads) (Olsson et al. 2022),这些可以被视为较为简单的关联回忆任务。MQAR 任务也与“电话簿查找”任务密切相关,后者已被证明对状态容量有限的循环模型(如 SSMs)具有挑战性 (De et al. 2024; Jelassi et al. 2024)。


Scaling Laws on The Pile (Sequence Length 8192)

序列长度为 8192 的 Pile 数据集上的扩展规律 (Scaling Laws on The Pile (Sequence Length 8192))

Figure 9: (Scaling Laws.) Models of size $\approx125M$ to $\approx1.3B$ parameters, trained on the Pile. Mamba-2 matches or exceeds the performance of Mamba as well as a strong “Transformer $++^{,\mathfrak{s}}$ recipe. Compared to our Transformer baseline, Mamba-2 is Pareto dominant on performance (perplexity), theoretical FLOPs, and actual wall-clock time.

图 9: (扩展定律)。模型规模从约 1.25 亿参数到约 13 亿参数,在 Pile 数据集上进行训练。Mamba-2 的性能匹配或超过了 Mamba 以及一个强大的 “Transformer $++^{,\mathfrak{s}}$” 方案。与我们的 Transformer 基线相比,Mamba-2 在性能(困惑度)、理论 FLOPs 和实际运行时间上都具有帕累托优势。

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

Table 1: (Zero-shot Evaluations.) Best results for each size in bold, second best unlined. We compare against open source LMs with various tokenizers, trained for up to 300B tokens. Pile refers to the validation split, comparing only against models trained on the same dataset and tokenizer (GPT-NeoX-20B). For each model size, Mamba-2 outperforms Mamba, and generally matches Pythia at twice the model size. Full results in Table 10.

表 1: (零样本评估)。每个规模的最佳结果用粗体表示,次佳结果不加下划线。我们将 Mamba-2 与各种分词器训练的开源大语言模型 (LLM) 进行比较,这些模型最多训练了 300B 个 Token。Pile 指的是验证集分割,仅与在同一数据集和分词器上训练的模型(如 GPT-NeoX-20B)进行比较。对于每个模型规模,Mamba-2 的表现优于 Mamba,并且通常与 Pythia 在两倍模型规模下的表现相当。完整结果见表 10。

| 模型 | Token. | PILE PPL↓ | LAMBADA ↑Tdd | LAMBADA ACC↑ | HELLASWAG ACC↑ | PIQA ACC ↑ | ARC-E ACC↑ | ARC-C ACC↑ | WINOGRANDE ACC↑ | OPENBOOKQA ACC↑ | AVERAGE ACC↑ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Pythia-1B | NeoX | 7.82 | 7.92 | 56.1 | 47.2 | 70.7 | 57.0 | 27.1 | 53.5 | 31.4 | 49.0 |
| Mamba-790M | NeoX | 7.33 | 6.02 | 62.7 | 55.1 | 72.1 | 61.2 | 29.5 | 56.1 | 34.2 | 53.0 |
| Mamba-2-780M | NeoX | 7.26 | 5.86 | 61.7 | 54.9 | 72.0 | 61.0 | 28.5 | 60.2 | 36.2 | 53.5 |
| Hybrid H3-1.3B | GPT2 | - | 11.25 | 49.6 | 52.6 | 71.3 | 59.2 | 28.1 | 56.9 | 34.4 | 50.3 |
| Pythia-1.4B | NeoX | 7.51 | 6.08 | 61.7 | 52.1 | 71.0 | 60.5 | 28.5 | 57.2 | 30.8 | 51.7 |
| RWKV4-1.5B | NeoX | 7.70 | 7.04 | 56.4 | 52.5 | 72.4 | 60.5 | 29.4 | 54.6 | 34.0 | 51.4 |
| Mamba-1.4B | NeoX | 6.80 | 5.04 | 65.0 | 59.1 | 74.2 | 65.5 | 32.8 | 61.5 | 36.4 | 56.4 |
| Mamba-2-1.3B | NeoX | 6.66 | 5.02 | 65.7 | 59.9 | 73.2 | 64.3 | 33.3 | 60.9 | 37.8 | 56.4 |
| Hybrid H3-2.7B | GPT2 | - | 7.92 | 55.7 | 59.7 | 73.3 | 65.6 | 32.3 | 61.4 | 33.6 | 54.5 |
| Pythia-2.8B | NeoX | 6.73 | 5.04 | 64.7 | 59.3 | 74.0 | 64.1 | 32.9 | 59.7 | 35.2 | 55.7 |
| RWKV4-3B | NeoX | 7.00 | 5.24 | 63.9 | 59.6 | 73.7 | 67.8 | 33.1 | 59.6 | 37.0 | 56.4 |
| Mamba-2.8B | NeoX | 6.22 | 4.23 | 69.2 | 66.1 | 75.2 | 69.7 | 36.3 | 63.5 | 39.6 | 59.9 |
| Mamba-2-2.7B | NeoX | 6.09 | 4.10 | 69.7 | 66.6 | 76.4 | 69.6 | 36.4 | 64.0 | 38.8 | 60.2 |

说明:

  • 表格中的“-”表示该位置没有数据。
  • “PILE PPL↓”表示 PILE 数据集上的困惑度(越低越好)。
  • “LAMBADA ↑Tdd” 和其他以“↑”结尾的指标表示准确率(越高越好)。


Figure 10: (Efficiency Benchmarks.) $(L e f t)$ Our SSD is $2-8\times$ faster than a Mamba fused scan for large state expansion $\left(N=64\right)$ ) and faster than Flash Attention-2 for sequence length 2k and above. $(R i g h t)$ Sequence length 4K: Increasing state expansion slows down the Mamba optimized scan implementation linearly. SSD can handle much larger state expansion factors without much slowdown.

![](https://u254848-88c6-e493554b.yza1.seetacloud.com:8443/miner/v2/analysis/pdf_img?as_attachment=False&pdf=6fe16414fa471a5ab150cddcba5840d3a94e347e96992c50c12c1fb8b919479f1735788830_2405.21060v1.pdf&filename=d8a50c2c3cabc11376ad87697747c8eab6ead23c19492b747e1da063c6070187.jpg)  
图 10: (效率基准测试。)  (左)  我们的 SSD 比 Mamba 融合扫描在大规模状态扩展 (N=64) 下快 2-8 倍,并且在序列长度为 2k 及以上时比 Flash Attention-2 更快。  (右)  序列长度为 4K:增加状态扩展会线性减慢 Mamba 优化的扫描实现。SSD 可以处理更大规模的状态扩展,而不会显著减速。

We compare on a challenging version of the MQAR setup from (Arora, Eyuboglu, Zhang, et al. 2024), using a harder task, longer sequences, and smaller models. Our baselines include standard multi-head softmax attention as well as the Based architecture which combines convolutions, local attention, and a linear attention variant.

我们使用来自 (Arora, Eyuboglu, Zhang, et al. 2024) 的更具挑战性的 MQAR 设置版本进行比较,任务更难、序列更长、模型更小。我们的基准模型包括标准的多头 softmax 注意力机制以及 Based 架构,后者结合了卷积、局部注意力和线性注意力变体。

Results are shown in Figure 8. While Mamba-1 struggles on this task, Mamba-2 performs well across all settings. Surprisingly, it is significantly better than Mamba-1 even when the state sizes are controlled $(\mathsf{N}=16)$ ). (We are not sure which aspect of the architecture is the predominant factor, which remains a question to explore in future work.) Additionally, this task validates the importance of state size: increasing from $\textsf{N}=16$ to $\mathsf{N},=,64$ and $\mathsf{N},=,256$ consistently improves performance on MQAR, as the larger state allows more information (key-value pairs) to be memorized.

结果如图 8 所示。虽然 Mamba-1 在此任务上表现不佳,但 Mamba-2 在所有设置下都表现出色。令人惊讶的是,即使在状态大小受控的情况下 (N=16),Mamba-2 的表现也显著优于 Mamba-1。(我们不确定架构的哪个方面是主要因素,这仍然是未来工作需要探索的问题。)此外,该任务验证了状态大小的重要性:从 N=16 增加到 N=64 和 N=256 时,MQAR 的性能持续提升,因为更大的状态允许记住更多的信息(键值对)。

图 8:

9.2 Language Modeling

9.2 语言模型 (Language Modeling)

语言模型是自然语言处理中的一个重要任务,旨在根据给定的上下文预测下一个词的概率分布。这一领域的研究不仅推动了文本生成、机器翻译等应用的发展,还为大语言模型 (LLM) 的训练提供了基础。在本节中,我们将详细介绍语言模型的基本概念、常见架构以及最新的进展。

Following standard protocols in LLMs, we train and evaluate the Mamba-2 architecture on standard auto regressive language modeling against other architectures. We compare both pre training metrics (perplexity) and zero-shot evaluations. The model sizes (depth and width) follow GPT3 specifications, from $125\mathrm{m}$ to 2.7B. We use the Pile dataset (L. Gao, Biderman, et al. 2020), and follow the training recipe described in Brown et al. (2020). This follows the same setup as reported in Mamba (Gu and Dao 2023); training details are in Appendix D.

遵循大语言模型 (LLM) 中的标准协议,我们在标准自回归语言建模任务上训练和评估 Mamba-2 架构,并与其他架构进行比较。我们比较了预训练指标(困惑度)和零样本评估。模型规模(深度和宽度)遵循 GPT3 的规范,从 125m 到 2.7B。我们使用 Pile 数据集 (L. Gao, Biderman 等, 2020),并遵循 Brown 等人 (2020) 描述的训练方案。这与 Mamba (Gu 和 Dao, 2023) 报告的设置相同;训练细节见附录 D。

遵循大语言模型 (LLM) 中的标准协议,我们在标准自回归语言建模任务上训练和评估 Mamba-2 架构,并与其他架构进行比较。我们比较了预训练指标(困惑度)和零样本评估。模型规模(深度和宽度)遵循 GPT3 的规范,从 125m 到 2.7B。我们使用 Pile 数据集 (L. Gao, Biderman 等, 2020),并遵循 Brown 等人 (2020) 描述的训练方案。这与 Mamba (Gu 和 Dao, 2023) 报告的设置相同;训练细节见附录 D。

9.2.1 Scaling Laws

9.2.1 扩展定律 (Scaling Laws)

scaling laws 在中文中通常翻译为“扩展定律”或“缩放定律”。根据要求,这里使用“扩展定律”。

For baselines, we compare against both Mamba and its Transformer $^{++}$ recipe (Gu and Dao 2023), which is based on the PaLM and LLaMa architectures (e.g. rotary embedding, SwiGLU MLP, RMSNorm instead of LayerNorm, no linear bias, and higher learning rates). As Mamba has already demonstrated that it outperforms the standard Transformer architecture (GPT3 architecture) as well as recent sub quadratic architectures (H3 (Dao, D. Y. Fu, et al. 2023), Hyena (Poli et al. 2023), RWKV-4 (B. Peng, Alcaide, et al. 2023), RetNet (Y. Sun et al. 2023)), we omit those in the plot for clarity (see Gu and Dao (2023) for comparisons).

对于基准测试,我们对比了 Mamba 以及其改进版 Transformer $^{++}$ 配方 (Gu and Dao 2023),该配方基于 PaLM 和 LLaMa 架构(例如旋转位置编码、SwiGLU MLP、使用 RMSNorm 而不是 LayerNorm、没有线性偏置以及更高的学习率)。由于 Mamba 已经证明其性能优于标准的 Transformer 架构(如 GPT-3 架构)以及最近的次二次架构(如 H3 (Dao, D. Y. Fu, et al. 2023)、Hyena (Poli et al. 2023)、RWKV-4 (B. Peng, Alcaide, et al. 2023)、RetNet (Y. Sun et al. 2023)),为了清晰起见,我们在图表中省略了这些对比(详见 Gu and Dao (2023) 的比较)。

Figure 9 shows scaling laws under the standard Chinchilla (Hoffmann et al. 2022) protocol, on models from $\approx,125M$ to $\approx1.3B$ parameters.

图 9: 展示了在标准 Chinchilla (Hoffmann et al. 2022) 协议下的扩展规律,模型参数规模从约 1.25 亿到约 13 亿。

注:$\approx,125M$ 表示约 1.25 亿参数,$\approx1.3B$ 表示约 13 亿参数。

9.2.2 Downstream Evaluations

9.2.2 下游评估

下游评估 (Downstream Evaluations) 是指在预训练模型完成后,使用特定任务来评估模型性能的过程。这些任务通常与模型的最终应用场景相关,能够更直接地反映模型的实际应用效果。下游评估可以帮助我们了解模型在不同任务上的表现,从而为模型的选择和优化提供依据。

Table 1 shows the performance of Mamba-2 on a range of popular downstream zero-shot evaluation tasks, compared to the most well-known open source models at these sizes, most importantly Pythia (Biderman et al. 2023) which were trained with the same tokenizer, dataset, and training length (300B tokens) as our models.

表 1: 显示了 Mamba-2 在一系列流行的下游零样本 (Zero-shot) 评估任务中的表现,并与这些规模下最著名的开源模型进行了比较,特别是 Pythia (Biderman et al. 2023),后者使用与我们模型相同的分词器、数据集和训练长度(300B Token)进行训练。

9.2.3 Hybrid Models: Combining SSD Layer with MLP and Attention

9.2.3 混合模型:将 SSD 层与 MLP 和注意力机制结合 (Hybrid Models: Combining SSD Layer with MLP and Attention)

混合模型通过将 SSD (Single Shot Detector) 层与多层感知机 (MLP) 和注意力机制 (Attention) 结合,旨在提升模型的性能和灵活性。SSD 层主要用于目标检测任务,而 MLP 和注意力机制则可以增强模型对复杂特征的学习能力。这种组合方式能够在保持 SSD 层高效性的同时,引入更强大的特征表示能力,从而提高模型在各种任务中的表现。

通过将 SSD 层与 MLP 和注意力机制结合,模型可以在处理图像或序列数据时,不仅能够快速定位目标,还能更好地理解上下文信息,进而提升整体的预测准确性。此外,注意力机制可以帮助模型聚焦于重要的特征区域,减少噪声干扰,进一步优化模型的性能。

Recent and concurrent work (Dao, D. Y. Fu, et al. 2023; De et al. 2024; Glorioso et al. 2024; Lieber et al. 2024) suggests that a hybrid architecture with both SSM layers and attention layers could improve the model quality over that of a Transformer, or a pure SSM (e.g., Mamba) model, especially for in-context learning. We explore the different ways that SSD layers can be combined with attention and MLP to understand the benefits of each. Empirically we find that having around $10%$ of the total number of layers being attention performs best. Combining SSD layers, attention layers, and MLP also works better than either pure Transformer $++$ or Mamba-2.

最近和同期的研究 (Dao, D. Y. Fu, et al. 2023; De et al. 2024; Glorioso et al. 2024; Lieber et al. 2024) 表明,结合 SSM 层和注意力层的混合架构可以提高模型的质量,优于纯 Transformer 或纯 SSM(例如 Mamba)模型,尤其是在上下文学习中。我们探索了 SSD 层与注意力层和 MLP 结合的不同方式,以理解每种组合的优势。实验结果表明,大约 10% 的总层数为注意力层时表现最佳。结合 SSD 层、注意力层和 MLP 的效果也优于纯 Transformer ++ 或 Mamba-2。

SSD and Attention We find that SSD and attention layers are complementary: by themselves (e.g. in the Mamba-2 architecture vs. Transformer $\mathrel{+{+}}$ ) their performance (measured by perplexity) is nearly the same, but a mixture of SSD and attention layers outperforms the pure Mamba-2 or Transformer $^{++}$ architecture. We show some results (Table 2) for the $350\mathrm{M}$ model (48 layers) trained to 7B tokens on the Pile with the GPT-2 tokenizer (same number of parameters, same hyper parameters, same training and validation set). Adding in just a few attention layers already yields notable improvement and strikes the best balance between quality and efficiency. We hypothesize that the SSM layers function well as a general sequence-to-sequence mapping, and attention layers act as a retrieval mechanism to quickly refer to previous tokens in the sequence instead of forcing the model to compress all the context to its memory (SSM states).

SSD 和注意力机制

我们发现 SSD 和注意力层是互补的:单独使用时(例如在 Mamba-2 架构与 Transformer + 中),它们的性能(通过困惑度衡量)几乎相同,但混合使用 SSD 和注意力层的模型优于纯 Mamba-2 或 Transformer ++ 架构。我们展示了一些 350M 模型(48 层)的结果(表 2),该模型使用 GPT-2 分词器在 Pile 数据集上训练到 7B 个 Token(参数数量相同,超参数相同,训练和验证集相同)。仅仅添加几个注意力层就已经带来了显著的改进,并在质量和效率之间达到了最佳平衡。我们推测 SSM 层作为通用的序列到序列映射功能良好,而注意力层则充当检索机制,可以快速引用序列中的先前 Token,而不是强迫模型将所有上下文压缩到其内存(SSM 状态)中。

表 2:

请注意,原文中的公式和特殊字符已按要求保留原样。

Table 2: (Combining SSD and Attention Blocks.) Perplexity of a $350\mathrm{M}$ model with 48 layers, with different number of attention layers. Having around a $10%$ ratio of attention layers performs best.

表 2: (结合 SSD 和 Attention Blocks) 模型困惑度,该模型有 350M 参数和 48 层,不同数量的注意力层。大约 10% 的注意力层比例表现最佳。

注意力层数 困惑度
0
5
10
15
20

注:表格中的具体数值未给出,因此保持空白。

| NUM.ATTNBLOCKS | 0 (Mamba-2) | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 6 | 11 | 15 | 24 | Transformer++ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| PERPLEXITY↓ | 8.60 | 8.38 | 8.32 | 8.29 | 8.29 | 8.28 | 8.26 | 8.27 | 8.28 | 8.30 | 8.34 | 8.50 | 8.68 |

Hybrid Models with SSD, MLP, and Attention We compare different ways that SSD can be combined with the (gated) MLP and attention layers, and evaluate at the 2.7B scale (64 layers), trained to 300B tokens on the Pile (same number of parameters, same hyper parameters, same training and validation set, same data order):

带有 SSD、MLP 和注意力机制的混合模型

我们比较了 SSD 与 (门控) MLP 和注意力层的不同组合方式,并在 2.7B 规模(64 层)上进行了评估,在 Pile 数据集上训练了 300B 个 Token(参数数量相同,超参数相同,训练和验证集相同,数据顺序相同):

We report the validation perplexity on the Pile, as well as zero-shot evaluation, in Table 3. In general, the quality of Transformer $\mathrel{\mathop{:}}++$ and Mamba-2 models are around the same. We see that adding just 6 attention layers noticeably improves over the pure Mamba-2 model (and over Transformer $^{++}$ ). Adding MLP layers reduces model quality, but can (i) speed up training and inference due to the simplicity and hardware-efficiency of the MLP layer (ii) be easier to up-cycle to MoE models by replacing MLP layers with mixture-of-experts.

我们在表 3 中报告了在 Pile 数据集上的验证困惑度以及零样本 (Zero-shot) 评估结果。总体而言,Transformer $\mathrel{\mathop{:}}++$ 和 Mamba-2 模型的质量大致相同。我们发现,仅添加 6 个注意力层就能显著提升纯 Mamba-2 模型的性能(也超过了 Transformer $^{++}$)。添加 MLP 层会降低模型质量,但可以 (i) 由于 MLP 层的简单性和硬件效率,加速训练和推理 (ii) 通过将 MLP 层替换为专家混合 (MoE) 模型,更容易进行模型升级。

表 3: 验证困惑度及零样本评估结果
模型 困惑度
Transformer $\mathrel{\mathop{:}}++$ ...
Mamba-2 ...
添加 6 个注意力层的 Mamba-2 ...
添加 MLP 层的 Mamba-2 ...

注:表格中的具体数值未给出,可根据实际情况填写。

Table 3: (Zero-shot Evaluations.) Best results for each size in bold. We compare different ways SSD, MLP, and attention layers can be combined, evaluated at 2.7B scale trained to 300B tokens on the Pile.

表 3: (零样本评估) 每个规模的最佳结果用粗体表示。我们比较了不同方式组合的 SSD、MLP 和注意力层,在 Pile 数据集上训练到 300B Token 的 2.7B 规模模型的表现。

模型 Token. PILE PPL↓ LAMBADA PPL↓ LAMBADA ACC↑ HELLASWAG ACC↑ PIQA ACC↑ ARC-E ACC↑ ARC-C ACC↑ WINOGRANDE ACC↑ OPENBOOKQA ACC↑ AVERAGE ACC↑
Transformer++ NeoX 6.13 3.99 70.3 66.4 75.2 67.7 37.8 63.9 40.4 60.2
Mamba-2 NeoX 6.09 4.10 69.7 66.6 76.4 69.6 36.4 64.0 38.8 60.2
Mamba-2-MLP NeoX 6.13 4.18 69.3 65.0 76.4 68.1 37.0 63.1 38.2 59.6
Mamba-2-Attention NeoX 5.95 3.85 71.1 67.8 75.8 69.9 37.8 65.3 39.0 61.0
Mamba-2-MLP-Attention NeoX 6.00 3.95 70.0 66.6 75.4 70.6 38.6 64.6 39.2 60.7

注释:

  • PPL 表示困惑度(Perplexity),数值越低越好,因此用 ↓ 表示。
  • ACC 表示准确率(Accuracy),数值越高越好,因此用 ↑ 表示。

9.3 Speed Benchmarks

9.3 速度基准测试

根据要求,仅提供上述标题的翻译,不添加额外内容。

We benchmark the speed of the SSD algorithm against Mamba’s scan implementation and Flash Attention-2 (Figure 10). SSD, thanks to its reformulation to use matrix multiplication as a subroutine, can exploit specialized matrix multiplication (matmul) units on GPUs, also known as tensor cores. As a result, it is $2{-}8\times$ faster than Mamba’s fused associative scan, which does not leverage matmul units. Due to its linear scaling in sequence length, SSD is faster than Flash Attention-2 starting at sequence length $2K$ .

我们对 SSD 算法的速度进行了基准测试,对比了 Mamba 的扫描实现和 Flash Attention-2 (图 10)。SSD 通过重新设计以使用矩阵乘法作为子程序,可以利用 GPU 上的专用矩阵乘法 (matmul) 单元,也称为张量核心 (tensor cores)。因此,它比 Mamba 的融合关联扫描快 2-8 倍,后者没有利用 matmul 单元。由于其在序列长度上的线性扩展特性,SSD 在序列长度达到 2K 时开始比 Flash Attention-2 更快。

图 10:

However, we note that the Mamba-2 model as a whole might not be as efficient to train as Transformer at short sequence length (e.g. at $2K$ ), since a Transformer with $L$ layers would have $\begin{array}{l}{{\frac{L}{2}}}\end{array}$ MLP layers and $\begin{array}{l}{{\frac{L}{2}}}\end{array}$ attention layers, while a Mamba-2 model would have $L$ SSD layers for the same number of parameters. Generally the MLP layers are very hardware efficient since they consist of simple matrix multiplication and pointwise linearity. As shown in Section 9.2.3, one can also combine $\begin{array}{l}{{\frac{L}{2}}}\end{array}$ SSD layers and $\begin{array}{l}{{\frac{L}{2}}}\end{array}$ MLP layers to speed up training at short sequence length.

然而,我们注意到 Mamba-2 模型整体上在短序列长度(例如 2K)时可能不如 Transformer 训练得高效。因为一个具有 ( L ) 层的 Transformer 会有 ( \frac{L}{2} ) 个 MLP 层和 ( \frac{L}{2} ) 个注意力层,而 Mamba-2 模型在同一参数数量下会有 ( L ) 个 SSD 层。通常,MLP 层非常硬件高效,因为它们由简单的矩阵乘法和逐点线性运算组成。如第 9.2.3 节所示,也可以将 ( \frac{L}{2} ) 个 SSD 层和 ( \frac{L}{2} ) 个 MLP 层结合起来,以加速短序列长度下的训练。

Table 4: (Ablations: Mamba-2 block.) We ablate the major differences between the Mamba-2 and Mamba-1 neural network blocks (Figure 6, Section 7.1). Note that these components are independent of the inner sequence mixing layer; in these ablations, we use SSD for the inner SSM layer (differing from the S6 layer of Mamba-1).

表 4: (消融实验:Mamba-2 块) 我们对 Mamba-2 和 Mamba-1 神经网络块之间的主要差异进行了消融实验 (图 6, 第 7.1 节)。请注意,这些组件与内部序列混合层无关;在这些消融实验中,我们使用 SSD 作为内部 SSM 层(不同于 Mamba-1 的 S6 层)。

ABCX 投影 额外归一化 参数 困惑度
Mamba-1 顺序 129.3M 11.76
顺序 129.3M 11.54
并行 126.5M 11.66
Mamba-2 并行 126.5M 11.49

9.4 Architecture Ablations

9.4 架构消融研究 (Architecture Ablations)

架构消融研究是指通过系统地移除或修改模型的某些组件,来分析这些组件对模型性能的影响。这种研究方法可以帮助我们更好地理解模型各个部分的作用,以及哪些设计选择对最终结果最为关键。在本节中,我们将详细探讨不同架构变体的实验结果,并分析它们对模型性能的具体影响。

9.4.1 Block Design

9.4.1 块设计 (Block Design)

块设计是一种实验设计方法,用于将实验单元分组,以减少实验误差并提高实验结果的精确度。在块设计中,实验单元被分为若干个“块”,每个块内的实验单元具有相似的特征,而不同块之间的特征可能存在差异。通过这种方式,可以更好地控制实验中的变异因素,从而获得更可靠的结果。

块设计常用于农业试验、医学研究等领域,特别是在需要考虑环境或个体差异的情况下。常见的块设计类型包括完全随机区组设计 (Completely Randomized Block Design, CRBD) 和拉丁方设计 (Latin Square Design) 等。

在实际应用中,块设计可以帮助研究人员更有效地分配资源,并确保实验结果的可重复性和可靠性。

Section 7.1 introduces the Mamba-2 block, which has small modifications to the Mamba-1 block which are partly motivated by the connection to attention and also to improve the s cal ability of Mamba-2. Table 4 ablates these architecture changes to the block, which occur outside of the core SSM layer.

7.1 节介绍了 Mamba-2 模块,该模块对 Mamba-1 模块进行了少量修改,这些修改部分是由于与注意力机制的联系,同时也为了提高 Mamba-2 的可扩展性。表 4 对这些发生在核心 SSM 层之外的架构更改进行了消融研究。

注意:在原文中,“s cal ability”可能是一个拼写错误,应为“scalability”。

The ablations validate that parallel projections to create $(A,B,C,X)$ saves parameters and performs slightly better than Mamba’s sequential projections. More importantly, this modification is amenable to tensor parallelism at larger model sizes (Section 8). Additionally, the extra normalization layer also slightly improves performance. More importantly, preliminary experiments at larger scales observed that it also helps with training stability.

消融实验验证了并行投影生成 $(A,B,C,X)$ 能够节省参数,并且性能略优于 Mamba 的顺序投影。更重要的是,这种修改在更大模型规模下有利于张量并行(第 8 节)。此外,额外的归一化层也略微提升了性能。更重要的是,初步的大规模实验观察到它还提高了训练稳定性。

9.4.2 Head Structure

9.4.2 头结构 (Head Structure)


由于该段落仅包含标题,没有更多内容需要翻译。根据规则,当仅有标题时,不需要询问提供更多内容。

Section 7.2 describes how the dimensions of the $B,C,X$ projections can be viewed as a hyper parameter analogous to notions of multi-head attention and multi-query attention. We also showed how the original Mamba architecture is analogous to multi-value attention (Proposition 7.2), which was a choice that naturally developed from the state-space model point of view and was not previously ablated.

7.2 节描述了如何将 $B,C,X$ 投影的维度视为类似于多头注意力 (multi-head attention) 和多查询注意力 (multi-query attention) 的超参数。我们还展示了原始 Mamba 架构如何类似于多值注意力 (multi-value attention)(命题 7.2),这是从状态空间模型的角度自然发展而来的选择,并且之前没有进行过消融实验。

Table 5 ablates choices of the multi-head structure for the Mamba-2 architecture. Strikingly, we find a large difference between multi-value and multi-query or multi-key head patterns, despite seeming very similar. Note that this is not explained by the total state size, which is the same for all of them (equal to HPN or the product of the number of heads, head dimension, and state dimension).

表 5: 对 Mamba-2 架构的多头结构选择进行了消融研究。令人惊讶的是,我们发现多值 (multi-value) 和多查询 (multi-query) 或多键 (multi-key) 头模式之间存在显著差异,尽管它们看起来非常相似。需要注意的是,这种差异不能用总状态大小来解释,因为所有情况下的总状态大小都是相同的(等于 HPN 或者是头数量、头维度和状态维度的乘积)。

注:HPN 是指 head projection number。

We also compare to multi-head patterns where the number of $C,B,X$ (analogous to $Q,K,V)$ heads is equal. We compare against the standard multi-head pattern, as well as one with aggressive sharing where they all have only 1 head. Note that in the latter case, the model still has H different sequence mixers $M$ , because each head still has a different $A$ . When parameter matched, these multi-head patterns perform similarly to each other, in between the MVA and MQA/MKA patterns.

我们还比较了多头模式,其中 $C, B, X$ (类似于 $Q, K, V$)的头数相等。我们与标准的多头模式进行了对比,同时也对比了一种极端共享的情况,即所有头都只有一个。需要注意的是,在后一种情况下,模型仍然有 H 个不同的序列混合器 $M$,因为每个头仍然有不同的 $A$。当参数匹配时,这些多头模式的表现彼此相似,介于 MVA 和 MQA/MKA 模式之间。

在多头模式中,$C, B, X$ 分别对应查询(Query)、键(Key)和值(Value),而 $Q, K, V$ 是 Transformer 中常用的术语。标准多头模式是指每个头都有独立的参数,而在极端共享的情况下,所有头共享相同的参数,但每个头仍然有自己的注意力矩阵 $A$。这种设计使得模型在参数量相同的情况下,表现介于多值注意力(MVA)和多查询/多键注意力(MQA/MKA)之间。

9.4.3 Attention Kernel Approximations

9.4.3 注意力核函数近似 (Attention Kernel Approximations)


根据要求,仅翻译了标题,并保留了原始 Markdown 格式。等待更多内容以继续翻译。

Section 7.3 noted how SSD can be combined with ideas from the linear attention literature, such as various forms of kernel approximations. We ablate several variants of these suggested by previous works in Table 6. These include the cosFormer (Qin, Weixuan Sun, et al. 2022), Random Feature Attention H. Peng et al. 2021, and Positive Random Features (Performer) (Cho roman ski et al. 2021).

第 7.3 节提到如何将 SSD 与线性注意力文献中的思想相结合,例如各种形式的核函数近似。我们在表 6 中对之前工作中提出的几种变体进行了消融实验。这些变体包括 cosFormer (Qin, Weixuan Sun, et al. 2022),随机特征注意力 (Random Feature Attention) (H. Peng et al. 2021),以及正随机特征 (Positive Random Features) (Performer) (Cho roman ski et al. 2021)。

We also ablate adding a normalization term, akin to the denominator of the softmax function in standard attention. We found that this introduced instabilities to most variants, but slightly improved performance for the ReLU activation function $\psi$ .

我们还测试了添加一个归一化项,类似于标准注意力机制中 softmax 函数的分母。我们发现这给大多数变体引入了不稳定性,但对于 ReLU 激活函数 $\psi$ 稍微提高了性能。

Table 7 also tests more recent proposals to improve linear attention that involve expanding the feature dimension (Based (Arora, Eyuboglu, Zhang, et al. 2024) and ReBased (Aksenov et al. 2024)). These linear attention extensions aim to appropriate the exp kernel with a quadratic approximation. ReBased also proposes to replace the QK activation function with a layer normalization; from an SSM-centric view we apply a normalization on top of $(B,C)$ before applying the SSM function.

表 7: 还测试了最近提出的改进线性注意力机制的方案,这些方案涉及扩展特征维度(Based (Arora, Eyuboglu, Zhang, et al. 2024) 和 ReBased (Aksenov et al. 2024))。这些线性注意力扩展的目标是通过二次近似来逼近 exp 核函数。ReBased 还提出用层归一化替换 QK 激活函数;从 SSM 为中心的角度来看,我们在应用 SSM 函数之前对 (B, C) 进行归一化。

表 7: 还测试了最近提出的改进线性注意力机制的方案,这些方案涉及扩展特征维度(Based (Arora, Eyuboglu, Zhang, et al. 2024) 和 ReBased (Aksenov et al. 2024))。这些线性注意力扩展的目标是通过二次近似来逼近 exp 核函数。ReBased 还提出用层归一化替换 QK 激活函数;从 SSM 为中心的角度来看,我们在应用 SSM 函数之前对 (B, C) 进行归一化。

Table 5: (Ablations: Multi-head structure.) All models have state expansion factor $N=64$ and head size $P,=,64$ and are trained to Chinchilla scaling law token counts. The number of $A$ heads is always equal to the total heads H, i.e. each head has a separate input-dependent $A$ decay factor. $(T o p)$ 125M models, 2.5B tokens (Bottom) 360M models, 7B tokens

表 5: (消融实验:多头结构) 所有模型的状态扩展因子 $N=64$ 和头大小 $P,=,64$,并按照 Chinchilla 缩放定律进行训练。$A$ 头的数量始终等于总头数 H,即每个头都有一个独立的输入依赖的 $A$ 衰减因子。 (上) 125M 模型,2.5B Token (下) 360M 模型,7B Token

SSM HEAD PATTERN ATTN.ANALOG A HEADS B HEADS C HEADS X HEADS LAYERS PARAMS PPL.
多输入 (MIS) 多值 (MVA) 24 1 1 24 24 126.5M 11.66
多收缩 (MCS) 多查询 (MQA) 24 1 24 1 24 126.5M 12.62
多扩展 (MES) 多键 (MKA) 24 24 1 1 24 126.5M 12.59
多头 (MHS) 多头 (MHA) 24 24 24 24 15 127.6M 12.06
多状态 (MSS) 24 1 1 1 36 129.6M 12.00
多输入 (MIS) 多值 (MVA) 32 1 1 32 48 361.8M 8.73
多收缩 (MCS) 多查询 (MQA) 32 1 32 1 48 361.8M 9.33
多扩展 (MES) 多键 (MKA) 32 32 1 1 48 361.8M 9.36
多头 (MHS) 多头 (MHA) 32 1 1 1 70 361.3M 9.01
多状态 (MSS) 32 32 32 32 29 357.3M 9.04

注释:

  • SSM HEAD PATTERN:SSM 头模式
  • ATTN.ANALOG:注意力模拟
  • A HEADS:A 头数量
  • B HEADS:B 头数量
  • C HEADS:C 头数量
  • X HEADS:X 头数量
  • LAYERS:层数
  • PARAMS:参数量
  • PPL.:困惑度

Table 6: (Ablations: Kernel approximations.) We test various proposals for the kernel activation function $\psi$ , including linear attention variants aiming to approximate the exp kernel from standard softmax attention.

表 6: (消融实验:核函数近似) 我们测试了多种用于核激活函数 $\psi$ 的方案,包括旨在近似标准 softmax 注意力中 exp 核的线性注意力变体。

核激活函数 困惑度
none 11.58
Swish 11.66
Exp 11.62
ReLU 11.73
ReLU+归一化 11.64
cosFormer 11.97
RandomFeatureAttention 11.57
Positive Random Features (Performer) 12.21

Table 7: (Ablations: Kernel approximations.) We test the (Re)Based methods for linear attention approximations, which involve expanded feature maps. (Top) 130M models. (Top) 380M models with $N=256$ .

表 7: (消融实验:核函数近似) 我们测试了用于线性注意力近似的 (Re)Based 方法,这些方法涉及扩展的特征映射。 (上) 130M 模型。 (下) 380M 模型,$N=256$ 。

  • 表 7: (Ablations: Kernel approximations.) 我们测试了用于线性注意力近似的 (Re)Based 方法,这些方法涉及扩展的特征映射。 (上) 130M 模型。 (下) 380M 模型,$N=256$ 。

注:为了保持一致性,建议使用第一种翻译方式。

| KERNEL ACTIVATION Q | PERPLEXITY |
| --- | --- |
| Swish | 11.67 |
| Swish + Taylor (Based) | 12.19 |
| LayerNorm | 11.50 |
| LayerNorm + Square (ReBased) | 11.84 |
| Swish | 8.58 |
| Swish + Taylor (Based) | 8.71 |
| LayerNorm | 8.61 |
| LayerNorm + Square (ReBased) | 8.63 |

We note that this technique has been independently proposed as the “QK-Norm” for softmax attention (Team 2024) and an “internal normalization” for Mamba (Lieber et al. 2024).

我们注意到,这种技术已被独立提出为 softmax 注意力的“QK-Norm” (Team 2024) 和 Mamba 的“内部归一化” (Lieber et al. 2024)。

Overall, Table 6 and Table 7 found that the kernel approximation methods we tried did not seem to improve over simple pointwise non-linear activation functions for $\psi$ . Thus our default settings for Mamba-2 used $\psi(x)=\mathsf{S w i s h}(x)$ to follow Mamba-1, but we suggest that removing this activation entirely may be a simpler choice that we did not extensively test.

总体而言,表 6 和表 7 的结果显示,我们尝试的核函数近似方法似乎并没有比简单的逐点非线性激活函数对 $\psi$ 有明显的改进。因此,Mamba-2 的默认设置采用了 $\psi(x)=\mathsf{S w i s h}(x)$,以保持与 Mamba-1 的一致性,但我们建议完全移除这个激活函数可能是一个更简单的选择,尽管我们没有对此进行广泛的测试。

We emphasize however that SSD and vanilla linear attention differ in the inclusion of the 1-semi separable mask $L$ , while the various linear attention methods in the literature were derived to approximate softmax attention without this term; thus, our negative results may be not unexpected.

然而,我们强调,SSD 和普通的线性注意力机制的主要区别在于是否包含 1-半可分离掩码 $L$,而文献中的各种线性注意力方法都是为了近似 softmax 注意力而不包含这个项;因此,我们的负面结果可能是意料之中的。

10 Related Work and Discussion

10 相关工作与讨论

根据上述要求,这是该标题的翻译结果。请提供接下来的内容以便继续翻译。

The state space duality framework bridges connections between SSMs, structured matrices, and attention. We discuss in more depth the relations between SSD and these concepts more broadly. Using ideas from each of the viewpoints, we also suggest some directions that the SSD framework can be extended in future work.

状态空间对偶性框架 (State Space Duality, SSD) 搭建了状态空间模型 (SSMs)、结构化矩阵和注意力机制之间的联系。我们将更深入地讨论 SSD 与这些概念之间的关系。借鉴每个视角的思想,我们还提出了一些未来工作中可以扩展 SSD 框架的方向。

10.1 State Space Models

10.1 状态空间模型 (State Space Models)

状态空间模型是一类用于描述系统动态行为的数学模型。这类模型通过定义系统的状态变量及其随时间的变化来刻画系统的演化过程。状态空间模型广泛应用于控制理论、信号处理、经济学等领域,能够有效地描述和预测复杂系统的动态特性。

Structured state space models can be characterized along the axes

结构化状态空间模型可以在多个维度上进行表征

Time Variance (Selectivity). The original structured SSMs (S4) were linear time-invariant (LTI) systems (Gu 2023; Gu, Goel, and Ré 2022) motivated by continuous-time online memorization (Gu, Dao, et al. 2020; Gu, Johnson, Goel, et al. 2021; Gu, Johnson, Timalsina, et al. 2023). Many variants of structured SSMs have been proposed (Dao, D. Y. Fu, et al. 2023; Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; Ma et al. 2023; J. T. Smith, Warrington, and Linderman 2023), including several that drop the recurrence and focus on the convolutional representation of LTI SSMs (D. Y. Fu et al. 2023; Y. Li et al. 2023; Poli et al. 2023; Qin, Han, Weixuan Sun, B. He, et al. 2023).

时间变化性 (Selectivity)。最初的结构化 SSMs (S4) 是线性时不变 (LTI) 系统 (Gu 2023; Gu, Goel, 和 Ré 2022),其灵感来源于连续时间在线记忆 (Gu, Dao, 等 2020; Gu, Johnson, Goel, 等 2021; Gu, Johnson, Timalsina, 等 2023)。许多结构化 SSMs 的变体已经被提出 (Dao, D. Y. Fu, 等 2023; Gu, Gupta, 等 2022; Gupta, Gu, 和 Berant 2022; Ma 等 2023; J. T. Smith, Warrington, 和 Linderman 2023),其中包括一些放弃递归并专注于 LTI SSMs 卷积表示的变体 (D. Y. Fu 等 2023; Y. Li 等 2023; Poli 等 2023; Qin, Han, Weixuan Sun, B. He, 等 2023)。

这些变体的研究表明,通过不同的设计选择,可以在保持 LTI 系统核心特性的同时,进一步提升模型的表达能力和效率。具体来说,一些研究者探索了如何通过卷积操作来替代传统的递归结构,从而简化模型的训练和推理过程,同时保持对长时间依赖性的有效捕捉。

SSD is a time-varying structured SSM, also known as a selective SSM introduced in Mamba (Gu and Dao 2023). Selective SSMs are closely related to gating mechanisms of RNNs, including classical RNNs such as the LSTM (Hochreiter and Schmid huber 1997) and GRU (J. Chung et al. 2014) as well as more modern variants such as the QRNN (Bradbury et al. 2016), SRU (Lei 2021; Lei et al. 2017), RWKV (B. Peng, Alcaide, et al. 2023), HGRN (Qin, Yang, and Zhong 2023), and Griffin (Botev et al. 2024; De et al. 2024). These RNNs differ in their parameter iz at ions in various ways, most importantly in the lack of a state expansion.

SSD 是一种时变结构化的 SSM,也称为选择性 SSM,由 Mamba (Gu and Dao 2023) 引入。选择性 SSM 与 RNN 的门控机制密切相关,包括经典的 RNN,如 LSTM (Hochreiter and Schmidhuber 1997) 和 GRU (J. Chung et al. 2014),以及更现代的变体,如 QRNN (Bradbury et al. 2016)、SRU (Lei 2021; Lei et al. 2017)、RWKV (B. Peng, Alcaide, et al. 2023)、HGRN (Qin, Yang, and Zhong 2023) 和 Griffin (Botev et al. 2024; De et al. 2024)。这些 RNN 在参数化方面有所不同,最重要的是它们缺乏状态扩展。

Dimensionality and State Expansion. An important characteristic of SSD, shared by previous SSMs in its lineage (S4, H3, Mamba), is that it is a single-input single-output (SISO) system where input channels are processed independently. This leads to a much larger effective state size of ND where N is the SSM state size (also called state expansion factor) and D is the standard model dimension. Traditional RNNs either have $\mathsf{N},=,1$ or are multi-input multi-output (MIMO) with dense $B,C$ matrices, either of which leads to a smaller state. While MIMO SSMs have been shown to work well in some domains (Lu et al. 2023; Orvieto et al. 2023; J. T. Smith, Warrington, and Linderman 2023), Mamba showed that state expansion is crucial for information-dense domains such as language. One of the main advantages of SSD is allowing for even larger state expansion factors without slowing down the model. Many subsequent works have since adopted state expansion (Section 10.4).

维度扩展与状态扩展。SSD 的一个重要特征,继承自其前身的 SSM(如 S4、H3、Mamba),是它是一个单输入单输出 (SISO) 系统,其中输入通道被独立处理。这导致了有效状态大小显著增大,达到 ND,其中 N 是 SSM 的状态大小(也称为状态扩展因子),D 是标准模型的维度。传统的递归神经网络 (RNN) 要么有 $\mathsf{N},=,1$,要么是多输入多输出 (MIMO) 系统,具有密集的 $B,C$ 矩阵,这两种情况都会导致较小的状态。虽然 MIMO SSM 在某些领域表现出色(Lu et al. 2023; Orvieto et al. 2023; J. T. Smith, Warrington, and Linderman 2023),但 Mamba 证明了在信息密集型领域(如语言)中,状态扩展是至关重要的。SSD 的主要优势之一是能够在不减慢模型速度的情况下实现更大的状态扩展因子。许多后续工作也采用了状态扩展(第 10.4 节)。

Structure. Compared to previous structured SSMs, the main restriction of SSD is on the expressivity of the state transitions $A_{t}$ . We note that more general SSMs, such as the case of diagonal $A_{t}$ , have the same theoretical efficiency as SSD, but are less hardware-friendly. This is because the dual quadratic form loses its attention-like interpretation and becomes more difficult to compute. Thus compared to Mamba, SSD differs only in a slightly more restrictive form of diagonal $A_{t}$ , and trades off this expressivity for improved hardware efficiency (and ease of implementation).

结构。与之前的结构化 SSM 相比,SSD 的主要限制在于状态转移矩阵 $A_{t}$ 的表达能力。我们注意到,更一般的 SSM,例如对角线形式的 $A_{t}$,在理论上具有与 SSD 相同的效率,但硬件友好性较差。这是因为双二次形式失去了类似注意力机制的解释,并且计算变得更加困难。因此,与 Mamba 相比,SSD 仅在对角线形式的 $A_{t}$ 上有稍微严格的限制,并通过这种限制换取了更好的硬件效率(和实现的简便性)。

We hypothesize that it may be possible to refine our structured matrix algorithms to improve to the general diagonal SSM case as well.

我们假设,可能可以改进我们的结构化矩阵算法,以提升到一般的对角线 SSM (diagonal SSM) 情况。

10.2 Structured Matrices

10.2 结构化矩阵

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

The first viewpoint of the state space duality adopts the viewpoint of these models as matrix sequence transformations or “matrix mixers”: sequence transformations (Definition 2.1) that can be represented as matrix multiplication (by a $\mathsf{T}\times\mathsf{T}$ matrix) along the sequence dimension T.

状态空间对偶的第一个观点将这些模型视为矩阵序列变换或“矩阵混合器”:序列变换(定义 2.1)可以表示为沿序列维度 T 的矩阵乘法(由一个 $\mathsf{T}\times\mathsf{T}$ 矩阵)。

根据这个观点,模型通过矩阵运算来处理和转换输入序列,其中每个输入序列被视为一个向量,并且变换过程可以通过矩阵乘法来实现。这种变换方式允许模型在序列的不同位置之间建立复杂的依赖关系。

Several such matrix mixers have been proposed before, where the primary axis of variation is the representation of the matrix. These include MLP-Mixer (Tolstikhin et al. 2021) (unstructured matrix), FNet (Lee-Thorp et al. 2021) (Fourier Transform matrix), M2 (Dao, B. Chen, et al. 2022; Dao, Gu, et al. 2019; Dao, Sohoni, et al. 2020; D. Fu et al. 2024) (butterfly/monarch matrix), Toeplitz matrices (Poli et al. 2023; Qin, Han, Weixuan Sun, B. He, et al. 2023), and even more exotic structures (De Sa et al. 2018; Thomas et al. 2018).

在此之前,已经提出了几种矩阵混合器,其中主要的变化轴是矩阵的表示形式。这些方法包括:

  • MLP-Mixer (Tolstikhin et al. 2021)(无结构矩阵)
  • FNet (Lee-Thorp et al. 2021)(傅里叶变换矩阵)
  • M2 (Dao, B. Chen, et al. 2022; Dao, Gu, et al. 2019; Dao, Sohoni, et al. 2020; D. Fu et al. 2024)(蝶形/君主矩阵)
  • 托普利茨矩阵 (Poli et al. 2023; Qin, Han, Weixuan Sun, B. He, et al. 2023)
  • 以及其他更奇特的结构 (De Sa et al. 2018; Thomas et al. 2018)。

An important characterization is that efficient (sub-quadratic) matrix sequence transformations are exactly those which have structured matrix mixers. A core result of the SSD framework is viewing SSMs as matrix mixers with a particular structure – semi separable matrices (Section 3). The linear vs. quadratic duality then takes the form of structured matrix multiplication vs. naive matrix multiplication.

一个重要特征是,高效的(次二次方)矩阵序列变换正是那些具有结构化矩阵混合器的变换。SSD框架的一个核心结果是将 SSM 视为具有特定结构——半分离矩阵(第 3 节)的矩阵混合器。线性与二次方的对偶性则表现为结构化矩阵乘法与朴素矩阵乘法之间的对比。

The structure matrix representation led to our efficient SSD algorithm through block decomposition s of particular semiseparable matrices (Section 6). We note that semi separable matrices are well-studied in the scientific computing literature, and incorporating those ideas may be a promising avenue for more improvements to state space models. We also suggest that focusing on the matrix mixer viewpoint can lead to more fruitful directions for sequence models, such as designing principled non-causal variants of Mamba, or finding ways to characterize and bridge the gap between softmax attention and sub-quadratic models through analyzing their matrix transformation structure.

结构矩阵表示通过特定的半可分矩阵 (semiseparable matrices) 的块分解,使我们能够实现高效的 SSD 算法(第 6 节)。我们注意到,半可分矩阵在科学计算文献中已有深入研究,结合这些研究成果可能会为状态空间模型带来更多改进。我们还建议,专注于矩阵混合器 (matrix mixer) 视角可以为序列模型带来更有成效的研究方向,例如设计有原则的非因果 Mamba 变体,或者通过分析它们的矩阵变换结构来表征并弥合 softmax 注意力机制和次二次模型之间的差距。

10.3 (Linear) Attention

10.3 (线性) Attention

线性 Attention 机制是一种特殊的注意力机制,它通过线性运算来计算注意力权重,而不是传统的 softmax 操作。这种机制可以显著减少计算复杂度,特别是在处理长序列时。线性 Attention 的主要优点是能够在保持模型性能的同时,提高计算效率和降低内存消耗。

在传统的 Attention 机制中,计算注意力权重通常需要进行矩阵乘法和 softmax 操作,这会导致计算复杂度为 O(n²),其中 n 是序列长度。而线性 Attention 通过使用特定的技巧(如局部窗口或稀疏连接)将复杂度降低到 O(n),从而使得模型能够更高效地处理长序列数据 [20]。

线性 Attention 的一个典型实现是 Transformer 模型中的局部敏感哈希 (LSH) 注意力机制,它通过将相似的 Token 映射到相同的桶中,减少了不必要的计算。此外,还有一些其他变体,如 Performer 模型中使用的随机特征方法,它们也能够在保持性能的同时减少计算开销。

总之,线性 Attention 为解决传统 Attention 机制在长序列上的计算瓶颈提供了一种有效的解决方案,使得大语言模型能够在更大规模的数据集上进行训练和推理。

Compared to standard (causal) attention, SSD has only two main differences.

与标准 (因果) 注意力机制相比,SSD 仅有两个主要区别。

First, SSD does not use the softmax activation of standard attention (Bahdanau, Cho, and Bengio 2015; Vaswani et al. 2017), which is what gives attention its quadratic complexity. When the softmax is dropped, the sequence can be computed with linear scaling through the linear attention framework (Katha ro poul os et al. 2020).

首先,SSD 不使用标准注意力机制中的 softmax 激活 (Bahdanau, Cho, 和 Bengio 2015; Vaswani 等 2017),这是导致注意力机制具有二次复杂度的原因。当去掉 softmax 后,序列可以通过线性注意力框架 (Katharopoulos 等 2020) 以线性缩放的方式进行计算。

Second, SSD multiplies the logits matrix by an input-dependent 1-semi separable mask. Thus this mask can be viewed as replacing the softmax in standard attention.

其次,SSD 将 logits 矩阵与一个依赖于输入的 1-半可分离掩码 (1-semi separable mask) 相乘。因此,这个掩码可以被视为替换了标准注意力机制中的 softmax。

This semi separable mask can also be viewed as providing positional information. The elements $a_{t}$ act as “gates” in the RNN sense, or a “selection” mechanism (see discussion in Mamba paper), and their cumulative products $a_{j:i}$ control how much interaction is allowed between positions 𝑖and $j$ . Positional embeddings (e.g. sinusoidal (Vaswani et al. 2017), AliBi (Press, N. Smith, and Lewis 2022), and RoPE (Su et al. 2021)) are an important component of Transformers that are often viewed as heuristics, and the 1-SS mask of SSD can be seen as a more principled form of relative positional embeddings. We note that this view was also posited concurrently by GateLoop (Katsch 2023).

这种半可分离掩码也可以被视为提供位置信息。元素 $a_{t}$ 在 RNN 的意义上充当“门”或“选择”机制(参见 Mamba 论文中的讨论),它们的累积乘积 $a_{j:i}$ 控制位置 𝑖 和 $j$ 之间的交互程度。

位置嵌入(例如正弦位置编码 (Vaswani et al. 2017)、AliBi (Press, N. Smith, and Lewis 2022) 和 RoPE (Su et al. 2021))是 Transformer 中的重要组成部分,通常被视为启发式方法。而 SSD 的 1-SS 掩码可以被看作是一种更原理性的相对位置嵌入形式。我们注意到,GateLoop (Katsch 2023) 也同时提出了这一观点。

The second viewpoint of state space duality is a special case of our more general structured masked attention (SMA) framework, where the duality is revealed as different contraction orderings on a simple 4-way tensor contraction. SMA is a strong generalization of linear attention that is much more general than SSD as well; other forms of structured masks may lead to more variants of efficient attention with different properties than SSD.

状态空间对偶的第二种观点是我们更为通用的结构化掩码注意力 (SMA) 框架的一个特例,在该框架中,对偶性表现为一个简单的四维张量收缩的不同收缩顺序。SMA 是线性注意力的一个强大泛化,比状态空间对偶 (SSD) 更加通用;其他形式的结构化掩码可能会导致具有不同性质的高效注意力机制的更多变体。

在 SMA 框架中,通过对张量收缩顺序的不同选择,可以揭示出不同的对偶关系。与 SSD 不同,SMA 提供了一个更广泛的视角,允许我们探索更多种类的高效注意力机制。这些机制不仅限于 SSD 的特性,还可以通过不同的结构化掩码设计来实现多样化的性能优化。

Beside leading to new models, these connections to attention can lead to other directions for understanding SSMs. For example, we are curious whether the phenomenon of attention sinks (Darcet et al. 2024; Xiao et al. 2024) exist for Mamba models, and more broadly whether interpret ability techniques can be transferred to SSMs (Ali, Zimerman, and Wolf 2024).

除了引导新的模型之外,这些与注意力机制的联系还可以为理解状态空间模型 (SSM) 提供其他方向。例如,我们好奇注意力汇现象 (attention sinks) 是否存在于 Mamba 模型中 (Darcet et al. 2024; Xiao et al. 2024),更广泛地说,是否可以将可解释性技术迁移到 SSM 中 (Ali, Zimerman, and Wolf 2024)。

Finally, many other variants of linear attention have been proposed (Arora, Eyuboglu, Timalsina, et al. 2024; Arora, Eyuboglu, Zhang, et al. 2024; Cho roman ski et al. 2021; H. Peng et al. 2021; Qin, Han, Weixuan Sun, Dongxu Li, et al. 2022; Qin, Weixuan Sun, et al. 2022; Schlag, Irie, and Schmid huber 2021; Zhang et al. 2024; Zheng, C. Wang, and Kong 2022) (see Section 4.1.3 for descriptions of several of these), and we expect that many techniques can be transferred to SSMs (e.g. Section 7.3).

最后,许多其他线性注意力机制的变体已经被提出 (Arora, Eyuboglu, Timalsina, et al. 2024; Arora, Eyuboglu, Zhang, et al. 2024; Cho roman ski et al. 2021; H. Peng et al. 2021; Qin, Han, Weixuan Sun, Dongxu Li, et al. 2022; Qin, Weixuan Sun, et al. 2022; Schlag, Irie, and Schmidhuber 2021; Zhang et al. 2024; Zheng, C. Wang, and Kong 2022)(详见第 4.1.3 节对其中一些变体的描述),我们预计许多技术可以迁移到 SSMs(例如,第 7.3 节)。

We emphasize that SSD does not generalize standard softmax attention, or any other transformation on the attention kernel matrix that does not have a finite feature map $\psi$ . Compared to general attention, SSD’s advantage is having a controllable state expansion factor N that compresses the history, compared to quadratic attention’s cache of the entire history scaling with sequence length ${\mathsf{T}}\gg{\mathsf{N}},$ . Concurrent work has starting studying the tradeoffs of these representations, for example on copying and in-context learning tasks (Akyürek et al. 2024; Grazzi et al. 2024; Jelassi et al. 2024; Park et al. 2024). We note that Mamba-2 significantly improves on Mamba on some of these capabilities (e.g. as demonstrated by MQAR results in Section 9.1), but more remains to be understood.

我们强调,SSD 并不泛化标准的 softmax 注意力机制,或任何其他对注意力核矩阵没有有限特征映射 $\psi$ 的变换。与一般的注意力机制相比,SSD 的优势在于具有可控的状态扩展因子 N,可以压缩历史信息,而二次注意力机制则需要缓存整个历史信息,其规模随序列长度 ${\mathsf{T}} \gg {\mathsf{N}}$ 增加。近期的研究已经开始探讨这些表示方法之间的权衡,例如在复制任务和上下文学习任务中的表现 (Akyürek et al. 2024; Grazzi et al. 2024; Jelassi et al. 2024; Park et al. 2024)。我们注意到,Mamba-2 在某些能力上显著优于 Mamba(例如,在第 9.1 节的 MQAR 结果中有所体现),但仍有更多问题需要进一步理解。

10.4 Related Models

10.4 相关模型

根据您的要求,我已经翻译了标题。请继续提供需要翻译的内容。

We finally highlight a growing body of recent and concurrent work that have developed sequence models very similar to Mamba and Mamba-2.

我们最后强调了一大批最近和同期的研究工作,这些工作开发了与 Mamba 和 Mamba-2 非常相似的序列模型。

这些研究展示了类似的技术路径和模型结构,表明在该领域内存在广泛的探索和共识 [20]。

• RetNet (Y. Sun et al. 2023) and Trans No rmer LL M (Qin, Dong Li, et al. 2023) generalize Linear Attention using decay terms instead of a cumulative sum, and propose dual parallel/recurrent algorithms as well as a hybrid “chunkwise” mode. These algorithms can be seen as an instantiation of SSD where $A_{t}$ is time-invariant (constant for all $t$ ); in the SMA interpretation, the mask matrix $L$ would be a decay matrix $L_{i,j}=\gamma^{i-j}$ . These models also differ architecturally in various ways. For example, since they were derived from an attention-centric perspective they preserve the multi-head attention (MHA) pattern; since Mamba-2 was derived from an SSM-centric pattern it preserves the multi-value attention (MVA) or multi-expand SSM (MES) pattern, which we show to be better (Section 9.4).

• RetNet (Y. Sun 等 2023) 和 Trans No rmer 大语言模型 (Qin, Dong Li 等 2023) 通过使用衰减项而不是累积和来推广线性注意力机制,并提出了双平行/递归算法以及混合的“分块”模式。这些算法可以被视为状态空间模型 (SSD) 的一种实例,其中 $A_{t}$ 是时间不变的(对所有 $t$ 均为常数);在滑动平均 (SMA) 解释中,掩码矩阵 $L$ 将是一个衰减矩阵 $L_{i,j} = \gamma^{i-j}$。这些模型在架构上也存在多种差异。例如,由于它们是从注意力机制的角度推导出来的,因此保留了多头注意力 (MHA) 模式;而 Mamba-2 是从状态空间模型 (SSM) 的角度推导出来的,因此保留了多值注意力 (MVA) 或多扩展状态空间模型 (MES) 模式,我们在第 9.4 节中展示了这种模式更为优越。

请注意,原文中的 "Trans No rmer LL M" 可能是拼写错误,通常应为 "Transformer LLM"。

11 Conclusion

We proposed a theoretical framework based on well-studied classes of structured matrices that bridges the conceptual gap between SSMs and attention variants. This framework yields insights on how recent SSMs (e.g. Mamba) perform as well as Transformers on language modeling. Moreover, our theoretical tools provide new ideas to improve SSMs (and potentially Transformers) by connecting the algorithmic and systems advances on both sides. As a demonstration, the framework guides our design of a new architecture (Mamba-2) at the intersection of SSMs and structured attention.

我们提出了一种基于结构化矩阵类的理论框架,该框架弥合了状态空间模型 (SSM) 和注意力机制变体之间的概念差距。这一框架揭示了最近的 SSM(例如 Mamba)在语言建模任务中表现得与 Transformer 一样出色的原因。此外,我们的理论工具通过连接双方的算法和系统进展,为改进 SSM(以及可能的 Transformer)提供了新的思路。作为演示,该框架指导我们设计了一种新的架构(Mamba-2),它位于 SSM 和结构化注意力的交界处。

我们提出了一种基于结构化矩阵类的理论框架,该框架弥合了状态空间模型 (SSM) 和注意力机制变体之间的概念差距。这一框架揭示了最近的 SSM(例如 Mamba)在语言建模任务中表现得与 Transformer 一样出色的原因。此外,我们的理论工具通过连接双方的算法和系统进展,为改进 SSM(以及可能的 Transformer)提供了新的思路。作为演示,该框架指导我们设计了一种新的架构(Mamba-2),它位于 SSM 和结构化注意力的交界处。

Acknowledgments

References

参考文献

根据您的要求,我已经完成了对标题的翻译。请提供需要翻译的具体内容,我将继续按照规则进行翻译。

A Glossary

术语表

相关说明

  • 本术语表旨在解释本文中使用的关键术语和概念。
  • 首次出现的专业术语会在括号内标注英文原文,例如:“生成式 AI (Generative AI)”,之后将只使用中文。
  • 保留原始段落格式、术语(如 FLAC,JPEG 等)、公司缩写(如 Microsoft, Amazon, OpenAI 等)以及引用格式(如 [20])。
  • 图表标题将翻译为中文,例如:“Figure 1: ”翻译为“图 1: ”,“Table 1: ”翻译为:“表 1: ”。

常见术语

  • Transformer:一种用于处理序列数据的神经网络架构,广泛应用于自然语言处理任务。
  • Token:在自然语言处理中,Token 是指文本的基本单位,可以是单词、字符或子词。
  • 大语言模型 (LLM/Large Language Model):指具有大量参数的语言模型,能够生成高质量的文本并执行各种语言任务。
  • 零样本 (Zero-shot):指模型在没有见过特定任务的训练数据的情况下,仍然能够完成该任务的能力。
  • 少样本 (Few-shot):指模型仅需少量示例即可学习并完成新任务的能力。
  • AI智能体 (AI Agent):能够感知环境并采取行动以实现目标的智能系统。
  • 通用人工智能 (AGI):指具有与人类智能相当的广泛认知能力的人工智能系统。
  • Python语言:一种高级编程语言,广泛用于数据分析、机器学习等领域。

Table 8: Glossary of notation and terminology; mnemonics bolded. (Top) Frequently used tensor dimensions. (Bottom) Matrices and tensors used in state space models or structured masked attention.

表 8: 符号和术语表;助记符加粗。 (上) 常用的张量维度。 (下) 状态空间模型或结构化掩码注意力中使用的矩阵和张量。

符号 含义
B 批次大小 (Batch size)
T 序列长度 (Sequence length)
D 隐藏层维度 (Hidden dimension)
H 头数 (Number of heads)
C 通道数 (Number of channels)
N 节点数 (Number of nodes)
K 关键值维度 (Key dimension)
V 值维度 (Value dimension)
符号 含义
X 输入张量 (Input tensor)
Y 输出张量 (Output tensor)
W 权重矩阵 (Weight matrix)
A 注意力矩阵 (Attention matrix)
M 掩码矩阵 (Mask matrix)
S 状态矩阵 (State matrix)
P 位置编码 (Positional encoding)
| 符号 | 描述 | 定义 |
| --- | --- | --- |
| T | 时间轴或目标序列轴 | 定义 2.1 |
| S | 源序列轴(在注意力机制中) | 公式 (9) |
| D | 模型维度或 d_model | 定义 7.1 |
| N | 状态/特征维度或 d_state | 公式 (2) 和 (9) |
| P | 头维度或 d_head | 定义 2.1 |
| H | 头的数量或 n_head | 定义 7.1 |
| M | 序列变换矩阵 | 定义 2.3 |
| A | 离散 SSM 递归(状态)矩阵 | 公式 (2) |
| B | 状态空间模型输入投影(扩展)矩阵 | 公式 (2) |
| C | 状态空间模型输出投影(收缩)矩阵 | 公式 (2) |
| X | 输入矩阵(形状 (T, P)) | 公式 (2) 和 (9) |
| Y | 输出矩阵(形状 (T, P)) | 公式 (2) 和 (9) |
| Q | 注意力查询矩阵 | 公式 (9) |
| K | 注意力键矩阵 | 公式 (9) |
| V | 注意力值矩阵 | 公式 (9) |
| G | 注意力 Gram 矩阵 | QK^T (或 CB^T)|
| L | (结构化)掩码矩阵(在因果设置中为下三角矩阵) | 定义 4.2 |

B Efficient Algorithms for the Scalar SSM Scan (1-SS Multiplication)

B 标量 SSM 扫描的高效算法 (1-SS 乘法)

根据上述要求,这是该标题的翻译结果。如果您有更多内容需要翻译,请继续提供。

In this section we flesh out various algorithms for computing the scalar SSM scan, through the lens of structured matrix decomposition s. The scalar SSM scan is defined as computing the recurrent part of the discrete SSM (7), in the case when $N=1$ (i.e. $A$ is a scalar). This is commonly used to compute SSMs re currently; in particular, the case of structured SSMs where $A$ is diagonally structured reduces down to this operation, such as in the S5 (J. T. Smith, Warrington, and Linderman 2023) and S6 (Gu and Dao 2023) models.

在本节中,我们将详细介绍用于计算标量 SSM (State Space Model) 扫描的各种算法,通过结构化矩阵分解的角度来探讨。标量 SSM 扫描的定义是计算离散 SSM (7) 的递归部分,在 $N=1$ 的情况下(即 $A$ 是一个标量)。这在当前计算 SSM 时非常常见;特别是当 $A$ 具有对角结构的结构化 SSM 时,该操作会简化为这种情况,例如在 S5 (J. T. Smith, Warrington, 和 Linderman 2023) 和 S6 (Gu 和 Dao 2023) 模型中。

在本节中,我们将详细介绍用于计算标量 SSM (State Space Model) 扫描的各种算法,通过结构化矩阵分解的角度来探讨。标量 SSM 扫描的定义是计算离散 SSM (7) 的递归部分,在 \$N=1\$ 的情况下(即 \$A\$ 是一个标量)。这在当前计算 SSM 时非常常见;特别是当 \$A\$ 具有对角结构的结构化 SSM 时,该操作会简化为这种情况,例如在 S5 (J. T. Smith, Warrington, 和 Linderman 2023) 和 S6 (Gu 和 Dao 2023) 模型中。

The goal of this section is to support a central theme of this paper that efficient algorithms for sequence models can be viewed as structured matrix multiplication algorithms. The various matrix decomposition ideas we show here are related to ideas used to derive fast SSM algorithms (Section 6), as well as directly used as a subroutine.

本节的目标是支持本文的一个核心主题,即序列模型的有效算法可以被视为结构化的矩阵乘法算法。我们在这里展示的各种矩阵分解思想与用于推导快速 SSM 算法 (第 6 节) 的思想相关,同时也被直接用作子程序。

B.1 Problem Definition

B.1 问题定义

问题定义部分通常会明确研究的目标和范围,介绍研究背景以及需要解决的具体问题。请提供后续内容以便继续翻译。

Let $a:(\mathsf{D},)$ and $b:(\mathsf{D},)$ be sequences of scalars. The scalar SSM scan is defined as

令 \$a:(\mathsf{D},)\$ 和 \$b:(\mathsf{D},)\$ 为标量序列。标量 SSM (State Space Model) 扫描定义为

$$
h_{t}=a_{t}h_{t-1}+b_{t}.
$$

$$
h_{t} = a_{t} h_{t-1} + b_{t}.
$$

该公式表示在时间步 t 时,隐藏状态 ( h_{t} ) 是由前一个时间步的隐藏状态 ( h_{t-1} ) 通过系数 ( a_{t} ) 加权后,再加上偏置项 ( b_{t} ) 得到的。

Here $h_{-1}$ can be an arbitrary value representing the previous hidden state to the SSM recurrence; unless otherwise specified, we assume $h_{-1}=0$ .

这里 $h_{-1}$ 可以是任意值,表示 SSM 递归的前一个隐藏状态;除非另有说明,我们假设 $h_{-1}=0$ 。

We also call equation (21) the cumprodsum (cumulative product sum). Note that the cumprodsum reduces to the cumprod (cumulative product) when $b=0$ is the additive identity and it reduces to the cumsum (cumulative sum) when $a=1$ is the multiplicative identity.

我们还称公式 (21) 为累乘和 (cumprodsum)。请注意,当加法单位元 $b=0$ 时,累乘和退化为累乘 (cumprod);当乘法单位元 $a=1$ 时,累乘和退化为累加 (cumsum)。

  • 累乘和 (cumprodsum):累积乘积和
  • 累乘 (cumprod):累积乘积
  • 累加 (cumsum):累积和

这里 $b=0$ 是加法单位元,意味着任何数与 0 相加结果不变;$a=1$ 是乘法单位元,意味着任何数与 1 相乘结果不变。

Finally, note that in vectorized form we can write

最后,注意在向量化形式下我们可以写为

请注意,由于原文内容较短且没有提供具体的数学表达式或上下文,这里仅翻译了句子的文本部分。如果需要完整的数学表达式或更多上下文,请提供完整的内容。

$$
\begin{array}{r}{M=M b}\ {M=\left[\begin{array}{c 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]}\end{array}
$$

$$
\begin{array}{r}
M = M b \
M = \left[
\begin{array}{cccccc}
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]
\end{array}
$$

该矩阵表示一个下三角矩阵,其中每个元素 $a_i$ 表示某个时间步的系数。矩阵的第一列是不同时间步的系数乘积,而对角线上的元素均为 1。这种结构常用于时间序列分析和递归计算中。

In other words, this is simply the matrix-vector product by a 1-SS matrix $M$ .

换句话说,这仅仅是 1-SS 矩阵 $M$ 的矩阵-向量乘积。

Therefore we have three ways of viewing this fundamental primitive operation that are all equivalent:

因此,我们有三种等价的方式来看待这个基本的原始操作:

(注:原文中未列出具体的三种方式,请根据上下文补充完整)

B.2 Classical Algorithms

B.2 经典算法

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

We first describe the two classical ways of computing the SSM scan (21), previously used by prior work.

我们首先描述两种经典的 SSM 扫描计算方法 (21),这些方法之前已被先前的研究使用过。

B.2.1 Sequential Recurrence

B.2.1 顺序递归 (Sequential Recurrence)

顺序递归是指在处理序列数据时,模型通过逐步处理每个时间步的数据来生成输出。这种机制常见于自然语言处理任务中,例如语言模型和机器翻译。在顺序递归中,模型的当前状态依赖于之前的时间步的状态,从而能够捕捉到序列中的长期依赖关系。

在大语言模型 (LLM) 中,顺序递归可以通过多种方式实现,例如使用循环神经网络 (RNN) 或 Transformer 模型。这些模型能够在处理长序列时保持上下文信息,确保生成的输出具有连贯性和一致性。

顺序递归的一个重要特点是它能够处理变长的输入序列,这使得它在处理自然语言、语音等非结构化数据时非常有效。然而,顺序递归也存在一些挑战,例如在处理非常长的序列时可能会出现梯度消失或爆炸的问题 [20]。

The recurrent mode simply computes (21) one timestep $t$ at a time. From the perspective of 1-SS multiplication, this was also described in Section 3.4.1.

递归模式 simply 计算 (21) 每个时间步 $t$。从 1-SS 乘法的角度来看,这也在第 3.4.1 节中进行了描述。

根据上下文,这里的一些术语可能需要进一步解释:

  • “递归模式”指的是在计算过程中逐个时间步进行处理的方式。
  • “1-SS 乘法”是指单步状态 (Single Step State) 乘法,这是一种特定的计算方法。

如果你对这些术语有更具体的理解或需要进一步解释,请参考相关文献或上下文。

B.2.2 Parallel Associative Scan

B.2.2 并行关联扫描 (Parallel Associative Scan)

并行关联扫描是一种用于在并行计算中高效处理数据的技术。它通过将数据划分为多个部分,并在每个部分上并行执行操作,从而加速计算过程。这种技术常用于图形处理、科学计算和大数据处理等领域。

并行关联扫描的核心思想是将一个全局的操作分解为多个局部操作,然后通过合并这些局部结果来得到最终的全局结果。常见的并行关联扫描算法包括前缀和 (Prefix Sum) 等。

在实现并行关联扫描时,通常需要考虑以下几个方面:

  • 数据划分:如何将数据有效地划分为多个子集,以便在不同的处理器或线程上并行处理。
  • 局部计算:在每个子集上执行的局部操作,通常是某种形式的累积或聚合操作。
  • 结果合并:如何将各个子集的局部结果合并为最终的全局结果,确保计算的正确性。

并行关联扫描的效率取决于数据的划分方式、局部操作的复杂度以及结果合并的策略。优化这些方面可以显著提高并行计算的性能。

Second, an important observation is that this recurrence can be turned into an associative scan (E. Martin and Cundy 2018; J. T. Smith, Warrington, and Linderman 2023). This fact is not completely obvious. For example, S5 defined the correct associative scan operator and then showed associativity of the operator through rote calculation.

其次,一个重要的观察是,这个递归可以被转换成一个结合扫描 (E. Martin 和 Cundy 2018; J. T. Smith, Warrington, 和 Linderman 2023)。这一事实并不完全显而易见。例如,S5 定义了正确的结合扫描操作符,然后通过机械计算证明了该操作符的结合性。

A slightly cleaner way to see that this is computable with an associative scan is to turn the multi-term recurrence into a single-term recurrence on a hidden state of size 2 instead of 1:

一种更清晰的方法是将多项递归转换为在大小为 2 的隐藏状态上的单项递归,而不是大小为 1 的状态,这样可以通过关联扫描 (associative scan) 来计算:

这种方式展示了如何通过使用大小为 2 的隐藏状态来简化多项递归关系,使其可以使用关联扫描进行计算。

$$
\begin{array}{c}{{h_{t}=a_{t}h_{t-1}+b_{t}}}\ {{\left[h_{t}\right]=\left[a_{t}\quad b_{t}\right]\left[h_{t-1}\right].}}\end{array}
$$

$$
\begin{array}{c}
h_{t} = a_{t} h_{t-1} + b_{t} \
\left[ h_{t} \right] = \left[ a_{t} \quad b_{t} \right] \left[ h_{t-1} \right].
\end{array}
$$

该公式表示当前状态 ( h_t ) 是由前一个状态 ( h_{t-1} ) 通过系数 ( a_t ) 和 ( b_t ) 线性组合得到的。具体来说,( h_t ) 可以看作是 ( a_t ) 与 ( h_{t-1} ) 的乘积加上 ( b_t ),或者用矩阵形式表示为 ( [h_t] = [a_t \quad b_t] [h_{t-1}] )。

Then computing all the $h_{t}$ is the same as taking the cumulative products of these $2\times2$ matrices. Since matrix multiplication is associative, this can be computed with an associative scan. The associative binary operator is simply matrix multiplication on these particular matrices:

然后计算所有的 $h_{t}$ 等同于对这些 $2\times2$ 矩阵进行累积乘积。由于矩阵乘法是结合的,这可以通过结合扫描(associative scan)来计算。这里的结合二元运算就是对这些特定矩阵进行矩阵乘法:

注意:在原文中,“associative scan” 是一种特定的计算方法,用于并行计算累积操作。为了保持专业性,这里保留了英文术语。

$$
\left[a_{t}\quad b_{t}\right]\left[a_{s}\quad b_{s}\right]=\left[a_{t}a_{s}\quad a_{t}b_{s}+b_{t}\right].
$$

$$
\left[ a_{t} \quad b_{t} \right] \left[ a_{s} \quad b_{s} \right] = \left[ a_{t} a_{s} \quad a_{t} b_{s} + b_{t} \right] .
$$

该公式表示两个向量的乘积运算,其中第一个向量为 ([a_t, b_t]),第二个向量为 ([a_s, b_s])。乘积结果为 ([a_t a_s, a_t b_s + b_t])。

Equating the top row yields the same associative scan operator as defined by S5:

将第一行相等,得到的关联扫描运算符与 S5 中定义的相同:

将第一行相等,得到的关联扫描运算符与 S5 中定义的相同:

$$
(a_{t},b_{t})\otimes(a_{s},b_{s})=(a_{t}a_{s},a_{t}b_{s}+b_{t}).
$$

$$
(a_{t}, b_{t}) \otimes (a_{s}, b_{s}) = (a_{t} a_{s}, a_{t} b_{s} + b_{t}).
$$

该公式表示两个元组的运算,其中第一个元素相乘,第二个元素按照特定规则相加。具体来说,给定两个元组 ((a_{t}, b_{t})) 和 ((a_{s}, b_{s})),它们的运算结果为 ((a_{t} a_{s}, a_{t} b_{s} + b_{t}))。

The reason why associative scans are important is that they can be parallel i zed using a divide-and-conquer algorithm (Blelloch 1990). We omit the details of this algorithm, and instead show that the entire associative SSM scan algorithm can be derived from scratch through matrix decomposition s (Appendix B.3.5).

关联扫描之所以重要,是因为它们可以使用分治算法 (Blelloch 1990) 进行并行化。我们省略了该算法的详细内容,而是通过矩阵分解从头推导出整个关联 SSM 扫描算法(附录 B.3.5)。

B.3 Efficient Algorithms via Structured Matrix Decomposition s

B.3 通过结构化矩阵分解实现高效算法

通过结构化矩阵分解 (Structured Matrix Decomposition) 可以显著提高算法的计算效率。这种方法利用了矩阵中的特定结构,使得在处理大规模数据时能够减少计算复杂度并加速运算过程。结构化矩阵分解在许多领域都有广泛应用,包括信号处理、图像压缩、机器学习等。

在本节中,我们将介绍几种常见的结构化矩阵分解方法及其应用场景。这些方法不仅能够提高计算效率,还能在某些情况下提供更好的数值稳定性。具体来说,我们将讨论以下内容:

  • 稀疏矩阵分解
  • 低秩矩阵近似
  • 快速傅里叶变换 (FFT) 和离散余弦变换 (DCT) 等变换域方法
  • 结构化矩阵在大语言模型 (LLM) 中的应用

通过这些技术,可以在保持高精度的同时,大幅减少计算资源的消耗,从而实现更高效的算法设计。

We discuss several algorithms for computing the SSM scan, all through the lens of finding structured matrix decompositions of the 1-SS matrix $M$ . These algorithms or computation modes include

我们讨论了多种用于计算 SSM 扫描的算法,所有这些算法都是通过寻找 1-SS 矩阵 $M$ 的结构化矩阵分解来实现的。这些算法或计算模式包括

B.3.1 Dilated Mode

B.3.1 稀疏模式 (Dilated Mode)

稀疏模式是一种特殊的网络结构设计,通过在卷积层中引入空洞(或称为膨胀)卷积来扩大感受野,而不需要增加参数量。这种设计使得模型能够在不显著增加计算成本的情况下捕获更大范围的上下文信息。稀疏模式常用于处理长序列数据或需要捕捉远距离依赖的任务。

This mode factors the 1-SS matrix in a particular way involving increasing “strides”. This is best illustrated through a concrete example:

这种模式以特定的方式对 1-SS 矩阵进行分解,涉及逐渐增大的“步幅”。通过一个具体的例子可以最好地说明这一点:

图 1:

表 1:

(如果原文中有图或表,请根据上下文补充)

Note that this closely resembles the computation of dilated convolutions.

请注意,这与空洞卷积 (dilated convolutions) 的计算非常相似。

We also note that this factorization shows that 1-SS matrices are a special case of butterfly matrices, another broad and fundamental type of structured matrix (Dao, Gu, et al. 2019; Dao, Sohoni, et al. 2020).

我们还注意到,这种分解表明 1-SS 矩阵是蝶形矩阵 (butterfly matrices) 的一个特例,而蝶形矩阵是另一种广泛且基础的结构化矩阵类型 (Dao, Gu, et al. 2019; Dao, Sohoni, et al. 2020)。

Remark 8. This algoritihm is sometimes described as a “work-inefficient but more parallel iz able” prefix sum algorithm (Hillis and Steele Jr 1986), becauses it uses $O(T\log(T))$ operations but has half the depth/span as the work-efficient associative scan algorithm.

备注 8. 这个算法有时被描述为“工作效率较低但更并行化”的前缀和算法 (Hillis and Steele Jr 1986),因为它使用了 $O(T\log(T))$ 操作,但其深度/跨度是工作高效的关联扫描算法的一半。

B.3.2 State-Passing (Chunkwise) Mode

B.3.2 状态传递(分块)模式 (State-Passing (Chunkwise) Mode)

在这种模式下,数据被分成多个块进行处理,每个块之间会传递状态信息。这种处理方式适用于需要保持上下文连续性的场景,例如在处理长文本或流式数据时。通过分块处理,可以有效减少内存占用,并提高处理效率。

在这种模式中,每个块的处理结果会影响后续块的处理,因此状态信息的传递至关重要。具体实现方式取决于应用场景和需求。

This mode can be viewed as a generalization of the standard recurrent mode where instead of passing forward the recurrent state $h$ one step at a time, we compute the answer on chunks of arbitrary length $k$ and pass the state through the chunk. This can also be derived from a simple block decomposition of the 1-SS matrix.

这种模式可以被视为标准递归模式的泛化。在标准递归模式中,我们是逐个时间步传递递归状态 $h$,而在这种模式下,我们可以在长度为任意 $k$ 的块上计算结果,并通过整个块传递状态。这也可以从 1-SS 矩阵的简单分块分解中推导出来。

Remark 9. While we call this “state-passing” to refer to how states are passed from one local segment to another, this is related to the “chunkwise” algorithms proposed by related models (Y. Sun et al. 2023; Yang et al. 2024).

备注 9. 尽管我们称其为“状态传递 (state-passing)”来描述状态如何从一个局部片段传递到另一个局部片段,这与相关模型提出的“分块 (chunkwise)”算法有关 (Y. Sun 等 2023; Yang 等 2024)。

Consider computing $h=M b$ in “chunks”: for some index $k,\in,[T]$ , we want to compute $h_{0:k}$ or the output up to index $k$ , and have a way to reduce the problem to a smaller problem on indices $[k:T]$ .

考虑分块计算 $h = M b$:对于某个索引 $k \in [T]$,我们希望计算 $h_{0:k}$ 或者直到索引 $k$ 的输出,并且有一种方法可以将问题简化为在索引 $[k:T]$ 上的较小问题。

这样做的目的是通过将计算过程分解为更小的部分,使得每个部分的计算更加高效和可控。通过递归地处理剩余的部分,我们可以逐步得到最终的结果。

We write $M$ as

我们用 $M$ 表示

注意:这里仅有一行内容,且没有更多的上下文信息。根据给定的规则,我保留了原始的特殊字符和公式,并将其直接翻译为中文。

$$
M=\left[\begin{array}{c c c c c c c c c}{a_{0;0}}&&&&&&&\ {a_{1;0}}&{a_{1;1}}&&&&&\ {\vdots}&&{\ddots}&&&&&\ {a_{k-1;0}}&{\ldots}&{\ldots}&{a_{k-1;k-1}}&&&&\ {a_{k;0}}&{\ldots}&{\ldots}&{a_{k;k-1}}&{a_{k;k}}&&&\ {\vdots}&&&{\vdots}&{\vdots}&{\ddots}&&\ {a_{T-1;0}}&{\ldots}&{\ldots}&{a_{T-1;k-1}}&{a_{T-1;k}}&{\ldots}&{a_{T-1;T-1}}\end{array}\right]
$$

$$
M=\left[\begin{array}{c c c c c c c c c}
a_{0;0} & & & & & & & \
a_{1;0} & a_{1;1} & & & & & \
\vdots & & \ddots & & & & & \
a_{k-1;0} & \ldots & \ldots & a_{k-1;k-1} & & & & \
a_{k;0} & \ldots & \ldots & a_{k;k-1} & a_{k;k} & & & \
\vdots & & & \vdots & \vdots & \ddots & & \
a_{T-1;0} & \ldots & \ldots & a_{T-1;k-1} & a_{T-1;k} & \ldots & a_{T-1;T-1}
\end{array}\right]
$$

矩阵 ( M ) 是一个 ( T \times T ) 的三角矩阵,其中元素 ( a_{i;j} ) 表示第 ( i ) 行第 ( j ) 列的值。矩阵的结构显示了从第 0 行到第 ( T-1 ) 行的元素分布,其中对角线上的元素为 ( a_{i;i} ),而其他位置的元素根据行和列的索引进行排列。

Let the upper-left triangle be $M_{L}$ , lower-right be $M_{R}$ (left and right sub problems), and lower-left be $M_{C}$ . Divide up $b$ into $b_{L}=b_{0:k}$ and $b_{R}=b_{k:T}$ in the same way. Note that

令左上三角为 \$M_{L}\$ ,右下为 \$M_{R}\$ (左子问题和右子问题),左下为 \$M_{C}\$ 。将 \$b\$ 按照相同的方式划分为 \$b_{L}=b_{0:k}\$ 和 \$b_{R}=b_{k:T}\$ 。注意,

$$
M b=\left[{\overset{M_{L}b_{L}}{M_{R}b_{R}+M_{C}b_{L}}}\right]
$$

$$
M b=\left[{\overset{M_{L}b_{L}}{M_{R}b_{R}+M_{C}b_{L}}}\right]
$$

公式中的各个符号和矩阵表示具体的数学运算关系,这里保持原样不变。如果您需要对公式的具体含义进行解释,请提供更多的上下文信息。

Also, $M_{C}$ has the rank-1 factorization (this is essentially the defining property of semi separable matrices)

此外,$M_{C}$ 具有秩为 1 的分解(这实际上是半可分矩阵的基本定义属性)[20]。

$$
M_{C}=\left[\begin{array}{c}{a_{k:k}}\ {\vdots}\ {a_{T-1:k}}\end{array}\right]a_{k},\left[a_{k-1:0}\quad\cdot\cdot\cdot\quad a_{k-1:k-1}\right]
$$

$$
M_{C}=\left[\begin{array}{c} a_{k:k} \ \vdots \ a_{T-1:k} \end{array}\right] a_{k} , \left[a_{k-1:0} \quad \cdot\cdot\cdot \quad a_{k-1:k-1}\right]
$$

该公式表示矩阵 ( M_C ) 的构造方式。具体来说,矩阵 ( M_C ) 由两部分组成:

  • 第一部分是一个列向量,包含从 ( a_{k:k} ) 到 ( a_{T-1:k} ) 的元素,这些元素按顺序排列。
  • 第二部分是一个行向量,包含从 ( a_{k-1:0} ) 到 ( a_{k-1:k-1} ) 的元素。

矩阵 ( M_C ) 的具体含义和应用场景可能需要结合上下文进一步解释。

Thus

因此

请注意,您提供的内容似乎不完整。如果您有更多内容需要翻译,请继续提供。根据您的要求,我将直接开始翻译并保持格式不变。如果这就是完整的句子,以上就是翻译结果。但通常情况下,这样的短句可能是一个更大段落的开头部分。

$$
M_{C}b_{L}=\left[\begin{array}{c}{{a_{k:k}}}\ {{\vdots}}\ {{a_{T-1:k}}}\end{array}\right]a_{k}\cdot(M b)_{k-1}.
$$

$$
M_{C}b_{L}= \left[ \begin{array}{c} a_{k:k} \ \vdots \ a_{T-1:k} \end{array} \right] a_{k} \cdot (M b)_{k-1}.
$$

该公式表示矩阵 ( M_C ) 与向量 ( b_L ) 的乘积,其中右边是一个由多个 ( a ) 向量组成的矩阵与 ( a_k ) 和 ( (M b){k-1} ) 的乘积。具体来说,矩阵的每一行是从 ( a{k:k} ) 到 ( a_{T-1:k} ) 的元素,最后乘以 ( a_k ) 和 ( (M b)_{k-1} )。

Here we think of $(M b){k-1},=,h{k-1}$ as the “final state” of the left chunk, because the row vector in $M_{C}$ ’s factorization is the same as the final row of $M_{L}$ . Furthermore, note that the column vector in $M_{C}$ ’s factorization is the same as the final column of $M_{R}$ .7 Thus

我们把 $(M b){k-1} = h{k-1}$ 视为左块的“最终状态”,因为 $M_{C}$ 分解中的行向量与 $M_{L}$ 的最后一行相同。此外,需要注意的是,$M_{C}$ 分解中的列向量与 $M_{R}$ 的最后一列相同 [7]。

因此

$$
M_{R}b_{R}+M_{C}b_{L}=M_{R}\left[\begin{array}{c}{{a_{k}h_{k-1}+b_{k}}}\ {{b_{k+1}}}\ {{\vdots}}\ {{b_{T-1}}}\end{array}\right]
$$

$$
M_{R}b_{R} + M_{C}b_{L} = M_{R} \left[ \begin{array}{c} a_{k}h_{k-1} + b_{k} \ b_{k+1} \ \vdots \ b_{T-1} \end{array} \right]
$$

该公式表示的是矩阵 ( M_{R} ) 与向量 ( b_{R} ) 的乘积加上矩阵 ( M_{C} ) 与向量 ( b_{L} ) 的乘积,等于矩阵 ( M_{R} ) 与一个列向量的乘积。这个列向量的元素从 ( a_{k}h_{k-1} + b_{k} ) 开始,后续元素依次为 ( b_{k+1} ),直到 ( b_{T-1} )。

Finally, we use the observation that $M_{L}$ and $M_{R}$ are self-similar to the original matrix $M$ ; the answers for these two smaller 1-SS matrix multiplications can be performed arbitrarily using any algorithm. In total, the algorithm proceeds as follows:

最后,我们观察到 $M_{L}$ 和 $M_{R}$ 与原始矩阵 $M$ 具有自相似性;这两个较小的 1-SS 矩阵乘法可以使用任意算法进行计算。总的来说,算法步骤如下:

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

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

翻译结果:

最后,我们观察到 $M_{L}$ 和 $M_{R}$ 与原始矩阵 $M$ 具有自相似性;这两个较小的 1-SS 矩阵乘法可以使用任意算法进行计算。总的来说,算法步骤如下:

In other words, we compute the left subproblem as a black box, pass its final state on to the right problem, and compute the right subproblem as a black box.

换句话说,我们将左子问题作为一个黑盒进行计算,将其最终状态传递给右子问题,然后将右子问题作为一个黑盒进行计算。

The utility of this method comes from more complicated settings, such as in the general $N{\cdot}$ -semi separable case, and when the input $b$ has an additional “batch” dimension (or in other words this is a matrix-matrix instead of matrix-vector multiplication). In this case, we can use an alternate algorithm for the chunks (corresponding to MM by $M_{L}$ and $M_{R}$ ) that does not materialize the full hidden states $h$ . Instead, we skip the hidden states and directly compute the final state $h_{k-1}$ in an alternate way, then “pass” the state to the next chunk.

这种方法的实用性体现在更复杂的场景中,例如在一般的 $N{\cdot}$ -半可分情况下,以及当输入 $b$ 具有额外的“批次”维度时(换句话说,这是矩阵-矩阵乘法而不是矩阵-向量乘法)。在这种情况下,我们可以使用一种替代算法来处理块(对应于由 $M_{L}$ 和 $M_{R}$ 执行的矩阵乘法),而不显式地生成完整的隐藏状态 $h$。相反,我们跳过隐藏状态,直接以另一种方式计算最终状态 $h_{k-1}$,然后将该状态“传递”到下一个块。

Complexity. This method can be very work-efficient because steps 2-3 takes only constant time. Therefore assuming the two sub problems (steps 1 and 4) are linear time, the whole method takes linear time.

复杂度。该方法可以非常高效,因为步骤 2-3 仅需常数时间。因此,假设两个子问题(步骤 1 和 4)是线性时间的,整个方法的时间复杂度为线性时间。

The downside is that this is also sequential.

缺点是这同样是一个顺序过程。

B.3.3 Fully Recurrent Mode

Note that the fully recurrent mode, where the recurrence is evolved one step at a time (21), is simply an instantiation of the state-passing mode with chunk size $k=1$ .

请注意,完全循环模式(fully recurrent mode),其中循环是一步一步演进的 (21),只是状态传递模式的一种实例,其块大小为 $k=1$ 。

B.3.4 (Parallel) Block Decomposition Mode

B.3.4 (并行) 块分解模式

在这种模式下,输入数据被划分为多个块,每个块可以独立处理。这种并行处理方式可以显著提高处理速度,特别是在处理大规模数据时。块分解模式适用于多种应用场景,例如图像压缩(如 JPEG)、音频编码(如 FLAC)等。

在并行块分解模式中,每个块可以由不同的处理器或线程同时处理,从而实现高效的并行计算。这种模式的优点是能够充分利用多核处理器的计算能力,减少处理时间。然而,需要注意的是,块与块之间的边界处理可能会引入一些额外的复杂性,需要特别注意数据的一致性和完整性。

引用的相关文献可以参考 [20]。

This uses the same matrix decomposition as the state-passing mode, but computes sub problems in a different order that trades off computation for parallel iz ation.

这使用了与状态传递模式相同的矩阵分解,但以不同的顺序计算子问题,这种顺序通过增加并行化来权衡计算量。

As usual, we write $M$ as

如往常一样,我们写作 $M$ 为

请注意:由于原文内容较少且没有完整表达,建议提供完整的句子或段落以便更准确地进行翻译。但如果这就是完整的句子,以上就是翻译结果。

$$
\boldsymbol{M}=\left[\begin{array}{c c c c c}{1}&{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]=\left[\begin{array}{c c c c c}{1}&&&&\ {-a_{1}}&{1}&&&\ {0}&{-a_{2}}&{1}&&\ {\vdots}&{\vdots}&{\ddots}&{\ddots}&\ {0}&{0}&{\ldots}&{-a_{T-1}}&{1}\end{array}\right]^{-1}\boldsymbol{\bar{G}}_{1}+\boldsymbol{D}_{2}.
$$

$$
\boldsymbol{M}=\left[\begin{array}{c c c c c} 1 & 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]=\left[\begin{array}{c c c c c} 1 & & & & \ -a_1 & 1 & & & \ 0 & -a_2 & 1 & & \ \vdots & \vdots & \ddots & \ddots & \ 0 & 0 & \ldots & -a_{T-1} & 1 \end{array}\right]^{-1} \boldsymbol{\bar{G}}_1 + \boldsymbol{D}_2.
$$

矩阵 (\boldsymbol{M}) 可以表示为两个矩阵的和,其中第一个矩阵是下三角矩阵的逆,第二个矩阵是 (\boldsymbol{D}_2)。具体形式如下:

  • 左边的矩阵是一个 (T \times T) 的矩阵,其主对角线元素均为 1,副对角线以下的元素由 (a_i) 组成。
  • 右边的矩阵是另一个 (T \times T) 的矩阵,其主对角线元素均为 1,副对角线以下的元素为 (-a_i),并且该矩阵的逆与 (\boldsymbol{\bar{G}}_1) 相乘后再加上 (\boldsymbol{D}_2)。

这个表达式展示了矩阵 (\boldsymbol{M}) 的分解形式,其中 (\boldsymbol{\bar{G}}_1) 和 (\boldsymbol{D}_2) 是给定的矩阵。

The key observation is again that the bottom-left quadrant of $M$ is rank-1. Aside from inspection, another way to see this is by using the RHS, observing that the bottom-left quadrant of it is a trivial rank-1 matrix (it is all 0 except the top-right corner is $-a_{T/2}$ ), and using the Woodbury inversion formula to see that the bottom-left corner of the LHS must also be rank 1. This also provides a way to deduce the rank-1 factorization, which can be verified through inspection:

关键观察是再次注意到矩阵 $M$ 的左下象限是秩为 1 的。除了通过直接检查外,另一种方法是使用右侧 (RHS),观察到其左下象限是一个平凡的秩为 1 的矩阵(它全部为 0,除了右上角为 $-a_{T/2}$),并使用 Woodbury 反演公式来推导出左侧 (LHS) 的左下角也必须是秩为 1 的。这还提供了一种推导秩为 1 分解的方法,该分解可以通过检查进行验证:

关键观察是再次注意到矩阵 \$M\$ 的左下象限是秩为 1 的。除了通过直接检查外,另一种方法是使用右侧 (RHS),观察到其左下象限是一个平凡的秩为 1 的矩阵(它全部为 0,除了右上角为 \$-a_{T/2}\$),并使用 Woodbury 反演公式来推导出左侧 (LHS) 的左下角也必须是秩为 1 的。这还提供了一种推导秩为 1 分解的方法,该分解可以通过检查进行验证:

$$
{\begin{array}{r l}&{M_{\mathrm{lower-left-quadrant}}=\left[\begin{array}{c c c c}{(a_{T/2}\ldots a_{1})}&{\ldots}&{a_{T/2}}\ {\vdots}&{\ddots}&{\vdots}\ {(a_{T-1}\ldots a_{T/2}a_{T/2-1}\ldots a_{1})}&{\ldots}&{(a_{T-1}\ldots a_{T/2})}\end{array}\right]}\ &{=\left[\begin{array}{c}{a_{T/2}}\ {\vdots}\ {a_{T-1}\ldots a_{T/2}}\end{array}\right]\left[(a_{T/2-1}\ldots a_{1})}&{\ldots}&{a_{T/2-1}}&{1\right],.}\end{array}}
$$

$$
{\begin{array}{r l}
& M_{\mathrm{lower-left-quadrant}} = \left[\begin{array}{c c c c}
{(a_{T/2} \ldots a_{1})} & {\ldots} & {a_{T/2}} \
{\vdots} & {\ddots} & {\vdots} \
{(a_{T-1} \ldots a_{T/2} a_{T/2-1} \ldots a_{1})} & {\ldots} & {(a_{T-1} \ldots a_{T/2})}
\end{array}\right] \
& = \left[\begin{array}{c}
{a_{T/2}} \
{\vdots} \
{a_{T-1} \ldots a_{T/2}}
\end{array}\right] \left[(a_{T/2-1} \ldots a_{1}) & {\ldots} & {a_{T/2-1}} & 1 \right] ,.
\end{array}}
$$

这段公式表示的是矩阵 ( M_{\mathrm{lower-left-quadrant}} ) 的构造方式。该矩阵的左下象限部分由两个矩阵相乘得到。第一个矩阵是一个列向量,第二个矩阵是一个行向量。具体来说:

  • 第一个矩阵的每一行是从 ( a_{T/2} ) 到 ( a_{T-1} ) 的序列。
  • 第二个矩阵的每一列是从 ( a_{T/2-1} ) 到 ( a_1 ) 的序列,并在最后添加了一个 1。

这种矩阵分解方式有助于简化计算和分析。

A second observation is that this matrix is self-similar: any principle submatrix has the same form. In particular, the top-left and bottom-right quadrants are both 1-SS matrices.

第二个观察是这个矩阵具有自相似性:任何主子矩阵都具有相同的结构。具体来说,左上角和右下角的四个分块都是 1-SS 矩阵。

This provides an easy way to perform the matrix multiplication by $M$ : recurse on the two halves (i.e. top-left and bottomright) in parallel, and then account for the bottom-left submatrix. This “combination” step in the divide-and-conquer algorithm is easy since the submatrix is rank 1. This leads to a parallel algorithm.

这提供了一种简单的方法来执行矩阵 $M$ 的乘法:并行递归处理两个半部分(即左上角和右下角),然后考虑左下角的子矩阵。在分治算法中,这个“组合”步骤非常简单,因为该子矩阵的秩为 1。这导致了一个并行算法。

这提供了一种简单的方法来执行矩阵 \$M\$ 的乘法:并行递归处理两个半部分(即左上角和右下角),然后考虑左下角的子矩阵。在分治算法中,这个“组合”步骤非常简单,因为该子矩阵的秩为 1。这导致了一个并行算法。

Complexity. Like the state-passing algorithm, this method uses the same block decomposition s of the rank-structured semi separable matrices. The difference is that we recurse on both sub problems in parallel, while the state-passing algorithm handles the left and then right sub problems. This lowers the depth/span of the algorithm from linear to $\log(T)$ . The tradeoff is that the combination step (accounting for the rank-1 bottom-left submatrix) requires linear instead of constant work, so the total work is $O(T\log(T))$ instead of linear.

复杂度。与状态传递算法类似,该方法使用相同的块分解 ( s ) 来处理秩结构的半可分离矩阵。不同之处在于,我们并行地递归处理两个子问题,而状态传递算法则是先处理左侧子问题再处理右侧子问题。这将算法的深度/跨度从线性降低到 ( \log(T) )。权衡是,组合步骤(考虑秩为 1 的左下角子矩阵)需要线性的工作量,而不是常数工作量,因此总的工作量为 ( O(T\log(T)) ),而不是线性。

Note also that in the recursion, we can stop at any time and compute the sub problems in any other way. This is a main idea behind the SSD algorithm (Section 6), where we switch to the dual quadratic attention formulation on small sub problems.

请注意,在递归过程中,我们可以在任何时候停止,并以其他任何方式计算子问题。这是 SSD 算法(第 6 节)背后的主要思想,在该算法中,我们在小的子问题上切换到对偶二次注意力公式。

B.3.5 Associative Scan Mode

B.3.5 关联扫描模式 (Associative Scan Mode)

关联扫描模式是一种用于处理数据的并行计算方法。在这种模式下,数据被划分为多个部分,每个部分可以独立进行处理,最后将结果合并。这种模式特别适用于需要对大量数据进行快速处理的应用场景。

关联扫描模式的核心思想是通过并行化操作来加速计算过程。具体来说,它可以通过以下步骤实现:

  1. 数据划分:将输入数据划分为多个子集,每个子集可以在不同的处理器或线程上并行处理。
  2. 局部计算:对每个子集进行独立的计算,生成中间结果。
  3. 结果合并:将所有中间结果合并,得到最终的输出。

这种模式在许多领域都有广泛的应用,例如图像处理、数据分析和机器学习等。通过利用并行计算的优势,关联扫描模式可以显著提高处理效率,尤其是在处理大规模数据时。

关联扫描模式的一个重要特点是它的可扩展性。随着数据量的增加,可以通过增加更多的处理器或线程来保持高效的处理速度。此外,关联扫描模式还可以与其他并行计算技术结合使用,进一步提升性能。

在实际应用中,关联扫描模式常常与特定的数据结构和算法相结合,以优化计算过程。例如,在图像处理中,关联扫描模式可以用于快速计算图像的累积和,从而加速后续的滤波和卷积操作 [20]。

图 1: 关联扫描模式的工作流程示例

表 1: 不同扫描模式的性能对比

模式 处理时间 内存占用 并行度
串行扫描 较长 较低
关联扫描 较短 较高

关联扫描模式的优点在于它能够充分利用现代多核处理器和分布式系统的计算能力,从而实现高效的数据处理。然而,它的实现也面临一些挑战,例如如何有效地划分数据以及如何最小化不同处理器之间的通信开销。未来的研究可能会集中在这些方面,以进一步优化关联扫描模式的性能。

The state passing (chunkwise) algorithm has linear work, but also involves sequential operations.

状态传递(chunkwise)算法的工作量是线性的,但也涉及顺序操作。

The block matrix reduction and dilated modes are parallel iz able: they have $\log(T)$ depth/span. However, they do extra work $(O(T\log(T))$ .

块矩阵约简和膨胀模式是可并行化的:它们具有 $\log(T)$ 的深度/跨度。然而,它们会进行额外的工作 $(O(T \log(T)))$ 。

注意:在原文中,$\log(T)$ 和 $O(T \log(T))$ 是数学表达式,因此保持原样不变。

As noted in Appendix B.2.2, there is an algorithm that achieves both ${\cal O}(\log{T})$ depth and $O(T)$ work by leveraging the associative scan (also called prefix scan) algorithm (Baker et al. 1996). This algorithm is most easily seen from the SSM scan or cumprodsum view, and even then is not obvious: it requires separately deriving an associative operator (22), and then leveraging the parallel/associative/prefix scan algorithm as a black box (Blelloch 1990).

如附录 B.2.2 所述,存在一种算法可以通过利用关联扫描(也称为前缀扫描)算法 (Baker et al. 1996) 实现 ${\cal O}(\log{T})$ 深度和 $O(T)$ 工作量。该算法从 SSM 扫描或 cumprodsum 视角最容易理解,但即便如此也不明显:它需要分别推导出一个关联操作符 (22),然后将并行/关联/前缀扫描算法作为黑箱使用 (Blelloch 1990)。

这种算法的关键在于如何有效地利用关联操作符来实现并行计算,从而在保持较低深度的同时完成大量的工作。具体来说,关联扫描算法通过将输入数据分解为多个子问题,并行处理这些子问题,最后将结果合并,从而实现了高效的计算。

Here we show that it is actually possible to derive this parallel scan from leveraging a different matrix decomposition:

这里我们展示实际上可以通过利用不同的矩阵分解来推导出这种并行扫描算法:

| M= | ao:0 a1:0 | a1:1 | a2:2 |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | a2:0 | a2:1 | a2:2 |  |  |  |  |  |  |  |
|  | a3:0 a4:0 | a3:1 a4:1 | a3:2 a4:2 | a3:3 a4:3 | a4:4 |  |  |  |  |  |
|  | a5:0 a6:0 | a5:1 a6:1 |  | a5:2 a6:2 | a5:3 a6:3 | a5:4 a6:4 | a5:5 a6:5 | a6:6 |  |  |
|  | a7:0 | a7:1 |  | a7:2 | a7:3 | a7:4 | a7:5 | a7:6 | a7:7 |  |
|  | a1:0 | ao:0 | a1:0 | a1:1 |  | a2:2 |  |  |  |  |
|  | a2:2 a3:2] a1:1] | a2:1 | a2:1 | a2:1 | a2:1 | a2:1 | a3:2 a3:3 | a4:4 |  |  |
|  | a4:4 a5:4 | a1:0 a4:1 a1:1 | a1:0 a4:1 a1:1 | a1:0 a4:1 a1:1 | a1:0 a4:1 a1:1 | a1:0 a4:1 a1:1 | a4:4 a3:2 a4:3 a3:3 | a5:4 @5:5 |  |  |
|  | a6:6 a7:6 |  | a6:1 | a1:0 a1:1 |  | a5:4 a6:6 a7:6 | a6:3 | a3:2 a3:3 | a6:6 a6:1 a7:6 | @5:4 a5:5 | a6:6 a7:6 a7:7 |

请注意,表格中的某些单元格包含重复或不完整的数据,可能需要进一步确认和清理。

Stage 1. First we compute the answers for each of the diagonal blocks in the multiplication $M b$ . This produces two numbers, but the first element is unchanged. For example, the second block is going to compute $b_{2}$ and $a_{3}b_{2}+b_{3}$

阶段 1

首先,我们计算乘法 $M b$ 中每个对角块的答案。这会产生两个数值,但第一个元素保持不变。例如,第二个块将计算 $b_{2}$ 和 $a_{3}b_{2}+b_{3}$。

Stage 2. Now consider each of the $2\times2$ blocks factored as a rank-1 matrix in the strictly lower triangular part of the matrix. Note that each of the right side row vectors is the same as the bottom row vector in the diagonal block in its column: in particular the $\left[a_{1:0}\ a_{1:1}\right],,\left[a_{3:2}\ a_{3:3}\right]$ , and $[a_{5:4}~a_{5:5}]$ rows.

阶段 2. 现在考虑每个 $2\times2$ 的块可以分解为矩阵的严格下三角部分中的秩-1 矩阵。注意,每个右侧的行向量与其所在列的对角块的底行向量相同:特别是 $\left[a_{1:0}\ a_{1:1}\right]$、$\left[a_{3:2}\ a_{3:3}\right]$ 和 $[a_{5:4}~a_{5:5}]$ 这几行。

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

Therefore we already have the answers to these from Stage 1, which is the second element of all $T/2$ sub problems in Stage 1. If we call this array of elements $b^{\prime}$ (of half the size of $b$ ), then we need to multiply $b^{\prime}$ by the 1-SS matrix generated by $a_{3:-1},a_{3:1},a_{5:3},a_{7:5}.$ .

因此我们已经从第一阶段得到了这些问题的答案,即第一阶段中所有 $T/2$ 子问题的第二个元素。如果我们称这个元素数组为 $b^{\prime}$ (其大小是 $b$ 的一半),那么我们需要将 $b^{\prime}$ 与由 $a_{3:-1}, a_{3:1}, a_{5:3}, a_{7:5}$ 生成的1-SS矩阵相乘。

注意:1-SS矩阵是指单移位 (Single Shift) 矩阵。

Stage 3. Finally, each of the answers to Stage 2 can be broadcast into two final answers by multiplying by the left-side column vectors: in particular the $\left[a_{2:2}\ a_{3:2}\right]^{\intercal}$ , $\left[a_{4:4}\ a_{5:4}\right]^{\top}$ , and $\left[a_{6:6}\ a_{7:6}\right]^{\top}$ vectors.

阶段 3. 最后,可以通过将第二阶段的每个答案与左侧列向量相乘,广播成两个最终答案:具体来说是 $\left[a_{2:2}\ a_{3:2}\right]^{\intercal}$ , $\left[a_{4:4}\ a_{5:4}\right]^{\top}$ 和 $\left[a_{6:6}\ a_{7:6}\right]^{\top}$ 向量。

注意:这里的向量符号和公式保持原样未做翻译。

Note that this can be slightly modified with some off-by-one shifting of the indices. An equivalent way to view this algorithm is as the three-step matrix factorization

请注意,这可以通过对索引进行一些偏差为一的调整来略微修改。另一种等价的理解该算法的方式是将其视为三步矩阵分解 (three-step matrix factorization)。

Note that Stage 1 and Stage 3 require $O(T)$ work, while Stage 2 reduces to a self-similar problem of half the size. It is easy to check that this requires $O(T)$ total work and ${\cal O}(\log{T})$ depth/span.

请注意,第 1 阶段和第 3 阶段需要 $O(T)$ 的工作量,而第 2 阶段缩减为一个规模减半的自相似问题。可以很容易地验证,这总共需要 $O(T)$ 的工作量和 ${\cal O}(\log{T})$ 的深度/跨度。

Remark 10. In fact, it is possible to see that the computation graph of this algorithm is identical to that of the associative scan algorithm described in Appendix B.2.2. The key takeaway is that instead of the steps of (1) recognizing that 𝑀defines a recurrence (2) observing that the recurrence can be defined with an associative binary operator; there is a completely different perspective of simply finding a structured matrix decomposition algorithm for 𝑀.

备注 10. 实际上,可以发现该算法的计算图与附录 B.2.2 中描述的关联扫描算法 (associative scan algorithm) 的计算图是相同的。关键在于,不同于传统的步骤(1)识别出 ( M ) 定义了一个递归关系(2)观察到该递归关系可以用一个结合二元运算符来定义;这里提供了一种完全不同的视角,即直接寻找 ( M ) 的结构化矩阵分解算法。

C Theory Details

C.1 Extras: Closure Properties of SSMs

C.1 附加内容:SSM 的闭包性质

根据要求,仅翻译以上标题,保留原始格式。

We present here some additional properties of semi separable matrices to illustrate their flexibility and utility. This section is not necessary to understand our core results.

我们在这里介绍半可分矩阵的一些附加性质,以展示它们的灵活性和实用性。本节内容不是理解我们核心结果所必需的。


我们在这里介绍半可分矩阵的一些附加性质,以展示它们的灵活性和实用性。本节内容不是理解我们核心结果所必需的。

(注:由于原文内容较少,已按照要求完整翻译,保留了原始段落格式。)

Proposition C.1 (Semi separable Closure Properties). Semi separable matrices are closed under several primitive operations.

命题 C.1 (半可分闭包性质)。半可分矩阵在几种基本操作下是封闭的。

The addition and multiplication properties are easily seen. The inverse property has many proofs; one approach follows immediately from the Woodbury inversion identity, which has also featured prominently in the structured SSM literature (Gu, Goel, and Ré 2022).

加法和乘法性质很容易理解。逆性质有多种证明方法;其中一种方法直接来源于 Woodbury 逆公式 (Woodbury inversion identity),该公式在结构化 SSM 文献中也占有重要地位 (Gu, Goel, 和 Ré 2022)。

注:SSM 在这里指的是 State Space Model (状态空间模型)。

In turn, these imply closure properties of state space models.

这些性质反过来意味着状态空间模型的封闭性属性。

For example, the addition property says that summing two parallel SSM models is still an SSM. The multiplication property says that sequentially composing or chaining two SSMs can still be viewed as an SSM, whose total state size is additive–a somewhat nontrivial fact.

例如,加法性质表明,将两个并行的 SSM 模型相加后仍然是一个 SSM。乘法性质表明,顺序组合或串联两个 SSM 仍然可以被视为一个 SSM,其总状态大小是累加的——这是一个相对不 trivial 的事实。

Finally, the inverse property can let us relate SSMs to other types of models. For example, one can notice that banded matrices are semi separable, so their inverses are semi separable. (In fact, the semi separable family of structure is often motivated by taking inverses of banded matrices (Vandebril et al. 2005)). Moreover, the fast recurrence properties of semi separable matrices can be viewed as a consequence of their inverse being banded.

最后,逆性质可以让我们将状态空间模型 (SSM) 与其他类型的模型联系起来。例如,可以注意到带状矩阵是半可分的,因此它们的逆也是半可分的。(事实上,半可分结构的家族通常是通过对带状矩阵求逆来定义的 (Vandebril et al. 2005))。此外,半可分矩阵的快速递归性质可以看作是其逆为带状矩阵的结果。

在上述内容中,“状态空间模型”对应英文中的“SSM”,第一次出现时已在括号中注明英文原文。

Remark 11. The fact that 1-SS matrices are simple recurrences (7) are equivalent to the fact that the inverse of a 1-SS matrix

备注 11. 1-SS 矩阵是简单递归 (7) 的事实等价于 1-SS 矩阵的逆矩阵的事实。

根据这一性质,1-SS 矩阵的结构可以通过简单的递推关系来描述,而这种递推关系与其逆矩阵的存在性是等价的。这意味着我们可以利用递推公式 (7) 来理解和计算 1-SS 矩阵及其逆矩阵的性质。

is a 2-banded matrix:

是一个双带矩阵:

根据上述规则,这段内容翻译完成。如果有更多内容需要翻译,请继续提供。

$$
\boldsymbol{M}=\left[\begin{array}{c c c c c c}{1}&&&&&\ {a_{1}}&&{1}&&&\ {a_{2}a_{1}}&&{a_{2}}&{1}&&\ &&{\vdots}&{\ddots}&{\ddots}&\ {\vdots}&&{\vdots}&{\ddots}&{a_{T-1}}&{1}\end{array}\right]=\left[\begin{array}{c c c c c c}{1}&&&&\ {-a_{1}}&{1}&&&\ {0}&{-a_{2}}&{1}&&\ {\vdots}&{\vdots}&{\ddots}&{\ddots}&\ {0}&{0}&{\cdots}&{-a_{T-1}}&{1}\end{array}\right]^{T}
$$

$$
\boldsymbol{M}=\left[\begin{array}{c c c c c c}
1 &&&&& \
a_1 && 1 &&& \
a_2 a_1 && a_2 & 1 && \
&& \vdots & \ddots & \ddots & \
\vdots && \vdots & \ddots & a_{T-1} & 1
\end{array}\right]=\left[\begin{array}{c c c c c c}
1 &&&& \
-a_1 & 1 &&& \
0 & -a_2 & 1 && \
\vdots & \vdots & \ddots & \ddots & \
0 & 0 & \cdots & -a_{T-1} & 1
\end{array}\right]^{T}
$$

矩阵 (\boldsymbol{M}) 的定义如上所示。该矩阵可以表示为两个矩阵的转置形式,其中第一个矩阵是下三角矩阵,第二个矩阵是单位矩阵与负对角线元素的组合。

Thus $y=M x\leftrightarrow M^{-1}y=x_{\mathrm{,}}$ , or

因此 $y = M x \leftrightarrow M^{-1} y = x_{\mathrm{,}}$ ,或

注意:这里保留了原始的数学公式,未进行翻译。

$$
\left[\begin{array}{c c c c c}{{1}}&{{}}&{{}}&{{}}&{{}}\ {{-a_{1}}}&{{1}}&{{}}&{{}}&{{}}\ {{0}}&{{-a_{2}}}&{{1}}&{{}}&{{}}\ {{\vdots}}&{{\vdots}}&{{\ddots}}&{{\ddots}}&{{}}\ {{0}}&{{0}}&{{\ldots}}&{{-a_{T-1}}}&{{1}}\end{array}\right]y=x.
$$

$$
\left[\begin{array}{c c c c c}
1 & & & & \
-a_1 & 1 & & & \
0 & -a_2 & 1 & & \
\vdots & \vdots & \ddots & \ddots & \
0 & 0 & \ldots & -a_{T-1} & 1
\end{array}\right] y = x.
$$

这个公式表示一个下三角矩阵与向量 ( y ) 的乘积等于向量 ( x )。矩阵的主对角线元素为 1,副对角线下方的元素为 (-a_i),其余元素为 0。

Or element wise,

或者按元素-wise,

请注意,“element wise”在中文中通常翻译为“按元素”或“逐元素”,根据上下文可以适当选择。这里我选择了“按元素”作为翻译。如果需要更具体的表达,请提供更多上下文。

$$
\begin{array}{r l}&{y_{t}-a_{t}y_{t-1}=x_{t}}\ &{y_{t}=a_{t}y_{t-1}+x_{t}.}\end{array}
$$

$$
\begin{array}{r l}
& y_{t} - a_{t} y_{t-1} = x_{t} \
& y_{t} = a_{t} y_{t-1} + x_{t}.
\end{array}
$$

该公式表示一个递推关系,其中 ( y_t ) 是当前时刻的输出,( a_t ) 是当前时刻的系数,( y_{t-1} ) 是前一时刻的输出,( x_t ) 是当前时刻的输入。第一个式子表示当前输出与前一时刻输出之间的差值等于当前输入,第二个式子则表示当前输出等于前一时刻输出乘以系数加上当前输入。

Conversely, we also use these closure results to prove that auto regressive structured attention (under certain assumptions) must be SSMs, allowing us to show that more general families of efficient sequence models including attention variants can be reduced to state space models (Appendix C.2).

相反,我们也利用这些闭包结果证明了自回归结构化注意力机制(在某些假设下)必须是状态空间模型 (SSMs),这使我们能够证明更广泛的高效序列模型家族(包括注意力机制的变体)可以被简化为状态空间模型(附录 C.2)。

C.2 Auto regressive Masked Attention is Semi separable-Structured Attention

C.2 自回归掩码注意力是半可分离结构化注意力 (Auto regressive Masked Attention is Semi separable-Structured Attention)

自回归掩码注意力机制是一种特殊的注意力机制,它结合了自回归模型的特点和掩码操作,使得模型在处理序列数据时能够有效地捕捉依赖关系。这种注意力机制可以被视为一种半可分离的结构化注意力,因为它在计算过程中既保留了部分依赖关系,又通过掩码操作限制了某些位置之间的直接交互。

在这种机制下,模型在生成每个输出 Token 时,只会依赖于之前的位置,而不会看到未来的信息。这种特性使得自回归掩码注意力在许多自然语言处理任务中表现出色,尤其是在需要顺序生成的任务中,如文本生成、机器翻译等。

通过引入掩码操作,自回归掩码注意力能够在保持一定结构化的同时,实现高效的并行计算。与完全可分离的注意力机制不同,自回归掩码注意力并不是完全独立地处理每个位置,而是通过掩码矩阵控制了信息流动的方向和范围,从而形成了半可分离的结构。

这种半可分离的结构化注意力机制在 Transformer 模型中得到了广泛应用,并且在大语言模型 (LLM) 中也发挥了重要作用。它不仅提高了模型的计算效率,还增强了模型对长距离依赖关系的建模能力。

We prove Theorem 5.2 from Section 5.2. In Section 4.3 we defined structured attention as a broad generalization of masked attention, where the property of efficiency (i.e. a linear-time form for the kernel attention) is abstracted into the efficiency of structured matrix multiplication. However, beyond computational efficiency, standard linear attention (Katha ro poul os et al. 2020) also has two important properties. First, it is causal, which is required for settings such as auto regressive modeling. Moreover, it has efficient auto regressive generation. In other words, the cost of an auto regressive step – i.e. the incremental cost of computing the output $y_{T}$ upon seeing $x_{T}$ , given that $x_{0:T}$ has already been seen and pre processed – requires only constant time.

我们证明了第 5.2 节中的定理 5.2。在第 4.3 节中,我们将结构化注意力定义为掩码注意力的广泛推广,其中效率属性(即核注意力的线性时间形式)被抽象为结构化矩阵乘法的效率。然而,除了计算效率之外,标准线性注意力 (Katha ro poul os et al. 2020) 还具有两个重要属性。首先,它是因果的,这是自回归建模等场景所必需的。此外,它还具有高效的自回归生成。换句话说,自回归步骤的成本——即在已经看到并预处理了 $x_{0:T}$ 的情况下,计算输出 $y_{T}$ 所需的增量成本——仅需要常数时间。

在自回归建模中,因果性确保了模型在生成序列时,每个时刻的输出只依赖于之前的输入,而不会受到未来输入的影响。这种特性对于许多自然语言处理任务非常重要,例如文本生成和语音合成。

Here we characterize which instances of SMA have efficient auto regression.

在这里,我们分析了哪些 SMA (Sequence Modeling with Attention) 实例具有高效的自回归。

In the framework of SMA, causality is equivalent to the constraint that the mask $L$ is a lower-triangular matrix.

在SMA框架中,因果性等价于约束掩码 $L$ 是一个下三角矩阵。

Characterizing the space of $L$ matrices that have efficient auto regression is more difficult. We will use a narrow technical definition of auto regressive processes, in the spirit of classical definitions from the time series literature (e.g. ARIMA processes (Box et al. 2015)).

表征具有高效自回归的 $L$ 矩阵空间更为困难。我们将使用一个狭义的技术定义来描述自回归过程,这一定义借鉴了时间序列文献中的经典定义(例如 ARIMA 过程 (Box et al. 2015))。

Definition C.2. We define an auto regressive transformation $x,\in,\mathbb{R}^{T},\mapsto,y,\in,\mathbb{R}^{T}$ of order $k$ as one where each output $y_{t}$ depends only on the current input and last $k$ outputs:

定义 C.2. 我们定义一个自回归变换 $x,\in,\mathbb{R}^{T},\mapsto,y,\in,\mathbb{R}^{T}$,其阶数为 $k$,即每个输出 $y_{t}$ 仅依赖于当前输入和最后的 $k$ 个输出:

  • 自回归变换 (auto regressive transformation) 是一种常见的序列建模方法,其中每个时间步的输出只与当前输入以及之前的若干个输出有关。在这里,$k$ 表示依赖的前 $k$ 个输出的数量。

$$
y_{t}=\mu_{t}x_{t}+\ell_{t1}y_{t-1}+\cdot\cdot\cdot+\ell_{t k}y_{t-k}.
$$

$$
y_{t} = \mu_{t} x_{t} + \ell_{t1} y_{t-1} + \cdot\cdot\cdot + \ell_{t k} y_{t-k}.
$$

该公式表示时间序列模型中的一个预测方程,其中:

  • ( y_{t} ) 是在时间 ( t ) 的预测值,
  • ( \mu_{t} ) 是时间 ( t ) 的系数,
  • ( x_{t} ) 是时间 ( t ) 的输入变量,
  • ( \ell_{t1}, \ell_{t2}, \cdot\cdot\cdot, \ell_{tk} ) 是滞后项的系数,
  • ( y_{t-1}, y_{t-2}, \cdot\cdot\cdot, y_{t-k} ) 是过去 ( k ) 个时间点的预测值。

Note that the case where $L$ is the cumsum matrix is a special case with $k,=,1$ and thus $y_{t},=,x_{t},+,y_{t_{1}}$ . With this definition, characterizing the space of efficient auto regressive linear transforms follows from the properties of semi separable matrices. Theorem C.3 formalizes and proves Theorem 5.2.

请注意,当 $L$ 是累积和矩阵时,这是一个特殊情况,其中 $k = 1$,因此 $y_{t} = x_{t} + y_{t-1}$。根据这个定义,描述高效的自回归线性变换的空间可以基于半可分矩阵的性质。定理 C.3 正式表述并证明了定理 5.2。

注意:在翻译过程中保留了原始的数学符号和公式格式。

Theorem C.3. Let $L\in\mathbb{R}^{T\times T}$ be an efficient auto regressive transformation of order $k$ . Then $L$ is a state space model of order $k+1$ .

定理 C.3. 设 $L\in\mathbb{R}^{T\times T}$ 是一个阶数为 $k$ 的高效自回归变换。那么 $L$ 是一个阶数为 $k+1$ 的状态空间模型。

注:在中文中,"auto regressive transformation" 通常翻译为“自回归变换”,而 "state space model" 翻译为“状态空间模型”。

Proof. Let $(x,y)$ be input and output sequences, so that $y=L x$ . Rearranging the definition (23),

证明。设 $(x, y)$ 为输入和输出序列,使得 $y = L x$ 。根据定义 (23) 重新排列,

注意:由于原文中没有更多的上下文信息,这里仅对给定的数学表达式和句子进行了直译。如果您有更多上下文或需要进一步解释,请提供更多信息。

$$
y_{t}-\ell_{t1}y_{t-1}-\cdot\cdot\cdot-\ell_{t k}y_{t-k}=\mu_{t}x_{t}.
$$

$$
y_{t} - \ell_{t1} y_{t-1} - \cdot \cdot \cdot - \ell_{t k} y_{t-k} = \mu_{t} x_{t}.
$$

该公式表示时间序列模型中的一个关系,其中 ( y_t ) 是当前时刻的输出,( \ell_{t1}, \cdot \cdot \cdot, \ell_{tk} ) 是过去的 ( k ) 个时刻的系数,( x_t ) 是当前时刻的输入,( \mu_t ) 是当前时刻的权重。

Vector i zing over $t$ , this can be expressed as a matrix transformation

将向量 $i$ 在 $t$ 上进行向量化,这可以表示为一个矩阵变换。

$$
\left[\begin{array}{c c c c c c c c}{1}&&&&&&&\ {-\ell_{t1}}&{1}&&&&&&\ {\vdots}&{\ddots}&{\ddots}&&&&&\ {-\ell_{t k}}&{\ldots}&{-\ell_{t1}}&{1}&&&&\ {\vdots}&{\ddots}&{\vdots}&{\ddots}&{\ddots}&&\ &{\ldots}&{-\ell_{T-1,k}}&{\ldots}&{-\ell_{T-1,1}}&{1}\end{array}\right]\left[\begin{array}{c}{y_{0}}\ {y_{1}}\ {\vdots}\ {y_{k}}\ {\vdots}\ {y_{k}}\ {\vdots}\ {y_{T-1}}\end{array}\right]=\left[\begin{array}{c c c c c c c c}{\mu_{0}}&&&&&&\ &{\mu_{1}}&&&&&\ &&{\ddots}&&&&\ &&&{\mu_{k}}&&&\ &&&&{\ddots}&&\ &&&&&&{\mu_{T-1}}\end{array}\right]\left[\begin{array}{c}{x_{0}}\ {x_{1}}\ {\vdots}\ {x_{k}}\ {x_{k}}\ {\vdots}\ {x_{T-1}}\end{array}\right].
$$

$$
\left[\begin{array}{c c c c c c c c}{1}&&&&&&&\ {-\ell_{t1}}&{1}&&&&&&\ {\vdots}&{\ddots}&{\ddots}&&&&&\ {-\ell_{t k}}&{\ldots}&{-\ell_{t1}}&{1}&&&&\ {\vdots}&{\ddots}&{\vdots}&{\ddots}&{\ddots}&&\ &{\ldots}&{-\ell_{T-1,k}}&{\ldots}&{-\ell_{T-1,1}}&{1}\end{array}\right]\left[\begin{array}{c}{y_{0}}\ {y_{1}}\ {\vdots}\ {y_{k}}\ {\vdots}\ {y_{k}}\ {\vdots}\ {y_{T-1}}\end{array}\right]=\left[\begin{array}{c c c c c c c c}{\mu_{0}}&&&&&&\ &{\mu_{1}}&&&&&\ &&{\ddots}&&&&\ &&&{\mu_{k}}&&&\ &&&&{\ddots}&&\ &&&&&&{\mu_{T-1}}\end{array}\right]\left[\begin{array}{c}{x_{0}}\ {x_{1}}\ {\vdots}\ {x_{k}}\ {x_{k}}\ {\vdots}\ {x_{T-1}}\end{array}\right].
$$

该公式表示一个矩阵方程,其中左边的矩阵是一个下三角矩阵,包含参数 (\ell_{ti}),右边的矩阵是一个对角矩阵,包含参数 (\mu_i)。向量 (y) 和 (x) 分别表示输出和输入序列。具体来说,该方程描述了如何通过输入序列 (x) 计算输出序列 (y),其中每个输出 (y_t) 依赖于当前和之前的输入 (x_t) 以及之前的输出 (y_{t-1}, \dots, y_1)。

The $\mu$ diagonal matrix can be moved to the left and folded into the matrix of $\ell$ coefficients, which remains a $k+1$ -band lower-triangular matrix. But we also have $L^{-1}y=x$ , so $L$ is the inverse of this matrix.

$\mu$ 对角矩阵可以移动到左边并与 $\ell$ 系数矩阵合并,仍然保持为一个 $k+1$ -带下三角矩阵。但我们也知道 $L^{-1}y=x$,因此 $L$ 是这个矩阵的逆矩阵。

Next, note that $k{+1}$ -band matrices are $k!+!1$ -semi separable by the rank characterization of semi se par ability (Definition 3.1). By Proposition C.1, the inverse $L$ is therefore at most $k{+2}$ -semi separable. A slightly stronger bound of $k{+1}$ can be obtained because of the additional structure of banded matrices. Finally, the characterization of $L$ as an order $.k!+!1$ state space model follows from Theorem 3.5. □

接下来,注意到 $k{+1}$ -band 矩阵根据半可分性的秩特征 (定义 3.1) 是 $k!+!1$ -半可分的。根据命题 C.1,逆矩阵 $L$ 因此最多是 $k{+2}$ -半可分的。由于带状矩阵的额外结构,可以得到稍强的界为 $k{+1}$ 。最后,根据定理 3.5,$L$ 可以被描述为一个阶数为 $k!+!1$ 的状态空间模型。□

接下来,注意到 \$k{+1}\$ -band 矩阵根据半可分性的秩特征 (定义 3.1) 是 \$k\!+\!1\$ -半可分的。根据命题 C.1,逆矩阵 \$L\$ 因此最多是 \$k{+2}\$ -半可分的。由于带状矩阵的额外结构,可以得到稍强的界为 \$k{+1}\$ 。最后,根据定理 3.5,\$L\$ 可以被描述为一个阶数为 \$k\!+\!1\$ 的状态空间模型。□

In other words, efficient auto regressive attention is semi separable SMA.

换句话说,高效的自回归注意力机制是半可分离的 SMA (Semi Separable Attention)。

D Experimental Details

D 实验细节

实验细节部分通常包含实验设置、参数配置、数据集信息等内容。请提供具体的实验细节内容,以便我继续翻译。

D.1 MQAR Details

D.1 MQAR 细节

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

We use a harder version of the task introduced in Based (Arora, Eyuboglu, Zhang, et al. 2024) where tokens that are not query/key/values are replaced with random tokens. We also use more key-value pairs, longer sequences, and smaller model sizes than the usual variant of MQAR used by prior work, all of which make the task harder.

我们使用了基于 Arora, Eyuboglu, Zhang 等 (2024) 引入的任务的更难版本,在该版本中,不是查询/键/值的 Token 被替换为随机 Token。我们还使用了更多的键值对、更长的序列以及更小的模型规模,这些都使得任务更加困难。与之前工作中常用的 MQAR 变体相比,这些改动进一步增加了任务的难度。

For each sequence length $T\in{256,512,1024}$ , we use $T/4$ key-value pairs. The total vocab size is 8192.

对于每个序列长度 $T \in {256, 512, 1024}$,我们使用 $T/4$ 个 key-value 对。总词汇量大小为 8192。

We use a form of curriculum training where training cycles through datasets using $(T/32,T/16,T/8,T/4)$ key-value pairs, where each dataset has $2^{18}\approx250000$ examples, for a total of 8 epochs through each dataset (total of $2^{28}\approx270M$ examples). The total batch size is $2^{18}\approx0.25M$ tokens (e.g. for $T=1024$ , the batch size is 256).

我们采用了一种课程训练 (curriculum training) 的形式,其中训练过程会循环使用包含 $(T/32, T/16, T/8, T/4)$ 个键值对的数据集。每个数据集包含 $2^{18} \approx 250,000$ 个样本,总共对每个数据集进行 8 个 epoch 的训练(总计 $2^{28} \approx 2.7亿$ 个样本)。总批量大小为 $2^{18} \approx 0.25M$ 个 Token(例如,当 $T=1024$ 时,批量大小为 256)。

All methods use 2 layers with default settings; the attention baseline additionally receives positional embeddings. For each method, we sweep over model dimensions $\mathsf{D}={32,64,128,256}$ and learning rates ${{10}^{-{\bar{3.5}}},{10}^{-2},{10}^{-2.5}}$ . We use a linear decay schedule that drops on every epoch (e.g. the last epoch would have a learning rate $1/8$ of the maximum/starting learning rate).

所有方法均使用 2 层,默认设置;注意力基线方法还接收位置嵌入。对于每种方法,我们在模型维度 $\mathsf{D}={32,64,128,256}$ 和学习率 ${{10}^{-3.5},{10}^{-2},{10}^{-2.5}}$ 上进行调参。我们使用线性衰减调度,在每个 epoch 都会降低学习率(例如,最后一个 epoch 的学习率将是最大/起始学习率的 $1/8$)。

D.2 Scaling Law Details

D.2 扩展定律细节 (Scaling Law Details)

由于原文内容较少,仅包含标题,因此翻译结果如上所示。如果有更多内容,请继续提供,我将按照要求进行翻译。

All models were trained on the Pile. For the scaling law experiments, we use the GPT2 tokenizer.

所有模型都在 Pile 上进行了训练。对于缩放定律实验,我们使用了 GPT2 分词器 (tokenizer)。

Model Sizes. Table 9 specifies the model sizes we use for scaling laws following GPT3 (Brown et al. 2020), First, we changed the batch size of the 1.3B model from 1M tokens to $0.5\mathrm{M}$ tokens for uniformity. Second, we changed the number of training steps and total tokens to roughly match Chinchilla scaling laws (Hoffmann et al. 2022), which specify that training tokens should increase proportionally to model size.

模型规模。表 9 规定了我们用于扩展定律的模型规模,这些扩展定律遵循 GPT3 (Brown et al. 2020)。首先,我们将 1.3B 模型的批量大小从 1M Token 减少到 0.5M Token,以保持一致性。其次,我们调整了训练步数和总 Token 数,使其大致符合 Chinchilla 扩展定律 (Hoffmann et al. 2022),该定律指出训练 Token 数应与模型规模成正比增加。

模型规模。表 9 规定了我们用于扩展定律的模型规模,这些扩展定律遵循 GPT3 (Brown et al. 2020)。首先,我们将 1.3B 模型的批量大小从 1M Token 减少到 0.5M Token,以保持一致性。其次,我们调整了训练步数和总 Token 数,使其大致符合 Chinchilla 扩展定律 (Hoffmann et al. 2022),该定律指出训练 Token 数应与模型规模成正比增加。

Training Recipes. All models used the AdamW optimizer with

训练配方 (Training Recipes)。所有模型使用了 AdamW 优化器

Table 9: (Scaling Law Model Sizes.) Our model sizes and hyper parameters for scaling experiments. (Model dimension and number of heads applies only to Transformer models.)

表 9: (扩展定律模型规模) 我们的模型规模和超参数用于扩展实验。(模型维度和头部数量仅适用于 Transformer 模型。)

PARAMS n_layers d_model n_heads / d_head TRAINING STEPS LEARNING RATE BATCH SIZE TOKENS
125M 12 768 12/64 4800 6e-4 0.5M tokens 2.5B
350M 24 1024 16/64 13500 3e-4 0.5M tokens 7B
760M 24 1536 16/96 29000 2.5e-4 0.5M tokens 15B
1.3B 24 2048 32/64 50000 2e-4 0.5M tokens 26B

注:PARAMS 表示模型参数量,n_layers 表示层数,d_model 表示模型维度,n_heads / d_head 表示头部数量/每个头部的维度,TRAINING STEPS 表示训练步数,LEARNING RATE 表示学习率,BATCH SIZE 表示批量大小,TOKENS 表示训练使用的 Token 数量。

By default, the peak learning rate is the GPT3 specification.

默认情况下,峰值学习率采用 GPT3 的规格。

Compared to GPT3 recipe, we use an “improved recipe”, inspired by changes adopted by popular large language models such as PaLM (Chowdhery et al. 2023) and LLaMa (Touvron, Lavril, et al. 2023). These include:

与 GPT3 的方案相比,我们采用了一种“改进方案”,该方案借鉴了流行的大语言模型(如 PaLM (Chowdhery et al. 2023) 和 LLaMa (Touvron, Lavril, et al. 2023))所采用的改进措施。这些改进包括:

  • 更好的数据预处理方法
  • 更高效的训练策略
  • 更优化的模型架构

这些改进使得我们的模型在性能和效率上都有显著提升。

D.3 Downstream Evaluation Details

D.3 下游评估细节

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

To evaluate downstream performance of fully trained, we train Mamba-2 on 300B tokens on the Pile, using the GPTNeoX (Black et al. 2022) tokenizer.

为了评估完全训练后的下游性能,我们在 Pile 数据集上使用 GPTNeoX (Black et al. 2022) 的分词器对 Mamba-2 进行了 3000 亿个 Token 的训练。

We use the same hyper parameters as the scaling experiments, except with batch size 1M for the 1.3B and 2.7B model. For the 2.7B model, we also follow GPT3 specification (32 layers, dimension 2560).

我们使用与扩展实验相同的超参数,除了对 1.3B 和 2.7B 模型使用 1M 的批量大小 (batch size)。对于 2.7B 模型,我们还遵循 GPT3 的配置 (32 层,维度 2560)。

For all models, we use 5x the learning rate of the corresponding GPT3 model.

对于所有模型,我们使用对应 GPT3 模型的学习率的 5 倍。

For downstream evaluation, we use the LM evaluation harness from EleutherAI (L. Gao, Tow, et al. 2021), on the same tasks as Mamba (Gu and Dao 2023) with one additional one:

对于下游评估,我们使用了来自 EleutherAI 的 LM 评估框架 (L. Gao, Tow, et al. 2021),在与 Mamba (Gu 和 Dao 2023) 相同的任务上进行评估,并增加了一个额外的任务:

  • 下游评估 (downstream evaluation) 指的是在特定任务上对模型性能的评估。
  • LM 评估框架 (LM evaluation harness) 是一个用于评估语言模型性能的工具。

• LAMBADA (Paperno et al. 2016) • HellaSwag (Zellers et al. 2019) • PIQA (Bisk et al. 2020) • ARC-challenge (P. Clark et al. 2018) • ARC-easy: an easy subset of ARC-challenge • WinoGrande (Sakaguchi et al. 2021) • OpenBookQA (Mihaylov et al. 2018)

• LAMBADA (Paperno 等 2016)
• HellaSwag (Zellers 等 2019)
• PIQA (Bisk 等 2020)
• ARC-challenge (P. Clark 等 2018)
• ARC-easy:ARC-challenge 的简单子集
• WinoGrande (Sakaguchi 等 2021)
• OpenBookQA (Mihaylov 等 2018)

这些是常见的自然语言处理基准数据集,用于评估大语言模型在不同任务上的表现。

D.4 Ablation Details

D.4 消融实验细节 (Ablation Details)

由于原文中只有标题而没有具体内容,根据规则,我不会询问提供更多内容。如果您有更多内容需要翻译,请直接提供。

(Re)Based Details. Our ablations in Section 9.4.3 considered the Based (Arora, Eyuboglu, Zhang, et al. 2024) and ReBased (Aksenov et al. 2024) models.

(Re)Based 模型细节。我们在第 9.4.3 节的消融实验中考虑了 Based (Arora, Eyuboglu, Zhang 等,2024) 和 ReBased (Aksenov 等,2024) 模型。

Based approximates the exp kernel with a quadratic Taylor expansion $\exp(x)\approx1+x+x^{2}/2$ , which can be accomplished by the feature map

基于二次泰勒展开 $\exp(x) \approx 1 + x + x^2 / 2$ 来近似指数核 (exp kernel),这可以通过特征映射来实现。

$$
\psi_{\mathrm{Taylor}}(x)=\mathsf{c o n c a t e n a t e}(1,x,1/\sqrt{2}x\otimes x).
$$

$$
\psi_{\mathrm{Taylor}}(x) = \mathsf{concatenate}(1, x, 1/\sqrt{2} x \otimes x).
$$

泰勒展开函数 (\psi_{\mathrm{Taylor}}(x)) 由连接操作 (\mathsf{concatenate}) 构成,具体为将常数 1、向量 (x) 以及 (x) 的二阶张量积 (x \otimes x)(乘以缩放因子 (1/\sqrt{2}))进行拼接。

ReBased proposes to use the simpler feature map $\psi_{\mathrm{Quadratic}}(x)=x\otimes x$ corresponding to the kernel transformation $x^{2}$ , but also applies a layer normalization beforehand. We view the layer normalization as an alternative non-linear activation to our default Swish activation, and ablate combinations of these.

ReBased 提议使用更简单的特征映射 $\psi_{\mathrm{Quadratic}}(x)=x\otimes x$,该映射对应于核变换 $x^{2}$,但在应用之前还进行了层归一化 (layer normalization)。我们将层归一化视为默认的 Swish 激活函数的一种替代非线性激活,并对这些组合进行了消融研究 (ablation)。

Table 10: (Zero-shot Evaluations.) Best results for each size in bold, second best unlined. We compare against open source LMs with various tokenizers, trained for up to 300B tokens. Pile refers to the validation split, comparing only against models trained on the same dataset and tokenizer (GPT-NeoX-20B). For each model size, Mamba-2 outperforms Mamba, and generally matches Pythia at twice the model size.

表 10: (零样本评估)。每个规模的最佳结果用粗体表示,次佳结果不加下划线。我们将 Mamba-2 与开源大语言模型 (LLM) 进行比较,这些模型使用不同的分词器训练,最多训练 300B 个 Token。Pile 指的是验证集的分割,仅与在同一数据集和分词器上训练的模型(如 GPT-NeoX-20B)进行比较。对于每个模型规模,Mamba-2 的表现优于 Mamba,并且通常在模型规模为两倍时与 Pythia 相当。

MODEL TOKEN. PILE PPL↓ LAMBADA PPL↓ LAMBADA ACC↑ HELLASWAG ACC↑ PIQA ACC ↑ ARC-E ACC ↑ ARC-C ACC↑ WINOGRANDE ACC↑ OPENBOOKQA ACC↑ AVERAGE ACC↑
Hybrid H3-130M GPT2 - 89.48 25.8 31.7 64.2 44.4 24.2 50.6 27.0 38.2
Pythia-160M NeoX 29.64 38.10 33.0 30.2 61.4 43.2 24.1 51.9 29.2 39.0
Mamba-130M NeoX 10.56 16.07 44.3 35.2 64.5 48.0 24.2 51.9 28.8 42.4
Mamba-2-130M NeoX 10.48 16.86 43.9 35.3 64.9 47.4 24.2 52.1 30.6 42.6
Hybrid H3-360M GPT2 - 12.58 48.0 41.5 68.1 51.4 24.7 54.1 31.6 45.6
Pythia-410M NeoX 9.95 10.84 51.4 40.6 66.9 52.1 24.6 53.8 30.0 45.6
Mamba-370M NeoX 8.28 8.14 55.6 46.5 69.5 55.1 28.0 55.3 30.8 48.7
Mamba-2-370M NeoX 8.21 8.02 55.8 46.9 70.5 54.9 26.9 55.7 32.4 49.0
Pythia-1B NeoX 7.82 7.92 56.1 47.2 70.7 57.0 27.1 53.5 31.4 49.0
Mamba-790M NeoX 7.33 6.02 62.7 55.1 72.1 61.2 29.5 56.1 34.2 53.0
Mamba-2-780M NeoX 7.26 5.86 61.7 54.9 72.0 61.0 28.5 60.2 36.2 53.5
GPT-Neo 1.3B GPT2 - 7.50 57.2 48.9 71.1 56.2 25.9 54.9 33.6 49.7
Hybrid H3-1.3B GPT2 - 11.25 49.6 52.6 71.3 59.2 28.1 56.9 34.4 50.3
OPT-1.3B OPT - 6.64 58.0 53.7 72.4 56.7 29.6 59.5 33.2 51.9
Pythia-1.4B NeoX 7.51 6.08 61.7 52.1 71.0 60.5 28.5 57.2 30.8 51.7
RWKV4-1.5B NeoX 7.70 7.04 56.4 52.5 72.4 60.5 29.4 54.6 34.0 51.4
Mamba-1.4B NeoX 6.80 5.04 65.0 59.1 74.2 65.5 32.8 61.5 36.4 56.4
Mamba-2-1.3B NeoX 6.66 5.02 65.7 59.9 73.2 64.3 33.3 60.9 37.8 56.4
GPT-Neo 2.7B GPT2 - 5.63 62.2 55.8 72.1 61.1 30.2 57.6 33.2 53.2
Hybrid H3-2.7B GPT2 - 7.92 55.7 59.7 73.3 65.6 32.3 61.4 33.6 54.5
OPT-2.7B OPT - 5.12 63.6 60.6 74.8 60.8 31.3 61.0 35.2 55.3
Pythia-2.8B NeoX 6.73 5.04 64.7 59.3 74.0 64.1 32.9 59.7 35.2 55.7
RWKV4-3B NeoX 7.00 5.24 63.9 59.6 73.7 67.8 33.1 59.6 37.0 56.4
Mamba-2.8B NeoX 6.22 4.23 69.2 66.1 75.2 69.7 36.3 63.5 39.6 59.9
Mamba-2-2.7B NeoX 6.09 4.10 69.7 66.6 76.4 69.6 36.4 64.0 38.8 60.2
GPT-J-6B GPT2 - 4.10 68.3 66.3 75.4 67.0 36.6 64.1 38.2 59.4
OPT-6.7B OPT - 4.25 67.7 67.2 76.3 65.6 34.9 65.5 37.4 59.2
Pythia-6.9B NeoX 6.51 4.45 67.1 64.0 75.2 67.3 35.5 61.3 38.0 58.3
RWKV4-7.4B NeoX 6.31 4.38 67.2 65.5 76.1 67.8 37.5 61.0 40.2 59.3

注:表中的“一”表示该模型在该指标上没有可用的数据。

Note that because these expand the feature dimension, we must project to smaller $B,C$ dimensions; in Table 7, use state size $N=64$ for $130\mathrm{M}$ models and $N=256$ for $380\mathrm{M}$ models. For the (Re)Based methods, we project to 8 and 16 dimensions respectively before applying their feature maps; this results in a total state size of $8^{2}=64$ for ReBased and $1+8+8^{2}=73$ for Based in the 130M model case. Because the $B$ and $C$ projections are smaller, these methods use fewer parameters, and we adjust the layer count appropriately.

请注意,由于这些方法扩展了特征维度,我们必须将维度投影到较小的 $B, C$ 维度;在表 7 中,对于 $130\mathrm{M}$ 模型使用状态大小 $N=64$,对于 $380\mathrm{M}$ 模型使用 $N=256$。对于 (Re)Based 方法,在应用它们的特征图之前,我们分别将维度投影到 8 和 16 维;这导致在 130M 模型情况下,ReBased 的总状态大小为 $8^{2}=64$,而 Based 的总状态大小为 $1+8+8^{2}=73$。由于 $B$ 和 $C$ 的投影维度较小,这些方法使用的参数较少,因此我们相应地调整了层数。

模型 状态大小 $N$ 特征维度投影 总状态大小
130M 64 ReBased: 8, Based: 16 ReBased: 64, Based: 73
380M 256 - -

因为 $B$ 和 $C$ 的投影维度较小,(Re)Based 方法使用的参数更少,因此我们适当减少了层数。

阅读全文(20积分)