Transformer 家族版本 2.0

The Transformer Family Version 2.0

翁荔 Lilian Weng · · 2023-01-27 · Article ↗

摘要 · Abstract

日期：2023 年 1 月 27 日 | 预计阅读时间：45 分钟 | 作者：Lilian Weng * 局部与全局上下文的结合 * 用于强化学习的 Transformer

Date: January 27, 2023 | Estimated Reading Time: 45 min | Author: Lilian Weng * Combination of Local and Global Context * Transformers for Reinforcement Learning

核心贡献 · Key contributions

对 2020 年以来 Transformer 架构变体与改进进行了全面、结构化的综述。
Provides a comprehensive structured survey of Transformer architecture variants and improvements since 2020.
系统梳理了注意力机制与位置编码，包括相对位置编码和旋转位置嵌入。
Systematically reviews attention mechanisms and positional encodings, including relative position encoding and rotary position embedding.
讨论了通过记忆与检索扩展上下文窗口的机制，如 Transformer-XL、Compressive Transformer、kNN-LM 和 Memorizing Transformer。
Discusses mechanisms for extending context window via memory and retrieval, such as Transformer-XL, Compressive Transformer, kNN-LM, and Memorizing Transformer.
涵盖了自适应计算方法，包括 Universal Transformer、自适应注意力跨度、深度自适应 Transformer 和 CALM。
Covers adaptive computation methods including Universal Transformer, adaptive attention span, depth-adaptive Transformer, and CALM.
综述了稀疏、局部敏感哈希、低秩和随机特征注意力等效率改进，以降低二次复杂度。
Surveys efficiency improvements like sparse, LSH, low-rank, and random feature attention to reduce quadratic complexity.
重点介绍了局部-全局混合注意力模型，包括 ETC、Longformer 和 Big Bird。
Highlights hybrid local-global attention models including ETC, Longformer, and Big Bird.

局限 · Limitations

该综述是一个快照；Transformer 研究发展迅速，其覆盖范围和框架可能很快过时。
The survey is a snapshot; Transformer research evolves rapidly, so coverage and frameworks may become outdated.
许多技术仅在特定任务和数据集上评估，其相对收益可能无法跨领域或跨模型规模迁移。
Many techniques are evaluated on specific tasks and datasets, so their relative gains may not transfer across domains or model scales.
高效注意力方法常常以精度换取速度和内存，引入近似误差或限制上下文范围。
Efficient attention methods often trade off accuracy for speed and memory, introducing approximation errors or restricted context.
ALiBi 等长上下文外推方法的泛化能力有限，在更长序列上可能性能下降。
Long-context extrapolation methods like ALiBi show limited extrapolation and may degrade on much longer sequences.
基于记忆和检索的增强方法需要额外的存储和计算，在极大规模下可能不切实际。
Memory and retrieval based augmentations require extra storage and computation, which may be impractical at very large scale.

论文章节 · Sections（共 25）

概述 Overview
Transformer 基础 Transformer Basics#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#transformer-basics)
注意力与自注意力 Attention and Self-Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#attention-and-self-attention)
多头自注意力 Multi-Head Self-Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#multi-head-self-attention)
编码器-解码器架构 Encoder-Decoder Architecture#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#encoder-decoder-architecture)
位置编码 Positional Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#positional-encoding)
正弦位置编码 Sinusoidal Positional Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#sinusoidal-positional-encoding)
可学习位置编码 Learned Positional Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#learned-positional-encoding)
相对位置编码 Relative Position Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#relative-position-encoding)
旋转位置嵌入 Rotary Position Embedding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#rotary-position-embedding)
更长上下文 Longer Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#longer-context)
上下文记忆 Context Memory#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#context-memory)
不可微外部记忆 Non-Differentiable External Memory#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#non-differentiable-external-memory)
距离增强注意力分数 Distance-Enhanced Attention Scores#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#distance-enhanced-attention-scores)
循环化 Make it Recurrent#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#make-it-recurrent)
自适应建模 Adaptive Modeling#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#adaptive-modeling)
自适应注意力跨度 Adaptive Attention Span#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#adaptive-attention-span)
深度自适应 Transformer Depth-Adaptive Transformer#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#depth-adaptive-transformer)
高效注意力 Efficient Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#efficient-attention)
固定局部上下文 Fixed Local Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#fixed-local-context)
跨步上下文 Strided Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#strided-context)
局部与全局上下文的组合 Combination of Local and Global Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#combination-of-local-and-global-context)
基于内容的注意力 Content-based Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#content-based-attention)
低秩注意力 Low-Rank Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#low-rank-attention)
强化学习中的 Transformer Transformers for Reinforcement Learning#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#transformers-for-reinforcement-learning)

全文 · Full text（逐段中英对照）

概述 Overview

日期：2023 年 1 月 27 日 | 预计阅读时间：45 分钟 | 作者：Lilian Weng

Date: January 27, 2023 | Estimated Reading Time: 45 min | Author: Lilian Weng

* 局部与全局上下文的结合

* Combination of Local and Global Context

* 用于强化学习的 Transformer

* Transformers for Reinforcement Learning

自大约三年前我发表关于“Transformer 家族”的文章以来，又涌现出许多新的 Transformer 架构改进。在此，我对 2020 年的那篇文章进行了大规模重构与扩充——重新组织章节层级，并用更多近期论文改进了许多部分。2.0 版本是旧版本的超集，篇幅约为原来的两倍。

Many new Transformer architecture improvements have been proposed since my last post on “The Transformer Family” about three years ago. Here I did a big refactoring and enrichment of that 2020 post — restructuring the hierarchy of sections and improving many sections with more recent papers. Version 2.0 is a superset of the old version, about twice the length.

Transformer 基础 Transformer Basics#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#transformer-basics)

Transformer（为与其他增强版本区分，常称为“vanilla Transformer”；Vaswani 等人，2017）模型采用编码器-解码器架构，这与许多神经机器翻译模型的做法一致。后续简化的 Transformer 在语言建模任务中表现出色，例如仅编码器的 BERT 或仅解码器的 GPT。

The Transformer (referred to as “vanilla Transformer” to distinguish it from other enhanced versions; Vaswani, et al., 2017) model has an encoder-decoder architecture, as commonly used in many NMT models. Later simplified Transformer was shown to achieve great performance in language modeling tasks, like in encoder-only BERT or decoder-only GPT.

注意力与自注意力 Attention and Self-Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#attention-and-self-attention)

注意力机制是神经网络中的一种机制，模型可以通过选择性地关注给定的数据集合来学习做出预测。注意力的大小由学习到的权重来量化，因此输出通常形成加权平均。

Attention is a mechanism in neural networks that enables a model to learn to make predictions by selectively attending to a given set of data. The amount of attention is quantified by learned weights, and thus the output is usually formed as a weighted average.

自注意力是一种注意力机制，模型利用同一数据样本中其他部分的观测来对该样本的某一部分进行预测。从概念上讲，它与非局部均值（non-local means）非常相似。另请注意，自注意力是排列不变的；换句话说，它是定义在集合上的操作。

Self-attention is a type of attention mechanism where the model makes predictions for one part of a data sample using other parts of the observation from the same sample. Conceptually, it feels quite similar to non-local means. Also note that self-attention is permutation-invariant; in other words, it is an operation on sets.

在多种注意力/自注意力形式中，Transformer（Vaswani 等，2017）采用*缩放点积注意力*：给定查询矩阵 $\mathbf{Q}$、键矩阵 $\mathbf{K}$ 和值矩阵 $\mathbf{V}$，输出是值向量的加权和，其中每个值槽位（value slot）的权重由查询与相应键的点积决定：

Among various forms of attention/self-attention, Transformer (Vaswani et al., 2017) relies on the *scaled dot-product attention*: given a query matrix $\mathbf{Q}$, a key matrix $\mathbf{K}$, and a value matrix $\mathbf{V}$, the output is a weighted sum of the value vectors, where the weight assigned to each value slot is determined by the dot-product of the query with the corresponding key:

\[ \text{attn}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}(\frac{\mathbf{Q} \mathbf{K}^{\top}}{\sqrt{d_k}})\mathbf{V} \]

对于查询向量和键向量 $\mathbf{q}_i, \mathbf{k}_j \in \mathbb{R}^d$（分别是查询矩阵和键矩阵中的行向量），我们得到一个标量分数：

And for a query and a key vector $\mathbf{q}_i, \mathbf{k}_j \in \mathbb{R}^d$ (row vectors in query and key matrices), we have a scalar score:

\[ a_{ij} = \text{softmax}(\frac{\mathbf{q}_i \mathbf{k}_j^\top}{\sqrt{d_k}}) = \frac{\exp(\frac{\mathbf{q}_i \mathbf{k}_j^\top}{\sqrt{d_k}})}{ \sum_{r \in \mathcal{S}_i} \exp(\frac{\mathbf{q}_i \mathbf{k}_r^\top}{\sqrt{d_k}}) } \]

其中 $\mathcal{S}_i$ 是第 $i$ 个查询需要关注的键位置集合。

where $\mathcal{S}_i$ is a collection of key positions for the $i$-th query to attend to.

如需了解其他类型的注意力机制,请参阅我之前的博文。

See my old post for other types of attention if interested.

多头自注意力 Multi-Head Self-Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#multi-head-self-attention)

多头自注意力模块是 Transformer 的关键组成部分。多头机制不是只计算一次注意力，而是将输入拆分成更小的块，然后并行地在每个子空间上计算缩放点积注意力。独立的注意力输出被简单地拼接，并线性变换为期望的维度。

The multi-head self-attention module is a key component in Transformer. Rather than only computing the attention once, the multi-head mechanism splits the inputs into smaller chunks and then computes the scaled dot-product attention over each subspace in parallel. The independent attention outputs are simply concatenated and linearly transformed into expected dimensions.

\[\begin{aligned} \text{MultiHeadAttn}(\mathbf{X}_q, \mathbf{X}_k, \mathbf{X}_v) &= [\text{head}_1; \dots; \text{head}_h] \mathbf{W}^o \\ \text{where head}_i &= \text{Attention}(\mathbf{X}_q\mathbf{W}^q_i, \mathbf{X}_k\mathbf{W}^k_i, \mathbf{X}_v\mathbf{W}^v_i) \end{aligned}\]

其中 $[.;.]$ 表示拼接操作。$\mathbf{W}^q_i, \mathbf{W}^k_i \in \mathbb{R}^{d \times d_k/h}, \mathbf{W}^v_i \in \mathbb{R}^{d \times d_v/h}$ 是将大小为 $L \times d$ 的输入嵌入映射为查询、键和值矩阵的权重矩阵；$\mathbf{W}^o \in \mathbb{R}^{d_v \times d}$ 是输出的线性变换。所有这些权重都应在训练过程中学习。

where $[.;.]$ is a concatenation operation. $\mathbf{W}^q_i, \mathbf{W}^k_i \in \mathbb{R}^{d \times d_k/h}, \mathbf{W}^v_i \in \mathbb{R}^{d \times d_v/h}$ are weight matrices to map input embeddings of size $L \times d$ into query, key and value matrices. And $\mathbf{W}^o \in \mathbb{R}^{d_v \times d}$ is the output linear transformation. All the weights should be learned during training.

多头缩放点积注意力机制示意图。（图片来源：Vaswani 等人，2017，图 2）

Illustration of the multi-head scaled dot-product attention mechanism. (Image source: Figure 2 in Vaswani, et al., 2017)

编码器-解码器架构 Encoder-Decoder Architecture#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#encoder-decoder-architecture)

编码器生成基于注意力机制的表示，能够从大上下文中定位特定信息。它由 6 个相同模块堆叠而成，每个模块包含两个子模块：一个*多头自注意力*层和一个*逐点*全连接前馈网络。逐点意味着对序列中每个元素应用相同的线性变换（具有相同权重）。这也可以视为滤波器大小为 1 的卷积层。每个子模块都有残差连接和层归一化。所有子模块输出相同维度$d$的数据。

The encoder generates an attention-based representation with capability to locate a specific piece of information from a large context. It consists of a stack of 6 identical modules, each containing two submodules, a *multi-head self-attention* layer and a *point-wise* fully connected feed-forward network. By point-wise, it means that it applies the same linear transformation (with same weights) to each element in the sequence. This can also be viewed as a convolutional layer with filter size 1. Each submodule has a residual connection and layer normalization. All the submodules output data of the same dimension $d$.

Transformer 解码器的功能是从编码表示中检索信息。其架构与编码器非常相似，不同之处在于每个相同的重复模块中，解码器包含两个多头注意力子模块而非一个。第一个多头注意力子模块被*掩蔽*，以防止位置关注未来信息。

The function of Transformer decoder is to retrieve information from the encoded representation. The architecture is quite similar to the encoder, except that the decoder contains two multi-head attention submodules instead of one in each identical repeating module. The first multi-head attention submodule is *masked* to prevent positions from attending to the future.

原始 Transformer 模型的架构。（图片来源：图 17）

The architecture of the vanilla Transformer model. (Image source: Figure 17)

位置编码 Positional Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#positional-encoding)

由于自注意力操作具有排列不变性，因此使用合适的位置编码向模型提供 *顺序信息* 非常重要。位置编码 $\mathbf{P} \in \mathbb{R}^{L \times d}$ 与输入嵌入具有相同的维度，因此可以直接加到输入上。经典 Transformer 考虑了两种编码方式：

Because self-attention operation is permutation invariant, it is important to use proper positional encoding to provide *order information* to the model. The positional encoding $\mathbf{P} \in \mathbb{R}^{L \times d}$ has the same dimension as the input embedding, so it can be added on the input directly. The vanilla Transformer considered two types of encodings:

正弦位置编码 Sinusoidal Positional Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#sinusoidal-positional-encoding)

正弦位置编码定义如下：给定词元位置 $i=1,\dots,L$ 和维度 $\delta=1,\dots,d$：

Sinusoidal positional encoding is defined as follows, given the token position $i=1,\dots,L$ and the dimension $\delta=1,\dots,d$:

