2.2 Linear Attention# 和诸多实用的序列建模方案一样，研究人员也十分关注线性模型的构建，以此提升长序列的处理速度。线性模型有多种定义形式，序列模型中常用的通用形式如下：
z i = f ( a ⋅ z i − 1 + b ⋅ s i ) (4) \mathbf{z}_i = f(a \cdot \mathbf{z}_{i-1} + b \cdot \mathbf{s}_i) \tag{4} z i = f ( a ⋅ z i − 1 + b ⋅ s i ) ( 4 ) 其中 s i \mathbf{s}_i s i i i i z i \mathbf{z}_i z i i i i i i i f ( ⋅ ) f(\cdot) f ( ⋅ ) f ( ⋅ ) f(\cdot) f ( ⋅ ) f ( ⋅ ) f(\cdot) f ( ⋅ )
只需对当前输入与上一时刻状态做线性变换，即可将 式 (4) z i = f ( z i − 1 W z + s i W s ) \mathbf{z}_i = f(\mathbf{z}_{i-1}\mathbf{W}_z + \mathbf{s}_i\mathbf{W}_s) z i = f ( z i − 1 W z + s i W s ) Transformer，构建混合模型。
实际上，研究人员更倾向于构建线性注意力模型，在保证全局序列建模能力的同时，提升整体运行效率。实现该目标的一大难点在于，标准自注意力本身并非线性形式。先回顾自注意力的表达式：
A t t s e l f = A ⋅ V = ψ ( Q ⋅ K ⊤ ) ⋅ V \begin{align*}
\mathrm{Att}_{\mathrm{self}} &= \mathbf{A} \cdot \mathbf{V} \\
&= \psi(\mathbf{Q} \cdot \mathbf{K}^\top) \cdot \mathbf{V} \tag{5}
\end{align*} Att self = A ⋅ V = ψ ( Q ⋅ K ⊤ ) ⋅ V ( 5 ) 其中 ψ ( ⋅ ) \psi(\cdot) ψ ( ⋅ ) ψ ( a ) = Normalize ( Mask ( exp ( a d ) ) ) \psi(a) = \text{Normalize}(\text{Mask}(\exp(\frac{a}{\sqrt{d}}))) ψ ( a ) = Normalize ( Mask ( exp ( d a ))) ψ ( ⋅ ) \psi(\cdot) ψ ( ⋅ ) ψ ( ⋅ ) \psi(\cdot) ψ ( ⋅ ) n n n
尽管自注意力中键与值是成对存在的，但二者在计算中是分步使用的。一种巧妙的解法是将查询与键值交互解耦，使上下文信息的编码独立于具体的查询。其核心思路是：通过对查询和键应用特征映射函数 ϕ ( ⋅ ) \phi(\cdot) ϕ ( ⋅ ) Softmax 运算。我们将 Q \mathbf{Q} Q K \mathbf{K} K Q ′ = ϕ ( Q ) ∈ R n × d ′ , K ′ = ϕ ( K ) ∈ R n × d ′ \mathbf{Q}' = \phi(\mathbf{Q}) \in \mathbb{R}^{n \times d'}, \quad \mathbf{K}' = \phi(\mathbf{K}) \in \mathbb{R}^{n \times d'} Q ′ = ϕ ( Q ) ∈ R n × d ′ , K ′ = ϕ ( K ) ∈ R n × d ′
A t t s e l f ≡ ψ ′ ( Q ′ ⋅ K ′ ⊤ ) ⋅ V = D − 1 ( Q ′ ⋅ K ′ ⊤ ) ⋅ V = D − 1 ( Q ′ ⋅ ( K ′ ⊤ ⋅ V ) ) \begin{align*}
\mathrm{Att}_{\mathrm{self}} &\equiv \psi'(\mathbf{Q}' \cdot \mathbf{K}'^\top) \cdot \mathbf{V} \\ &= \mathbf{D}^{-1} (\mathbf{Q}' \cdot \mathbf{K}'^\top) \cdot \mathbf{V} \\
&= \mathbf{D}^{-1} \left( \mathbf{Q}' \cdot (\mathbf{K}'^\top \cdot \mathbf{V}) \right) \tag{6}
\end{align*} Att self ≡ ψ ′ ( Q ′ ⋅ K ′⊤ ) ⋅ V = D − 1 ( Q ′ ⋅ K ′⊤ ) ⋅ V = D − 1 ( Q ′ ⋅ ( K ′⊤ ⋅ V ) ) ( 6 ) 其中 D \mathbf{D} D ψ ′ ( Z ) = D − 1 Z \psi'(\mathbf{Z}) = \mathbf{D}^{-1}\mathbf{Z} ψ ′ ( Z ) = D − 1 Z Softmax 归一化，仅需乘以 D − 1 \mathbf{D}^{-1} D − 1
由此可以得到一种高效的自回归形式：首先通过 K ′ ⊤ ⋅ V \mathbf{K}'^\top \cdot \mathbf{V} K ′⊤ ⋅ V K ′ ⊤ ⋅ V = ∑ j = 1 n k ′ j ⊤ ⋅ v j \mathbf{K}'^\top \cdot \mathbf{V} = \sum_{j=1}^{n} \mathbf{k}'{_j}^\top \cdot \mathbf{v}_j K ′⊤ ⋅ V = ∑ j = 1 n k ′ j ⊤ ⋅ v j K ′ ⊤ ⋅ V \mathbf{K}'^\top \cdot \mathbf{V} K ′⊤ ⋅ V 式 (4)
μ j = μ j − 1 + k ′ j ⊤ ⋅ v j (7) \mu_j = \mu_{j-1} + \mathbf{k}'{_j}^\top \cdot \mathbf{v}_j \tag{7} μ j = μ j − 1 + k ′ j ⊤ ⋅ v j ( 7 ) 其中 μ j ∈ R d ′ × d \mu_j \in \mathbb{R}^{d' \times d} μ j ∈ R d ′ × d k ′ j ⊤ ⋅ v j \mathbf{k}'{_j}^\top \cdot \mathbf{v}_j k ′ j ⊤ ⋅ v j ν j ∈ R d ′ \nu_j \in \mathbb{R}^{d'} ν j ∈ R d ′
ν j = ν j − 1 + k ′ j ⊤ (8) \nu_j = \nu_{j-1} + \mathbf{k}'{_j}^\top \tag{8} ν j = ν j − 1 + k ′ j ⊤ ( 8 ) 据此，第 j j j
A t t s e l f , j = q j ′ ⋅ μ n q j ′ ⋅ ν n (9) \mathrm{Att}_{\mathrm{self},j} = \frac{\mathbf{q}'_j \cdot \mu_n}{\mathbf{q}'_j \cdot \nu_n} \tag{9} Att self , j = q j ′ ⋅ ν n q j ′ ⋅ μ n ( 9 ) 显然，该模型实现了线性时间复杂度，因为 μ n \mu_n μ n ν n \nu_n ν n n n n μ j \mu_j μ j ν j \nu_j ν j μ j \mu_j μ j ν j \nu_j ν j A t t s e l f , j = q j ′ ⋅ μ j q j ′ ⋅ ν j \mathrm{Att}_{\mathrm{self},j} = \frac{\mathbf{q}'_j \cdot \mu_j}{\mathbf{q}'_j \cdot \nu_j} Att self , j = q j ′ ⋅ ν j q j ′ ⋅ μ j
对线性注意力模型的一个直接扩展，是让式 (7), (8) μ j \mu_j μ j ν j \nu_j ν j
μ j = a ⋅ μ j − 1 + ( 1 − a ) ⋅ k j ′ ⊤ ⋅ v j ν j = a ⋅ ν j − 1 + ( 1 − a ) ⋅ k j ′ ⊤ \begin{align*}
\mu_j &= a \cdot \mu_{j-1} + (1 - a) \cdot \mathbf{k}'^\top_j \cdot \mathbf{v}_j \tag{10} \\
\nu_j &= a \cdot \nu_{j-1} + (1 - a) \cdot \mathbf{k}'^\top_j \tag{11}
\end{align*} μ j ν j = a ⋅ μ j − 1 + ( 1 − a ) ⋅ k j ′⊤ ⋅ v j = a ⋅ ν j − 1 + ( 1 − a ) ⋅ k j ′⊤ ( 10 ) ( 11 ) 并照常对参数 a a a a a a
我们已经了解了为注意力机制设计线性模型的通用思路。这类模型的核心设计选择是移除基于 Softmax 的归一化，从而得到基于模型各类中间状态的线性表示形式。这一思路催生了多种近年提出的自注意力替代方案。尽管这些系统的架构各不相同，但其底层模型均具有式 (4) QKV 注意力机制；相反，我们可以为查询、键和值赋予新的含义与形式。
此处的讨论也与之前介绍的记忆模型相关。从记忆的角度来看，键和值可以被视为上下文的编码。因此，在上述线性注意力模型中，我们构建了一个记忆系统，其中仅用两个简单变量 μ j \mu_j μ j ν j \nu_j ν j j j j
2.3 State-Space Models# 在控制系统中，状态空间模型 (SSMs )
d z ( t ) d t = z ( t ) ⋅ A + s ( t ) ⋅ B o ( t ) = z ( t ) ⋅ C + s ( t ) ⋅ D \begin{align*} \frac{d\mathbf{z}(t)}{dt}
&= \mathbf{z}(t) \cdot \mathbf{A} + \mathbf{s}(t) \cdot \mathbf{B} \tag{12} \\
\mathbf{o}(t) &= \mathbf{z}(t) \cdot \mathbf{C} + \mathbf{s}(t) \cdot \mathbf{D} \tag{13}
\end{align*} d t d z ( t ) o ( t ) = z ( t ) ⋅ A + s ( t ) ⋅ B = z ( t ) ⋅ C + s ( t ) ⋅ D ( 12 ) ( 13 ) 其中 s ( t ) \mathbf{s}(t) s ( t ) o ( t ) \mathbf{o}(t) o ( t ) z ( t ) \mathbf{z}(t) z ( t ) t t t s ( t ) , o ( t ) ∈ R d \mathbf{s}(t), \mathbf{o}(t) \in \mathbb{R}^d s ( t ) , o ( t ) ∈ R d z ( t ) ∈ R d z \mathbf{z}(t) \in \mathbb{R}^{d_z} z ( t ) ∈ R d z 式 (12) A ∈ R d z × d z \mathbf{A} \in \mathbb{R}^{d_z \times d_z} A ∈ R d z × d z B ∈ R d × d z \mathbf{B} \in \mathbb{R}^{d \times d_z} B ∈ R d × d z 式 (13) C ∈ R d z × d \mathbf{C} \in \mathbb{R}^{d_z \times d} C ∈ R d z × d D ∈ R d × d \mathbf{D} \in \mathbb{R}^{d \times d} D ∈ R d × d
这些方程描述了从输入 s ( t ) \mathbf{s}(t) s ( t ) o ( t ) \mathbf{o}(t) o ( t ) { s 0 , s 1 , … , s n } \{\mathbf{s}_0, \mathbf{s}_1, \dots, \mathbf{s}_n\} { s 0 , s 1 , … , s n } s ( t ) \mathbf{s}(t) s ( t ) Δ t \Delta t Δ t { z 0 , z 1 , … , z n } \{\mathbf{z}_0, \mathbf{z}_1, \dots, \mathbf{z}_n\} { z 0 , z 1 , … , z n } { o 0 , o 1 , … , o n } \{\mathbf{o}_0, \mathbf{o}_1, \dots, \mathbf{o}_n\} { o 0 , o 1 , … , o n } SSM 的方程：
z t = z t − 1 ⋅ A ‾ + s t ⋅ B ‾ o t = z t ⋅ C ‾ + s t ⋅ D ‾ \begin{align*}
\mathbf{z}_t &= \mathbf{z}_{t-1} \cdot \overline{\mathbf{A}} + \mathbf{s}_t \cdot \overline{\mathbf{B}} \tag{14} \\
\mathbf{o}_t &= \mathbf{z}_t \cdot \overline{\mathbf{C}} + \mathbf{s}_t \cdot \overline{\mathbf{D}} \tag{15}
\end{align*} z t o t = z t − 1 ⋅ A + s t ⋅ B = z t ⋅ C + s t ⋅ D ( 14 ) ( 15 ) 这种形式的 SSM 定义了一个带有残差连接的 RNN。式 (14) t t t t − 1 t-1 t − 1 式 (15) z t \mathbf{z}_t z t s t \mathbf{s}_t s t
参数 A ‾ \overline{\mathbf{A}} A B ‾ \overline{\mathbf{B}} B C ‾ \overline{\mathbf{C}} C D ‾ \overline{\mathbf{D}} D A \mathbf{A} A B \mathbf{B} B C \mathbf{C} C D \mathbf{D} D 式 (12) 式 (14) 双线性变换 零阶保持 (ZOH )
式 (14) z t \mathbf{z}_t z t o t \mathbf{o}_t o t
很容易得到：
z t = ∑ i = 0 t s i ⋅ B ‾ ⋅ A ‾ t − i o t = ∑ i = 0 t s i ⋅ B ‾ ⋅ A ‾ t − i ⋅ C ‾ + s t ⋅ D ‾ \begin{align*}
\mathbf{z}_t &= \sum_{i=0}^{t} \mathbf{s}_i \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-i} \tag{16} \\
\mathbf{o}_t &= \sum_{i=0}^{t} \mathbf{s}_i \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-i} \cdot \overline{\mathbf{C}} + \mathbf{s}_t \cdot \overline{\mathbf{D}} \tag{17}
\end{align*} z t o t = i = 0 ∑ t s i ⋅ B ⋅ A t − i = i = 0 ∑ t s i ⋅ B ⋅ A t − i ⋅ C + s t ⋅ D ( 16 ) ( 17 ) 显然，式 (17)
∑ i = 0 t s i ⋅ B ‾ ⋅ A ‾ t − i ⋅ C ‾ = [ s 0 s 1 … s t ] ⋅ [ B ‾ ⋅ A ‾ t ⋅ C ‾ B ‾ ⋅ A ‾ t − 1 ⋅ C ‾ … B ‾ ⋅ C ‾ ] ⊤ \begin{align*}
\sum_{i=0}^{t} \mathbf{s}_i \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-i} \cdot \overline{\mathbf{C}}
={}&
\begin{bmatrix} \mathbf{s}_0 & \mathbf{s}_1 & \dots & \mathbf{s}_t \end{bmatrix} \cdot \\
&
\begin{bmatrix} \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t} \cdot \overline{\mathbf{C}} & \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-1} \cdot \overline{\mathbf{C}} & \dots & \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} \end{bmatrix}^\top \tag{18}
\end{align*} i = 0 ∑ t s i ⋅ B ⋅ A t − i ⋅ C = [ s 0 s 1 … s t ] ⋅ [ B ⋅ A t ⋅ C B ⋅ A t − 1 ⋅ C … B ⋅ C ] ⊤ ( 18 ) 我们定义一个参数化的卷积滤波器：
W ssm = [ B ‾ ⋅ A ‾ n max ⋅ C ‾ B ‾ ⋅ A ‾ n max − 1 ⋅ C ‾ … B ‾ ⋅ C ‾ ] (19) \mathbf{W}_{\text{ssm}} = \begin{bmatrix} \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{n_{\max}} \cdot \overline{\mathbf{C}} & \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{n_{\max}-1} \cdot \overline{\mathbf{C}} & \ldots & \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} \end{bmatrix} \tag{19} W ssm = [ B ⋅ A n m a x ⋅ C B ⋅ A n m a x − 1 ⋅ C … B ⋅ C ] ( 19 ) 其中 n max n_{\max} n m a x S = [ s 0 ⋮ s n ] \mathbf{S} = \begin{bmatrix} \mathbf{s}_0 \\ \vdots \\ \mathbf{s}_n \end{bmatrix} S = s 0 ⋮ s n
O = Conv ( S , W ssm ) + Linear ( S , D ‾ ) (20) \mathbf{O} = \text{Conv}(\mathbf{S}, \mathbf{W}_{\text{ssm}}) + \text{Linear}(\mathbf{S}, \overline{\mathbf{D}}) \tag{20} O = Conv ( S , W ssm ) + Linear ( S , D ) ( 20 ) 其中 Conv ( ⋅ ) \text{Conv}(\cdot) Conv ( ⋅ ) Linear ( ⋅ ) \text{Linear}(\cdot) Linear ( ⋅ )
另一个计算瓶颈在于求解 A ‾ n \overline{\mathbf{A}}^n A n n n n 对角化 A \mathbf{A} A A ‾ \overline{\mathbf{A}} A ( A , B , C , D ) (\mathbf{A},\mathbf{B},\mathbf{C},\mathbf{D}) ( A , B , C , D ) U \mathbf{U} U ( U A U − 1 , B U − 1 , U C , D ) (\mathbf{U}\mathbf{A}\mathbf{U}^{-1},\mathbf{B}\mathbf{U}^{-1},\mathbf{U}\mathbf{C},\mathbf{D}) ( UA U − 1 , B U − 1 , UC , D )
若对矩阵执行上述变换，且假设 A \mathbf{A} A P − 1 Λ P \mathbf{P}^{-1}\Lambda\mathbf{P} P − 1 Λ P Λ \boldsymbol{\Lambda} Λ A = P − 1 Λ P \mathbf{A} = \mathbf{P}^{-1}\Lambda\mathbf{P} A = P − 1 Λ P B ⋅ A ‾ t ⋅ C \mathbf{B} \cdot \overline{\mathbf{A}}^t \cdot \mathbf{C} B ⋅ A t ⋅ C
B ‾ ⋅ A ‾ t ⋅ C ‾ = B ‾ ⋅ ( P − 1 Λ P ) t ⋅ C ‾ = B ‾ ⋅ ( P − 1 Λ P ) ⋅ ( P − 1 Λ P ) ⋯ ( P − 1 Λ P ) ⋅ C ‾ = ( B ‾ ⋅ P − 1 ) ⋅ Λ t ⋅ ( P ⋅ C ‾ ) \begin{align*}
\overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^t \cdot \overline{\mathbf{C}} &= \overline{\mathbf{B}} \cdot (\mathbf{P}^{-1}\Lambda\mathbf{P})^t \cdot \overline{\mathbf{C}} \\
&= \overline{\mathbf{B}} \cdot (\mathbf{P}^{-1}\Lambda\mathbf{P}) \cdot (\mathbf{P}^{-1}\Lambda\mathbf{P}) \cdots (\mathbf{P}^{-1}\Lambda\mathbf{P}) \cdot \overline{\mathbf{C}} \\
&= (\overline{\mathbf{B}} \cdot \mathbf{P}^{-1}) \cdot \Lambda^t \cdot (\mathbf{P} \cdot \overline{\mathbf{C}}) \tag{21}
\end{align*} B ⋅ A t ⋅ C = B ⋅ ( P − 1 Λ P ) t ⋅ C = B ⋅ ( P − 1 Λ P ) ⋅ ( P − 1 Λ P ) ⋯ ( P − 1 Λ P ) ⋅ C = ( B ⋅ P − 1 ) ⋅ Λ t ⋅ ( P ⋅ C ) ( 21 ) 由于 Λ \Lambda Λ t t t Λ t \Lambda^t Λ t
A ‾ ′ = Λ B ‾ ′ = B ‾ ⋅ P − 1 C ‾ ′ = P ⋅ C ‾ D ‾ ′ = D ‾ \begin{align*}
\overline{\mathbf{A}}' &= \Lambda \tag{22} \\
\overline{\mathbf{B}}' &= \overline{\mathbf{B}} \cdot \mathbf{P}^{-1} \tag{23} \\
\overline{\mathbf{C}}' &= \mathbf{P} \cdot \overline{\mathbf{C}} \tag{24} \\
\overline{\mathbf{D}}' &= \overline{\mathbf{D}} \tag{25}
\end{align*} A ′ B ′ C ′ D ′ = Λ = B ⋅ P − 1 = P ⋅ C = D ( 22 ) ( 23 ) ( 24 ) ( 25 ) 将 SSM 应用于 Transformer 的方式十分简便：只需按照式 (14), (15)
SSM 的形式体系为序列建模搭建了统一框架。借助并行与串行两种互补视角，我们可针对不同硬件条件与延迟要求进行优化。近期多款模型均受这种双重特性启发，成功融合了 Transformer 的并行计算能力与循环神经网络的高效推理优势。