\[ \text{PE}(i,\delta) = \begin{cases} \sin(\frac{i}{10000^{2\delta'/d}}) & \text{if } \delta = 2\delta' \\ \cos(\frac{i}{10000^{2\delta'/d}}) & \text{if } \delta = 2\delta' + 1 \end{cases} \]

这样，位置编码的每个维度都对应一个正弦波，不同维度具有不同的波长，范围从 $2\pi$ 到 $10000 \cdot 2\pi$。

In this way each dimension of the positional encoding corresponds to a sinusoid of different wavelengths in different dimensions, from $2\pi$ to $10000 \cdot 2\pi$.

$L=32$ 且 $d=128$ 时的正弦位置编码。数值介于 -1（黑色）和 1（白色）之间，0 为灰色。

Sinusoidal positional encoding with $L=32$ and $d=128$. The value is between -1 (black) and 1 (white) and the value 0 is in gray.

可学习位置编码 Learned Positional Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#learned-positional-encoding)

可学习位置编码为每个元素分配一个可学习的列向量，该向量编码其绝对位置（Gehring 等，2017）；此外，这种编码可以在每一层以不同方式学习（Al-Rfou 等，2018）。

Learned positional encoding assigns each element a learned column vector that encodes its absolute position (Gehring et al., 2017); furthermore, this encoding can be learned differently per layer (Al-Rfou et al., 2018).

相对位置编码 Relative Position Encoding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#relative-position-encoding)

Shaw 等人（2018）将相对位置信息纳入$\mathbf{W}^k$和$\mathbf{W}^v$。最大相对位置被裁剪为最大绝对值$k$，这种裁剪操作使模型能够泛化到未见过的序列长度。因此，需要考虑$2k + 1$个独特的边标签，我们记$\mathbf{P}^k, \mathbf{P}^v \in \mathbb{R}^{2k+1}$为可学习的相对位置表示。

Shaw et al. (2018) incorporated relative positional information into $\mathbf{W}^k$ and $\mathbf{W}^v$. Maximum relative position is clipped to a maximum absolute value of $k$ and this clipping operation enables the model to generalize to unseen sequence lengths. Therefore, $2k + 1$ unique edge labels are considered and let us denote $\mathbf{P}^k, \mathbf{P}^v \in \mathbb{R}^{2k+1}$ as learnable relative position representations.

\[ A_{ij}^k = P^k_{\text{clip}(j - i, k)} \quad A_{ij}^v = P^v_{\text{clip}(j - i, k)} \quad \text{其中 }\text{clip}(x, k) = \text{clip}(x, -k, k) \]

\[ A_{ij}^k = P^k_{\text{clip}(j - i, k)} \quad A_{ij}^v = P^v_{\text{clip}(j - i, k)} \quad \text{where }\text{clip}(x, k) = \text{clip}(x, -k, k) \]

如果省略标量$1/\sqrt{d_k}$和 softmax 中的归一化项，但包含位置编码，则可以将位置$i$处的查询与位置$j$处的键之间的注意力分数写为：

If omitting the scalar $1/\sqrt{d_k}$ and the normalizing term in softmax but including positional encodings, we can write the attention score between query at position $i$ and key at position $j$ as:

\[ \begin{aligned} a_{ij} &= \mathbf{q}_i {\mathbf{k}_j}^\top = (\mathbf{x}_i + \mathbf{p}_i)\mathbf{W}^q ((\mathbf{x}_j + \mathbf{p}_j)\mathbf{W}^k)^\top \\ &= \mathbf{x}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{x}_j^\top + \mathbf{x}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{p}_j^\top + \mathbf{p}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{x}_j^\top + \mathbf{p}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{p}_j^\top \end{aligned} \]

Transformer-XL 将上述四项重新参数化如下：

Transformer-XL reparameterizes the above four terms as follows:

\[ a_{ij}^{\text{rel}} = \underbrace{ \mathbf{x}_i\mathbf{W}^q \color{blue}{ {\mathbf{W}_E^k}^{\top} } \mathbf{x}_j^{\top} }_{\text{content-based addressing}} + \underbrace{ \mathbf{x}_i\mathbf{W}^q \color{blue}{ {\mathbf{W}_R^k}^{\top} } \color{green}{\mathbf{r}_{i-j}^{\top}} }_{\text{content-dependent positional bias}} + \underbrace{ \color{red}{\mathbf{u}} \color{blue}{ {\mathbf{W}_E^k}^{\top} } \mathbf{x}_j^{\top} }_{\text{global content bias}} + \underbrace{ \color{red}{\mathbf{v}} \color{blue}{ {\mathbf{W}_R^k}^{\top} } \color{green}{\mathbf{r}_{i-j}^{\top}} }_{\text{global positional bias}} \]

将 $\mathbf{p}_j$ 替换为相对位置编码 $\mathbf{r}_{i-j} \in \mathbf{R}^{d}$;

- Replace $\mathbf{p}_j$ with relative positional encoding $\mathbf{r}_{i-j} \in \mathbf{R}^{d}$;

将 $\mathbf{p}_i\mathbf{W}^q$ 替换为两个可训练参数 $\mathbf{u}$(用于内容)和 $\mathbf{v}$(用于位置),分别位于两个不同的项中;

- Replace $\mathbf{p}_i\mathbf{W}^q$ with two trainable parameters $\mathbf{u}$ (for content) and $\mathbf{v}$ (for location) in two different terms;

将 $\mathbf{W}^k$ 拆分为两个矩阵:$\mathbf{W}_E^k$ 用于内容信息,$\mathbf{W}_R^k$ 用于位置信息。

- Split $\mathbf{W}^k$ into two matrices, $\mathbf{W}_E^k$ for content information and $\mathbf{W}_R^k$ for location information.

旋转位置嵌入 Rotary Position Embedding#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#rotary-position-embedding)

旋转位置嵌入（_RoPE_；Su 等，2021）通过旋转矩阵对绝对位置进行编码，并将它与每个注意力层的键和值矩阵相乘，从而在每一层注入相对位置信息。

Rotary position embedding (_RoPE_; Su et al. 2021) encodes the absolute position with a rotation matrix and multiplies key and value matrices of every attention layer with it to inject relative positional information at every layer.

在将相对位置信息编码进第 $i$ 个键与第 $j$ 个查询的内积时，我们希望构造一个函数，使得内积只与相对位置 $i-j$ 有关。旋转位置嵌入（RoPE）利用欧几里得空间中的旋转操作，将相对位置嵌入建模为：以与其位置索引成比例的角度简单地旋转特征矩阵。

When encoding relative positional information into the inner product of the $i$-th key and the $j$-th query, we would like to formulate the function in a way that the inner product is only about the relative position $i-j$. Rotary Position Embedding (RoPE) makes use of the rotation operation in Euclidean space and frames the relative position embedding as simply rotating feature matrix by an angle proportional to its position index.

给定一个向量 $\mathbf{z}$，如果我们想将它逆时针旋转 $\theta$，可以将其乘以一个旋转矩阵得到 $R\mathbf{z}$，其中旋转矩阵 $R$ 定义为：

Given a vector $\mathbf{z}$, if we want to rotate it counterclockwise by $\theta$, we can multiply it by a rotation matrix to get $R\mathbf{z}$ where the rotation matrix $R$ is defined as:

\[ R = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} \]

当推广到更高维空间时，RoPE 将 $d$ 维空间划分为 $d/2$ 个子空间，并为位于位置 $i$ 的词元构建一个大小为 $d \times d$ 的旋转矩阵 $R$：

When generalizing to higher-dimensional space, RoPE divides the $d$-dimensional space into $d/2$ subspaces and constructs a rotation matrix $R$ of size $d \times d$ for a token at position $i$:

\[ R^d_{\Theta, i} = \begin{bmatrix} \cos i\theta_1 & -\sin i\theta_1 & 0 & 0 & \dots & 0 & 0 \\ \sin i\theta_1 & \cos i\theta_1 & 0 & 0 & \dots & 0 & 0 \\ 0 & 0 & \cos i\theta_2 & -\sin i\theta_2 & \dots & 0 & 0 \\ 0 & 0 & \sin i\theta_2 & \cos i\theta_2 & \dots & 0 & 0 \\ \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & 0 & \dots & \cos i\theta_{d/2} & -\sin i\theta_{d/2} \\ 0 & 0 & 0 & 0 & \dots & \sin i\theta_{d/2} & \cos i\theta_{d/2} \end{bmatrix} \]

在论文中，$ \Theta = { \theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, \ldots, d/2] } $。注意，这本质上等同于正弦位置编码，但以旋转矩阵的形式表述。

where in the paper we have $ \Theta = { \theta_i = 10000^{-2(i-1)/d}, i \in [1, 2, \ldots, d/2] } $. Note that this is essentially equivalent to sinusoidal positional encoding but formulated as a rotation matrix.

然后，键矩阵和查询矩阵都通过与这个旋转矩阵相乘来融入位置信息：

Then both key and query matrices incorporate the positional information by multiplying with this rotation matrix:

\[ \begin{aligned} & \mathbf{q}_i^\top \mathbf{k}_j = (R^d_{\Theta, i} \mathbf{W}^q\mathbf{x}_i)^\top (R^d_{\Theta, j} \mathbf{W}^k\mathbf{x}_j) = \mathbf{x}_i^\top\mathbf{W}^q R^d_{\Theta, j-i}\mathbf{W}^k\mathbf{x}_j \\ & \text{ 其中 } R^d_{\Theta, j-i} = (R^d_{\Theta, i})^\top R^d_{\Theta, j} \end{aligned} \]

\[ \begin{aligned} & \mathbf{q}_i^\top \mathbf{k}_j = (R^d_{\Theta, i} \mathbf{W}^q\mathbf{x}_i)^\top (R^d_{\Theta, j} \mathbf{W}^k\mathbf{x}_j) = \mathbf{x}_i^\top\mathbf{W}^q R^d_{\Theta, j-i}\mathbf{W}^k\mathbf{x}_j \\ & \text{ where } R^d_{\Theta, j-i} = (R^d_{\Theta, i})^\top R^d_{\Theta, j} \end{aligned} \]

旋转位置编码实现的示意图。（图片来源：Su 等人，2021）注意：与论文原图相比，我使用 $i$ 而非 $m$ 来表示位置索引。

Visual illustration of how rotary position embedding is implemented. (Image source: Su et al., 2021) Note: I used $i$ instead of $m$ to represent the position index compared to the original figure in the paper.

更长上下文 Longer Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#longer-context)

在推理时，Transformer 模型的输入序列长度受限于训练时所使用的上下文长度。单纯增加上下文长度会导致高消耗：时间上为 $\mathcal{O}(L^2 d)$，内存上为 $\mathcal{O}(L^2)$，并且可能因硬件约束而无法支持。

The length of an input sequence for Transformer models at inference time is upper-bounded by the context length used for training. Naively increasing context length leads to high consumption in both time $\mathcal{O}(L^2 d)$ and memory $\mathcal{O}(L^2)$ and may not be supported due to hardware constraints.

本节介绍 Transformer 架构中的几项改进，以更好地支持推理时的长上下文；例如，使用附加内存、设计更好的上下文外推，或采用循环机制。

This section introduces several improvements in Transformer architecture to better support long context at inference; e.g., using additional memory, designing for better context extrapolation, or employing recurrence mechanisms.

上下文记忆 Context Memory#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#context-memory)

原始 Transformer 具有固定且有限的注意力范围。在每次更新步骤中，模型只能关注同一分段内的其他元素，信息无法跨隔离的固定长度分段流动。这种*上下文分割*导致了几个问题：

The vanilla Transformer has a fixed and limited attention span. The model can only attend to other elements in the same segments during each update step and no information can flow across separated fixed-length segments. This *context segmentation* causes several issues:

- 模型无法捕捉长期依赖关系。

- The model cannot capture very long-term dependencies.

- 在缺乏或仅有少量上下文的情况下，很难预测每个分段的前几个词元。

- It is hard to predict the first few tokens in each segment given no or thin context.

- 评估成本很高。每当分段向右移动一位时，新分段都要从头重新处理，尽管存在大量重叠的词元。

- The evaluation is expensive. Whenever the segment is shifted to the right by one, the new segment is re-processed from scratch, although there are a lot of overlapped tokens.

原始 Transformer 与 Transformer-XL 在分段长度为 4 时的训练阶段比较。（图片来源：Dai 等人，2019，图 2 左侧）。

A comparison between the training phase of vanilla Transformer and Transformer-XL with a segment length of 4. (Image source: left part of Figure 2 in Dai et al., 2019).

我们将模型中第 $n$ 层在第 $\tau+1$ 段的隐藏状态记为 $\mathbf{h}_{\tau+1}^{(n)} \in \mathbb{R}^{L \times d}$。除了同一片段上一层的隐藏状态 $\mathbf{h}_{\tau+1}^{(n-1)}$ 之外，它还依赖于上一片段中同一层的隐藏状态 $\mathbf{h}_{\tau}^{(n)}$。通过纳入先前隐藏状态的信息，模型将注意力跨度扩展到过去多个片段，显著延长了可关注的过去范围。

Let’s label the hidden state of the $n$-th layer for the $\tau + 1$-th segment in the model as $\mathbf{h}_{\tau+1}^{(n)} \in \mathbb{R}^{L \times d}$. In addition to the hidden state of the previous layer for the same segment $\mathbf{h}_{\tau+1}^{(n-1)}$, it also depends on the hidden state of the same layer for the previous segment $\mathbf{h}_{\tau}^{(n)}$. By incorporating information from the previous hidden states, the model extends the attention span much longer in the past, over multiple segments.

其中，红色部分 $\widetilde{\mathbf{h}}_{\tau+1}^{(n-1)}$ 定义为对前一段隐藏状态 $\mathbf{h}_{\tau}^{(n-1)}$ 取停止梯度后与当前段隐藏状态 $\mathbf{h}_{\tau+1}^{(n-1)}$ 沿序列长度维度拼接。具体地，查询 $\mathbf{Q}_{\tau+1}^{(n)}$ 仅由当前段隐藏状态计算，而键 $\mathbf{K}_{\tau+1}^{(n)}$ 和值 $\mathbf{V}_{\tau+1}^{(n)}$ 都使用扩展后的隐藏状态；最后经 Transformer 层得到输出 $\mathbf{h}_{\tau+1}^{(n)}$。

\[ \begin{aligned} \color{red}{\widetilde{\mathbf{h}}_{\tau+1}^{(n-1)}} &= [\text{stop-gradient}(\mathbf{h}_{\tau}^{(n-1)}) \circ \mathbf{h}_{\tau+1}^{(n-1)}] \\ \mathbf{Q}_{\tau+1}^{(n)} &= \mathbf{h}_{\tau+1}^{(n-1)}\mathbf{W}^q \\ \mathbf{K}_{\tau+1}^{(n)} &= \color{red}{\widetilde{\mathbf{h}}_{\tau+1}^{(n-1)}} \mathbf{W}^k \\ \mathbf{V}_{\tau+1}^{(n)} &= \color{red}{\widetilde{\mathbf{h}}_{\tau+1}^{(n-1)}} \mathbf{W}^v \\ \mathbf{h}_{\tau+1}^{(n)} &= \text{transformer-layer}(\mathbf{Q}_{\tau+1}^{(n)}, \mathbf{K}_{\tau+1}^{(n)}, \mathbf{V}_{\tau+1}^{(n)}) \end{aligned} \]

注意，键和值都依赖于扩展后的隐藏状态，而查询只使用当前步的隐藏状态。拼接操作 $[. \circ .]$ 是沿序列长度维度进行的。Transformer-XL 需要使用相对位置编码，因为如果对前一个片段和当前片段使用绝对位置编码，它们会被赋予相同的位置编码，这是我们不希望的。

Note that both keys and values rely on extended hidden states, while queries only consume hidden states at the current step. The concatenation operation $[. \circ .]$ is along the sequence length dimension. And Transformer-XL needs to use relative positional encoding because previous and current segments would be assigned with the same encoding if we encode absolute positions, which is undesired.

压缩 Transformer（Rae 等人，2019）通过压缩过去的记忆来扩展 Transformer-XL，以支持更长的序列。它明确地在每一层添加大小为 $m_m$ 的 _memory_ 槽，用于存储该层过去的激活值，以保留长上下文。当某些过去的激活值足够旧时，它们会被压缩并保存到每一层额外的 _compressed memory_ 中，其大小为 $m_{cm}$。

Compressive Transformer (Rae et al. 2019) extends Transformer-XL by compressing past memories to support longer sequences. It explicitly adds _memory_ slots of size $m_m$ per layer for storing past activations of this layer to preserve long context. When some past activations become old enough, they are compressed and saved in an additional _compressed memory_ of size $m_{cm}$ per layer.

压缩 Transformer 维护两种类型的记忆槽：memory 和 compressed memory，以支持长上下文。（图片来源：Rae 等人，2019。）

Compressive Transformer maintains two types of memory slots, memory and compressed memory, to support long context. (Image source: Rae et al. 2019).

记忆与压缩记忆都是 FIFO 队列。给定模型上下文长度 $L$，压缩率为 $c$ 的压缩函数定义为 $f_c: \mathbb{R}^{L \times d} \to \mathbb{R}^{[\frac{L}{c}] \times d}$，它将 $L$ 个最旧的激活映射到 $[\frac{L}{c}]$ 个压缩记忆元素。压缩函数有以下几种选择：

Both memory and compressed memory are FIFO queues. Given the model context length $L$, the compression function of compression rate $c$ is defined as $f_c: \mathbb{R}^{L \times d} \to \mathbb{R}^{[\frac{L}{c}] \times d}$, mapping $L$ oldest activations to $[\frac{L}{c}]$ compressed memory elements. There are several choices of compression functions:

1. 最大/平均池化，核大小和步长均为 $c$；

1. Max/mean pooling of kernel and stride size $c$;

2. 一维卷积，核大小和步长均为 $c$（需要学习额外参数）；

2. 1D convolution with kernel and stride size $c$ (need to learn additional parameters);

3. 空洞卷积（需要学习额外参数）。在他们的实验中，卷积压缩在 EnWik8 数据集上效果最好；

3. Dilated convolution (need to learn additional parameters). In their experiments, convolution compression works out the best on EnWik8 dataset;

Compressive Transformer 还有两个额外的训练损失：

Compressive transformer has two additional training losses:

1. 自编码损失（无损压缩目标）衡量我们从压缩记忆中重建原始记忆的效果。

1. Auto-encoding loss (lossless compression objective) measures how well we can reconstruct the original memories from compressed memories.

\[ \mathcal{L}_{ac} = \| \textbf{old_mem}^{(i)} - g(\textbf{new_cm}^{(i)}) \|_2 \] 其中 $ g: \mathbb{R}^{[\frac{L}{c}] \times d} \to \mathbb{R}^{L \times d} $ 是压缩函数 $ f $ 的逆操作。

\[ \mathcal{L}_{ac} = \| \textbf{old_mem}^{(i)} - g(\textbf{new_cm}^{(i)}) \|_2 \] where $ g: \mathbb{R}^{[\frac{L}{c}] \times d} \to \mathbb{R}^{L \times d} $ reverses the compression function $ f $.

2. 注意力重建损失（有损目标）重建内容注意力在原始记忆与压缩记忆上的差异，并最小化该差异。

2. Attention-reconstruction loss (lossy objective) reconstructs content-based attention over memory vs compressed memory and minimizes the difference.

\[ \mathcal{L}_{ar} = \| \text{attn}(\mathbf{h}^{(i)}, \textbf{old_mem}^{(i)}) − \text{attn}(\mathbf{h}^{(i)}, \textbf{new_cm}^{(i)})\|_2 \]

大小为 $m$ 的 Transformer-XL 记忆的最大时间范围为 $m \times N$，其中 $N$ 是模型的层数，注意力成本为 $\mathcal{O}(L^2 + Lm)$。相比之下，压缩 Transformer 的时间范围为 $(m_m + c \cdot m_{cm}) \times N$，注意力成本为 $\mathcal{O}(L^2 + L(m_m + m_{cm}))$。较大的压缩率 $c$ 在时间范围长度与注意力成本之间提供了更好的权衡。

Transformer-XL with a memory of size $m$ has a maximum temporal range of $m \times N$, where $N$ is the number of layers in the model, and attention cost $\mathcal{O}(L^2 + Lm)$. In comparison, compressed transformer has a temporal range of $(m_m + c \cdot m_{cm}) \times N$ and attention cost $\mathcal{O}(L^2 + L(m_m + m_{cm}))$. A larger compression rate $c$ gives better tradeoff between temporal range length and attention cost.

注意力权重按从最旧到最新的顺序存储于三个位置：压缩记忆→记忆→因果掩码序列。在实验中，他们观察到，从存储在常规记忆中的最旧激活到存储在压缩记忆中的激活，注意力权重有所增加，这表明网络正在学习保留重要信息。

Attention weights, from oldest to newest, are stored in three locations: compressed memory → memory → causally masked sequence. In the experiments, they observed an increase in attention weights from oldest activations stored in the regular memory to activations stored in the compressed memory, implying that the network is learning to preserve salient information.

注意力权重（误差棒为一个标准差）随记忆位置变化，从最旧（左）到最新（右）。（图片来源：Rae 等，2019。）

Attention weights, with one standard deviation as error bars, plotted against memory positions, from oldest (left) to newest (right). (Image source: Rae et al. 2019).

不可微外部记忆 Non-Differentiable External Memory#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#non-differentiable-external-memory)

$k$NN-LM（Khandelwal 等，2020）通过线性插值两个模型预测的下一个词概率，用独立的 $k$NN 模型增强预训练 LM。$k$NN 模型建立在外部键值存储之上，该存储可以存储任意大规模预训练数据集或 OOD 新数据集。该数据存储经过预处理，保存大量对（上下文的 LM 嵌入表示，下一个词），最近邻检索在 LM 嵌入空间中进行。由于数据存储可能非常庞大，我们需要借助诸如 FAISS 或 ScaNN 之类的快速密集向量搜索库。索引过程只需一次，推理时易于实现并行。

$k$NN-LM (Khandelwal et al. 2020) enhances a pretrained LM with a separate $k$NN model by linearly interpolating the next token probabilities predicted by both models. The $k$NN model is built upon an external key-value store which can store any large pre-training dataset or OOD new dataset. This datastore is preprocessed to save a large number of pairs (LM embedding representation of context, next token), and the nearest neighbor retrieval happens in the LM embedding space. Because the datastore can be gigantic, we need to rely on libraries for fast dense vector search such as FAISS or ScaNN. The indexing process only happens once and parallelism is easy to implement at inference time.

在推理时，下一个词概率是两种预测的加权和：

At inference time, the next token probability is a weighted sum of two predictions:

\[ \begin{aligned} p(y \vert \mathbf{x}) &= \lambda \; p_\text{kNN}(y \vert \mathbf{x}) + (1- \lambda) \; p_\text{LM}(y \vert \mathbf{x}) \\ p_\text{kNN}(y \vert \mathbf{x}) &\propto \sum_{(k_i, w_i) \in \mathcal{N}} \mathbb{1}[y = w_i] \exp(-d(k_i, f(\mathbf{x}))) \end{aligned} \]

其中 $\mathcal{N}$ 包含由 $k$NN 检索到的一组最近邻数据点；$d(., .)$ 是诸如 L2 距离之类的距离函数。

where $\mathcal{N}$ contains a set of nearest neighbor data points retrieved by $k$NN; $d(., .)$ is a distance function such as L2 distance.

根据实验，更大的数据存储规模或更大的 $k$ 与更好的困惑度相关。权重标量 $\lambda$ 需要调节，但一般来说，与域内数据相比，域外数据期望的 $\lambda$ 更大，而且更大的数据存储可以承受更大的 $\lambda$。

According to the experiments, larger datastore size or larger $k$ is correlated with better perplexity. The weighting scalar $\lambda$ should be tuned, but in general it is expected to be larger for out-of-domain data compared to in-domain data and larger datastore can afford a larger $\lambda$.

SPALM（*自适应半参数语言模型*；Yogatama 等人，2021）同时引入了（1）Transformer-XL 式记忆，将外部上下文中的隐藏状态作为短期记忆；（2）$k$NN-LM 式键值存储，作为长期记忆。

SPALM (*Adaptive semiparametric language models*; Yogatama et al. 2021) incorporates both (1) Transformer-XL style memory for hidden states from external context as short-term memory and (2) $k$NN-LM style key-value store as long memory.

图示 SPALM 如何将过去隐藏状态的上下文记忆（短期记忆）与外部键值数据存储（长期记忆）相结合，以支持更长的上下文。（图片来源：Yogatama 等人，2021。）

Illustration of how SPALM combines context memory of past hidden states (short term memory) with an external key-value datastore (long term memory) to support longer context. (Image source: Yogatama et al. 2021).

SPALM 运行 $k$NN 搜索，获取上下文最相关的 $k$ 个词符。对于每个词符，我们都可以获得预训练语言模型提供的相同嵌入表示，记为 $\{\mathbf{y}_i\}_{i=1}^k$。门控机制首先使用一个简单的注意力层，以 $\mathbf{h}^R_t$（词符 $x_t$ 在第 $R$ 层的隐藏状态）作为查询来聚合检索到的词符嵌入，然后学习一个门控参数 $\mathbf{g}_t$，以平衡局部信息 $\mathbf{h}^R_t$ 与长期信息 $\mathbf{m}_t$。

SPALM runs $k$NN search to fetch $k$ tokens with most relevant context. For each token we can get the same embedding representation provided by a pretrained LM, denoted as $\{\mathbf{y}_i\}_{i=1}^k$. The gating mechanism first aggregates the retrieved token embeddings with a simple attention layer using $\mathbf{h}^R_t$ (the hidden state for token $x_t$ at layer $R$) as a query and then learns a gating parameter $\mathbf{g}_t$ to balance between local information $\mathbf{h}^R_t$ and long-term information $\mathbf{m}_t$.

\[ \begin{aligned} \mathbf{m}_t &= \sum_{i=1}^k \frac{\exp(\mathbf{y}_i^\top \mathbf{h}^R_t)}{\sum_{j=1}^k \exp(\mathbf{y}_j^\top \mathbf{h}^R_t)} \cdot \mathbf{y}_i \\ \mathbf{g}_t &= \sigma(\mathbf{w}_g^\top \mathbf{h}_t^R) \\ \mathbf{z}_t &= (1 - \mathbf{g}_t) \odot \mathbf{m}_t + \mathbf{g}_t \odot \mathbf{h}^R_t \\ p(x_{t+1}\mid \mathbf{x}_{\leq t}) &= \text{softmax}(\mathbf{z}_t; \mathbf{W}) \end{aligned} \]

其中 $\mathbf{w}_g$ 为待学习的参数向量；$\sigma(\cdot)$ 为 sigmoid 函数；$\mathbf{W}$ 是输入和输出词符共享的词嵌入矩阵。与 $k$NN-LM 不同，他们没有发现最近邻距离在检索词符的聚合中有帮助。

where $\mathbf{w}_g$ is a parameter vector to learn; $\sigma(\cdot)$ is sigmoid; $\mathbf{W}$ is the word embedding matrix shared between both input and output tokens. Different from $k$NN-LM, they didn’t find the nearest neighbor distance to be helpful in the aggregation of retrieved tokens.

在训练期间，长期记忆中的键表示由预训练语言模型生成并保持恒定，而值编码器（即词嵌入矩阵）会被更新。

During training, the key representations in the long-term memory stay constant, produced by a pretrained LM, but the value encoder, aka the word embedding matrix, gets updated.

Memorizing Transformer（Wu 等人，2022）在仅解码器 Transformer 的顶部堆栈附近添加了一个由 $k$NN 增强的注意力层。该特殊层维护一个 Transformer-XL 风格的 FIFO 缓存，用于存储过去的键值对。

Memorizing Transformer (Wu et al. 2022) adds a $k$NN-augmented attention layer near the top stack of a decoder-only Transformer. This special layer maintains a Transformer-XL style FIFO cache of past key-value pairs.

相同的 QKV 值同时用于局部注意力和 $k$NN 机制。$k$NN 查找为输入序列中的每个查询返回前 $k$ 个（键，值）对，然后通过自注意力堆栈处理它们，计算检索值的加权平均。两种注意力通过一个可学习的逐头门控参数进行组合。为防止值幅度出现大的分布偏移，缓存中的键和值均被归一化。

The same QKV values are used for both local attention and $k$NN mechanisms. The $k$NN lookup returns top-$k$ (key, value) pairs for each query in the input sequence and then they are processed through the self-attention stack to compute a weighted average of retrieved values. Two types of attention are combined with a learnable per-head gating parameter. To prevent large distributional shifts in value magnitude, both keys and values in the cache are normalized.

他们在 Memorizing Transformer 实验中的发现： * 一些实验观察到，先使用小记忆训练模型，再用更大记忆微调，效果优于从头开始使用大记忆训练。

What they found during experiments with Memorizing Transformer: * It is observed in some experiments that training models with a small memory and then finetuned with a larger memory works better than training with a large memory from scratch.

During training, the key representations in the long-term memory stay constant, produced by a pretrained LM, but the value encoder, aka the word embedding matrix, gets updated.

* 较小的 Memorizing Transformer 只需在内存中保存 8k 个 token，即可匹配具有 5 倍可训练参数的更大规模 vanilla Transformer 的困惑度。

* The smaller Memorizing Transformer with just 8k tokens in memory can match the perplexity of a larger vanilla Transformer with 5× more trainable parameters.

* 将外部内存的大小增加到 262K 时，性能持续获得提升。

* Increasing the size of external memory provided consistent gains up to a size of 262K.

* 无内存 Transformer 可以通过微调来使用内存。

* A non-memory Transformer can be fine-tuned to use memory.

使用键值内存微调 vanilla Transformer，可以达到与从头训练 Memorizing Transformer 相当的性能（图片来源：Wu et al., 2022）。

Fine-tuning a vanilla Transformer with a key-value memory can achieve similar performance to training a Memorizing Transformer from scratch (Image source: Wu et al., 2022).

距离增强注意力分数 Distance-Enhanced Attention Scores#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#distance-enhanced-attention-scores)

Distance-Aware Transformer（DA-Transformer；Wu 等，2021）和 Attention with Linear Biases（ALiBi；Press 等，2022）的动机类似：为了鼓励模型在比训练时更长的上下文上进行外推，我们可以根据键（key）和查询（query）token 之间的距离，将位置信息显式附加到每一对注意力分数上。

Distance-Aware Transformer (DA-Transformer; Wu et al., 2021) and Attention with Linear Biases (ALiBi; Press et al., 2022) are motivated by similar ideas: to encourage the model to extrapolate over longer contexts than it was trained on, we can explicitly attach positional information to every pair of attention scores based on the distance between key and query tokens.

需要注意的是，原始 Transformer 中的默认位置编码只向输入序列添加位置信息，而后来改进的编码机制（如旋转位置嵌入）会修改每一层的注意力分数，其形式与距离增强的注意力分数非常相似。

Note that the default positional encoding in the vanilla Transformer only adds positional information to the input sequence, whereas later improved encoding mechanisms—such as rotary position embedding—alter the attention scores of every layer and take a form very similar to distance-enhanced attention scores.

DA-Transformer（Wu 等，2021）将每一层的注意力分数乘以一个可学习的偏置，该偏置被建模为键与查询之间距离的函数。不同的注意力头使用不同的参数，以区分对短期上下文与长期上下文的多样化偏好。给定两个位置 $i, j$，DA-Transformer 使用以下加权函数来修正自注意力分数：

_DA-Transformer_ (Wu et al., 2021) multiplies the attention scores at each layer by a learnable bias formulated as a function of the distance between key and query. Different attention heads use different parameters to distinguish their preferences for short-term versus long-term context. Given two positions, $i, j$, DA-Transformer uses the following weighting function to alter the self-attention score:

\[ \begin{aligned} \mathbf{R}^{(i)} &= \alpha_i \mathbf{R} \quad \text{where }R_{ij} = \vert i-j \vert\\ f(\mathbf{R}^{(i)}; \beta_i) &= \frac{1 + \exp(\beta_i)}{1 + \exp(\beta_i - \mathbf{R}^{(i)})} \\ \text{attn}(\mathbf{Q}^{(i)}, \mathbf{K}^{(i)}, \mathbf{V}^{(i)}) &= \text{row-softmax}\Big(\frac{\text{ReLU}(\mathbf{Q}^{(i)}\mathbf{K}^{(i)\top})f(\mathbf{R}^{(i)})}{\sqrt{d}}\Big) \mathbf{V}^{(i)} \end{aligned} \]

其中 $\alpha_i$ 是一个可学习参数，用于按头对相对距离进行不同加权，头用上标 $^{(i)}$ 索引；$\beta_i$ 是一个可学习参数，用于控制第 $i$ 个注意力头相对于距离的上界和上升斜率。加权函数 $f(.)$ 的设计满足：(1) $f(0)=1$；(2) 当 $\mathbf{R}^{(i)} \to -\infty$ 时 $f(\mathbf{R}^{(i)}) = 0$；(3) 当 $\mathbf{R}^{(i)} \to +\infty$ 时 $f(\mathbf{R}^{(i)})$ 有界；(4) 尺度可调；(5) 函数单调。由 $f(\mathbf{R}^{(i)})$ 带来的额外时间复杂度为 $\mathcal{O}(L^2)$，相对于自注意力时间复杂度 $\mathcal{O}(L^2 d)$ 较小。额外内存消耗极小，约为 $\mathcal{O}(2h)$。

where $\alpha_i$ is a learnable parameter that weights relative distance differently per head, with the head indexed by superscript $^{(i)}$; $\beta_i$ is a learnable parameter that controls the upper bound and ascending slope with respect to distance for the $i$-th attention head. The weighting function $f(.)$ is designed such that: (1) $f(0)=1$; (2) $f(\mathbf{R}^{(i)}) = 0$ when $\mathbf{R}^{(i)} \to -\infty$; (3) $f(\mathbf{R}^{(i)})$ is bounded when $\mathbf{R}^{(i)} \to +\infty$; (4) the scale is tunable; and (5) the function is monotonic. The extra time complexity introduced by $f(\mathbf{R}^{(i)})$ is $\mathcal{O}(L^2)$, which is small relative to the self-attention time complexity $\mathcal{O}(L^2 d)$. The extra memory consumption is minimal, approximately $\mathcal{O}(2h)$.

与乘法方法不同，ALiBi（Press 等人，2022）在查询-键注意力分数上添加了一个常数偏置项，该项与两两距离成比例。该偏置引入了强烈的近因偏好，并对距离过远的键进行惩罚。不同头以不同速率增加惩罚：\[ \text{softmax}( \mathbf{q}_i \mathbf{K}^\top + \alpha_i \cdot [0, -1, -2, \dots, -(i-1)] ) \] 其中 $\alpha_i$ 是头特定的加权标量。与 DA-transformer 不同，$\alpha_i$ 不是学习得到的，而是固定为几何序列；例如，对于 8 个头，$\alpha_i = \{\frac{1}{2}, \frac{1}{2^2}, \dots, \frac{1}{2^8}\}$。总体思路与相对位置编码所旨在解决的问题非常相似。

Instead of multipliers, ALiBi (Press et al. 2022) adds a constant bias term on query-key attention scores, proportional to pairwise distances. The bias introduces a strong recency preference and penalizes keys that are too far away. The penalties are increased at different rates within different heads: \[ \text{softmax}( \mathbf{q}_i \mathbf{K}^\top + \alpha_i \cdot [0, -1, -2, \dots, -(i-1)] ) \] where $\alpha_i$ is a head-specific weighting scalar. Different from DA-transformer, $\alpha_i$ is not learned but fixed as a geometric sequence; for example, for 8 heads, $\alpha_i = \{\frac{1}{2}, \frac{1}{2^2}, \dots, \frac{1}{2^8}\}$. The overall idea is very much similar to what relative positional encoding aims to solve.

ALiBi 如何通过位置偏置项增强注意力分数的示意图。（图片来源：Press 等人，2021 年。）

Illustration of how ALiBi enhances attention scores with a positional bias term. (Image source: Press et al. 2021).

使用 ALiBi，Press 等人（2022）在训练时以上下文长度 1024 训练了一个 1.3B 模型，并在推理时外推到 2046。

With ALiBi, Press et al. (2022) trained a 1.3B model on context length 1024 during training and extrapolated to 2046 at inference time.

不同配置的 Transformer 模型运行推理的外推实验，包括正弦位置编码、旋转位置编码、T5 中的简化相对位置编码和 ALiBi。所有模型都以较小的上下文长度训练，但推理时使用了更长的上下文。（图片来源：Press 等人，2021 年。）

Extrapolation experiments for running inference with Transformers of different configurations, including sinusoidal positional encoding, rotary positional encoding, simplified relative positional encoding in T5, and ALiBi. All models were trained with small context lengths but inference ran with much longer contexts. (Image source: Press et al., 2021).

循环化 Make it Recurrent#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#make-it-recurrent)

Universal Transformer（Dehghani 等人，2019）将 Transformer 中的自注意力机制与 RNN 中的循环机制相结合，旨在同时利用 Transformer 的长期全局感受野和 RNN 的学习归纳偏置。Universal Transformer 不是固定通过一定数量的层，而是使用自适应计算时间动态调整步数。如果固定步数，Universal Transformer 就相当于一个跨层共享参数的多层 Transformer。

Universal Transformer (Dehghani, et al. 2019) combines self-attention in Transformer with the recurrent mechanism in RNN, aiming to benefit from both a long-term global receptive field of Transformer and learned inductive biases of RNN. Rather than going through a fixed number of layers, Universal Transformer dynamically adjusts the number of steps using adaptive computation time. If we fix the number of steps, a Universal Transformer is equivalent to a multi-layer Transformer with shared parameters across layers.

从高层次来看，Universal Transformer 可以看作一个用于学习每个 token 的隐藏状态表示的循环函数。该循环函数在各 token 位置上并行演化，位置之间的信息通过注意力机制共享。

On a high level, the universal transformer can be viewed as a recurrent function for learning the hidden state representation per token. The recurrent function evolves in parallel across token positions and the information between positions is shared through self-attention.

Universal Transformer 如何对每个位置并行地重复精炼一组隐藏状态表示。（图片来源：Dehghani 等人，2019 的图 1。）

How the Universal Transformer refines a set of hidden state representations repeatedly for every position in parallel. (Image source: Figure 1 in Dehghani, et al. 2019).

给定长度为 $L$ 的输入序列，Universal Transformer 在步骤 $t$ 迭代更新表示 $\mathbf{h}^t \in \mathbb{R}^{L \times d}$，步数可调。在步骤 0，$\mathbf{h}^0$ 被初始化为与输入嵌入矩阵相同。所有位置在多头注意力机制中并行处理，然后经过一个循环转换函数。

Given an input sequence of length $L$, Universal Transformer iteratively updates the representation $\mathbf{h}^t \in \mathbb{R}^{L \times d}$ at step $t$ for an adjustable number of steps. At step 0, $\mathbf{h}^0$ is initialized to be same as the input embedding matrix. All the positions are processed in parallel in the multi-head self-attention mechanism and then go through a recurrent transition function.

\[ \begin{aligned} \mathbf{A}^t &= \text{LayerNorm}(\mathbf{h}^{t-1} + \text{MultiHeadAttention}(\mathbf{h}^{t-1} + \mathbf{P}^t) \\ \mathbf{h}^t &= \text{LayerNorm}(\mathbf{A}^{t-1} + \text{Transition}(\mathbf{A}^t)) \end{aligned} \]

其中 $\text{Transition}(.)$ 要么是可分离卷积，要么是一个全连接神经网络，由两个逐位置的（即分别应用于 $\mathbf{A}^t$ 的每一行）仿射变换加一个 ReLU 组成。

where $\text{Transition}(.)$ is either a separable convolution or a fully-connected neural network that consists of two position-wise (i.e. applied to each row of $\mathbf{A}^t$ individually) affine transformations plus one ReLU.

位置编码 $\mathbf{P}^t$ 使用正弦位置信号，但额外增加了一个时间维度：

The positional encoding $\mathbf{P}^t$ uses sinusoidal position signal but with an additional time dimension:

\[ \text{PE}(i, t, \delta) = \begin{cases} \sin(\frac{i}{10000^{2\delta'/d}}) \oplus \sin(\frac{t}{10000^{2\delta'/d}}) & \text{if } \delta = 2\delta'\\ \cos(\frac{i}{10000^{2\delta'/d}}) \oplus \cos(\frac{t}{10000^{2\delta'/d}}) & \text{if } \delta = 2\delta' + 1\\ \end{cases} \]

Universal Transformer 的简化示意图。编码器和解码器共享相同的基本循环结构。但解码器还会关注最终的编码器表示 $\mathbf{h}^T$。（图片来源：Dehghani 等人，2019，图 2）

A simplified illustration of Universal Transformer. The encoder and decoder share the same basic recurrent structure. But the decoder also attends to the final encoder representation $\mathbf{h}^T$. (Image source: Figure 2 in Dehghani et al., 2019)

在 Universal Transformer 的自适应版本中，循环步数 $T$ 由 ACT 动态确定。每个位置都配备了一个动态的 ACT 停止机制。一旦某个逐 token 的循环模块停止，它就不再接受更多的循环更新，而只是将当前值复制到下一步，直到所有模块都停止，或直到模型达到最大步数限制。

In the adaptive version of Universal Transformer, the number of recurrent steps $T$ is dynamically determined by ACT. Each position is equipped with a dynamic ACT halting mechanism. Once a per-token recurrent block halts, it stops taking more recurrent updates but simply copies the current value to the next step until all the blocks halt or until the model reaches a maximum step limit.

自适应建模 Adaptive Modeling#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#adaptive-modeling)

自适应建模指的是一种能够根据不同输入调整计算量的机制。例如，一些词元可能只需要局部信息，因此需要较短的注意力跨度；或者一些词元相对更容易预测，无需经过整个注意力堆栈处理。

Adaptive modeling refers to a mechanism that can adjust the amount of computation according to different inputs. For example, some tokens may only need local information and thus demand a shorter attention span; or some tokens are relatively easier to predict and do not need to be processed through the entire attention stack.

自适应注意力跨度 Adaptive Attention Span#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#adaptive-attention-span)

Transformer 的一个关键优势是能够捕获长期依赖关系。根据上下文的不同，模型有时可能更倾向于关注更远的内容；或者一个注意力头可能与其他注意力头具有不同的注意力模式。如果注意力跨度能够灵活调整其长度，并且仅在需要时才关注更远的内容，这将有助于减少计算和内存开销，从而支持模型中更长的最大上下文长度。

One key advantage of Transformer is the capability of capturing long-term dependencies. Depending on the context, the model may prefer to attend further sometime than others; or one attention head may had different attention pattern from the other. If the attention span could adapt its length flexibly and only attend further back when needed, it would help reduce both computation and memory cost to support longer maximum context size in the model.

这是自适应注意力跨度（Adaptive Attention Span）的动机。Sukhbaatar 等人（2019）提出了一种自我注意力机制，旨在寻找最优的注意力跨度。他们假设不同的注意力头在同一上下文窗口内可能分配不同的分数（见图 14），因此每个头的最优跨度应分别训练。

This is the motivation for Adaptive Attention Span. Sukhbaatar et al (2019) proposed a self-attention mechanism that seeks an optimal attention span. They hypothesized that different attention heads might assign scores differently within the same context window (See Fig. 14) and thus the optimal span would be trained separately per head.

同一模型中的两个注意力头 A 和 B 在同一上下文窗口内分配注意力的方式不同。头 A 更多关注最近的词元，而头 B 则均匀地回顾更远的过去。（图片来源：Sukhbaatar 等人，2019）

Two attention heads in the same model, A & B, assign attention differently within the same context window. Head A attends more to the recent tokens, while head B look further back into the past uniformly. (Image source: Sukhbaatar, et al. 2019)

给定第 $i$ 个词元，我们需要计算该词元与其注意力跨度 $s$ 内的其他键之间的注意力权重：

Given the $i$-th token, we need to compute the attention weights between this token and other keys within its attention span of size $s$:

$$ \begin{aligned} e_{ij} &= \mathbf{q}_i {\mathbf{k}_j}^\top \\ a_{ij} &= \text{softmax}(e_{ij}) = \frac{\exp(e_{ij})}{\sum_{r=i-s}^{i-1} \exp(e_{ir})} \\ \mathbf{y}_i &= \sum_{r=i-s}^{i-1}a_{ir}\mathbf{v}_r = \sum_{r=i-s}^{i-1}a_{ir}\mathbf{x}_r\mathbf{W}^v \end{aligned} $$

增加了一个_软掩码函数_$m_z$ 来控制有效的可调节注意力跨度，该函数将查询与键之间的距离映射为 [0, 1] 区间内的值。$m_z$ 由参数 $z \in [0, s]$ 参数化，且 $z$ 是可学习的：

A _soft mask function_ $m_z$ is added to control an effective adjustable attention span, which maps the distance between query and key into a [0, 1] value. $m_z$ is parameterized by $z \in [0, s]$ and $z$ is to be learned:

\[ m_z(x) = \text{clip}(\frac{1}{R}(R+z-x), 0, 1) \]

其中 $R$ 是定义 $m_z$ 软度的超参数。

where $R$ is a hyper-parameter which defines the softness of $m_z$.

自适应注意力跨度中使用的软掩码函数。（图片来源：Sukhbaatar 等人，2019。）

The soft masking function used in the adaptive attention span. (Image source: Sukhbaatar, et al. 2019.)

软掩码函数应用于注意力权重中的 softmax 元素：

The soft mask function is applied to the softmax elements in the attention weights:

\[ a_{ij} = \frac{m_z(i-j)\exp(s_{ij})}{\sum_{r=i-s}^{i-1} m_z(i-r) \exp(s_{ir})} \]

在上述方程中，$z$ 是可微的，因此它与模型的其他部分联合训练。参数 $z^{(i)}, i=1, \dots, h$ 按注意力头分别学习。此外，损失函数在 $\sum_{i=1}^h z^{(i)}$ 上有一个额外的 L1 惩罚项。

In the above equation, $z$ is differentiable, so it is trained jointly with other parts of the model. Parameters $z^{(i)}, i=1, \dots, h$ are learned *separately per head*. Moreover, the loss function has an extra L1 penalty on $\sum_{i=1}^h z^{(i)}$.

利用自适应计算时间（Adaptive Computation Time），该方法可以进一步增强，使注意力跨度长度具有灵活性，并动态适应当前输入。时刻 $t$ 的注意力头的跨度参数 $z_t$ 是一个 sigmoid 函数（S 形函数），$z_t = S \sigma(\mathbf{v} \cdot \mathbf{x}_t + b)$，其中向量 $\mathbf{v}$ 和偏置标量 $b$ 与其他参数联合学习。

Using Adaptive Computation Time, the approach can be further enhanced to have flexible attention span length, adapting dynamically to the current input. The span parameter $z_t$ of an attention head at time $t$ is a sigmoidal function, $z_t = S \sigma(\mathbf{v} \cdot \mathbf{x}_t + b)$, where the vector $\mathbf{v}$ and the bias scalar $b$ are learned jointly with other parameters.

在使用自适应注意力跨度的 Transformer 实验中，Sukhbaatar 等人（2019）发现一个总体趋势：较低层不需要很长的注意力跨度，而较高层中的少数注意力头可能会使用异常长的跨度。自适应注意力跨度还有助于大幅减少 FLOPS 的数量，尤其是在具有许多注意力层和大上下文长度的大模型中。

In experiments with the Transformer using adaptive attention spans, Sukhbaatar et al. (2019) found a general tendency that lower layers do not require very long attention spans, while a few attention heads in higher layers may use exceptionally long spans. Adaptive attention spans also help greatly reduce the number of FLOPS, especially in a large model with many attention layers and a large context length.

深度自适应 Transformer Depth-Adaptive Transformer#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#depth-adaptive-transformer)

在推理时，我们自然会假设某些词元比其他词元更容易预测，因此不需要同样多的计算量。于是，我们可以只让这些词元的预测经过有限数量的层，从而在速度与性能之间取得良好的平衡。

At inference time, it is natural to assume that some tokens are easier to predict and thus do not require as much computation as others. Therefore, we may only process their predictions through a limited number of layers to achieve a good balance between speed and performance.

深度自适应 Transformer（Elabyad 等人，2020）和置信自适应语言模型（CALM；Schuster 等人，2022）都受这一思路启发，学习为不同的输入词元预测所需的最优层数。

Both Depth-Adaptive Transformer (Elabyad et al., 2020) and Confident Adaptive Language Model (CALM; Schuster et al., 2022) are motivated by this idea and learn to predict optimal numbers of layers needed for different input tokens.

深度自适应 Transformer（Elabyad 等人，2020）在每一层都附加了一个输出分类器，基于该层的激活产生退出预测。分类器权重矩阵可以逐层不同，也可以跨层共享。训练时，模型会对不同的退出序列进行采样，从而利用不同层的隐藏状态来优化模型。学习目标融合了在不同层 $n = 1, \dots, N$ 预测出的似然概率：

_Depth-adaptive transformer_ (Elabyad et al., 2020) attaches an output classifier to every layer to produce exit predictions based on activations of that layer. The classifier weight matrices can be different per layer or shared across layers. During training, the model samples different sequences of exits such that the model is optimized with hidden states of different layers. The learning objective incorporates likelihood probabilities predicted at different layers, $n = 1, \dots, N$:

\[ \text{LL}^n_t = \log p(y_t \vert \mathbf{h}^n_{t-1}) \quad \text{LL}^n = \sum_{t=1}^{\vert \mathbf{y} \vert} \text{LL}^n_t \]

自适应深度分类器输出参数化分布 $q_t$。它使用交叉熵损失，针对一个 oracle 分布 $q^*_t$ 进行训练。论文探索了学习这种分类器 $q_t$ 的三种配置。

Adaptive depth classifiers output a parametric distribution $q_t$. It is trained with cross-entropy loss against an oracle distribution $q^*_t$. The paper explored three configurations for how to learn such a classifier $q_t$.

三种自适应深度分类器的示意图。

Illustration of three types of adaptive depth classifiers.

1. **序列特定深度分类器**：同一序列的所有词元共享相同的退出块。它取决于该序列编码器表示的平均值。给定长度为 $L$ 的输入序列 $\mathbf{x}$，分类器以 $\bar{\mathbf{x}} = \frac{1}{L} \sum_{t=1}^L \mathbf{x}_t$ 为输入，输出一个对应 $N$ 层的 $N$ 维多项分布。

1. *Sequence-specific depth classifier*: All tokens of the same sequence share the same exit block. It depends on the average of the encoder representation of the sequence. Given an input sequence $\mathbf{x}$ of length $L$, the classifier takes $\bar{\mathbf{x}} = \frac{1}{L} \sum_{t=1}^L \mathbf{x}_t$ as input and outputs a multinomial distribution of $N$ dimensions, corresponding to $N$ layers.

\[\begin{aligned} q(n \vert \mathbf{x}) & = \text{softmax}(\mathbf{W}_n \bar{\mathbf{x}} + b_n) \in \mathbb{R}^N \\ q_\text{lik}^*(\mathbf{x}, \mathbf{y}) & = \delta(\arg\max_n \text{LL}^n - \lambda n) \\ \text{or }q_\text{corr}^*(\mathbf{x}, \mathbf{y}) & = \delta(\arg\max_n C^n - \lambda n) \text{ where }C^n = \vert\{ t \vert y_t = \arg\max_y p(y \vert \mathbf{h}^n_{t-1}) \}\vert \\ \end{aligned}\]

其中 $\delta$ 是狄拉克 δ 函数（单位冲激函数），$-\lambda n$ 是鼓励在较低层退出的正则化项。真实分布 $q^*$ 可以通过两种方式构造：基于最大似然的 $q_\text{lik}^*$ 或基于正确性的 $q_\text{corr}^*$。

where $\delta$ is Dirac delta (unit impulse) function and $-\lambda n$ is a regularization term to encourage lower layer exits. The ground truth $q^*$ can be prepared in two ways, based on maximum likelihood $q_\text{lik}^*$ or correctness $q_\text{corr}^*$.

2. **词元特定深度分类器（多项式）**：每个词元使用不同的退出块进行解码，其预测以第一个解码器隐藏状态 $\mathbf{h}^1_t$ 为条件：

2. *Token-specific depth classifier (multinomial)*: Each token is decoded with a different exit block, predicted conditioned on the first decoder hidden state $\mathbf{h}^1_t$:

\[ q_t(n \vert \mathbf{x}, \mathbf{y}_{< t}) = \text{softmax}(\mathbf{W}_n \mathbf{h}^1_t + b_n) \]

3. **Token-specific 深度分类器（geometric-like）**：为每个 token 的每一层生成一个二值退出预测分布 $\mathcal{X}^n_t$。RBF 核 $\kappa(t, t') = \exp\frac{\vert t - t' \vert^2}{\sigma}$ 用于平滑预测，从而将当前决策对未来时间步的影响纳入考虑。

3. _Token-specific depth classifier (geometric-like)_: A binary exit prediction distribution is produced for each layer and each token, $\mathcal{X}^n_t$. The RBF kernel $\kappa(t, t') = \exp\frac{\vert t - t' \vert^2}{\sigma}$ is used to smooth the predictions, incorporating the impact of the current decision on future time steps.

\[ \begin{aligned} \mathcal{X}^n_t &= \text{sigmoid}(\mathbf{w}_n^\top \mathbf{h}^n_t + b_n)\quad \forall n \in [1, \dots, N-1] \\ q_t(n \vert \mathbf{x}, \mathbf{y}_{< t}) &= \begin{cases} \mathcal{X}^n_t \prod_{n' < n} (1 - \mathcal{X}^{n'}_t) & \text{if } n < N \\ \prod_{n' < N} (1 - \mathcal{X}^{n'}_t) & \text{otherwise} \end{cases} \\ q_\text{lik}^*(\mathbf{x}, \mathbf{y}) &= \delta(\arg\max_n \widetilde{\text{LL}}^n_t - \lambda n) \text{ where } \widetilde{\text{LL}}^n_t = \sum_{t'=1}^{\vert\mathbf{y}\vert}\kappa(t, t') LL^n_{t'} \\ \text{or }q_\text{cor}^*(\mathbf{x}, \mathbf{y}) &= \delta(\arg\max_n \tilde{C}_t^n - \lambda n) \text{ where }C_t^n = \mathbb{1}[y_t = \arg\max_y p(y \vert \mathbf{h}^n_{t-1})],\; \tilde{C}^n_t = \sum_{t'=1}^{\vert\mathbf{y}\vert}\kappa(t, t') C^n_{t'} \\ \end{aligned} \]

在推理时，用于做出退出决策的置信度阈值需要校准。深度自适应 Transformer 通过网格搜索在验证集上找到这样的阈值。_CALM_ (Schuster et al. 2022) 应用 Learn then Test (LTT) 框架 (Angelopoulos et al. 2021) 来识别一组有效阈值，并选择最小值作为推理阈值。除了训练逐层退出分类器外，CALM 还探索了其他用于自适应深度预测的方法，包括 softmax 响应（即前两个 softmax 输出之差）和隐藏状态饱和（即 $\cos(\mathbf{h}^n_t, \mathbf{h}^{n+1}_t)$）作为退出决策的置信度分数。他们发现 softmax 响应能带来最佳的推理加速。

At inference time, the confidence threshold for making an exit decision needs to be calibrated. Depth-adaptive transformer finds such a threshold on a validation set via grid search. _CALM_ (Schuster et al. 2022) applied the Learn then Test (LTT) framework (Angelopoulos et al. 2021) to identify a subset of valid thresholds and chose the minimum value as the threshold for inference. Besides training a per-layer exit classifier, CALM also explored other methods for adaptive depth prediction, including softmax responses (i.e., the difference between the top two softmax outputs) and hidden state saturation (i.e., $\cos(\mathbf{h}^n_t, \mathbf{h}^{n+1}_t)$) as confidence scores for exit decisions. They found that softmax responses yield the best inference speedup.

高效注意力 Efficient Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#efficient-attention)

vanilla Transformer 的计算和内存成本随序列长度呈二次方增长，因此很难应用于非常长的序列。针对 Transformer 架构的许多效率改进都与自注意力模块有关——使其运行更廉价、更小巧或更快速。参见关于高效 Transformer 的综述论文 (Tay et al. 2020)。

The computation and memory cost of the vanilla Transformer grows quadratically with sequence length, and hence it is hard to apply to very long sequences. Many efficiency improvements for the Transformer architecture have something to do with the self-attention module—making it cheaper, smaller, or faster to run. See the survey paper on Efficient Transformers (Tay et al. 2020).

固定局部上下文 Fixed Local Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#fixed-local-context)

一种降低自注意力成本的简单改动是，将每个 token 的注意力范围限制为仅局部上下文，从而使得自注意力随序列长度线性增长。

A simple alteration to make self-attention less expensive is to restrict the attention span of each token to local context only, so that self-attention grows linearly with the sequence length.

该想法由 Image Transformer（Parmar 等人，2018）提出，它将图像生成建模为使用编码器-解码器 Transformer 架构的序列建模：

The idea was introduced by Image Transformer (Parmar, et al. 2018), which formulates image generation as sequence modeling using an encoder-decoder transformer architecture:

* 编码器生成源图像的上下文相关的逐像素通道表示；

* The encoder generates a contextualized, per-pixel-channel representation of the source image;

* 然后解码器自回归地生成输出图像，每一步生成一个像素的一个通道。

* Then the decoder autoregressively generates an output image, one channel per pixel at each time step.

我们将当前待生成像素的表示记为查询 $\mathbf{q}$。用于计算 $\mathbf{q}$ 的其他位置的表示为键向量 $\mathbf{k}_1, \mathbf{k}_2, \dots$，它们共同构成记忆矩阵 $\mathbf{M}$。$\mathbf{M}$ 的范围定义了像素查询 $\mathbf{q}$ 的上下文窗口。

Let’s label the representation of the current pixel to be generated as the query $\mathbf{q}$. Other positions whose representations will be used for computing $\mathbf{q}$ are key vector $\mathbf{k}_1, \mathbf{k}_2, \dots$ and they together form a memory matrix $\mathbf{M}$. The scope of $\mathbf{M}$ defines the context window for pixel query $\mathbf{q}$.

Image Transformer 引入了两种局部化的 $\mathbf{M}$，如下图所示。

Image Transformer introduced two types of localized $\mathbf{M}$, as illustrated below.

Image Transformer 中视觉输入的 1D 和 2D 注意力范围示意图。黑线标记一个查询块，青色勾勒出像素 q 的实际注意力范围。（图片来源：Parmar 等人 2018 年的论文图 2）

Illustration of 1D and 2D attention span for visual inputs in Image Transformer. The black line marks a query block and the cyan outlines the actual attention span for pixel q. (Image source: Figure 2 in Parmar et al, 2018)

1. **1D 局部注意力**：输入图像按光栅扫描顺序展开，即从左到右、从上到下。线性化后的图像被划分为多个不重叠的查询块。上下文窗口由与 $\mathbf{q}$ 同一查询块中的像素以及该查询块之前生成的固定数量的额外像素组成。

1. **1D Local Attention**: The input image is flattened in the raster scanning order, that is, from left to right and top to bottom. The linearized image is then partitioned into non-overlapping query blocks. The context window consists of pixels in the same query block as $\mathbf{q}$ and a fixed number of additional pixels generated before this query block.

2. **2D 局部注意力**：图像被划分为多个不重叠的矩形查询块。查询像素可以关注同一记忆块中的所有其他像素。为确保左上角的像素也有有效的上下文窗口，记忆块分别向上、左、右扩展固定量。

2. **2D Local Attention**: The image is partitioned into multiple non-overlapping rectangular query blocks. The query pixel can attend to all others in the same memory blocks. To make sure the pixel at the top-left corner can also have a valid context window, the memory block is extended to the top, left and right by a fixed amount, respectively.

跨步上下文 Strided Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#strided-context)

Sparse Transformer（Child 等人，2019）引入了**因子化自注意力**机制，通过稀疏矩阵分解，使得在高达 16,384 的序列长度上训练具有数百层的稠密注意力网络成为可能，而这一做法在现代硬件上原本是不可行的。

Sparse Transformer (Child et al., 2019) introduced *factorized self-attention* via sparse matrix factorization, making it possible to train dense attention networks with hundreds of layers on sequence lengths up to 16,384, which would otherwise be infeasible on modern hardware.

给定一组注意力连接模式 $\mathcal{S} = \{S_1, \dots, S_n\}$，其中每个 $S_i$ 表示第 $i$ 个查询向量所关注的键位置集合。

Given a set of attention connectivity patterns $\mathcal{S} = \{S_1, \dots, S_n\}$, where each $S_i$ denotes the set of key positions that the $i$-th query vector attends to.

\[\begin{aligned} \text{Attend}(\mathbf{X}, \mathcal{S}) &= \Big( a(\mathbf{x}_i, S_i) \Big)_{i \in \{1, \dots, L\}} \\ \text{where } a(\mathbf{x}_i, S_i) &= \text{softmax}\Big( \frac{(\mathbf{x}_i \mathbf{W}^q)(\mathbf{x}_j \mathbf{W}^k)_{j \in S_i}^{\top}}{\sqrt{d_k}} \Big) (\mathbf{x}_j \mathbf{W}^v)_{j \in S_i} \end{aligned}\]

注意，尽管 $S_i$ 的大小不固定，$a(\mathbf{x}_i, S_i)$ 的维度始终为 $d_v$，因此 $\text{Attend}(\mathbf{X}, \mathcal{S}) \in \mathbb{R}^{L \times d_v}$。

Note that although the size of $S_i$ is not fixed, $a(\mathbf{x}_i, S_i)$ always has dimension $d_v$, and therefore $\text{Attend}(\mathbf{X}, \mathcal{S}) \in \mathbb{R}^{L \times d_v}$.

在自回归模型中，注意力跨度定义为 $S_i = \{j : j \leq i\}$，因为这允许每个词元关注过去的所有位置。

In autoregressive models, an attention span is defined as $S_i = \{j : j \leq i\}$, since it allows each token to attend to all previous positions.

在分解式自注意力中，集合 $S_i$ 被分解为一棵 *依赖树*，使得对于每一对满足 $j \leq i$ 的 $i, j$，都存在一条从 $i$ 回到 $j$ 的路径，并且 $i$ 可以直接或间接地关注到 $j$。

In factorized self-attention, the set $S_i$ is decomposed into a *tree* of dependencies, such that for every pair $i, j$ with $j \leq i$, there is a path connecting $i$ back to $j$, and $i$ can attend to $j$ either directly or indirectly.

更准确地说，集合 $S_i$ 被划分为 $p$ 个 *非重叠* 子集，其中第 $m$ 个子集记为 $A_i^{(m)} \subset S_i, m = 1, ..., p$。因此，输出位置 $i$ 与任意 $j$ 之间的路径最大长度为 $p + 1$。例如，如果 $j, a, b, c, ..., i$ 是 $i$ 和 $j$ 之间的一条索引路径，那么会有 $j \in A_a^{(1)}$、$a \in A_b^{(2)}$、$b \in A_c^{(3)}$，依此类推。

More precisely, the set $S_i$ is divided into $p$ *non-overlapping* subsets, where the $m$-th subset is denoted as $A_i^{(m)} \subset S_i, m = 1, ..., p$. Therefore, the path between the output position $i$ and any $j$ has a maximum length $p + 1$. For example, if $j, a, b, c, ..., i$ is a path of indices between $i$ and $j$, then $j \in A_a^{(1)}$, $a \in A_b^{(2)}$, $b \in A_c^{(3)}$, and so on.

Sparse Transformer 提出了两种分解式注意力。通过图 10 中以 2D 图像输入为例，可以更容易地理解这些概念。

Sparse Transformer proposed two types of factorized attention. It is easier to understand the concepts as illustrated in Fig. 10 with 2D image inputs as examples.

顶行展示了 (a) Transformer、(b) 使用跨步注意力的 Sparse Transformer 以及 (c) 使用固定注意力的 Sparse Transformer 中的注意力连接模式。底行包含对应的自注意力连接矩阵。请注意，顶行和底行不在同一尺度上。（图片来源：Child 等人，2019，另加少量注释。）

The top row illustrates the attention connectivity patterns in (a) Transformer, (b) Sparse Transformer with strided attention, and (c) Sparse Transformer with fixed attention. The bottom row contains the corresponding self-attention connectivity matrices. Note that the top and bottom rows are not in the same scale. (Image source: Child et al., 2019 + a few extra annotations.)

1. **跨步注意力**，步长为 $\ell \sim \sqrt{n}$。这种方式对图像数据很有效，因为图像的结构与步长对齐。在图像场景中，每个像素会按光栅扫描顺序关注其前面的 $\ell$ 个像素（自然覆盖整个图像宽度），然后这些像素再关注同一列中的其他像素（由另一个注意力连接子集定义）。

1. *Strided* attention with stride $\ell \sim \sqrt{n}$. This works well with image data, as the structure is aligned with strides. In the image case, each pixel would attend to all the previous $\ell$ pixels in raster scanning order (naturally covering the entire width of the image), and then those pixels would attend to other pixels in the same column (defined by another attention connectivity subset).

\[ \begin{aligned} A_i^{(1)} &= \{ t, t+1, \dots, i\} \text{，其中 } t = \max(0, i - \ell) \\ A_i^{(2)} &= \{j: (i-j) \mod \ell = 0\} \end{aligned} \]

\[ \begin{aligned} A_i^{(1)} &= \{ t, t+1, \dots, i\} \text{, where } t = \max(0, i - \ell) \\ A_i^{(2)} &= \{j: (i-j) \mod \ell = 0\} \end{aligned} \]

2. **Fixed（固定）注意力机制**。一小部分词元汇总先前位置的信息，并将该信息传播到所有未来位置。

2. _Fixed_ attention. A small set of tokens summarize previous locations and propagate that information to all future locations.

\[ \begin{aligned} A_i^{(1)} &= \{j: \lfloor \frac{j}{\ell} \rfloor = \lfloor \frac{i}{\ell} \rfloor \} \\ A_i^{(2)} &= \{j: j \mod \ell \in \{\ell-c, \dots, \ell-1\} \} \end{aligned} \]

其中 $c$ 是一个超参数。如果 $c=1$，它会限制表示，而许多位置依赖于少数几个位置。该论文为 $\ell \in \{ 128, 256 \}$ 选择了 $c \in \{ 8, 16, 32 \}$。

where $c$ is a hyperparameter. If $c=1$, it restricts the representation whereas many depend on a few positions. The paper chose $c \in \{ 8, 16, 32 \}$ for $\ell \in \{ 128, 256 \}$.

在 Transformer 中使用因子化自注意力机制

Use Factorized Self-Attention in Transformer

在 Transformer 架构中，使用稀疏因子化注意力模式有三种方式：

There are three ways to use sparse factorized attention patterns in Transformer architecture:

1. 每个残差块使用一种注意力类型，然后交替使用这些类型，

1. One attention type per residual block and then interleave them,

$\text{attn}(\mathbf{X}) = \text{Attend}(\mathbf{X}, A^{(n \mod p)}) \mathbf{W}^o$，其中 n 是当前残差块的索引。

$\text{attn}(\mathbf{X}) = \text{Attend}(\mathbf{X}, A^{(n \mod p)}) \mathbf{W}^o$, where n is the index of the current residual block.

2. 设置一个单一头，使其关注所有因子化头所关注的位置，

2. Set up a single head which attends to locations that all the factorized heads attend to,

$\text{attn}(\mathbf{X}) = \text{Attend}(\mathbf{X}, \cup_{m=1}^p A^{(m)}) \mathbf{W}^o$。

$\text{attn}(\mathbf{X}) = \text{Attend}(\mathbf{X}, \cup_{m=1}^p A^{(m)}) \mathbf{W}^o$.

3. Use a multi-head attention mechanism, but different from vanilla Transformer, each head might adopt a pattern presented above, 1 or 2. $\rightarrow$ This option often performs the best.

Sparse Transformer also proposed a set of changes so as to train the Transformer up to hundreds of layers, including gradient checkpointing, recomputing attention & FF layers during the backward pass, mixed precision training, efficient block-sparse implementation, etc. Please check the paper for more details or my previous post on techniques for scaling up model training.

Blockwise Attention (Qiu et al. 2019) introduces a _sparse block matrix_ to only allow each token to attend to a small set of other tokens. Each attention matrix of size $L \times L$ is partitioned into $n \times n$ smaller blocks of size $\frac{L}{n}\times\frac{L}{n}$ and a sparse block matrix $\mathbf{M} \in \{0, 1\}^{L \times L}$ is defined by a permutation $\pi$ of ${1, \dots, n}$, which records the column index per row in the block matrix.

$$ \begin{aligned} \text{attn}(\mathbf{Q}, \mathbf{K}, \mathbf{V}, \mathbf{M}) &= \text{softmax}\Big(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}} \odot \mathbf{M}\Big)\mathbf{V} \\ (\mathbf{A} \odot \mathbf{M})_{ij} &= \begin{cases} A_{ij} & \text{if }M_{ij} = 1 \\ -\infty & \text{if }M_{ij} = 0 \\ \end{cases} \\ \text{where } M_{ij} &= \begin{cases} 1 & \text{if }\pi\big(\lfloor\frac{(i-1)n}{L} + 1\rfloor\big) = \lfloor\frac{(j-1)n}{L} + 1\rfloor \\ 0 & \text{otherwise} \end{cases} \end{aligned} $$

The actual implementation of Blockwise Attention only stores QKV as block matrices, each of size $n\times n$:

\[ \text{Blockwise-attn}(\mathbf{Q}, \mathbf{K}, \mathbf{V}, \mathbf{M}) = \begin{bmatrix} \text{softmax}\big(\frac{\hat{\mathbf{q}}_1 \hat{\mathbf{k}}_{\pi(1)}^\top}{\sqrt{d}} \odot \mathbf{M}\Big) \hat{\mathbf{v}}_{\pi(1)} \\ \vdots \\ \text{softmax}\big(\frac{\hat{\mathbf{q}}_n \hat{\mathbf{k}}_{\pi(n)}^\top}{\sqrt{d}} \odot \mathbf{M}\Big) \hat{\mathbf{v}}_{\pi(n)} \end{bmatrix} \]

其中 $\hat{\mathbf{q}}_i$、$\hat{\mathbf{k}}_i$ 和 $\hat{\mathbf{v}}_i$ 分别是 QKV 分块矩阵中的第 $i$ 行。每个 $\hat{\mathbf{q}}_i \hat{\mathbf{k}}_{\pi(i)}^\top$（$i = 1, \dots, n$）的大小为 $\frac{N}{n} \times \frac{N}{n}$，因此 Blockwise Attention（分块注意力）能够将注意力矩阵的内存复杂度从 $\mathcal{O}(L^2)$ 降低到 $\mathcal{O}(\frac{L}{n} \times \frac{L}{n} \times n) = \mathcal{O}(L^2/n)$。

where $\hat{\mathbf{q}}_i$, $\hat{\mathbf{k}}_i$, and $\hat{\mathbf{v}}_i$ are the $i$-th rows in the QKV block matrices, respectively. Each $\hat{\mathbf{q}}_i \hat{\mathbf{k}}_{\pi(i)}^\top$, $\forall i = 1, \dots, n$, is of size $\frac{N}{n} \times \frac{N}{n}$, and therefore Blockwise Attention is able to reduce the memory complexity of the attention matrix from $\mathcal{O}(L^2)$ to $\mathcal{O}(\frac{L}{n} \times \frac{L}{n} \times n) = \mathcal{O}(L^2/n)$.

局部与全局上下文的组合 Combination of Local and Global Context#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#combination-of-local-and-global-context)

ETC（“扩展 Transformer 构造”；Ainslie 等，2019）、Longformer（Beltagy 等，2020）和 Big Bird（Zaheer 等，2020）模型在构建注意力矩阵时同时结合局部与全局上下文。所有这些模型都可以从现有预训练模型初始化。

ETC (_Extended Transformer Construction_; Ainslie et al. 2019), Longformer (Beltagy et al. 2020) and Big Bird (Zaheer et al. 2020) models combine both local and global context when building an attention matrix. All these models can be initialized from existing pretrained models.

ETC 的全局-局部注意力（Ainslie 等，2019）接收两个输入：（1）长度 $n_l$ 的长输入 $\mathbf{x}^l$，即常规输入序列；（2）长度 $n_g$ 的全局输入 $\mathbf{x}^g$，其中包含少量辅助 token，$n_g \ll n_l$。因此，注意力根据这两个输入之间的方向性注意力拆分为四个部分：g2g、g2l、l2g 和 l2l。由于 l2l 注意力块可能非常大，它被限制为半径 $w$ 的固定大小注意力跨度（即局部注意力跨度），l2l 矩阵可重塑为 $n_l \times (2w+1)$。

Global-Local Attention of _ETC_ (Ainslie et al. 2019) takes two inputs: (1) the long input $\mathbf{x}^l$ of size $n_l$, which is the regular input sequence, and (2) the global input $\mathbf{x}^g$ of size $n_g$, which contains a smaller number of auxiliary tokens, $n_g \ll n_l$. Attention is thus split into four components based on directional attention across these two inputs: g2g, g2l, l2g, and l2l. Because the l2l attention piece can be very large, it is restricted to a fixed-size attention span of radius $w$ (i.e., local attention span), and the l2l matrix can be reshaped to $n_l \times (2w+1)$.

ETC 使用四个二值矩阵来处理结构化输入：$\mathbf{M}^{g2g}$、$\mathbf{M}^{g2l}$、$\mathbf{M}^{l2g}$ 和 $\mathbf{M}^{l2l}$。例如，g2g 注意力部分中注意力输出 $z^g = (z^g_1, \dots, z^g_{n_g})$ 的每个元素 $z^g_i \in \mathbb{R}^d$ 格式化为：

ETC utilizes four binary matrices to handle structured inputs, $\mathbf{M}^{g2g}$, $\mathbf{M}^{g2l}$, $\mathbf{M}^{l2g}$, and $\mathbf{M}^{l2l}$. For example, each element $z^g_i \in \mathbb{R}^d$ in the attention output $z^g = (z^g_1, \dots, z^g_{n_g})$ for the g2g attention piece is formatted as:

\[ \begin{aligned} a^{g2g}_{ij} = \frac{1}{\sqrt{d}} x^g_i \mathbf{W}^Q (x^g_j \mathbf{W}^K + P^K_{ij})^\top - (1- M^{g2g}_{ij})C \\ A^{g2g}_{ij} = \frac{\exp(a^{g2g}_{ij})}{\sum_{k=1}^{n_g} \exp(a^{g2g}_{ik})} \quad z^g_i = \sum^{n_g}_{j=1} A^{g2g}_{ij} x^g_j \mathbf{W}^V \end{aligned} \]

其中 $P^K_{ij}$ 是用于相对位置编码的可学习向量，$C$ 是一个非常大的常数（论文中 $C=10000$），用于在掩码关闭时抵消任意注意力权重。

where $P^K_{ij}$ is a learnable vector for relative position encoding, and $C$ is a very large constant ($C=10000$ in the paper) to offset any attention weights when the mask is off.

ETC、Longformer 和 Big Bird 的注意力模式。

Attention patterns of ETC, Longformer, and Big Bird.

ETC 的另一个更新是在预训练阶段除了 MLM 任务之外还引入了使用 NCE 损失的 CPC（对比预测编码）任务：当一个句子被掩码时，该句子的表示应与其周围上下文的表示相似。

Another update in ETC is to incorporate a CPC (contrastive predictive coding) task using NCE loss into the pretraining stage, in addition to the MLM task: The representation of one sentence should be similar to the representation of the context around it when that sentence is masked.

ETC 的全局输入 $\mathbf{x}^g$ 构造如下：假设长输入中存在一些分段（例如按句子划分），每个分段附加一个辅助 token 来学习全局输入。使用相对位置编码，用 token 位置标记全局分段 token。单向硬掩码（即对位于其前和其后的 token 分别做不同标记）在一些数据集上能带来性能提升。

The global input $\mathbf{x}^g$ for ETC is constructed as follows: Assuming there are some segments within the long inputs (e.g., by sentence), each segment is attached with an auxiliary token to learn global inputs. Relative position encoding is used to mark the global segment tokens with their token positions. One-directional hard masking (i.e., labeling tokens before vs. after differently) is found to bring performance gains in some datasets.

Longformer 中的注意力模式包含三个组成部分：

The attention pattern in Longformer contains three components:

1. **局部注意力**：与 ETC 类似，局部注意力由固定大小 $w$ 的滑动窗口控制；

1. *Local attention*: Similar to ETC, local attention is controlled by a sliding window of fixed size $w$;

2. *预选词的全局注意力*：Longformer 中有少量预选词（如 [CLS] token）被赋予全局注意力范围，即它们会关注输入序列中的所有其他 token。

2. *Global attention of preselected tokens*: Longformer has a few pre-selected tokens (e.g. [CLS] token) assigned with global attention span, that is, attending to all other tokens in the input sequence.

3. *扩张注意力*：固定大小 $r$ 的扩张滑动窗口，扩张间隙大小为 $d$，类似于 Sparse Transformer；

3. *Dilated attention*: Dilated sliding window of fixed size $r$ and gaps of dilation size $d$, similar to Sparse Transformer;

*Big Bird* 与 Longformer 非常相似，同时具备局部注意力和少量具有全局注意力范围的预选词，但 Big Bird 用一种新机制取代了扩张注意力：所有 token 都关注一组随机 token。这一设计的动机是，注意力模式可以被看作一个有向图，而随机图具有信息能够在任意两个节点之间快速流动的特性。

*Big Bird* is quite similar to Longformer, equipped with both local attention and a few preselected tokens with global attention span, but Big Bird replaces dilated attention with a new mechanism where all tokens attend to a set of random tokens. The design is motivated by the fact that attention pattern can be viewed as a directed graph and a random graph has the property that information is able to rapidly flow between any pair of nodes.

*Longformer* 在较低层使用较小的窗口大小，在较高层使用较大的窗口大小。消融研究表明，这种设置比相反或固定大小的配置效果更好。较低层不使用扩张滑动窗口，以便更好地学习利用直接局部上下文。Longformer 还采用分阶段训练过程：最初用较小的窗口大小训练模型以学习局部上下文，随后各训练阶段会增大窗口大小并降低学习率。

*Longformer* uses smaller window size at lower layers and larger window sizes at higher layers. Ablation studies showed that this setup works better than reversed or fixed size config. Lower layers do not have dilated sliding windows to better learn to use immediate local context. Longformer also has a staged training procedure where initially the model is trained with small window size to learn from local context and then subsequent stages of training have window sizes increased and learning rate decreased.

基于内容的注意力 Content-based Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#content-based-attention)

Reformer（Kitaev 等人，2020）提出的改进旨在解决 vanilla Transformer 中的以下痛点：

The improvements proposed by Reformer (Kitaev et al., 2020) aim to solve the following pain points in the vanilla Transformer:

自注意力模块中的二次时间复杂度和内存复杂度。

Quadratic time and memory complexity within the self-attention module.

在具有 $N$ 层的模型中，内存占用是单层模型的 $N$ 倍，因为需要存储用于反向传播的激活值。

Memory in a model with $N$ layers is $N$ times larger than in a single-layer model because we need to store activations for back-propagation.

中间的 FF 层通常相当大。

The intermediate FF layers are often quite large.

用**局部敏感哈希（LSH）注意力**替换点积注意力，将复杂度从 $\mathcal{O}(L^2)$ 降低到 $\mathcal{O}(L \log L)$。

Replace the dot-product attention with *locality-sensitive hashing (LSH) attention*, reducing the complexity from $\mathcal{O}(L^2)$ to $\mathcal{O}(L \log L)$.

2. 用*可逆残差层*替代标准残差块，这样在训练时只需存储一次激活，而不是$N$次（即与层数成正比）。

2. Replace the standard residual blocks with *reversible residual layers*, which allows storing activations only once during training instead of $N$ times (i.e. proportional to the number of layers).

在注意力公式的$\mathbf{Q}\mathbf{K}^\top$部分，我们只对最大的元素感兴趣，因为只有大的元素在 softmax 之后贡献很大。对于每个查询$\mathbf{q}_i \in \mathbf{Q}$，我们寻找$\mathbf{K}$中与$\mathbf{q}_i$最近的行向量。为了在高维空间中快速找到最近邻居，Reformer 将局部敏感哈希（LSH）引入其注意力机制。

In $\mathbf{Q}\mathbf{K}^\top$ part of the attention formula, we are only interested in the largest elements as only large elements contribute a lot after softmax. For each query $\mathbf{q}_i \in \mathbf{Q}$, we are looking for row vectors in $\mathbf{K}$ closest to $\mathbf{q}_i$. In order to find nearest neighbors quickly in high-dimensional space, Reformer incorporates Locality-Sensitive Hashing (LSH) into its attention mechanism.

如果哈希方案$x \mapsto h(x)$保持数据点之间的距离信息，使得相近的向量获得相似的哈希值，而遥远的向量获得非常不同的哈希值，则它是*局部敏感*的。Reformer 采用了这样的哈希方案：给定一个固定的随机矩阵$\mathbf{R} \in \mathbb{R}^{d \times b/2}$（其中$b$是超参数），哈希函数为$h(x) = \arg\max([xR; −xR])$。

A hashing scheme $x \mapsto h(x)$ is *locality-sensitive* if it preserves the distancing information between data points, such that close vectors obtain similar hashes while distant vectors have very different ones. The Reformer adopts a hashing scheme as such, given a fixed random matrix $\mathbf{R} \in \mathbb{R}^{d \times b/2}$ (where $b$ is a hyperparam), the hash function is $h(x) = \arg\max([xR; −xR])$.

局部敏感哈希（LSH）注意力示意图。（图片来源：Kitaev 等人，2020 年图 1 的右半部分。）

Illustration of Locality-Sensitive Hashing (LSH) attention. (Image source: right part of Figure 1 in Kitaev, et al. 2020).

在 LSH 注意力中，查询只能关注同一哈希桶中的位置，$S_i = \{j: h(\mathbf{q}_i) = h(\mathbf{k}_j)\}$。这个过程如图 20 所示：

In LSH attention, a query can only attend to positions in the same hashing bucket, $S_i = \{j: h(\mathbf{q}_i) = h(\mathbf{k}_j)\}$. It is carried out in the following process, as illustrated in Fig. 20:

* (a) 完整注意力的注意力矩阵通常是稀疏的。

* (a) The attention matrix for full attention is often sparse.

* (b) 使用 LSH，我们可以将键和查询按照它们的哈希桶进行排序对齐。

* (b) Using LSH, we can sort the keys and queries to be aligned according to their hash buckets.

* (c) 设置 $\mathbf{Q} = \mathbf{K}$（准确地说 $\mathbf{k}_j = \mathbf{q}_j / |\mathbf{q}_j|$），这样每个桶中键和查询的数量相等，便于批处理。有趣的是，这种“共享-QK”配置并不影响 Transformer 的性能。

* (c) Set $\mathbf{Q} = \mathbf{K}$ (precisely $\mathbf{k}_j = \mathbf{q}_j / |\mathbf{q}_j|$), so that there are equal numbers of keys and queries in one bucket, easier for batching. Interestingly, this “shared-QK” config does not affect the performance of the Transformer.

* (d) 应用批处理，将连续的 $m$ 个查询分为一组。

* (d) Apply batching where chunks of $m$ consecutive queries are grouped together.

LSH 注意力包含 4 个步骤：分桶、排序、分块和注意力计算。（图片来源：Kitaev 等人 2020 年论文中图 1 的左侧部分。）

The LSH attention consists of 4 steps: bucketing, sorting, chunking, and attention computation. (Image source: left part of Figure 1 in Kitaev, et al. 2020).

Reformer 的另一项改进是使用 **可逆残差层**（Gomez 等人，2017）。可逆残差网络的动机是设计这样一种架构：给定某一层的激活，可以利用模型参数从下一层的激活中恢复出该层的激活。因此，我们可以在反向传播时重新计算激活值，而不是存储所有激活值，从而节省内存。

Another improvement by Reformer is to use **reversible residual layers** (Gomez et al. 2017). The motivation for reversible residual network is to design the architecture in a way that activations at any given layer can be recovered from the activations at the following layer, using only the model parameters. Hence, we can save memory by recomputing the activation during backprop rather than storing all the activations.

给定一层 $x \mapsto y$，普通残差层执行 $y = x + F(x)$，而可逆层将输入和输出都拆分为成对形式 $(x_1, x_2) \mapsto (y_1, y_2)$，然后执行以下操作：

Given a layer $x \mapsto y$, the normal residual layer does $y = x + F(x)$, but the reversible layer splits both input and output into pairs $(x_1, x_2) \mapsto (y_1, y_2)$ and then executes the following:

\[ y_1 = x_1 + F(x_2), \quad y_2 = x_2 + G(y_1) \]

\[ x_2 = y_2 - G(y_1), \quad x_1 = y_1 - F(x_2) \]

Reformer 将同一思想应用于 Transformer，在一个可逆网络块中组合注意力（\F\）和前馈层（\G\）：

Reformer applies the same idea to Transformer by combining attention (\F\) and feed-forward layers (\G\) within a reversible net block:

\[ Y_1 = X_1 + \text{Attention}(X_2), \; Y_2 = X_2 + \text{FeedForward}(Y_1) \]

通过将前馈计算分块，可以进一步减少内存使用：

The memory can be further reduced by chunking the feed-forward computation:

\[ Y_2 = [Y_2^{(1)}; \dots; Y_2^{(c)}] = [X_2^{(1)} + \text{FeedForward}(Y_1^{(1)}); \dots; X_2^{(c)} + \text{FeedForward}(Y_1^{(c)})] \]

由此得到的可逆 Transformer 无需在每一层存储激活值。

The resulting reversible Transformer does not need to store activation in every layer.

Routing Transformer（Roy 等人，2021）同样采用基于内容的键和查询聚类。它没有使用像 LSH 这样的静态哈希函数，而是利用在线 $k$-均值聚类，并将其与局部时间稀疏注意力相结合，从而将注意力复杂度从 $O(L^2)$ 降低到 $O(L^{1.5})$。

Routing Transformer (Roy et al. 2021) is also built on content-based clustering of keys and queries. Instead of using a static hashing function like LSH, it utilizes online $k$-means clustering and combines it with local, temporal sparse attention to reduce the attention complexity from $O(L^2)$ to $O(L^{1.5})$.

在路由注意力中，键和查询都使用\k\均值聚类方法以及同一组质心$\boldsymbol{\mu} = (\mu_1, \dots, \mu_k) \in \mathbb{R}^{k \times d}$进行聚类。查询被路由到分配给同一质心的键。总复杂度为\O(Lkd + L^2d/k)\，其中\O(Lkd)\用于运行聚类分配，\O(L^2d/k)\用于注意力计算。聚类质心通过 EMA（指数移动平均）使用所有关联的键和查询进行更新。

Within routing attention, both keys and queries are clustered using the \k\-means clustering method and the same set of centroids $\boldsymbol{\mu} = (\mu_1, \dots, \mu_k) \in \mathbb{R}^{k \times d}$. Queries are routed to keys that get assigned to the same centroid. The total complexity is \O(Lkd + L^2d/k)\, where \O(Lkd)\ is for running clustering assignments and \O(L^2d/k)\ is for attention computation. The cluster centroids are updated by EMA (exponential moving average) using all associated keys and queries.

在 Routing Transformer 的实验中，一些最佳配置仅在模型的最后两层和一半的注意力头中启用路由注意力，而另一半则使用局部注意力。他们还观察到，局部注意力是一个相当强的基线，并且更大的注意力窗口总是会带来更好的结果。

In the experiments for Routing Transformer, some of the best configurations only have routing attention enabled in the last two layers of the model and in half of the attention heads, while the other half utilize local attention. They also observed that local attention is a pretty strong baseline and that a larger attention window always leads to better results.

低秩注意力 Low-Rank Attention#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#low-rank-attention)

Linformer（Wang 等，2020）用 *低秩* 矩阵近似完整的注意力矩阵，将时间与空间复杂度降低为 *线性*。Linformer 没有使用昂贵的 SVD 来识别低秩分解，而是为键矩阵和值矩阵分别添加两个线性投影 $\mathbf{E}_i, \mathbf{F}_i \in \mathbb{R}^{L \times k}$，将其维度从 $L \times d$ 降到 $k \times d$。只要 $k \ll L$，注意力内存就能大大减少。

Linformer (Wang et al., 2020) approximates the full attention matrix with a *low-rank* matrix, reducing the time and space complexity to be *linear*. Instead of using expensive SVD to find the low-rank decomposition, Linformer adds two linear projections $\mathbf{E}_i, \mathbf{F}_i \in \mathbb{R}^{L \times k}$ for the key and value matrices, respectively, reducing their dimensions from $L \times d$ to $k \times d$. As long as $k \ll L$, the attention memory can be greatly reduced.

\[ \begin{aligned} \overline{\text{head}}_i &= \text{attn}(\mathbf{X}_q\mathbf{W}^q_i, \mathbf{E}_i\mathbf{X}_k\mathbf{W}^k_i, \mathbf{F}_i\mathbf{X}_v\mathbf{W}^v_i) \\ &= \underbrace{\text{softmax}\Big( \frac{\mathbf{X}_q\mathbf{W}^q_i (\mathbf{E}_i \mathbf{X}_k\mathbf{W}^k_i)^\top}{\sqrt{d}} \Big)}_{\text{low rank attention matrix }\bar{A} \in \mathbb{R}^{k \times d}} \mathbf{F}_i \mathbf{X}_v\mathbf{W}^v_i \end{aligned} \]

可以应用其他技术进一步提升 Linformer 的效率： * 在投影层之间共享参数，例如按头共享、键值共享和逐层（跨所有层）共享。 * 在不同层使用不同的 $k$，因为较高层的头往往具有更偏斜的分布（更低秩），因此可以在较高层使用较小的 $k$。

Additional techniques can be applied to further improve the efficiency of Linformer: * Parameter sharing between projection layers, such as head-wise, key-value, and layer-wise (across all layers) sharing. * Use different $k$ at different layers, as heads in higher layers tend to have a more skewed distribution (lower rank) and thus we can use smaller $k$ at higher layers.

Linformer (Wang et al. 2020) approximates the full attention matrix with a _low rank_ matrix, reducing the time & space complexity to be _linear_. Instead of using expensive SVD to identify low rank decomposition, Linformer adds two linear projections $\mathbf{E}_i, \mathbf{F}_i \in \mathbb{R}^{L \times k}$ for key and value matrices, respectively, reducing their dimensions from $L \times d$ to $k \times d$. As long as $k \ll L$, the attention memory can be greatly reduced.

$$ \begin{aligned} \overline{\text{head}}_i &= \text{attn}(\mathbf{X}_q\mathbf{W}^q_i, \mathbf{E}_i\mathbf{X}_k\mathbf{W}^k_i, \mathbf{F}_i\mathbf{X}_v\mathbf{W}^v_i) \\ &= \underbrace{\text{softmax}\Big( \frac{\mathbf{X}_q\mathbf{W}^q_i (\mathbf{E}_i \mathbf{X}_k\mathbf{W}^k_i)^\top}{\sqrt{d}} \Big)}_{\text{low rank attention matrix }\bar{A} \in \mathbb{R}^{k \times d}} \mathbf{F}_i \mathbf{X}_v\mathbf{W}^v_i \end{aligned} $$

使用不同类型的投影；例如，均值/最大池化、核与步长为 \frac{L}{k}\ 的卷积层。

- Use different types of projections; e.g., mean/max pooling, convolution layer with kernel and stride \frac{L}{k}\.

（左）Informer 为键和值添加了两个投影层。（右）推理时间随序列长度变化的曲线图。（图片来源：Wang et al., 2020）。

(Left) Informer has two projection layers added for keys and values. (Right) Plot of inference time as a function of sequence length. (Image source: Wang et al., 2020).

随机特征注意力（RFA；Peng et al., 2021）依赖 _随机特征方法_（Rahimi & Recht, 2007）来近似自注意力中的 softmax 操作，使用低秩特征映射以实现线性时间和空间复杂度。Performers（Choromanski et al., 2021）也采用了随机特征注意力，并改进了核构造以进一步减小核近似误差。

Random Feature Attention (RFA; Peng et al., 2021) relies on _random feature methods_ (Rahimi & Recht, 2007) to approximate the softmax operation in self-attention with low-rank feature maps, achieving linear time and space complexity. Performers (Choromanski et al., 2021) also adopt random feature attention, with improvements to the kernel construction that further reduce the kernel approximation error.

RFA 背后的主要定理来自 Rahimi 和 Recht（2007）：

The main theorem behind RFA is from Rahimi & Recht, 2007:

\exp(\mathbf{x} \cdot \mathbf{y})\ 的一个无偏估计量是：

An unbiased estimator of \exp(\mathbf{x} \cdot \mathbf{y})\ is:

\[ \begin{aligned} \exp(\mathbf{x} \cdot \mathbf{y} / \sigma^2) &= \exp(\frac{1}{2\sigma^2}(\|\mathbf{x}\|^2 + \|\mathbf{y}\|^2 - \|\mathbf{x} - \mathbf{y}\|^2)) \\ &= \exp(\frac{\|\mathbf{x}\|^2}{2\sigma^2}) \exp(\frac{\|\mathbf{y}\|^2}{2\sigma^2}) \exp(-\frac{\|\mathbf{x} - \mathbf{y}\|^2}{2\sigma^2}) \\ &\approx \exp(\frac{\|\mathbf{x}\|^2}{2\sigma^2}) \exp(\frac{\|\mathbf{y}\|^2}{2\sigma^2}) \phi(\mathbf{x})\cdot\phi(\mathbf{y}) \\ &= \exp(\frac{1}{\sigma^2}) \phi(\mathbf{x})\cdot\phi(\mathbf{y}) & \text{; unit vectors} \end{aligned} \]

然后我们可以将注意力函数写成如下形式，其中 $\otimes$ 是外积运算，$\sigma^2$ 是温度：

Then we can write the attention function as follows, where $\otimes$ is the outer product operation and $\sigma^2$ is the temperature:

\[ \begin{aligned} \text{attn}(\mathbf{q}_t, {\mathbf{k}_i}, {\mathbf{v}_i}) &= \sum_i \frac{\exp(\mathbf{q}_t\cdot\mathbf{k}_i/\sigma^2)}{\sum_j \exp(\mathbf{q}_t\cdot\mathbf{k}_j/\sigma^2)}\mathbf{v}_i^\top \approx \sum_i \frac{\phi(\mathbf{q}_t)\phi(\mathbf{k}_i)\mathbf{v}_i^\top}{\sum_j \phi(\mathbf{q}_t)\phi(\mathbf{k}_j)} \\ &= \frac{\phi(\mathbf{q}_t)^\top \sum_i \phi(\mathbf{k}_i)\otimes\mathbf{v}_i}{\phi(\mathbf{q}_t)^\top \sum_j \phi(\mathbf{k}_j)} = \text{RFA}(\mathbf{q}_t, {\mathbf{k}_i}, {\mathbf{v}_i}) \end{aligned} \]

（左）默认 softmax 操作的计算顺序。（右）使用随机特征注意力时的计算顺序，比默认 softmax 便宜得多。（图片来源：Peng et al. 2021。）

(Left) The order of computation for the default softmax operation. (Right) The order of computation when using random feature attention, which is much cheaper than the default softmax. (Image source: Peng et al. 2021).

在因果注意力 RFA 中，时间步 $t$ 的词元只能关注更早的键和值 ${\mathbf{k}_i}_{i \leq t}, {\mathbf{v}_i}_{i \leq t}$。让我们使用一个变量元组 $(\mathbf{S}_t \in \mathbb{R}^{2D \times d}, \mathbf{z} \in \mathbb{R}^{2D})$ 来跟踪时间步 $t$ 的隐状态历史，类似于 RNN：

In Causal Attention RFA, a token at time step $t$ only attends to earlier keys and values ${\mathbf{k}_i}_{i \leq t}, {\mathbf{v}_i}_{i \leq t}$. Let us use a tuple of variables, $(\mathbf{S}_t \in \mathbb{R}^{2D \times d}, \mathbf{z} \in \mathbb{R}^{2D})$, to track the hidden state history at time step $t$, similar to RNNs:

\[ \begin{aligned} &\text{causal-RFA}(\mathbf{q}_t, \{\mathbf{k}_i\}_{i \leq t}, \{\mathbf{v}_i\}_{i \leq t}) = \frac{\phi(\mathbf{q}_t)^\top \mathbf{S}_t}{\phi(\mathbf{q}_t) \cdot \mathbf{z}_t} \\ &\text{where } \mathbf{S}_t = \mathbf{S}_{t-1} + \phi(\mathbf{k}_t)\otimes\mathbf{v}_t, \quad \mathbf{z}_t = \mathbf{z}_{t-1} + \phi(\mathbf{k}_t) \end{aligned} \]

其中 $2D$ 是 $\phi(.)$ 的大小，为了合理的近似，$D$ 不应小于模型大小 $d$。

where $2D$ is the size of $\phi(.)$ and $D$ should be no less than the model size $d$ for reasonable approximation.

RFA 在自回归解码中显著加速，内存复杂度主要取决于构造核 $\phi(.)$ 时对 $D$ 的选择。

RFA leads to significant speedup in autoregressive decoding and the memory complexity mainly depends on the choice of $D$ when constructing the kernel $\phi(.)$.

Performer 用正随机特征映射修改随机特征注意力，以减少估计误差。它还使随机采样的 $\mathbf{w}_1, \dots, \mathbf{w}_D$ 保持正交，以进一步降低估计器的方差。

Performer modifies the random feature attention with positive random feature maps to reduce the estimation error. It also keeps the randomly sampled $\mathbf{w}_1, \dots, \mathbf{w}_D$ to be orthogonal to further reduce the variance of the estimator.

使用（左）i.i.d. 与正交特征以及（右）sin/cos 与正随机特征时的近似误差比较。（图片来源：Choromanski 等人，2021。）

Comparison of approximation error when using (Left) i.i.d vs orthogonal features and (Right) sin/cos vs positive random features. (Image source: Choromanski et al. 2021).

强化学习中的 Transformer Transformers for Reinforcement Learning#(https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/#transformers-for-reinforcement-learning)

自注意力机制避免了将整个历史压缩到固定大小的隐藏状态中，并且与 RNN 相比，它不那么容易出现梯度消失或梯度爆炸。强化学习任务当然可以从这些特性中受益。然而，即使在监督学习中训练 Transformer 就已经相当困难，更不用说在 RL 环境中了。毕竟，仅靠自己来稳定地训练一个 LSTM 智能体本身就很有挑战性。

The self-attention mechanism avoids compressing the whole past into a fixed-size hidden state and does not suffer from vanishing or exploding gradients as much as RNNs. Reinforcement learning tasks can for sure benefit from these traits. However, it is quite difficult to train Transformer even in supervised learning, let alone in the RL context. It could be quite challenging to stabilize and train a LSTM agent by itself, after all.

门控 Transformer-XL（GTrXL；Parisotto 等人，2019）是将 Transformer 用于 RL 的一次尝试。GTrXL 通过在 Transformer-XL 的基础上做出两项修改，成功稳定了训练：

The Gated Transformer-XL (GTrXL; Parisotto, et al. 2019) is one attempt to use Transformer for RL. GTrXL succeeded in stabilizing training with two changes on top of Transformer-XL:

1. 层归一化仅应用于残差模块中的输入流，而非快捷流。这种重排的一个关键好处是允许原始输入从第一层流到最后一层。

1. The layer normalization is only applied on the input stream in a residual module, but NOT on the shortcut stream. A key benefit to this reordering is to allow the original input to flow from the first to last layer.

2. 残差连接被替换为 GRU 风格（门控循环单元；Chung 等人，2014）的门控机制。

2. The residual connection is replaced with a GRU-style (Gated Recurrent Unit; Chung et al., 2014) gating mechanism.

\[\begin{aligned} r &= \sigma(W_r^{(l)} y + U_r^{(l)} x) \\ z &= \sigma(W_z^{(l)} y + U_z^{(l)} x - b_g^{(l)}) \\ \hat{h} &= \tanh(W_g^{(l)} y + U_g^{(l)} (r \odot x)) \\ g^{(l)}(x, y) &= (1-z)\odot x + z\odot \hat{h} \end{aligned}\]

门控函数的参数被显式地初始化为接近恒等映射——这就是存在 $ b_g $ 项的原因。$ b_g > 0 $ 对加速学习大有帮助。

The gating function parameters are explicitly initialized to be close to an identity map — this is why there is a $ b_g $ term. A $ b_g > 0 $ greatly helps with the learning speedup.

Transformer-XL、调整层归一化顺序的 Transformer-XL 和门控 Transformer-XL 的模型架构比较。（图片来源：Parisotto 等人 2019 年的论文图 1）

Comparison of the model architecture of Transformer-XL, Transformer-XL with the layer norm reordered, and Gated Transformer-XL. (Image source: Figure 1 in Parisotto, et al. 2019)

Decision Transformer（DT；Chen 等人，2021）将强化学习问题表述为一种 *条件序列建模* 过程，根据期望回报、过去的状态和动作输出最优动作。因此，使用 Transformer 架构变得直接。Decision Transformer 适用于离策略强化学习，即模型只能访问由其他策略收集的固定轨迹集合。

Decision Transformer (DT; Chen et al 2021) formulates Reinforcement Learning problems as a process of *conditional sequence modeling*, outputting the optimal actions conditioned on the desired return, past states and actions. It therefore becomes straightforward to use Transformer architecture. Decision Transformer is for off-policy RL, where the model only has access to a fixed collection of trajectories collected by other policies.

为了鼓励模型学会如何行动以实现期望回报，它向模型输入期望的未来回报 $ \hat{R} = \sum_{t'=t}^T r_{t'} $，而不是当前奖励。轨迹由一组三元组（return-to-go $ \hat{R}_t $、状态 $ s_t $、动作 $ a_t $）组成，并作为 Transformer 的输入序列：

To encourage the model to learn how to act in order to achieve a desired return, it feeds the model with desired future return $ \hat{R} = \sum_{t'=t}^T r_{t'} $ instead of the current reward. The trajectory consists of a list of triplets, (return-to-go $ \hat{R}_t $, state $ s_t $, action $ a_t $), and it is used as an input sequence for Transformer:

\[ \tau = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \dots, \hat{R}_T, s_T, a_T) \]

添加并训练了三个线性层，分别用于 return-to-go、状态和动作，以提取 token 嵌入。预测头学习预测与输入 token $s_t$ 对应的动作 $a_t$。训练时，离散动作使用交叉熵损失，连续动作使用均方误差（MSE）。在他们的实验中，预测状态或 return-to-go 并未提高性能。

Three linear layers are added and trained for return-to-go, state and action respectively to extract token embeddings. The prediction head learns to predict $a_t$ corresponding to the input token $s_t$. The training uses cross-entropy loss for discrete actions or MSE for continuous actions. Predicting the states or return-to-go was not found to help improve the performance in their experiments.

实验将 DT 与多种无模型 RL 算法基线进行了比较，结果表明：

The experiments compared DT with several model-free RL algorithm baselines and showed that:

* 在低数据规模场景下，DT 比行为克隆更高效；

* DT is more efficient than behavior cloning in low data regime;

* DT 能很好地建模回报分布；

* DT can model the distribution of returns very well;

* 拥有较长的上下文对于获得良好结果至关重要；

* Having a long context is crucial for obtaining good results;

互动版：图/公式 + 针对本篇提问 →