UTF8gbsn

1 Background↩︎

Transformers are a type of neural network [1]. They were originally known for their strong performance in machine translation, and are now a de facto standard for building large-scale self-supervised learning systems [2], [3]. The past few years have seen the rise of Transformers not only in natural language processing (NLP) but also in several other fields, such as computer vision and multi-modal processing. As Transformers continue to mature, these models are playing an increasingly important role in the research and application of artificial intelligence (AI).

Looking back at the history of neural networks, Transformers have not been around for a long time. While Transformers are “newcomers” in NLP, they were developed on top of several ideas, the origins of which can be traced back to earlier work, such as word embedding [4], [5] and attention mechanisms [6], [7]. As a result, Transformers can benefit from the advancements of different sub-fields of deep learning, and provide an elegant way to combine these neural models. On the other hand, Transformers are unique, and differ from previous models in several ways. First, they do not depend on recurrent or convolutional neural networks for modeling sequences of words, but use only attention mechanisms and feed-forward neural networks. Second, the use of self-attention in Transformers makes it easier to deal with global contexts and dependencies among words. Third, Transformers are very flexible architectures and can be easily modified to accommodate different tasks.

The widespread use of Transformers motivates the development of cutting-edge techniques in deep learning. For example, there are significant refinements in self-attention mechanisms, which have been incorporated into many state-of-the-art NLP systems. The resulting techniques, together with the progress in self-supervised learning, have led us to a new era of AI: we are beginning to obtain models of universal language understanding, generation and reasoning. This has been evidenced by recent Transformer-based large language models (LLMs) which demonstrate amazing performance across a broad variety of tasks [8].

This paper provides an introduction to Transformers while reflecting the recent developments in applying these models to different problems. However, Transformers are so successful that there have been numerous related studies and we cannot give a full description of them. Therefore, we focus this work on the core ideas of Transformers, and present a basic description of the common techniques. We also discuss some recent advances in Transformers, such as model improvements for efficiency and accuracy considerations. Because the field is very active and new techniques are coming out every day, it is impossible to survey all the latest literature and we are not attempting to do so. Instead, we focus on just those concepts and algorithms most relevant to Transformers, aimed at the people who wish to get a general understanding of these models.

2 The Basic Model↩︎

Here we consider the model presented in [1]’s work. We start by considering the Transformer architecture and then discuss its components in detail.

2.1 The Transformer Architecture↩︎

Figure 1 shows the standard Transformer model which follows the standard encoder-decoder framework. A Transformer encoder comprises a number of stacked encoding layers (or encoding blocks). Each encoding layer has two different sub-layers (or sub-blocks), called the self-attention sub-layer and the feed-forward neural network (FFN) sub-layer. Suppose we have a source-side sequence \(\mathbf{x}=x_1 \cdots x_m\) and a target-side sequence \(\mathbf{y}=y_1 \cdots y_n\). The input of an encoding layer is a sequence of \(m\) vectors \(\mathbf{h}_{1} \cdots \mathbf{h}_{m}\), each having \(d_{\mathrm{model}}\) dimensions (or \(d\) dimensions for simplicity). Here we use \(\mathbf{H} \in \mathbb{R}^{m \times d}\) to denote these input vectors¹. The self-attention sub-layer first performs a self-attention operation \(\mathrm{Att}_{\mathrm{self}}(\cdot)\) on \(\mathbf{H}\) to generate an output \(\mathbf{C}\): \[\begin{align} \mathbf{C}& = & \mathrm{Att}_{\mathrm{self}}(\mathbf{H}). \label{eq:self-attention-basic} \end{align}\tag{1}\]

Here \(\mathbf{C}\) is of the same size as \(\mathbf{H}\), and can thus be viewed as a new representation of the inputs. Then, a residual connection and a layer normalization unit are added to the output so that the resulting model is easier to optimize.

The original Transformer model employs the post-norm structure where a residual connection is created before layer normalization is performed, as follows: \[\begin{align} \mathbf{H}_{\mathrm{self}} & = & \mathrm{LNorm}(\mathbf{C} + \mathbf{H}), \label{eq:post-norm-basic} \end{align}\tag{2}\]

where the addition of \(\mathbf{H}\) denotes the residual connection, and \(\mathrm{LNorm}(\cdot)\) denotes the layer normalization function. Substituting Eq. (1 ) into Eq. (2 ), we obtain the form of the self-attention sub-layer \[\begin{align} \mathrm{Layer}_{\mathrm{self}}(\mathbf{H}) & = & \mathbf{H}_{\mathrm{self}} \nonumber \\ & = & \mathrm{LNorm}(\mathrm{Att}_{\mathrm{self}}(\mathbf{H}) + \mathbf{H}). \label{eq:self-attention-layer} \end{align}\tag{3}\]

The definitions of \(\mathrm{LNorm}(\cdot)\) and \(\mathrm{Att}_{\mathrm{self}}(\cdot)\) have been given previously, and we will also discuss them later in the section.

The FFN sub-layer takes \(\mathbf{H}_{\mathrm{self}}\) and outputs a new representation \(\mathbf{H}_{\mathrm{ffn}} \in \mathbb{R}^{m \times d}\). It has the same form as the self-attention sub-layer, with the attention function replaced by the FFN function, given by \[\begin{align} \mathrm{Layer}_{\mathrm{ffn}}(\mathbf{H}_{\mathrm{self}}) & = & \mathbf{H}_{\mathrm{ffn}} \nonumber \\ & = & \mathrm{LNorm}(\mathrm{FFN}(\mathbf{H}_{\mathrm{self}}) + \mathbf{H}_{\mathrm{self}}). \label{eq:ffn-layer} \end{align}\tag{4}\]

Here \(\mathrm{FFN}(\cdot)\) can be any feed-forward neural network with non-linear activation functions. The most common structure of \(\mathrm{FFN}(\cdot)\) is a two-layer network involving two linear transformations and a ReLU activation function between them.

For deep models, we can stack the above neural networks. Let \(\mathbf{H}^{l}\) be the output of layer \(l\). Then, we can express \(\mathbf{H}^{l}\) as a function of \(\mathbf{H}^{l-1}\). We write this as a composition of two sub-layers \[\begin{align} \mathbf{H}^{l} & = & \mathrm{Layer}_{\mathrm{ffn}}(\mathbf{H}_{\mathrm{self}}^{l}), \\ \mathbf{H}_{\mathrm{self}}^{l} & = & \mathrm{Layer}_{\mathrm{self}}(\mathbf{H}^{l-1}). \end{align}\]

If there are \(L\) encoding layers, then \(\mathbf{H}^{L}\) serves as the final output of the encoder. In this case, \(\mathbf{H}^{L}\) can be viewed as a learned contextual representation of the input sequence. \(\mathbf{H}^{0}\) denotes the initial input to the encoder. In recurrent and convolutional models, \(\mathbf{H}^{0}\) typically consists simply of the word embeddings of the input sequence. The Transformer takes a different approach to representing input words and explicitly encodes positional information. In Section 2.2, we will discuss the embedding model used in the Transformer.

The Transformer decoder has a similar structure to the Transformer encoder, comprising \(L\) stacked decoding layers (or decoding blocks). Let \(\mathbf{S}^{l}\) be the output of the \(l\)-th decoding layer. We can formulate a decoding layer using the following equations:

\[\begin{align} \mathbf{S}^{l} & = & \mathrm{Layer}_{\mathrm{ffn}}(\mathbf{S}_{\mathrm{cross}}^{l}), \\ \mathbf{S}_{\mathrm{cross}}^{l} & = & \mathrm{Layer}_{\mathrm{cross}}(\mathbf{H}^{L},\mathbf{S}_{\mathrm{self}}^{l-1}) \tag{5}, \\ \mathbf{S}_{\mathrm{self}}^{l} & = & \mathrm{Layer}_{\mathrm{self}}(\mathbf{S}^{l-1}). \tag{6} \end{align}\]

Here there are three decoder sub-layers. The self-attention and FFN sub-layers are the same as those used in the encoder. \(\mathrm{Layer}_{\mathrm{cross}}(\cdot)\) denotes a cross attention sub-layer (or encoder-decoder sub-layer) which models the transformation from the source-side to the target-side. In Section 2.6we will see that \(\mathrm{Layer}_{\mathrm{cross}}(\cdot)\) can be implemented using the same function as \(\mathrm{Layer}_{\mathrm{self}}(\cdot)\).

The Transformer decoder outputs a distribution over a vocabulary \(V_\mathrm{y}\) at each target-side position. This is achieved by using a softmax layer that normalizes a linear transformation of \(\mathbf{S}^{L}\) to distributions of target-side words. To do this, we map \(\mathbf{S}^{L}\) to an \(n \times |V_{\mathrm{y}}|\) matrix \(\mathbf{O}\) by \[\begin{align} \mathbf{O} & = & \mathbf{S}^{L} \cdot \mathbf{W}_{\mathrm{o}}, \end{align}\]

where \(\mathbf{W}_{\mathrm{o}} \in \mathbb{R}^{d \times |V_{\mathrm{y}}|}\) is the parameter matrix of the linear transformation.

Then, the output of the Transformer decoder is given in the form \[\begin{align} \begin{bmatrix} \mathop{\mathrm{Pr}}(\cdot|y_0,\mathbf{x})\\ \vdots \\ \mathop{\mathrm{Pr}}(\cdot|y_0 \cdots y_{n-1},\mathbf{x}) \end{bmatrix} & = & \mathrm{Softmax}(\mathbf{O}) \nonumber \\ & = & \begin{bmatrix} \mathrm{Softmax}(\mathbf{o}_{1})\\ \vdots\\ \mathrm{Softmax}(\mathbf{o}_{n}) \end{bmatrix}, \end{align}\]

where \(\mathbf{o}_{i}\) denotes the \(i\)-th row vector of \(\mathbf{O}\), and \(y_0\) denotes the start symbol \(\langle \mathrm{SOS} \rangle\). Under this model, the probability of \(\mathbf{y}\) given \(\mathbf{x}\) can be defined as usual, \[\begin{align} \log \mathop{\mathrm{Pr}}(\mathbf{y}|\mathbf{x}) & = & \sum_{i=1}^{n} \log \mathop{\mathrm{Pr}}(y_{i}|y_0 \cdots y_{i-1},\mathbf{x}). \end{align}\]

This equation resembles the general form of language modeling: we predict the word at position \(i\) given all the words up to position \(i-1\). Therefore, the target output sequence is shifted one step forward in time relative to the input sequence, that is, the decoder takes \(y_0 \cdots y_{n-1}\) as input to predict the output sequence \(y_1 \cdots y_{n}\).

The Transformer architecture discussed above has several variants which have been successfully used in different fields of NLP. For example, we can use a Transformer encoder to represent text (referred to as the encoder-only architecture), a Transformer decoder to generate text (referred to as the decoder-only architecture), or a standard encoder-decoder Transformer model to transform an input sequence to an output sequence. In the rest of this chapter, much of the discussion is independent of the particular choice of application, and will be primarily focused on the encoder-decoder architecture. In Section 6, we will see applications of the encoder-only and decoder-only architectures.

2.2 Positional Encoding↩︎

In their original form, both FFNs and attention models used in Transformers ignore an important aspect of sequence modeling, namely that the order of the words plays a crucial role in expressing the meaning of a sequence. This means that the encoder and decoder are insensitive to the positional information of the input words. A simple approach to overcoming this problem is to add positional encodings to the representation of each word in the sequence. More formally, a word \(x_j\) can be represented as a \(d\)-dimensional vector \[\begin{align} \mathbf{e}_j & = & \mathbf{x}_j + \mathrm{PE}(j). \end{align}\]

Here \(\mathbf{x}_j \in \mathbb{R}^{d}\) is the word embedding, which can be obtained using standard word embedding models. \(\mathrm{PE}(j) \in \mathbb{R}^{d}\) is the representation of the position \(j\). The vanilla Transformer employs a sinusoidal positional encoding model, which we write in the form \[\begin{align} \mathrm{PE}(j,2k) & = & \sin(j \cdot \frac{1}{10000^{2k/d}}) \tag{7}, \\ \mathrm{PE}(j,2k+1) & = & \cos(j \cdot \frac{1}{10000^{2k/d}}) \tag{8}, \end{align}\]

where \(\mathrm{PE}(j,k)\) denotes the \(k\)-th entry of \(\mathrm{PE}(j)\). The idea of positional encoding is to distinguish different positions using continuous functions. Here we use the sine and cosine functions with different frequencies. Such a method can be interpreted as a positional numeral system. Because the encoding is based on individual positions, it is also called absolute positional encoding. In Section 4.1, we will see an improvement to this method.

Once we obtain these embeddings, the sequence \(\mathbf{e}_1 \cdots \mathbf{e}_m\) is used as input to the Transformer encoder, that is, \[\begin{align} \mathbf{H}^0 & = & \begin{bmatrix} \mathbf{e}_1 \\ \vdots \\ \mathbf{e}_m \end{bmatrix}. \end{align}\]

Similarly, we can define the input on the decoder side.

2.3 Multi-head Self-attention↩︎

The use of self-attention is perhaps one of the most significant advances in sequence-to-sequence models. It attempts to learn and make use of direct interactions between all pairs of inputs. From a representation learning perspective, self-attention models assume that the learned representation at position \(i\) (denoted by \(\mathbf{c}_{i}\)) is a weighted sum of the inputs over the sequence. The output \(\mathbf{c}_{i}\) is thus given by \[\begin{align} \mathbf{c}_i & = & \sum_{j=1}^{m} \alpha_{i,j} \mathbf{h}_{j} \label{eq:transformer-basic-attention-model}, \end{align}\tag{9}\]

where \(\alpha_{i,j}\) represents the attention weight assigned to \(\mathbf{h}_j\) when computing the representation at position \(i\). We can thus view \(\mathbf{c}_i\) as a representation of the global context at position \(i\). \(\alpha_{i,j}\) can be defined in different ways depending on the specific attention model considered. Here we use the scaled dot-product attention function to compute \(\alpha_{i,j}\), as follows \[\begin{align} \alpha_{i,j} & = & \mathrm{Softmax}\left( \mathbf{h}_i \mathbf{h}_j^\top/\beta \right) \nonumber \\ & = & \frac{\exp(\mathbf{h}_i \mathbf{h}_j^\top/\beta)}{\sum_{k=1}^m \exp(\mathbf{h}_i \mathbf{h}_k^\top/\beta)}, \end{align}\]

where the scaling factor \(\beta\) is typically set to \(\sqrt{d}\).

Compared with conventional recurrent and convolutional models, an advantage of self-attention models is that they shorten the computational “distance” between two inputs. Figure 2 illustrates the information flow in these models. We can see that, given the input at position \(i\), self-attention models can directly access any other input. By contrast, recurrent and convolutional models might need two or more jumps to see the entire sequence.

We can adopt a more general perspective on self-attention by using the QKV attention model. Suppose we have a sequence of \(\kappa\) queries \(\mathbf{Q} = \begin{bmatrix} \mathbf{q}_1 \\ \vdots \\ \mathbf{q}_{\kappa} \end{bmatrix}\), and a sequence of \(\psi\) key-value pairs \((\mathbf{K} = \begin{bmatrix} \mathbf{k}_1 \\ \vdots \\ \mathbf{k}_{\psi} \end{bmatrix}, \mathbf{V} = \begin{bmatrix} \mathbf{v}_1 \\ \vdots \\ \mathbf{v}_{\psi} \end{bmatrix})\). The output of the model is a sequence of vectors, each corresponding to a query. The form of the QKV attention is given by \[\begin{align} \mathrm{Att}_{\mathrm{qkv}}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) & = & \mathrm{Softmax}(\frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d}}) \mathbf{V}. \end{align}\]

We can write the output of the QKV attention model as a sequence of row vectors \[\begin{align} \mathbf{C} & = & \begin{bmatrix} \mathbf{c}_1 \\ \vdots \\ \mathbf{c}_{\kappa} \end{bmatrix} \nonumber \\ & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) \label{eq:transformer-qkv-attention}. \end{align}\tag{10}\]

To apply this equation to self-attention, we simply have \[\begin{align} \mathbf{H}^{q} & = & \mathbf{H} \mathbf{W}^{q}, \\ \mathbf{H}^{k} & = & \mathbf{H} \mathbf{W}^{k}, \\ \mathbf{H}^{v} & = & \mathbf{H} \mathbf{W}^{v}, \end{align}\]

where \(\mathbf{W}^{q}, \mathbf{W}^{k}, \mathbf{W}^{v} \in \mathbb{R}^{d \times d}\) represent linear transformations of \(\mathbf{H}\).

By considering Eq. (1 ), we then obtain \[\begin{align} \mathbf{C} & = & \mathrm{Att}_{\mathrm{self}}(\mathbf{H}) \nonumber \\ & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{H}^{q}, \mathbf{H}^{k}, \mathbf{H}^{v}) \nonumber \\ & = & \mathrm{Softmax}(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}}) \mathbf{H}^{v}. \end{align}\]

Here \(\mathrm{Softmax}(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}})\) is an \(m \times m\) matrix in which each row represents a distribution over all input positions, that is \[\begin{align} \textrm{row i} & = & \begin{bmatrix} \alpha_{i,1} & \cdots & \alpha_{i,m} \end{bmatrix}. \end{align}\]

We can improve the above self-attention model by using a technique called multi-head attention. This method can be motivated from the perspective of learning from multiple lower-dimensional feature subspaces. It projects the input into multiple subspaces and learns separate representations in each. Specifically, we project the entire input space into \(\tau\) subspaces (referred to as heads). For example, we transform \(\mathbf{H} \in \mathbb{R}^{m \times d}\) into \(\tau\) matrices of size \(m \times \frac{d}{\tau}\), denoted by \(\{\mathbf{H}_{1}^{\mathrm{head}}, \dots, \mathbf{H}_{\tau}^{\mathrm{head}}\}\). The attention model is then applied \(\tau\) times, once for each head. Finally, the outputs of these model runs are concatenated and transformed by a linear projection. This procedure can be expressed by \[\begin{align} \mathbf{C} & = & \mathrm{Merge}(\mathbf{C}_{1}^{\mathrm{head}},\cdots,\mathbf{C}_{\tau}^{\mathrm{head}}) \mathbf{W}_{c}. \end{align}\]

For each head \(h\), \[\begin{align} \mathbf{C}_{h}^{\mathrm{head}} & = & \mathrm{Softmax} \left( \frac{\mathbf{H}_{h}^{q} [\mathbf{H}_{h}^{k}]^\top}{\sqrt{d}} \right) \mathbf{H}_{h}^{v}, \\ \mathbf{H}_{h}^{q} & = & \mathbf{H} \mathbf{W}_{h}^{q}, \\ \mathbf{H}_{h}^{k} & = & \mathbf{H} \mathbf{W}_{h}^{k}, \\ \mathbf{H}_{h}^{v} & = & \mathbf{H} \mathbf{W}_{h}^{v}. \end{align}\]

Here \(\mathrm{Merge}(\cdot)\) is the concatenation function, and \(\mathrm{Att}_{\mathrm{qkv}}(\cdot)\) is the attention function described in Eq. (10 ). \(\mathbf{W}_{h}^{q}, \mathbf{W}_{h}^{k}, \mathbf{W}_{h}^{v} \in \mathbb{R}^{d \times \frac{d}{\tau}}\) are the parameters of the projections from a \(d\)-dimensional space to a \(\frac{d}{\tau}\)-dimensional space for the queries, keys, and values. Thus, \(\mathbf{H}_{h}^{q}\), \(\mathbf{H}_{h}^{k}\), \(\mathbf{H}_{h}^{v}\), and \(\mathbf{C}_{h}^{\mathrm{head}}\) are all \(m \times \frac{d}{\tau}\) matrices. \(\mathrm{Merge}(\mathbf{C}_{1}^{\mathrm{head}}, \cdots ,\mathbf{C}_{\tau}^{\mathrm{head}})\) produces an \(m \times d\) matrix. It is then transformed by a linear mapping \(\mathbf{W}_{c} \in \mathbb{R}^{d \times d}\), leading to the final result \(\mathbf{C} \in \mathbb{R}^{m \times d}\).

Although the notation may seem tedious, implementing multi-head models is straightforward using modern deep learning toolkits. A common method in Transformer-based systems is to store inputs from all the heads in data structures called tensors, enabling us to leverage parallel computing resources for greater efficiency. A more general discussion of the QKV attention and multi-head attention models can be found in the literature on encoder-decoder models.

2.4 Layer Normalization↩︎

Layer normalization provides a simple and effective means to make the training of neural networks more stable by standardizing the activations of the hidden layers in a layer-wise manner. As introduced by [9], given a layer’s output \(\mathbf{h} \in \mathbb{R}^{d}\), the layer normalization method computes a standardized output \(\mathrm{LNorm}(\mathbf{h}) \in \mathbb{R}^{d}\) by \[\begin{align} \mathrm{LNorm}(\mathbf{h}) & = & \boldsymbol{\alpha} \odot \frac{\mathbf{h} - \mu}{\sigma + \epsilon} + \boldsymbol{\beta}. \label{eq:layer-normalization} \end{align}\tag{11}\]

Here \(\mu \in \mathbb{R}\) and \(\sigma \in \mathbb{R}\) are the scalar mean and standard deviation of the activations, respectively. Let \(h_k\) be the \(k\)-th dimension of \(\mathbf{h}\). The scalars \(\mu\) and \(\sigma\) are defined as \[\begin{align} \mu & = & \frac{1}{d} \sum_{k=1}^{d} h_k, \\ \sigma & = & \sqrt{\frac{1}{d} \sum_{k=1}^{d} (h_k-\mu)^2 }. \end{align}\]

Here \(\boldsymbol{\alpha} \in \mathbb{R}^{d}\) and \(\boldsymbol{\beta} \in \mathbb{R}^{d}\) are the scale and shift parameters. They can be treated as parameters of layer normalization, whose values are learned jointly with other parameters of the Transformer model. The addition of \(\epsilon\) to \(\sigma\) is used for numerical stability. In general, \(\epsilon\) is chosen to be a very small number.

We illustrate the layer normalization method for the hidden states of an encoder in the following example (assume that \(m=4\), \(d=3\), \(\boldsymbol{\alpha}=\mathbf{1}\), \(\boldsymbol{\beta}=\mathbf{0}\), and \(\epsilon=0.1\)). \[\begin{matrix} \mathbf{h}_1 \\ \mathbf{h}_2 \\ \mathbf{h}_3 \\ \mathbf{h}_4 \end{matrix} \begin{bmatrix} 1 & 1 & 2 \\ 0.9 & 0.9 & 0 \\ 0.7 & 0.8 & 0 \\ 3 & 1 & 7 \end{bmatrix} \;\;\; \begin{matrix} \mu = 1.3,\;\sigma=0.5 \\ \mu = 0.6,\;\sigma=0.4 \\ \mu = 0.5,\;\sigma=0.4 \\ \mu = 3.7,\;\sigma=2.5 \end{matrix} \;\;\;\implies\;\;\; \begin{bmatrix} \frac{{\color{red} 1}-1.3}{0.5 + 0.1} & \frac{{\color{red} 1}-1.3}{0.5 + 0.1} & \frac{{\color{red} 2}-1.3}{0.5 + 0.1} \\ \frac{{\color{red} 0.9}-0.6}{0.4 + 0.1} & \frac{{\color{red} 0.9}-0.6}{0.4 + 0.1} & \frac{{\color{red} 0}-0.6}{0.4 + 0.1} \\ \frac{{\color{red} 0.7}-0.5}{0.4 + 0.1} & \frac{{\color{red} 0.8}-0.5}{0.4 + 0.1} & \frac{{\color{red} 0}-0.5}{0.4 + 0.1} \\ \frac{{\color{red} 3}-3.7}{2.5 + 0.1} & \frac{{\color{red} 1}-3.7}{2.5 + 0.1} & \frac{{\color{red} 7}-3.7}{2.5 + 0.1} \end{bmatrix} \nonumber\]

As discussed in Section 2.1, the layer normalization unit in each sub-layer is used to standardize the output of a residual block. Here we describe a more general formulation for this structure. Suppose that \(F(\cdot)\) is a neural network sub-layer. Then, the post-norm structure of \(F(\cdot)\) is given by: \[\begin{align} \mathbf{H}_{\mathrm{out}} & = & \mathrm{LNorm}(F(\mathbf{H}_{\mathrm{in}}) + \mathbf{H}_{\mathrm{in}}), \label{eq:post-norm-simple} \end{align}\tag{12}\]

where \(\mathbf{H}_{\mathrm{in}}\) and \(\mathbf{H}_{\mathrm{out}}\) are the input and output of this sub-layer. Clearly, Eq. (3 ) is an instance of this equation.

An alternative approach to incorporating layer normalization and residual connections into the modeling is to execute the \(\mathrm{LNorm}(\cdot)\) function right before the \(F(\cdot)\) function, and to establish an identity mapping from the input to the output of the entire sub-layer. This structure, known as the pre-norm structure, can be expressed in the form \[\begin{align} \mathbf{H}_{\mathrm{out}} & = & F(\mathrm{LNorm}(\mathbf{H}_{\mathrm{in}})) + \mathbf{H}_{\mathrm{in}}. \end{align}\]

Both post-norm and pre-norm Transformer models are widely used in NLP systems. See Figure 3 for a comparison of these two structures. In general, residual connections are considered an effective means to make the training of multi-layer neural networks easier. In this sense, the pre-norm Transformer seems promising because it follows the convention that a residual connection is created to bypass the whole network and that the identity mapping from the input to the output leads to easier optimization of deep models. However, when considering the expressive power of a model, there may be modeling advantages to using the post-norm Transformer, because it does not rely as heavily on residual connections and enforces more sophisticated modeling for representation learning. In Section 4.2, we will present a further discussion on this issue.

2.5 Feed-forward Neural Networks↩︎

The use of FFNs in Transformers is inspired in part by the fact that complex outputs can be formed by transforming the inputs through non-linearities. While the self-attention model itself has some non-linearity (in \(\mathrm{Softmax}(\cdot)\)), a more common way to introduce non-linearity is to consider additional layers with non-linear activation functions and linear transformations. Given an input \(\mathbf{H}_{\mathrm{in}} \in \mathbb{R}^{m \times d}\) and an output \(\mathbf{H}_{\mathrm{out}} \in \mathbb{R}^{m \times d}\), the \(\mathbf{H}_{\mathrm{out}} = \mathrm{FFN}(\mathbf{H}_{\mathrm{in}})\) function in the Transformer has the following form \[\begin{align} \mathbf{H}_{\mathrm{out}} & = & \mathbf{H}_{\mathrm{hidden}} \mathbf{W}_f + \mathbf{b}_f, \\ \mathbf{H}_{\mathrm{hidden}} & = & \mathrm{ReLU}(\mathbf{H}_{\mathrm{in}}\mathbf{W}_h + \mathbf{b}_h), \end{align}\]

where \(\mathbf{H}_{\mathrm{hidden}} \in \mathbb{R}^{m \times d_{\mathrm{ffn}}}\) represents the hidden states, and \(\mathbf{W}_h \in \mathbb{R}^{d \times d_{\mathrm{ffn}}}\), \(\mathbf{b}_h \in \mathbb{R}^{d_{\mathrm{ffn}}}\), \(\mathbf{W}_f \in \mathbb{R}^{d_{\mathrm{ffn}} \times d}\) and \(\mathbf{b}_f \in \mathbb{R}^{d}\) are the parameters. This is a two-layer FFN in which the first layer (or hidden layer) introduces a non-linearity through \(\mathrm{ReLU}(\cdot)\)² and the second layer involves only a linear transformation. It is common practice in the Transformer to use a larger hidden layer size. For example, a common choice is \(d_{\mathrm{ffn}} = 4d\), that is, the size of each hidden representation is 4 times that of the input.

Note that using a wide FFN sub-layer has proven to be of great practical value in many state-of-the-art systems. However, a consequence of this is that the model’s parameter count is dominated by the FFN. Table 1 shows the number of parameters and time complexities for different modules of a standard Transformer-based system. We see that FFNs dominate the model size when \(d_{\mathrm{ffn}}\) is large, though they are not the most time-consuming components. In the case of very large Transformer models, we therefore wish to address this problem to build more efficient systems.

2.6 Decoder-side Attention↩︎

A decoder layer involves two attention sub-layers: the first is a self-attention sub-layer, and the second is a cross-attention sub-layer. These sub-layers are based on either the post-norm or the pre-norm structure, but differ in how the attention functions are defined. Consider, for example, the post-norm structure described in Eq. (12 ). We can define the cross-attention and self-attention sub-layers for a decoding layer to be \[\begin{align} \mathbf{S}_{\mathrm{cross}} & = & \mathrm{Layer}_{\mathrm{cross}}(\mathbf{H}_{\mathrm{enc}},\mathbf{S}_{\mathrm{self}}) \nonumber \\ & = & \mathrm{LNorm}(\mathrm{Att}_{\mathrm{cross}}(\mathbf{H}_{\mathrm{enc}},\mathbf{S}_{\mathrm{self}}) + \mathbf{S}_{\mathrm{self}}), \\ \mathbf{S}_{\mathrm{self}} & = & \mathrm{Layer}_{\mathrm{self}}(\mathbf{S}) \nonumber \\ & = & \mathrm{LNorm}(\mathrm{Att}_{\mathrm{self}}(\mathbf{S}) + \mathbf{S}), \end{align}\]

where \(\mathbf{S} \in \mathbb{R}^{n \times d}\) is the input to the self-attention sub-layer, \(\mathbf{S}_{\mathrm{cross}} \in \mathbb{R}^{n \times d}\) and \(\mathbf{S}_{\mathrm{self}} \in \mathbb{R}^{n \times d}\) are the outputs of the sub-layers, and \(\mathbf{H}_{\mathrm{enc}} \in \mathbb{R}^{m \times d}\) is the output of the encoder ³.

As with conventional attention models, cross-attention is primarily used to model the correspondence between the source-side and target-side sequences. The \(\mathrm{Att}_{\mathrm{cross}}(\cdot)\) function is based on the QKV attention model, which generates an output by querying a collection of key-value pairs. More specifically, we define the queries, keys, and values as linear mappings of \(\mathbf{S}_{\mathrm{self}}\) and \(\mathbf{H}_{\mathrm{enc}}\), as follows \[\begin{align} \mathbf{S}_{\mathrm{self}}^{q} & = & \mathbf{S}_{\mathrm{self}} \mathbf{W}_{\mathrm{cross}}^{q}, \\ \mathbf{H}_{\mathrm{enc}}^{k} & = & \mathbf{H}_{\mathrm{enc}} \mathbf{W}_{\mathrm{enc}}^{k}, \\ \mathbf{H}_{\mathrm{enc}}^{v} & = & \mathbf{H}_{\mathrm{enc}} \mathbf{W}_{\mathrm{enc}}^{v}, \end{align}\]

where \(\mathbf{W}_{\mathrm{cross}}^{q}, \mathbf{W}_{\mathrm{enc}}^{k}, \mathbf{W}_{\mathrm{enc}}^{v} \in \mathbb{R}^{d \times d}\) are the parameters of the mappings. In other words, the queries are defined based on \(\mathbf{S}_{\mathrm{self}}\), while the keys and values are defined based on \(\mathbf{H}_{\mathrm{enc}}\).

\(\mathrm{Att}_{\mathrm{cross}}(\cdot)\) is then defined as \[\begin{align} \mathrm{Att}_{\mathrm{cross}}(\mathbf{H}_{\mathrm{enc}},\mathbf{S}_{\mathrm{self}}) & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{S}_{\mathrm{self}}^{q}, \mathbf{H}_{\mathrm{enc}}^{k}, \mathbf{H}_{\mathrm{enc}}^{v}) \nonumber \\ & = & \mathrm{Softmax} \left(\frac{\mathbf{S}_{\mathrm{self}}^{q} [\mathbf{H}_{\mathrm{enc}}^{k}]^\top}{\sqrt{d}} \right) \mathbf{H}_{\mathrm{enc}}^{v}. \end{align}\]

The \(\mathrm{Att}_{\mathrm{self}}(\cdot)\) function has a similar form to \(\mathrm{Att}_{\mathrm{cross}}(\cdot)\), with linear mappings of \(\mathbf{S}\) taken as the queries, keys, and values, like this \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{S}^{q},\mathbf{S}^{k},\mathbf{S}^{v}) \nonumber \\ & = & \mathrm{Softmax}(\frac{\mathbf{S}^{q} [\mathbf{S}^{k}]^\top}{\sqrt{d}} + \mathbf{M}) \mathbf{S}^{v}, \end{align}\]

where \(\mathbf{S}^{q} = \mathbf{S} \mathbf{W}_{\mathrm{dec}}^{q}\), \(\mathbf{S}^{k} = \mathbf{S} \mathbf{W}_{\mathrm{dec}}^{k}\), and \(\mathbf{S}^{v} = \mathbf{S} \mathbf{W}_{\mathrm{dec}}^{v}\) are linear mappings of \(\mathbf{S}\) with parameters \(\mathbf{W}_{\mathrm{dec}}^{q}, \mathbf{W}_{\mathrm{dec}}^{k}, \mathbf{W}_{\mathrm{dec}}^{v} \in \mathbb{R}^{d \times d}\).

This form is similar to that of Eq. (10 ). However, a key difference compared to encoder self-attention is that the model here must follow the rule of left-to-right generation (see Figure 4). That is, given a target-side word at position \(i\), the model can only see the target-side words in the left context \(y_1, \dots, y_{i-1}\). To enforce this, we add a masking variable \(\mathbf{M}\) to the unnormalized weight matrix \(\frac{\mathbf{S}^{q} [\mathbf{S}^{k}]^\top}{\sqrt{d}}\). Both \(\mathbf{M}\) and the weight matrix are of size \(n \times n\), so large negative values in \(\mathbf{M}\) suppress the corresponding attention scores. To prevent attention to the right context (future words) at step \(i\), \(\mathbf{M}\) is defined as: \[\begin{align} M(i,j) & = & \begin{cases} 0 & i \ge j \\ -\infty & i < j \end{cases}, \end{align}\]

where \(M(i,j)\) indicates a bias term for the alignment score between positions \(i\) and \(j\). Below we show an example of how the masking variable is applied (assume \(n=4\)). \[\begin{align} & & \mathrm{Softmax}(\frac{\mathbf{S}^{q} [\mathbf{S}^{k}]^\top}{\sqrt{d}} + \mathbf{M}) \nonumber \\ & = & \mathrm{Softmax}( \begin{bmatrix} 2 & 0.1 & 1 & 1 \\ 0 & 0.9 & 0.9 & 0.9 \\ 0.2 & 0.8 & 0.7 & 2 \\ 0.3 & 1 & 0.3 & 3 \end{bmatrix} + \begin{bmatrix} 0 & {\color{red} -\infty} & {\color{red} -\infty} & {\color{red} -\infty} \\ 0 & 0 & {\color{red} -\infty} & {\color{red} -\infty} \\ 0 & 0 & 0 & {\color{red} -\infty} \\ 0 & 0 & 0 & 0 \\ \end{bmatrix} ) \nonumber \\ & = & \mathrm{Softmax}( \begin{bmatrix} 2 & {\color{red} -\infty} & {\color{red} -\infty} & {\color{red} -\infty} \\ 0 & 0.9 & {\color{red} -\infty} & {\color{red} -\infty} \\ 0.2 & 0.8 & 0.7 & {\color{red} -\infty} \\ 0.3 & 1 & 0.3 & 3 \\ \end{bmatrix} ) \nonumber \\ & = & \begin{bmatrix} 1 & {\color{red} 0} & {\color{red} 0} & {\color{red} 0} \\ 0.3 & 0.7 & {\color{red} 0} & {\color{red} 0} \\ 0.2 & 0.4 & 0.4 & {\color{red} 0} \\ 0.05 & 0.1 & 0.05 & 0.8 \\ \end{bmatrix}. \end{align}\]

As noted in Section 2.3, it is straightforward to improve these models by using the multi-head attention mechanism. In addition, because decoders are typically the most time-consuming components in practical applications, the overall computational efficiency of these systems depends heavily on optimizing decoder-side attention.

2.7 Training and Inference↩︎

Transformers can be trained and used in a standard manner. For example, we can train a Transformer model by performing gradient descent to minimize a loss function on the training data, and test the trained model by performing beam search on unseen data. Below we present some of the techniques that are typically used in the training and inference of Transformer models.

• Learning Rate Scheduling. Like standard neural networks, Transformers can be directly trained using back-propagation. The training process is generally iterated many times to make the model fit the training data well. In each training step, we update the weights of the neural networks by moving them a small step in the direction of the negative gradient of the loss function. There are many ways to design the update rule for training. A popular choice is the Adam optimization method [10]. To adjust the learning rate during training, [1] present a learning rate scheduling strategy that increases the learning rate linearly for a number of steps and then decays it gradually. They design a learning rate of the form \[\begin{align} \eta & = & \eta_{0} \cdot \min \left\{ n_{\mathrm{step}}^{-0.5},\;n_{\mathrm{step}} \cdot (n_{\rm{warmup}})^{-1.5} \right\}, \label{eq:transformer-learning-rate} \end{align}\tag{13}\]

  where \(\eta_{0}\) denotes the initial learning rate, \(n_{\mathrm{step}}\) denotes the number of training steps we have executed, and \(n_{\mathrm{warmup}}\) denotes the number of warmup steps. In the first \(n_{\mathrm{warmup}}\) steps, the learning rate \(\eta\) grows larger as training proceeds. It reaches its highest value at the point where \(n_{\mathrm{step}} = n_{\mathrm{warmup}}\), and then decreases as an inverse square root function (i.e., \(\eta_{0} \cdot n_{\mathrm{step}}^{-0.5}\)).

• Batching and Padding. To make a trade-off between global optimization and training convergence, it is common practice to update the weights using a relatively small collection of samples, called a minibatch. Therefore, we can consider a batch version of the forward and backward computation processes in which the entire minibatch is used together to obtain the gradient information. One advantage of batching is that it allows the system to leverage efficient tensor operations to process multiple sequences in a single run. This requires all input sequences in a minibatch to be stored in a single memory block so they can be read and processed simultaneously. To illustrate this idea, consider a minibatch containing four samples whose source-sides are:

  A	B	C	D	E	F
  M	N
  R	S	T
  W	X	Y	Z

  We can store these sequences in a \(4 \times 6\) continuous block where each “row” represents a sequence, like this

  A	B	C	D	E	F
  M	N	\(\square\)	\(\square\)	\(\square\)	\(\square\)
  R	S	T	\(\square\)	\(\square\)	\(\square\)
  W	X	Y	Z	\(\square\)	\(\square\)

  Here, padding words \(\square\) are appended to the shorter sequences so that all sequences are properly aligned in memory. Typically, we do not want padding to affect the operation of the system, so we can simply define \(\square\) as a zero vector (referred to as zero padding). On the other hand, in some cases we are interested in using padding to describe something that is not covered by the input sequences. For example, we can replace padding words with context words from the left (or right) of a sequence, though this may require modifications to the system to ensure that the newly added context words do not cause unintended content to appear in the output.

• Search and Caching. At test time, we need to search the space of candidate hypotheses (or candidate target-side sequences) to identify the hypothesis with the highest score \[\begin{align} \hat{\mathbf{y}} & = & \mathop{\mathrm{arg\,max}}_{\mathbf{y}}\;\mathrm{score}(\mathbf{x}, \mathbf{y}), \end{align}\]

  where \(\mathrm{score}(\mathbf{x}, \mathbf{y})\) is the model score of the target-side sequence \(\mathbf{y}\) given the source-side sequence \(\mathbf{x}\). While there are many search algorithms designed to achieve this, most share a similar structure: the search program operates by extending a pool of candidate target-side sequences one token at a time. In this way, the resulting algorithm can be viewed as a left-to-right generation procedure. For a more detailed discussion of search algorithms and model scores of general sequence-to-sequence models, we refer the reader to the literature. Note that all designs for \(\mathrm{score}(\mathbf{x}, \mathbf{y})\), no matter how complex, are based on computing \(\mathop{\mathrm{Pr}}(\mathbf{y}|\mathbf{x})\). Because the attention models used in Transformers require computing the dot-product of each pair of input vectors in a layer, the time complexity of the search algorithm is a quadratic function of the length of \(\mathbf{y}\). It is therefore inefficient to repeatedly compute the outputs of the attention models for previously processed positions. This problem can be addressed by caching the states of each layer for the words we have already seen. Figure 5 illustrates the use of the caching mechanism in a search step. All states for positions \(< i\) are maintained and easily accessed in a cache. At position \(i\), all we need is to compute the states for the newly added word and then update the cache.

  

3 Syntax-aware Models↩︎

Although Transformers are deep learning models that do not explicitly make use of linguistic structures or assumptions, it may be necessary to incorporate our prior knowledge into such systems. This is in part because NLP researchers have long believed that a higher level of abstraction in data representation is needed to develop ideal NLP systems, and there have been many systems that use structural priors. However, structure is a broad topic. For a discussion of the various types of structure one might consider, see the work of [11]. For example, the inductive biases used in our model design can be thought of as some structural prior, and NLP models can also learn the underlying structure of problems by themselves. In this subsection we will discuss some of these issues. We will focus on the methods of introducing linguistic structure into Transformer models. As Transformers can be applied to many NLP tasks, which differ much in their input and output formats, we will primarily discuss modifications to Transformer encoders (call them syntax-aware Transformer encoders). Our discussion, however, is general, and these methods can easily be extended to Transformer decoders.

3.1 Syntax-aware Input and Output↩︎

One of the simplest methods of incorporating structure into NLP systems is to modify the input sequence, leaving the system unchanged. As a simple example, consider a sentence where each word \(x_j\) is assigned a set of \(\kappa\) syntactic labels \(\{\mathrm{tag}_j^{1}, \cdots ,\mathrm{tag}_j^{\kappa}\}\) (e.g., POS labels and dependency labels). We can concatenate these symbols to define a new “word” \[x_j/\mathrm{tag}_j^{1}/\cdots/\mathrm{tag}_j^{\kappa} \nonumber\]

Then, the embedding of this word is given by \[\begin{align} \mathbf{e}_j & = & e(x_j/\mathrm{tag}_j^{1}/\cdots/\mathrm{tag}_j^{\kappa}) + \mathrm{PE}(j), \end{align}\]

where \(e(x_j/\mathrm{tag}_j^{1}/ \cdots /\mathrm{tag}_j^{\kappa}) \in \mathbb{R}^{d}\) is the embedding of \(x_j/\mathrm{tag}_j^{1}/\cdots/\mathrm{tag}_j^{\kappa}\). Since \(x_j/\mathrm{tag}_j^{1}/ \cdots /\mathrm{tag}_j^{\kappa}\) is a complex symbol, we decompose the learning problem of \(e(x_j/\mathrm{tag}_j^{1}/ \cdots /\mathrm{tag}_j^{\kappa})\) into easier problems. For example, we can develop \(\kappa\) embedding models, each producing an embedding for a given tag. Then, we write \(e(x_j/\mathrm{tag}_j^{1}/ \cdots /\mathrm{tag}_j^{\kappa})\) as a sum of the word embedding and tag embeddings \[\begin{align} e(x_j/\mathrm{tag}_j^{1}/ \cdots /\mathrm{tag}_j^{\kappa}) & = & \mathbf{x}_j + e(\mathrm{tag}_j^{1}) + \cdots + e(\mathrm{tag}_j^{\kappa}), \end{align}\]

where \(\{e(\mathrm{tag}_j^{1}), \cdots ,e(\mathrm{tag}_j^{\kappa})\}\) are the embeddings of the tags. Alternatively, we can combine these embeddings via a neural network in the form \[\begin{align} e(x_j/\mathrm{tag}_j^{1}/ \cdots /\mathrm{tag}_j^{\kappa}) & = & \mathrm{FFN}_{\mathrm{embed}}(\mathbf{x}_j,e(\mathrm{tag}_j^{1}), \cdots ,e(\mathrm{tag}_j^{\kappa})), \end{align}\]

where \(\mathrm{FFN}_{\mathrm{embed}}(\cdot)\) is a feed-forward neural network that has one or two layers.

We can apply a similar technique on the decoder side, treating \(y_i/\mathrm{tag}_i^{1}/\cdots/\mathrm{tag}_i^{\kappa}\) as a syntax-augmented target word. However, doing so may drastically increase the size of the target-side vocabulary, posing a computational challenge for both training and inference.

Another common method for representing a sentence is via a syntax tree. In linguistics, the syntax of a sentence can be interpreted in many different ways, resulting in various grammars and corresponding tree-based (or graph-based) representations. While these representations differ in their exact syntactic forms, a general approach for utilizing them in sequence modeling is tree linearization. Consider the following sentence annotated with a constituency-based parse tree:

We can write this tree structure as a sequence of words, syntactic labels and brackets via a tree traversal algorithm, as follows

(S	(NP	(PRP	It	)\(_{\textrm{PRP}}\)	)\(_{\textrm{NP}}\)	(VP	(VBZ	’s	)\(_{\textrm{VBZ}}\)	(ADJP	(JJ
interesting		)\(_{\textrm{JJ}}\)	)\(_{\textrm{ADJP}}\)	)\(_{\textrm{VP}}\)	(.	!	)\(_{\textrm{.}}\)	)\(_{\textrm{S}}\)

This sequence of syntactic tokens can be used directly as input to the encoder, that is, each token is represented by the sum of its word and positional embeddings, and this combined embedding is treated as a standard input to the encoder. One application of linearized trees is tree-to-string machine translation, in which a syntax tree in one language is translated into a string in another language [12], [13]. Linearized trees can also be utilized for tree generation. For example, parsing tasks can be framed as sequence-to-sequence problems, mapping an input text to a sequential representation of its corresponding syntax tree [14], [15]. Figure 7 illustrates these models. It should be noted that the methods described here are not specific to the Transformer and can be applied to many sequence-based architectures, such as RNN-based models.

3.2 Syntax-aware Attention Models↩︎

For Transformer models, it is also natural to make use of syntax trees to guide the process of learning sequence representations. In the previous section we saw how representations of a sequence can be computed by relating different positions within that sequence. This allows us to impose structure on the attention distributions over positions. To do this, we use the encoder self-attention with an additive mask \[\begin{align} \mathrm{AttSyn}_{\mathrm{self}}(\mathbf{H}) & = & \mathrm{Softmax}(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}} + \mathbf{M}) \mathbf{H}^{v} \label{eq:self-attention-syn-add-masking} \end{align}\tag{14}\]

or alternatively with a multiplicative mask \[\begin{align} \mathrm{AttSyn}_{\mathrm{self}}(\mathbf{H}) & = & \mathrm{Softmax}(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}} \odot \mathbf{M}) \mathbf{H}^{v}, \label{eq:self-attention-syn-mul-masking} \end{align}\tag{15}\]

where \(\mathbf{M} \in \mathbb{R}^{m \times m}\) is a matrix of masking variables in which a larger value of \(M(i,j)\) indicates a stronger syntactic correlation between positions \(i\) and \(j\). In the following description we choose Eq. (15 ) as the basic form.

One common way to design \(\mathbf{M}\) is to project syntactic relations of the input tree structure into constraints over the sequence. Here we consider constituency parse trees and dependency parse trees for illustration. Generally, two types of masking methods are employed.

• 0-1 Masking. This method assigns \(M(i,j)=1\) if the words at positions \(i\) and \(j\) are considered syntactically correlated and \(M(i,j)=0\) otherwise [16], [17]. To model the relation between two words in a syntax tree, we can consider the distance between their corresponding nodes. One of the simplest forms is given by \[\begin{align} M(i,j) & = & \begin{cases} 1 & \omega(i,j) \le \omega_{\mathrm{max}} \\ 0 & \textrm{otherwise} \end{cases}, \label{eq:0-1-masking-syntactic-attention} \end{align}\tag{16}\]

  where \(\omega(i,j)\) is the shortest path length between the nodes corresponding to the words at positions \(i\) and \(j\). For example, given a dependency parse tree, \(\omega(i,j)\) is the number of dependency edges in the path between the two words. For a constituency parse tree, all words are leaf nodes. Thus, \(\omega(i,j)\) represents the tree distance between the two leaves within the same branch of the tree. \(\omega_{\mathrm{max}}\) is a parameter used to control the maximum distance between two nodes that can still be considered syntactically correlated. For example, assuming a dependency parse tree with \(\omega_{\mathrm{max}}=1\), Eq. (16 ) enforces the constraint that the attention score between positions \(i\) and \(j\) is computed only if they have a parent-dependent relation⁴.

• Soft Masking. Instead of treating \(\mathbf{M}\) as a hard constraint, we can use it as a soft constraint that scales the attention weight between positions \(i\) and \(j\) based on the degree to which the corresponding words are correlated. One approach is to reduce the attention weight as \(\omega(i,j)\) becomes larger. A simple method is to transform \(\omega(i,j)\) such that \(M(i,j)\) is negatively correlated with \(\omega(i,j)\) and falls within the interval \([0,1]\): \[\begin{align} M(i,j) & = & \mathrm{DNorm}(\omega(i,j)). \end{align}\]

  There are several alternative designs for \(\mathrm{DNorm}(\cdot)\). For example, one can compute a standardized score of \(-\omega(i,j)\) by subtracting its mean and dividing by its standard deviation [18], or normalize \(1/\omega(i,j)\) over all possible \(j\) in the sequence [19]. In cases where parsers can output a score between positions \(i\) and \(j\), this score can also be used to compute \(M(i,j)\). For example, a dependency parser can produce the probability of the word at position \(i\) being the parent of the word at position \(j\) [20]. We can then write \(M(i,j)\) as \[\begin{align} M(i,j) & = & \mathrm{Pr}_{\mathrm{parent}}(i|j) \end{align}\]

  or alternatively \[\begin{align} M(i,j) & = & \max\{\mathrm{Pr}_{\mathrm{parent}}(i|j), \mathrm{Pr}_{\mathrm{parent}}(j|i)\}, \end{align}\]

  where \(\mathrm{Pr}_{\mathrm{parent}}(i|j)\) and \(\mathrm{Pr}_{\mathrm{parent}}(j|i)\) are the probabilities given by the parser. See Figure 8 for an example of inducing a soft masking variable from a dependency parse tree.

3.3 Multi-branch Models↩︎

Integrating syntax into NLP systems presents several challenges. This is partly because automatic parse trees may contain errors, and partly because the use of syntax may lead to strong assumptions about the underlying structure of a sentence. Rather than combining syntactic and word information into a single monolithic model, it may be more flexible and effective to build one model to encode syntactic information and another to encode word sequences. One way to achieve this is through the use of multiple neural networks (referred to as branches or paths), each dealing with one type of input. The outputs of these branches are then combined to produce a final output [21]–[23]. Various methods have therefore been used to combine different types of input for neural models like Transformers.

One commonly used approach is to build two separate encoders, in which one model is trained to encode the syntactic input (denoted by \(\mathbf{t}\)), and the other is trained to encode the standard textual input (denoted by \(\mathbf{x}\)). Figure 9 (a) illustrates this multi-encoder architecture. The syntactic encoder \(\mathrm{Encode}_{\mathrm{syn}}(\mathbf{t})\) is based on models presented in Sections 3.1 and 3.2, and the text encoder \(\mathrm{Encode}_{\mathrm{text}}(\mathbf{x})\) is a standard Transformer encoder. The representations generated by these encoders are then fed into the combination model as input and combined into a hybrid representation, given by \[\begin{align} \mathbf{H}_{\mathrm{hybrid}} & = & \mathrm{Combine}(\mathbf{H}_{\mathrm{syn}},\mathbf{H}_{\mathrm{text}}) \nonumber \\ & = & \mathrm{Combine}(\mathrm{Encode}_{\mathrm{syn}}(\mathbf{t}),\mathrm{Encode}_{\mathrm{text}}(\mathbf{x})). \label{eq:transformer-multi-encoders} \end{align}\tag{17}\]

There are several designs for \(\mathrm{Combine}(\cdot)\), depending on the tasks to which the encoders are applied. For example, if we want to develop a text classifier, \(\mathrm{Combine}(\cdot)\) can be a simple pooling network. For more complicated tasks, such as machine translation, \(\mathrm{Combine}(\cdot)\) can also be a Transformer encoder, and we can fuse information from different sources by performing self-attention on \([\mathbf{H}_{\mathrm{syn}},\mathbf{H}_{\mathrm{text}}]\).

Although we focus primarily on syntactic models in this section, the general multi-encoder architecture can be used in many tasks where inputs from additional sources are required. For example, one can use one encoder to represent a sentence and another to represent the previous sentence in the same document. Combining the two encoders thus yields a context-aware model [24], [25]. Furthermore, the architectures of the encoders do not need to be restricted to Transformers. We can choose different models for different branches. For instance, in a widely used two-branch encoding architecture, we can employ a CNN-based encoder to model local context and a Transformer encoder to model global context [26].

Sub-models of Transformers can also be multi-branch neural networks. See Figure 9 (b) for an example involving two self-attention branches. One is the standard self-attention network \(\mathrm{Att}_{\mathrm{self}}(\mathbf{H})\), and the other is the syntax-aware self-attention network \(\mathrm{AttSyn}_{\mathrm{self}}(\mathbf{H})\). The output of the self-attention model is a linear combination of the outputs of these two branches [19], given by \[\begin{align} \mathbf{H}_{\mathrm{self}} & = & \alpha \cdot \mathrm{Att}_{\mathrm{self}}(\mathbf{H}) + (1 - \alpha) \cdot \mathrm{AttSyn}_{\mathrm{self}}(\mathbf{H}), \end{align}\]

where \(\alpha\) is a combination coefficient. \(\mathbf{H}_{\mathrm{self}}\) can be used as usual by applying a layer normalization function and adding a residual connection, thus keeping the overall architecture identical to standard Transformer models.

Multi-head attention networks can also be viewed as a form of multi-branch models. Therefore, we can provide guidance from syntax to only some of the heads while keeping the rest unchanged [20]. This approach is illustrated in Figure 9 (c) where only one head of the self-attention sub-layer makes use of syntax trees for computing attention weights.

3.4 Multi-scale Models↩︎

In linguistics, syntax studies how sentences are built from smaller constituents. Different levels of these constituents are generally organized in a hierarchical structure, called a syntactic hierarchy. It is therefore possible to use multiple levels of syntactic constituents to explain the same sentence. For example, words capture how the sentence is constructed from small meaningful units, and phrases capture how the sentence is constructed from larger linguistic units.

Multi-scale Transformers leverage different levels of abstraction in data to represent a sentence using features at multiple scales. A common approach is to write a sentence in multiple different forms and then combine them using a multi-branch network [27]. For example, consider the following sentence:

We can tokenize it into a sequence of words, denoted by

Alternatively, we can write it as a sequence of phrases by using a parser, denoted by

The simplest way to build a multi-scale model is to encode \(\mathbf{x}_{\mathrm{words}}\) and \(\mathbf{x}_{\mathrm{phrases}}\) using two separate Transformer encoders. The outputs of these encoders are then combined in some way. This leads to the same form as Eq. (17 ), and we can view this model as an instance of the general multi-encoder architecture.

Both \(\mathbf{x}_{\mathrm{words}}\) and \(\mathbf{x}_{\mathrm{phrases}}\) can be viewed as sequences of tokens. For example, \(\mathbf{x}_{\mathrm{words}}\) has nine word-based tokens, and \(\mathbf{x}_{\mathrm{phrases}}\) has three phrase-based tokens⁵. However, enumerating all possible phrases would result in a huge vocabulary. We therefore need a method to represent each phrase as an embedding in a computationally efficient manner. By treating phrase embedding as a sequence modeling problem, it becomes straightforward to learn sub-sequence representations by employing various sequence models. We now have a two-stage learning process. In the first stage, we learn the embeddings of input units at different scales using separate models. In the second stage, we learn to encode sequences at different scales using a multi-branch model.

More generally, we do not need to restrict ourselves to linguistically meaningful units in multi-scale representation learning. For example, we can learn sub-word segmentations from data and represent an input sentence as a sequence of sub-words. This results in a hierarchical representation of the sentence (e.g., sub-words \(\to\) words \(\to\) phrases). While the learned sub-words may not have linguistic meanings, they provide new insights into modeling words and phrases, as well as a new scale of features. Also, we do not strictly need to develop multiple encoders for multi-scale modeling. An alternative approach is to incorporate representations at different scales directly into the multi-head self-attention modules, which makes it easier to model the interactions among different scales [28], [29].

A problem with the approaches described above, however, is that the representations (or attention weight matrices) learned at different scales are of different sizes. For example, in the scenarios above, the representation learned from \(\mathbf{x}_{\mathrm{words}}\) is a \(9 \times d\) matrix, while the representation learned from \(\mathbf{x}_{\mathrm{phrases}}\) is a \(3 \times d\) matrix. A simple solution to this problem is to perform upsampling on the phrase-based representation to expand it to a \(9 \times d\) matrix. Likewise, we can perform downsampling on the word-based representation to shrink it to a \(3 \times d\) matrix. The combination model \(\mathrm{Combine}(\cdot)\) can then be implemented as described in Section 3.3.

It is worth noting that multi-scale modeling is widely discussed across several fields. In computer vision, for example, multi-scale modeling is often referred to as the process of learning a series of feature maps from the input image [30], [31]. Unlike the multi-branch models presented here, multi-scale vision Transformer models make use of the hierarchical nature of features when representing images. Systems of this kind are typically based on a stack of layers in which each layer learns features at a larger scale (e.g., a higher channel capacity) from the features produced by the previous layer.

3.5 Transformers as Syntax Learners↩︎

So far we have discussed syntax trees as constraints or priors on the encoding process so that we can make use of linguistic representations in learning neural networks. It is natural to wonder whether these neural models can learn some knowledge of linguistic structure directly from data, without human-designed linguistic annotations. This reflects a fundamental goal in NLP: to learn linguistic knowledge from data and encode it implicitly within models.

To explore the linguistic properties learned by NLP systems, a straightforward method is to examine the syntactic behaviors of the systems’ outputs. For example, we can examine whether the outputs of language generation systems contain grammatical errors. Another approach is to ask these systems to accomplish tasks that are linguistically meaningful, even though they were not explicitly trained to do so [3]. However, merely examining how model predictions exhibit syntactic abilities is not sufficient to fully answer the question. It is also possible that neural networks learn linguistic knowledge but do not utilize it during prediction [32]. Therefore, we need to investigate what is modeled and learned inside the neural networks themselves.

One approach to examining the latent linguistic structure in Transformer models is to develop probes to determine whether and to what extent these models capture linguistic notions such as dependency relations and parts-of-speech. A general probing approach involves extracting the internal representations of the models and probing them for linguistic phenomena. For Transformers, this is usually achieved by examining the attention maps and/or the outputs of attention layers. We then construct a probing predictor (or probing classifier) that takes these internal representations as input and outputs the corresponding linguistic notions [33]. The probing predictor can be based on either simple heuristics or parameterized models optimized specifically for the probing task. Recent work has shown that large-scale Transformer-based language models exhibit strong capabilities, referred to as emergent abilities, across various probing tasks. Figure 10 illustrates this probing process.

Many probing methods have been used in recent work on analyzing and understanding what is learned in neural encoders. Here we describe some of the popular ones.

• Trees. Given a trained Transformer encoder, it is easy to determine how likely it is that two words in a sentence share a linguistic relationship by computing the attention weight between them. We can use this quantity to define a metric measuring the syntactic distance between two words at positions \(i\) and \(j\) \[\begin{align} d_{s}(i,j) & = & 1 - \alpha(i,j). \end{align}\] By using this metric, it is straightforward to construct the minimum-spanning tree for the sentence. That is, we connect all the words to form a tree structure with the minimum total distance. This tree structure can be viewed as a latent tree representation of the sentence induced from the neural network. While this dependency-tree-like structure can serve as a source of learned syntactic information for downstream tasks, it says little about the model’s actual knowledge of syntax. An approach to aligning the representations in the encoder with linguistic structure is to learn to produce syntax trees that are consistent with human annotations. To achieve this, we need to develop a probing predictor trained on tree-annotated data. Suppose we have a human-annotated dependency tree for a given sentence. For each pair of words, we can obtain a distance \(\omega(i,j)\) by counting the number of edges between them. Then, we can learn a distance metric based on the internal representations of the encoder to approximate \(\omega(i,j)\). A simple form of such a metric relies on the Euclidean distance [34]. Let \(\mathbf{A} \in \mathbb{R}^{d \times k_s}\) be a parameter matrix. The squared Euclidean distance is given by \[\begin{align} d_{s}^2(i,j) & = & \left\| (\mathbf{h}_i - \mathbf{h}_j) \mathbf{A} \right\|_2^{2}, \end{align}\] where \(\mathbf{h}_i\) and \(\mathbf{h}_j\) are the representations produced by an encoding layer at positions \(i\) and \(j\)⁶. Given a set of tree-annotated sentences \(S\), we can optimize the model by \[\begin{align} \hat{\mathbf{A}} & = & \mathop{\mathrm{arg\,min}}_{\mathbf{A}} \sum_{s \in S} \frac{1}{|s|^2} \sum_{i \in s, j \in s} |\omega(i,j) - d_{s}^2(i,j)|, \end{align}\]

  where \(|s|\) is the length of the sentence \(s\), and \((i,j)\) indicates a pair of words in \(s\). The optimized model is then used to parse test sentences via the minimum-spanning tree algorithm, and we can compare the parse trees against the human-annotated trees. To obtain directed trees, which are standard forms of dependency syntax, one can update the above model by considering the relative distance of a word to the root. More details can be found in [34]’s work. Here the probing predictor functions similarly to a neural parser, trained to predict a syntax tree based on a representation of the input sentence. This idea can be extended to other forms of syntactic structure, such as phrase structure trees [35].

• Syntactic and Semantic Labels. Many syntactic and semantic parsing tasks can be framed as problems of predicting linguistic labels given a sentence or its segments. A simple example is part-of-speech tagging in which each word of a sentence is labeled with a word class. A probe for part-of-speech tagging can be a classifier that takes a representation \(\mathbf{h}_j\) each time and outputs the corresponding word class. One general probing approach to these problems is edge probing [36], [37]. Given a sentence, a labeled edge is defined as a tuple \[\begin{align} (\mathrm{span}_1,\mathrm{span}_2,\mathrm{label}) \nonumber, \end{align}\]

  where \(\mathrm{span}_1\) is a span \([i_1,j_1]\), and \(\mathrm{span}_2\) is an optional second span \([i_2,j_2]\), and \(\mathrm{label}\) is the corresponding label. Our goal is to learn a probe to predict \(\mathrm{label}\) given \(\mathrm{span}_1\) and \(\mathrm{span}_2\). For example, for part-of-speech tagging, \(\mathrm{span}_1\) is a unit span \([j,j]\) for each position \(j\), \(\mathrm{span}_2\) is an empty span, and \(\mathrm{label}\) is the part-of-speech tag corresponding to the \(j\)-th word of the sentence. For dependency parsing and coreference resolution, \(\mathrm{span}_1\) and \(\mathrm{span}_2\) are two words or entities, and \(\mathrm{label}\) is the relationship between them. For constituency parsing, \(\mathrm{span}_1\) is a span of words, \(\mathrm{span}_2\) is an empty span, and \(\mathrm{label}\) is the syntactic category of the tree node yielding \(\mathrm{span}_1\). In simple cases, the probing model can be a multi-layer feed-forward neural network with a softmax output layer. As usual, this model is trained on labeled data, and then tested on new data.

• Surface Forms of Words and Sentences. Probing tasks can also be designed to examine whether the representations encode the surface information of sentences or words [38], [39]. A simple sentence-level probing task is sentence length prediction. To do this, we first represent the sentence as a single vector \(\mathbf{h}\)⁷, and then build a classifier to categorize \(\mathbf{h}\) into the corresponding length bin. Similarly, probes can be built to predict whether two words at positions \(i\) and \(j\) are reordered in the sentence given \(\mathbf{h}_i\) and \(\mathbf{h}_j\). Also, we can develop probes to address conventional problems in morphology. For example, we reconstruct the word at position \(j\) or predict its sense with the representation \(\mathbf{h}_j\). In addition, probing tasks can be focused on particular linguistic problems, for example, numeracy [40] and function words [41].

• Cloze. Of course, we can probe neural models for problems beyond syntax and morphology. One perspective on large-scale pre-trained Transformer models is to view them as knowledge bases containing facts about the world. It is therefore tempting to see if we can apply them to test factual knowledge. A simple method is to ask a probe to recover the missing item of a sentence [42]. For example, if we have a cloze test

  Shiji was written by .

  we wish the probe to give an answer Sima Qian because there is a subject-relation-object fact (Shiji, written-by, Sima Qian). This probe can simply be a masked language model that is widely used in self-supervised learning of Transformer encoders.

In NLP, probing is closely related to the pre-training of large language models. In general, we can view probing tasks as downstream applications of these pre-trained models, though probing is ordinarily used as an analytical tool to quickly test the capabilities of the models. Ideally, we want to develop a probe that makes the best use of the internal representations to solve a given task. However, when a probe is complex and sufficiently well-trained, it becomes difficult to determine whether the problem was solved by the rich information in the representations or by the modeling capacity of the probe itself. A common way to isolate the contribution of the representations is to compare the probes against reasonable baselines or to conduct comparisons using control tasks [33], [43].

4 Improved Architectures↩︎

In this section we present several improvements to the vanilla Transformer model. Unlike the previous section, most of the improvements are from the perspective of machine learning, rather than linguistics.

4.1 Locally Attentive Models↩︎

The self-attention methods discussed in Section 2.3 can be viewed as learning representations over the entire input sequence. While this global attention mechanism provides a strong ability to capture long-distance dependencies, it does not explicitly model local information. Here, we consider a few techniques that attempt to explicitly capture locality in representations.

4.1.1 Priors for Local Modeling↩︎

One of the simplest ways to introduce local context modeling into Transformers is to add a penalty term to the attention function, thereby discouraging large attention weights between distant positions. On the encoder side, this leads to a form that we have already encountered several times in this chapter: \[\begin{align} \mathrm{AttLocal}_{\mathrm{self}}(\mathbf{H}) & = & \mathrm{Softmax} \left(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}} - \gamma \cdot \mathbf{G} \right) \mathbf{H}^{v}, \label{eq:self-attention-local} \end{align}\tag{18}\]

where \(\gamma\) is the weight (or temperature) of the penalty term, and \(\mathbf{G} \in \mathbb{R}^{m \times m}\) is the penalty matrix. Each entry \(G(i,j)\) indicates how heavily we penalize the model for attending to position \(j\) from position \(i\). A simple choice for \(G(i,j)\) is the absolute distance between \(i\) and \(j\), for example \[\begin{align} G(i,j) & = & |i - j|. \end{align}\]

Alternatively, \(G(i,j)\) can be defined using a Gaussian penalty function [44] \[\begin{align} G(i,j) & = & \frac{(i - j)^2}{2 \sigma_i^2}, \end{align}\]

where \(\sigma_i\) is the standard deviation of the Gaussian distribution. For a given \(i\), both penalty terms increase (either linearly or exponentially) as the distance \(|i-j|\) grows.

This method can be extended to the cross-attention model, as follows \[\begin{align} \mathrm{AttLocal}_{\mathrm{cross}}(\mathbf{H},\mathbf{S}) & = & \mathrm{Softmax} \left(\frac{\mathbf{S}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}} - \gamma \cdot \mathbf{G} \right) \mathbf{H}^{v}, \label{eq:self-attention-cross} \end{align}\tag{19}\]

where \(\mathbf{G}\) is an \(n \times m\) matrix. Each entry of \(\mathbf{G}\) can be defined as \[\begin{align} G(i,j) & = & \frac{(\mu_i - j)^2}{2 \sigma_i^2}, \end{align}\]

where \(\mu_i\) is the mean of the Gaussian distribution over the source-side positions. Both \(\mu_i\) and \(\sigma_i\) can be determined using heuristics. Alternatively, we can develop additional neural networks to model them, learning their parameters jointly with the rest of the model. For example, we could use a feed-forward neural network to predict \(\mu_i\) given \(\mathbf{s}_i\).

An alternative to Eq. (18 ) (or Eq. (19 )) is to treat the penalty term as a separate distribution and linearly combine it with the original attention model. For example, we can define the self-attention model as \[\begin{align} \mathrm{AttLocal}_{\mathrm{self}}(\mathbf{H}) & = & \left((1-\beta) \cdot \mathrm{Softmax} \left(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}} \right) + \beta \cdot \mathrm{Softmax} \left( - \gamma \cdot \mathbf{G} \right) \right) \mathbf{H}^{v}, \end{align}\]

where \(\beta \in [0,1]\) is the linear combination coefficient. Note that, to avoid manually tuning the hyperparameter \(\beta\), we can use a gating network to dynamically predict \(\beta\) and train it end-to-end.

Another alternative is to use a multiplicative mask to incorporate the prior into the modeling, similar to Eq. (15 ). This is given by \[\begin{align} \mathrm{AttLocal}_{\mathrm{self}}(\mathbf{H}) & = & \mathrm{Softmax} \left(\frac{\mathbf{H}^{q} [\mathbf{H}^{k}]^\top}{\sqrt{d}} \odot \mathbf{G}' \right) \mathbf{H}^{v}. \label{eq:self-attention-local-multiplicative} \end{align}\tag{20}\]

Here, \(\mathbf{G}' \in [0,1]^{m \times m}\) is a matrix with entries in \([0,1]\). The scalar \(G'(i,j)\) takes a maximum value of 1 when \(i=j\) and decays as \(j\) moves further away from \(i\). \(G'(i,j)\) can be obtained by normalizing \(-G(i,j)\) over all \(j\), or by using other decaying functions.

4.1.2 Local Attention↩︎

The term local attention has been used broadly to describe a wide range of problems and to refer to many different models in the NLP literature. The methods discussed above are those that impose soft constraints on attention models. In fact, local attention has its origins in attempts to restrict the scope of attention models for modeling and computational considerations [7]. Research in this area often focuses on introducing hard constraints, so that the resulting models can focus on parts of the input and ignore the rest. For example, we can predict a span of source-side positions to attend to, given a target-side position [44]–[46]. Also, attention spans can be induced from syntax trees, for example, syntactic subtree structures can help narrow the attention scope. Thus, many syntax-constrained models can be viewed as instances of local attention (see Section 3.4). In addition, the concept of local attention can be generalized to a rich set of models, such as sparse attention models, although these models are often discussed in the context of efficient machine learning methods. We will see a few examples in Section 5.

In deep learning, one of the most widely used architectures for learning features from a restricted region of the input is CNNs. It is thus interesting to consider methods of combining CNNs and Transformer models to obtain the benefits of both approaches, for example, CNNs deal with short-term dependencies, and self-attention models deal with long-term dependencies. One approach is to build a two-branch sequence model where one branch is based on CNNs and the other is based on self-attention models [26]. Another approach is to incorporate CNN layers into Transformer blocks in a way that enables the model to learn both local and global representations [47], [48].

4.1.3 Relative Positional Embedding↩︎

Relative positional embedding, also known as relative positional representation (RPR), is an improvement to the absolute positional embedding method used in standard Transformers[49], [50]. The idea of RPR is that we model the distance between two positions in a sequence rather than giving each position a fixed representation. As a result, we have a pair-wise representation \(\mathrm{PE}(i,j)\) for two positions \(i\) and \(j\). One simple way to define \(\mathrm{PE}(i,j)\) is to consider it as a lookup table for all pairs of \(i\) and \(j\). More specifically, let \(\mathbf{u}_\pi\) be a \(d\)-dimensional representation for a given distance \(\pi\). The form of \(\mathrm{PE}(i,j)\) in the vanilla RPR method is given by \[\begin{align} \mathrm{PE}(i,j) & = & \mathbf{u}_{\mathrm{clip}(j-i,k_{\mathrm{rpr}})}, \end{align}\]

where \(\mathrm{clip}(x,k_{\mathrm{rpr}})\) is a function that clips \(x\) in the interval \([-k_{\mathrm{rpr}},k_{\mathrm{rpr}}]\) \[\begin{align} \mathrm{clip}(x,k_{\mathrm{rpr}}) & = & \max\{-k_{\mathrm{rpr}},\min\{x,k_{\mathrm{rpr}}\}\}. \label{eq:transformer-rpr-clip} \end{align}\tag{21}\]

Thus, we have a model with parameters \[\begin{align} \mathbf{U}_{\mathrm{rpr}} & = & \begin{bmatrix} \mathbf{u}_{-k_{\mathrm{rpr}}} \\ \vdots \\ \mathbf{u}_{0} \\ \vdots \\ \mathbf{u}_{k_{\mathrm{rpr}}} \end{bmatrix}. \end{align}\]

While this notation is somewhat informal, we can view \(\mathbf{U}_{\mathrm{rpr}}\) as a matrix \(\in \mathbb{R}^{(2k_{\mathrm{rpr}}+1) \times d}\), and select a row corresponding to \(\mathrm{clip}(j-i,k_{\mathrm{rpr}})\) when computing RPR for given \(i\) and \(j\).

Using the above method, we can define three RPR models \(\mathrm{PE}^{q}(i,j)\), \(\mathrm{PE}^{k}(i,j)\) and \(\mathrm{PE}^{v}(i,j)\) for queries, keys, and values, respectively. Then, following the form of Eq. (9 ), the output of the self-attention model at position \(i\) can be written as \[\begin{align} \mathbf{c}_i & = & \sum_{j=1}^{m} \alpha_{i,j} \left[\mathbf{h}_{j}^{v} + \mathrm{PE}^{v}(i,j) \right] \nonumber \\ & = & \sum_{j=1}^{m} \alpha_{i,j} \mathbf{h}_{j}^{v} + \sum_{j=1}^{m} \alpha_{i,j} \mathrm{PE}^{v}(i,j), \label{eq:transformer-rpr-representation} \end{align}\tag{22}\]

where \(\mathbf{h}_{j}^{v}\) is the \(j\)-th row vector of \(\mathbf{H}^{v}\). This representation comprises two components: \(\sum_{j=1}^{m} \alpha_{i,j} \mathbf{h}_{j}^{v}\) is the basic representation, and \(\sum_{j=1}^{m} \alpha_{i,j} \mathrm{PE}^{v}(i,j)\) is the positional representation.

The attention weight \(\alpha_{i,j}\) is computed in a regular way, but with additional terms \(\mathrm{PE}^{q}(i,j)\) and \(\mathrm{PE}^{k}(i,j)\) added to each query and key: \[\begin{align} \alpha_{i,j} & = & \mathrm{Softmax}(\frac{[\mathbf{h}_{i}^{q} + \mathrm{PE}^{q}(i,j)] [\mathbf{h}_{j}^{k} + \mathrm{PE}^{k}(i,j)]^\top}{\sqrt{d}}). \label{eq:transformer-rpr-alpha} \end{align}\tag{23}\]

Figure 11 shows the Transformer encoder architectures with and without RPR. When RPR is adopted, \(\mathrm{PE}^{q}(i,j)\), \(\mathrm{PE}^{k}(i,j)\), \(\mathrm{PE}^{v}(i,j)\) are directly fed to each self-attention sub-layer, and so we can make better use of positional information for sequence modeling. Note that the use of the clipping function (see Eq. (21 )) makes the modeling simple because we do not need to distinguish the relative distances for cases \(|j-i| \ge k_{\mathrm{rpr}}\). This clipped distance-based model can lead to better modeling in local context windows.

Eqs. (22 ) and (23 ) provide a general approach to position-sensitive sequence modeling. There are many variants of this model. In the early work on RPR by [49], the positional representations for queries are removed, and the model works only with \(\mathrm{PE}^{k}(i,j)\) and \(\mathrm{PE}^{v}(i,j)\), like this \[\begin{align} \alpha_{i,j} & = & \mathrm{Softmax}(\frac{\mathbf{h}_{i}^{q} [\mathbf{h}_{j}^{k} + \mathrm{PE}^{k}(i,j)]^\top}{\sqrt{d}}). \label{eq:transformer-rpr-alpha-simple} \end{align}\tag{24}\]

By contrast, some variants attempt to improve the RPR model when computing attention weights but ignore \(\mathrm{PE}^{v}(i,j)\) when computing value representations [51], [52]. Instead of treating RPR as an additive term to each representation, researchers have also explored other ways of introducing RPR into the Transformer [53], [54]. We refer interested readers to these papers for more details.

4.2 Deep Models↩︎

Many state-of-the-art NLP systems are based on deep Transformer models. For example, recent large language models generally comprise tens of Transformer layers (or more precisely, hundreds of layers), demonstrating strong performance on many tasks [55], [56]. By stacking Transformer layers, it is straightforward to construct a deep model. However, as is often the case, training very deep neural networks is challenging. One difficulty arises from the fact that the error surfaces of deep neural networks are highly non-convex and contain many local optima, making the training process likely to get stuck. While there are optimization algorithms that can help alleviate this problem, most practical approaches rely on gradient-based methods for optimizing deep neural networks. As a result, training a model with many Transformer layers becomes challenging due to vanishing and exploding gradients during back-propagation. Here, we consider several techniques for training deep Transformer models.

4.2.1 Re-thinking the Pre-Norm and Post-Norm Architectures↩︎

As introduced previously, a Transformer sub-layer is a residual network where a shortcut is created to add the input of the network directly to the output of the sub-layer. This allows gradients to flow more directly from the output back to the input, mitigating the vanishing gradient problem. In general, a residual connection in a Transformer is used together with a layer normalization unit to form a sub-layer. This leads to two types of architectures, called post-norm and pre-norm. Specifically, recall from Section 2.4 that the post-norm architecture can be expressed as: \[\begin{align} \mathbf{z}^{l} & = & \mathrm{LNorm}(F^{l}(\mathbf{z}^{l-1}) + \mathbf{z}^{l-1}), \label{eq:post-norm-general} \end{align}\tag{25}\]

where \(\mathbf{z}^{l}\) and \(\mathbf{z}^{l-1}\) are the output and input of the sub-layer \(l\), respectively, and \(F^{l}(\cdot)\) is the core function of this sub-layer. The pre-norm architecture takes the identity mapping \(\mathbf{z}^{l-1}\) outside the layer normalization function, given in the form \[\begin{align} \mathbf{z}^{l} & = & F^{l}(\mathrm{LNorm}(\mathbf{z}^{l-1})) + \mathbf{z}^{l-1}. \label{eq:pre-norm-general} \end{align}\tag{26}\]

Consider the difference between the information flow in these two architectures:

• The post-norm architecture prevents the identity mapping of the input from being directly added to the final output of the sub-layer. This is not a true residual network because all the information is passed through a non-linear function (i.e., the layer normalization unit). Thus, the post-norm architecture is less effective for back-propagation. [57] show that the gradient of the loss of an \(L\)-sub-layer Transformer network with respect to \(\mathbf{z}^{l}\) is given by \[\begin{align} \frac{\partial E}{\partial \mathbf{z}^{l}} & = & \frac{\partial E}{\partial \mathbf{z}^{L}} \cdot \prod_{k=l}^{L-1} \frac{\partial \mathrm{LNorm}(\mathbf{v}^k)}{\partial \mathbf{v}^k} \cdot \prod_{k=l}^{L-1} \left( 1 + \frac{\partial F^{k}(\mathbf{z}^{k})}{\partial \mathbf{z}^{k}} \right), \end{align}\]

  where \(\mathbf{z}^{L}\) is the output of the last layer, \(\mathbf{v}^k\) is short for \(F^{k}(\mathbf{z}^{k-1})\), and \(E\) is the error measured by some loss function. \(\frac{\partial \mathrm{LNorm}(\mathbf{v}^k)}{\partial \mathbf{v}^k}\) and \(\frac{\partial F^{k}(\mathbf{z}^{k})}{\partial \mathbf{z}^{k}}\) are the gradients of the layer normalization function and the core function, respectively. Although the equation here appears complex, we see that \(\prod_{k=l}^{L-1} \frac{\partial \mathrm{LNorm}(\mathbf{v}^k)}{\partial \mathbf{v}^k}\) is simply a product of \(L-l\) factors. This means that the error gradient will be rescaled repeatedly as \(L\) becomes larger, leading to a higher risk of vanishing or exploding gradients in deeper models.

• The pre-norm architecture describes a standard residual neural network where the input of the entire network is added directly to its output. We can write the gradient of the error at \(\mathbf{z}^{l}\) as: \[\begin{align} \frac{\partial E}{\partial \mathbf{z}^{l}} & = & \frac{\partial E}{\partial \mathbf{z}^{L}} \cdot \left( 1 + \sum_{k=l}^{L-1} \frac{\partial F^{k}(\mathrm{LNorm}(\mathbf{z}^{k}))}{\partial \mathbf{z}^{k}} \right) \nonumber \\ & = & \frac{\partial E}{\partial \mathbf{z}^{L}} + \frac{\partial E}{\partial \mathbf{z}^{L}} \cdot \sum_{k=l}^{L-1} \frac{\partial F^{k}(\mathrm{LNorm}(\mathbf{z}^{k}))}{\partial \mathbf{z}^{k}}. \end{align}\]

  It is easy to see that \(\frac{\partial E}{\partial \mathbf{z}^{l}}\) receives direct feedback regarding the error, because the first term on the right-hand side (i.e., \(\frac{\partial E}{\partial \mathbf{z}^{L}}\)) is the gradient of the model output, which is independent of the network depth.

The use of the pre-norm architecture also helps optimization during early gradient descent steps. For example, it has been found that pre-norm Transformer models can be trained using a larger learning rate in the early stages of training, instead of requiring a gradual learning rate warmup from a small value [58].

While the pre-norm architecture facilitates easier optimization of deep Transformer models, we cannot simply conclude that it is a universally better choice compared to the post-norm architecture. In fact, both post-norm and pre-norm Transformer models have been successfully used in many applications. For example, the post-norm architecture is widely used in BERT-like models, while the pre-norm architecture is the more popular choice for recent generative large language models. Broadly, these two architectures provide different paradigms for designing deep Transformer models, each with distinct advantages and disadvantages. The post-norm architecture forces the representation to be learned through more non-linear functions, resulting in a highly expressive but complicated model that is relatively hard to train. By contrast, the pre-norm architecture makes the training of Transformer models easier, but it may be less expressive than its post-norm counterpart if the learned model becomes overly dependent on the shortcut paths.

An improvement to these architectures is to control the extent to which we want to “skip” a sub-layer. A simple way to do this is to weight different paths rather than treating them equally. For example, a scalar factor of a residual connection can be introduced to determine how heavily we weight this residual connection relative to the path of the core function [59]–[61]. A more general form of this model is given by \[\begin{align} \mathbf{z}^{l} & = & \mathrm{LNorm}(F^{l}(\mathbf{z}^{l-1}) + \beta \cdot \mathbf{z}^{l-1}) + \gamma \cdot \mathbf{z}^{l-1}, \end{align}\]

where \(\beta\) is the weight of the identity mapping inside the layer normalization function, and \(\gamma\) is the weight of the identity mapping outside the layer normalization function. Clearly, both the post-norm and pre-norm architectures can be seen as special cases of this equation. That is, if \(\beta=1\) and \(\gamma=0\), then it will become Eq. (25 ); if \(\beta=0\) and \(\gamma=1\), it will become Eq. (26 ). This model provides a multi-branch view of building residual blocks. The input to this block can be computed through multiple paths with different modeling complexities. When \(\beta\) and \(\gamma\) are small, the representation is forced to be learned through a “deep” model with multiple layers of cascaded non-linear units. In contrast, when \(\beta\) and \(\gamma\) are large, the representation is more likely to be learned using a “shallow” model with fewer layers. To determine the optimal choices of \(\beta\) and \(\gamma\), one can give them fixed values by considering some theoretical properties or system performance on validation sets, or compute these values by using additional functions that can be trained to do so [62]. It should be emphasized that many other types of architecture can be considered in the design of a Transformer sub-layer. It is possible, for instance, to introduce more layer normalization units into a sub-layer [63], [64], or, conversely, to simply remove them from a sub-layer [65].

4.2.2 Parameter Initialization↩︎

As with other deep neural networks, there is considerable interest in developing parameter initialization methods for deep Transformer models in order to facilitate optimization in a more favorable region of the parameter space. However, initialization is a broad topic in the optimization of machine learning models, and a general discussion of this topic lies beyond the scope of this section. Here, we will focus on specific parameter initialization methods used in Transformer-based systems rather than general optimization problems.

While the parameters of a neural network can be initialized in various ways, most practical systems adopt simple techniques to provide appropriate initial values. Consider, for example, the Xavier initialization for a parameter matrix \(\mathbf{W} \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}}\) [66]. We define a variable \(\eta\) as follows \[\begin{align} \eta & = & \mathrm{gain} \cdot \sqrt{\frac{6}{d_{\rm{in}} + d_{\rm{out}}}}, \end{align}\]

where \(\mathrm{gain}\) is a hyper-parameter which equals 1 by default. Then, each entry of \(\mathbf{W}\) can be initialized by using a uniform distribution \[\begin{align} W & \sim & U \left( -\eta,\eta \right) \end{align}\]

or, alternatively, using a Gaussian distribution \[\begin{align} W & \sim & \mathcal{N} \left(0, \eta^2 \right). \end{align}\]

This method can be easily adapted to initialize Transformer models with a large number of layers. One common approach is to find a more suitable value for \(\mathrm{gain}\) by taking into account the fact that the optimal initial conditions for optimization might differ across neural networks of varying depths. For example, one can increase the value of \(\mathrm{gain}\) as the depth of the model grows. Thus, \(\mathrm{gain}\) can be defined as a function of the network depth in the following form \[\begin{align} \mathrm{gain} & = & a \cdot L^{b}, \label{eq:transformer-init-gain-general} \end{align}\tag{27}\]

where \(a\) is a scalar, and \(L^{b}\) is the network depth raised to the power of \(b\). Typically, both \(a\) and \(b\) are positive numbers, indicating a preference for larger initial parameter values in deeper models. For example, [67] demonstrate that by choosing appropriate values for \(a\) and \(b\), a very deep Transformer model can be successfully trained.

Eq. (27 ) assigns \(\mathrm{gain}\) the same value for all of the sub-layers. However, it is found that the norm of gradients becomes smaller as a sub-layer is located farther from the output layer. This consistent application of \(\mathrm{gain}\) across the entire model could result in under-training of the lower layers due to the gradient vanishing problem. For this reason, one can develop methods that are sensitive to the position of a sub-layer in the neural network. The general form of such methods is given by \[\begin{align} \mathrm{gain} & = & \frac{a}{l^{b}}. \label{eq:transformer-init-gain-sub-layer-depth} \end{align}\tag{28}\]

Here \(l\) denotes the depth of a sub-layer. If \(l\) is larger (i.e., the sub-layer is closer to the output), \(\mathrm{gain}\) will be smaller and the corresponding parameters will be set to smaller values. An example of this method can be found in the work of [68].

It is also, of course, straightforward to apply general methods of initializing deep multi-layer neural networks to Transformer models. An example is to consider the Lipschitz constant in parameter initialization, which has been shown to help improve the stability of training deep models [69], [70]. Another approach is to use second-order methods to estimate the proper values of the parameters. For example, one can compute the Hessian of each parameter matrix to model its curvature [71].

For models with a large number of layers, it is also possible to pre-train some of the layers via smaller models and use their trained parameters to initialize bigger models [72]. That is, we first obtain a rough estimation of the parameters in a cheap way, and then continue the training process on the whole model as usual. These methods fall into a class of training methods, called model growth or depth growth.

As a simple example, consider a Transformer model (e.g., a Transformer encoder) of \(2L\) sub-layers. We can train this model by using the shallow-to-deep training method [73]. First, we train an \(L\)-sub-layer model (call it the shallow model) in a regular way. Then, we create a \(2L\)-sub-layer model (call it the deep model) by stacking the shallow model twice, and further train this deep model. To construct deeper models, this procedure can be repeated multiple times, say, we start with a model of \(L\) sub-layers, and obtain a model of \(L \cdot 2^{I}\) after \(I\) iterations. Note that many pre-trained models are used in a similar manner. For example, for BERT-like methods, a Transformer encoder is trained on large-scale data, and the optimized parameters are then used to initialize downstream systems.

4.2.3 Layer Fusion↩︎

Another problem with training a deep Transformer model is that the final prediction is conditioned on the last layer of the neural network. While the use of residual connections enables direct access to lower-level layers from a higher-level layer, there is still a “long” path for passing information from the bottom to the top. One simple way to address this is to create residual connections that skip more layers. For example, consider a group of \(L\) Transformer sub-layers. For the sub-layer at depth \(l\), we can introduce \(l-1\) residual connections, each connecting this sub-layer to a previous sub-layer. In this way, we develop a densely connected network where each sub-layer takes as input the outputs of all previous sub-layers [74]. The output of the last sub-layer can be seen as a combination of representations at different levels.

Following the notation used in the previous subsections, we denote the output of the sub-layer at depth \(l\) by \(\mathbf{z}^{l}\), and denote the function of the sub-layer by \(\mathrm{Layer}^{l}(\cdot)\). Then, \(\mathbf{z}^{l}\) can be expressed as \[\begin{align} \mathbf{z}^{l} & = & \mathrm{Layer}^{l}(\mathbf{z}^{1}, \cdots ,\mathbf{z}^{l-1}). \end{align}\]

We can simply view \(\mathrm{Layer}^{l}(\cdot)\) as a function that fuses the information from \(\{\mathbf{z}^{1}, \cdots ,\mathbf{z}^{l-1}\}\). There are many possible choices for \(\mathrm{Layer}^{l}(\cdot)\). For example, a simple form of \(\mathrm{Layer}^{l}(\cdot)\) is given by \[\begin{align} \mathrm{Layer}^{l}(\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}) & = & \mathrm{LNorm}(F^{l}(\mathbf{Z}^{l})), \\ \mathbf{Z}^{l} & = & \phi(\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}). \end{align}\]

Here \(\phi(\cdot)\) takes the layer outputs \(\{\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}\}\) and fuses them into a single representation \(\mathbf{Z}^{l}\). A simple instance of \(\phi(\cdot)\) is average pooling which computes the sum of \(\{\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}\}\) divided by \(l-1\). See Table 2 for more examples of \(\phi(\cdot)\).

Adopting a similar architecture for a Transformer sub-layer, we can also consider a post-norm form \[\begin{align} \mathrm{Layer}^{l}(\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}) & = & \mathrm{LNorm}(\mathbf{Z}^{l}), \tag{29} \\ \mathbf{Z}^{l} & = & \phi(F^{l}(\mathbf{z}^{l-1}),\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}), \tag{30} \end{align}\]

or a pre-norm form \[\begin{align} \mathrm{Layer}^{l}(\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}) & = & \mathbf{Z}^{l} \tag{31}, \\ \mathbf{Z}^{l} & = & \phi(\mathrm{LNorm}(F^{l}(\mathbf{z}^{l-1})),\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}). \tag{32} \end{align}\]

These models are quite general. For example, a standard post-norm encoder sub-layer can be recovered as a special case of Eqs. (29 30 ), if we remove the dependencies on sub-layers \(1\) to \(l-2\), and define \(\phi(\cdot)\) to be \[\begin{align} \phi(F^{l}(\mathbf{z}^{l-1}),\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}) & = & F^{l}(\mathbf{z}^{l-1}) + \mathbf{z}^{l-1}. \end{align}\]

Densely connected networks make it easier for information to flow through direct connections between sub-layers, but the resulting models are slightly more complex, especially when we use parameterized fusion functions. In practice, we typically add dense connections only to a subset of the sub-layers, so the overall network is not overly dense. For example, we might only add connections from lower sub-layers to the final few sub-layers. Thus, predictions can be made with direct access to different levels of representation [75].

4.2.4 Regularization↩︎

In machine learning, regularization is used to mitigate overfitting when training deep neural networks. It is therefore straightforward to apply regularization techniques to Transformer models. Since the topic of regularization is well studied, here we consider several methods that are particularly useful for training deep Transformer models.

One approach to regularizing deep Transformer models is to randomly skip sub-layers or layers during training [76], [77]. In each run of the model, we retain each sub-layer with a probability \(\rho\) and stack the selected sub-layers to form a new model. Thus, we essentially train an ensemble of different neural networks with shared architectures and parameters on the same dataset. In this way, each sub-layer learns to operate somewhat independently, and overfitting is reduced by preventing the co-adaptation of sub-layers. In fact, dropping out sub-layers (or layers) and dropping out neurons are essentially variations on the same theme. Sometimes, the method described here is referred to as sub-layer dropout or layer dropout.

At test time, we need to combine all possible networks to make predictions. A simple method to achieve this is to rescale the outputs of the stochastic components of the model [78]. As an example, suppose each sub-layer has a pre-norm architecture. Then, the output of the sub-layer at depth \(l\) is given by \[\begin{align} \mathbf{z}^{l} & = & \rho \cdot \mathrm{LNorm}(F^{l}(\mathbf{z}^{l-1})) + \mathbf{z}^{l-1}. \end{align}\]

Another idea is to share parameters across sub-layers. One of the simplest methods is to use the exact same parameters for all corresponding sub-layers [79]. For example, all the FFN sub-layers could be based on the same feed-forward network. This approach has a regularizing effect similar to methods that add the norms of parameter matrices to the loss function to penalize model complexity. For practical systems, there can be significant benefits to adopting a shared architecture, as it allows us to reuse the same sub-model to build a deep neural network, thereby substantially reducing the overall memory footprint. We will see further discussion regarding efficiency issues in Section 5.4.

4.3 Numerical Method-Inspired Models↩︎

A residual network computes its output as the sum of the identity mapping and a transformation of the input. Such a model can be interpreted as an Euler discretization of ordinary differential equations (ODEs) [80], [81]. To illustrate this idea, we consider a general form of residual networks \[\begin{align} \mathbf{z}^{l} & = & f^l \big(\mathbf{z}^{l-1} \big) + \mathbf{z}^{l-1}, \label{eq:pre-norm-ode-form} \end{align}\tag{33}\]

where \(f^l(\mathbf{z}^{l-1})\) denotes a function that takes an input variable \(\mathbf{z}^{l-1}\) and produces an output variable in the same space. Clearly, a Transformer sub-layer is a special case of this equation. For example, for a pre-norm Transformer, we have \(f^l(\cdot)=\mathrm{LNorm}(F^{l}(\cdot))\).

For notational simplicity, we rewrite the above equation in an equivalent form \[\begin{align} \mathbf{z}(l) & = & f \big(\mathbf{z}(l-1),l \big) + \mathbf{z}(l-1). \label{eq:pre-norm-ode-form-t} \end{align}\tag{34}\]

We use the notations \(\mathbf{z}(l)\) and \(f(\mathbf{z}(\cdot,l))\) to emphasize that \(\mathbf{z}(\cdot)\) and \(f(\cdot)\) are functions of \(l\). Here we assume that \(l\) is a discrete variable. If we relax \(l\) to a continuous variable and \(\mathbf{z}(l)\) to a continuous function of \(l\), then we can express Eq. (34 ) as \[\begin{align} \mathbf{z}(l) & = & \triangle l \cdot f \big(\mathbf{z}(l - \triangle l), l \big) + \mathbf{z}(l - \triangle l). \label{eq:pre-norm-ode-form-delta-f} \end{align}\tag{35}\]

This can be further written as \[\begin{align} \frac{\mathbf{z}(l) - \mathbf{z}(l - \triangle l)}{\triangle l } & = & f \big( \mathbf{z}(l - \triangle l),l \big). \label{eq:pre-norm-ode-form-delta-f-2} \end{align}\tag{36}\]

Taking the limit \(\triangle l \rightarrow 0\), we have an ODE \[\begin{align} \frac{\mathrm{d} \mathbf{z}(l)}{\mathrm{d} l} & = & f \big(\mathbf{z}(l),l \big). \label{eq:transformer-to-ode} \end{align}\tag{37}\]

We can therefore say that a pre-norm Transformer sub-layer (i.e., Eqs. (34 ) and (33 )) is an Euler discretization of the above ODE. This is an interesting result! A sub-layer essentially serves as an ODE solver.

Eqs. (34 ) and (33 ) represent standard forms of the Euler method. This method computes a new estimate of the solution by moving one step forward along \(l\) from the previous estimation. In general, two aspects can be considered in the design of numerical methods for ODEs.

• Linear Multi-step Methods. A linear multi-step method computes the current estimate of the solution by taking the estimations and derivative information from multiple previous steps. A general formulation of \(p\)-step methods can be expressed as \[\begin{align} \mathbf{z}(l) & = & \sum_{i=1}^{p} a_{i} \cdot \mathbf{z}(l-i) + h \sum_{i=1}^{p+1} b_{i} \cdot f \big(\mathbf{z}(l-i),\;l-i+1 \big), \label{eq:linear-multi-step-methods-general} \end{align}\tag{38}\] where \(h\) is the size of the step we move each time⁸, that is, \(\triangle l\) in Eqs. (35 ) and (36 ). \(\{a_i\}\) and \(\{b_i\}\) are coefficients of the solution points and derivatives in the linear combination. Given this definition, we can think of the Euler method as a single-step, low-order method of solving ODEs⁹.

• (Higher-order) Runge-Kutta Methods. Runge-Kutta (RK) methods and their variants provide ways to compute the next step solution by taking intermediate results in solving an ODE. As a result, we obtain higher-order methods but still follow the form of single-step methods, that is, the estimated solution is dependent only on \(\mathbf{z}(l-1)\) rather than on the outputs at multiple previous steps.

In fact, linear multi-step methods, though not explicitly mentioned, have been used in the layer fusion methods discussed in Section 4.2. For example, taking Eqs. (31 ) and (32 ) and a linear fusion function, a pre-norm sub-layer with dense connections to all previous sub-layers can be expressed as \[\begin{align} \mathrm{Layer}^{l}(\mathbf{z}^{1},\cdots,\mathbf{z}^{l-1}) & = & a_1 \cdot \mathbf{z}^{l-1} + \cdots + a_{l-1} \cdot \mathbf{z}^{1} + b_1 \cdot \mathrm{LNorm}(F^{l}(\mathbf{z}^{l-1})). \end{align}\]

This equation is an instance of Eq. (38 ) where we set \(h=1\) and remove some of the terms on the right-hand side.

It is also straightforward to apply Runge-Kutta methods to Transformer models [82]. Given an ODE as described in Eq. (37 ), an explicit \(p\)-th order Runge-Kutta solution is given by \[\begin{align} \mathbf{z}(l) & = & \mathbf{z}(l-1) + \sum_{i=1}^{p} \gamma_i \cdot \mathbf{g}_i, \\ \mathbf{g}_i & = & h \cdot f \big( \mathbf{z}(l-1) + \sum_{j=1}^{i-1} \beta_{i,j} \cdot \mathbf{g}_j,\;l - 1 + \lambda_i \cdot h \big). \label{eq:runge-kutta-general-g} \end{align}\tag{39}\]

Here \(\mathbf{g}_i\) represents an intermediate step which is present only during the above process. \(\{\gamma_i\}\), \(\{\beta_{i,j}\}\) and \(\{\lambda_i\}\) are coefficients that are determined by using the Taylor series of \(\mathbf{z}(l)\). To simplify the model, we assume that the same function \(f\) is used for all \(\{\mathbf{g}_i\}\). Then, we remove the dependency on the term \(l - 1 + \lambda_i \cdot h\) in \(f\), and rewrite Eq. (39 ) as \[\begin{align} \mathbf{g}_i & = & h \cdot f \big( \mathbf{z}(l-1) + \sum_{j=1}^{i-1} \beta_{i,j} \cdot \mathbf{g}_j\big), \label{eq:runge-kutta-general-g-new} \end{align}\tag{40}\]

where \(f(\cdot)\) is a function that is independent of \(i\).

As an example, consider the 4th-order Runge-Kutta (RK4) solution \[\begin{align} \mathbf{z}(l) & = & \mathbf{z}(l-1) + \frac{1}{6}(\mathbf{g}_1 + 2 \mathbf{g}_2 + 2 \mathbf{g}_3 + \mathbf{g}_4), \\ \mathbf{g}_1 & = & h \cdot f(\mathbf{z}(l-1)), \\ \mathbf{g}_2 & = & h \cdot f(\mathbf{z}(l-1) + \frac{1}{2} \mathbf{g}_1), \\ \mathbf{g}_3 & = & h \cdot f(\mathbf{z}(l-1) + \frac{1}{2} \mathbf{g}_2), \\ \mathbf{g}_4 & = & h \cdot f(\mathbf{z}(l-1) + \mathbf{g}_3). \end{align}\]

These equations define a new architecture of sub-layer. For example, by setting \(h=1\) and \(f(\cdot) = \mathrm{LNorm}(F^{l}(\cdot))\), we obtain an RK4 Transformer sub-layer, as shown in Figure 12. This method leads to a deep model because each sub-layer involves four evaluations of \(f(\cdot)\) in sequence. On the other hand, the resulting model is parameter efficient because we reuse the same function \(f(\cdot)\) within the sub-layer, without introducing additional parameters.

So far in this subsection, our discussion has focused on designing Transformer architectures inspired by dynamical systems. While the basic ODE model is continuous with respect to the depth \(l\), these methods still follow the general framework of neural networks in which \(l\) is treated as a discrete variable, and the representational power of the models is largely determined by this hyper-parameter. An alternative approach is to use neural ODEs to relax the “depth” into a continuous variable. In this way, we can obtain a continuous-depth model for computing the solution to ODEs. However, as a detailed discussion of neural ODEs lies beyond the scope of this chapter, we refer interested readers to related papers for more details [83]–[85].

4.4 Wide Models↩︎

Most of the methods that we have studied so far in this section are examples of learning and using deep models. Another design choice we generally face is determining the width of a neural network. Typically, the width of a Transformer model is defined as the dimensionality of a representation at some position in the input sequence, that is, the parameter \(d\). Increasing this width is a common method for obtaining a more complex and powerful model. For example, in the work of [1], a wide model (referred to as Transformer-Big) leads to significant improvements in translation quality for machine translation systems. Such wide models have been widely used to boost system performance on large-scale tasks [86], [87].

However, developing very wide Transformer models is difficult. One major challenge is that training such systems is computationally expensive. While the number of model parameters (the model size) grows linearly with \(d\), the time complexity of the model grows quadratically with \(d\) (see Table 1). In some NLP tasks, it has been found empirically that the training effort required to obtain satisfactory performance follows a scaling law [88]. These results emphasize the need for ways to improve training efficiency when enlarging \(d\).

One simple method is to incrementally grow the model along the dimension of \(d\), rather than training the model from scratch. Suppose we have an initial model with a \(d_1 \times d_1\) parameter matrix \(\mathbf{W}_1\), for example, the linear transformation of each query or key in some layer. We can train this model to obtain optimized \(\mathbf{W}_1\) in a regular way. Then, we want to extend this model to a wider model where \(\mathbf{W}_1\) is replaced by a \(d_2 \times d_2\) parameter matrix \(\mathbf{W}_2\). Let us assume for simplicity that \(d_2 = k d_1\). There are several ways to expand a \(d_1 \times d_1\) matrix to a \(k d_1 \times k d_1\) matrix. The simplest of these may be to use \(\mathbf{W}_1\) to fill \(\mathbf{W}_2\). We can write \(\mathbf{W}_2\) in the form \[\begin{align} \mathbf{W}_2 & = & \begin{matrix} k\;\textrm{times} & \;\\ \begin{bmatrix} \frac{\mathbf{W}_1}{\rho} & \cdots & \frac{\mathbf{W}_1}{\rho} \\ \vdots & & \vdots \\ \frac{\mathbf{W}_1}{\rho} & \cdots & \frac{\mathbf{W}_1}{\rho} \end{bmatrix} & \begin{rotate}{90}\textrm{ k times}\end{rotate} \\ \end{matrix}, \end{align}\]

where \(\rho\) is a hyper-parameter that is used to control the norm of \(\mathbf{W}_2\). For example, if \(\rho = k\), \(\mathbf{W}_2\) will have the same \(l_1\) norm as \(\mathbf{W}_1\). The above equation provides a good starting point for training the wide model, and we can train \(\mathbf{W}_2\) as usual after initialization. The procedure can be repeated a number of times for constructing a model with arbitrary width. Both this method and the depth growth method described in Section 4.2 are instances of the general method of model growth. In other words, we can obtain a larger model by extending a small model either vertically or horizontally, or both. Alternative methods for transforming \(\mathbf{W}_1\) to \(\mathbf{W}_2\) involve those considering other mathematical properties of the transformation [72]. These models can fall under the reusable neural networks where we are concerned with models and algorithms for transferring parameters from small models to (significantly) larger models [89].

A second difficulty in building a wide Transformer model is the large memory requirement. Since the feedforward network generally has a larger hidden layer than other parts of the model, it demands relatively more memory as the model becomes wider. Consider the feedforward network described in Section 2.5 \[\begin{align} \mathbf{H}_{\mathrm{out}} & = & \mathrm{FFN}(\mathbf{H}_{\mathrm{in}}) \nonumber \\ & = & \mathrm{ReLU}(\mathbf{H}_{\mathrm{in}} \cdot \mathbf{W}_h + \mathbf{b}_h) \cdot \mathbf{W}_f + \mathbf{b}_f, \label{eq:transformer-ffn-moe} \end{align}\tag{41}\]

where \(\mathbf{W}_h \in \mathbb{R}^{d \times d_{\mathrm{ffn}}}\) and \(\mathbf{W}_f \in \mathbb{R}^{d_{\mathrm{ffn}} \times d}\) are the parameters of the linear transformations. \(d_{\mathrm{ffn}}\) is typically several times larger than \(d\). Therefore, \(\mathbf{W}_h\) and \(\mathbf{W}_f\) will dominate the model size if \(d\) and \(d_{\mathrm{ffn}}\) have very large values.

In some cases, the size of the feedforward network may exceed the memory capacity of a single device. This problem can be addressed by using the mixture-of-experts (MoE) models [90]. An MoE model consists of \(M\) expert models \(\{e_1(\cdot),\cdots,e_M(\cdot)\}\). Given an input \(\mathbf{h}_{\mathrm{in}} \in \mathbb{R}^{d}\), each expert model produces an output \(e_k(\mathbf{h}_{\mathrm{in}})\). The output of the MoE model is a linear combination of \(\{e_1(\mathbf{h}_{\mathrm{in}}),\cdots,e_M(\mathbf{h}_{\mathrm{in}})\}\), given by \[\begin{align} \mathbf{h}_{\mathrm{out}} & = & \sum_{i=1}^{M} g_i(\mathbf{h}_{\mathrm{in}}) \cdot e_i(\mathbf{h}_{\mathrm{in}}), \end{align}\]

where \(g(\cdot)\) is a gating model (also called routing model). Its output is a vector \(g(\mathbf{h}_{\mathrm{in}}) = \begin{bmatrix} g_1(\mathbf{h}_{\mathrm{in}}) & \cdots & g_M(\mathbf{h}_{\mathrm{in}})\end{bmatrix}\) in which each entry \(g_i(\mathbf{h}_{\mathrm{in}})\) indicates the weight of the corresponding expert model. In many applications, it is assumed that \(g(\mathbf{h}_{\mathrm{in}})\) is a sparse vector. This means that only a small number of expert models are involved in computing the output. A widely-used form of \(g(\mathbf{h}_{\mathrm{in}})\) is given by using the softmax layer \[\begin{align} g(\mathbf{h}_{\mathrm{in}}) & = & \mathrm{Softmax}(\mathbf{h}_{\mathrm{in}} \cdot \mathbf{W}_{g}), \end{align}\]

where \(\mathbf{W}_{g} \in \mathbb{R}^{d \times M}\) is the parameter matrix of the layer. To enforce sparsity on \(g(\mathbf{h}_{\mathrm{in}})\), we can select the top-\(k\) entries of \(g(\mathbf{h}_{\mathrm{in}})\), that is, we set non-top-\(k\) entries to 0. An alternative method is to first perform top-\(k\) selection on \(\mathbf{h}_{\mathrm{in}} \cdot \mathbf{W}_{g}\) and then normalize the top-\(k\) entries using the softmax function.

Let \(\pi\) be the set of the indices of the top-\(k\) expert models. The MoE model with top-\(k\) routing has the following form \[\begin{align} \mathbf{h}_{\mathrm{out}} & = & \sum_{i \in \pi} g_i(\mathbf{h}_{\mathrm{in}}) \cdot e_i(\mathbf{h}_{\mathrm{in}}). \end{align}\]

An advantage of this approach is that we can distribute different expert models to different processors, making it possible to execute these models on parallel computing machines. In each run of the MoE model, either during training or inference, we only need to activate and use \(k\) expert models rather than all of the expert models. In this way, the MoE approach is automatically learning a sparse model by limiting the number of active expert models each time in training and inference. The sparsity is determined by the hyperparameter \(k\), e.g., a small value of \(k\) leads to a sparse model, and a large value of \(k\) leads to a dense model.

Let us return to the discussion of Eq. (41 ). It is straightforward to apply the MoE approach to feedforward neural networks. To simplify the discussion, consider the linear transformation of the first layer as shown in Eq. (41 ), that is, \(\mathbf{H}_{\mathrm{in}} \cdot \mathbf{W}_h\). We can approximate \(\mathbf{H}_{\mathrm{in}} \cdot \mathbf{W}_h\) in an MoE form \[\begin{align} \mathbf{H}_{\mathrm{in}} \cdot \mathbf{W}_h & \approx & \sum_{i \in \pi} g_i(\mathbf{H}_{\mathrm{in}}) \cdot e_i(\mathbf{H}_{\mathrm{in}}) \nonumber \\ & = & \sum_{i \in \pi} g_i(\mathbf{H}_{\mathrm{in}}) \cdot [\mathbf{H}_{\mathrm{in}} \cdot \mathbf{W}_h^{i}]. \label{eq:moe-ffn-w-h} \end{align}\tag{42}\]

Here \(\mathbf{W}_h\) is divided into \(M\) slices (or sub-matrices) \(\{\mathbf{W}_h^{1},\cdots,\mathbf{W}_h^{M}\}\), written as \[\begin{align} \mathbf{W}_h & = & \begin{bmatrix} \mathbf{W}_h^{1} & \cdots & \mathbf{W}_h^{M} \end{bmatrix}. \end{align}\]

Hence each expert model \(e_i(\mathbf{H}_{\mathrm{in}}) = \mathbf{H}_{\mathrm{in}} \cdot \mathbf{W}_h^{i}\) solves a sub-problem of the original linear mapping, and Eq. (42 ) can be thought of as a divide-and-conquer solution to the matrix multiplication problem.

We can, of course, treat any feedforward neural network as an expert model, resulting in the following model \[\begin{align} \mathbf{H}_{\mathrm{out}} & = & \sum_{i \in \pi} g_i(\mathbf{H}_{\mathrm{in}}) \cdot \mathrm{FFN}_i(\mathbf{H}_{\mathrm{in}}), \end{align}\]

where \(\mathrm{FFN}_i(\cdot)\) is a “small” feedforward neural network that has the same form as Eq. (41 ). This model is illustrated with an example in Figure 13. In practical implementations, all these expert models can be run in parallel on different devices, and so the resulting system is efficient.

Note that, from a perspective of machine learning, MoE is a general approach to combining different neural networks, each of which is developed to address a different aspect of the problem [91], [92]. The application here is just a special instance of the general framework of MoE. The approach is also often used to improve the overall performance of predictors, which can be discussed in the field of ensemble learning [93].

Another difficulty in developing massive Transformer models is training instability. As with many other exceptionally large neural networks, the straightforward optimization of a Transformer model with billions of parameters may lead to the model getting trapped in local minima and, occasionally, to large spikes in the loss during training [86], [87], [94]. Even with careful tuning of hyperparameters, training strategies, and initial parameter values, it is common to encounter situations where the training must be restarted from an earlier checkpoint to escape unstable optimization regions. One of the reasons for this training difficulty is that the standard implementations of linear algebra operations, such as matrix multiplication, can become numerically unstable when they operate on extremely large vectors and matrices. It is therefore crucial to improve training stability by employing numerically robust methods and mixed-precision techniques.

5 Efficient Models↩︎

Efficiency is an important consideration for many practical applications of Transformer models. For example, we may wish to run and/or train a Transformer model given memory and time constraints. Efficiency is not a single issue, but covers a wide range of problems. While these problems can be categorized in several different ways, there are two fundamental aspects one may consider in an efficiency problem.

• Time and Space Efficiencies. For a given problem, we wish the model to be small and fast, while remaining as accurate as possible in solving the problem. For example, in some machine translation applications, we may develop a model with a small number of parameters to fit into limited memory, and may develop a fast search algorithm to achieve low-latency translation. A practical difficulty here is that improving efficiency often leads to worse predictions. In many cases, we need to seek a trade-off between efficiency and accuracy.

• Scalability. When the problem is scaled up, we wish that the additional computational overhead required is minimized. For example, the training of a neural network is called efficient if it takes a reasonably short time to optimize it as the volume of training data increases. Another metric of efficiency is the amount of resources consumed in processing larger inputs. For instance, a machine translation system is inefficient in translating long sentences if the memory footprint and latency grow exponentially with the number of input words.

In this section, we will not discuss all the issues related to efficiency, which is a very broad topic. We instead consider the widely-used efficient approaches to Transformer-based sequence modeling and generation, some of which are refinements of model architectures, and some of which are architecture-agnostic approaches and could be used in other systems as well. Most of the discussions here are focused on developing lightweight and fast Transformer models that are relatively robust to long input and output sequences.

In general, the same optimization method can be applied to different modules of a Transformer-based system. To simplify the discussion, we will mostly consider self-attention sub-layers and FFN sub-layers in this section. Our discussion, however, is general and the methods presented here can be applied to other parts of a Transformer model, for example, cross-attention sub-layers.

5.1 Sparse Attention↩︎

In practice, the attention mechanisms used in Transformers are time-consuming, especially when the input sequences are long. To illustrate, consider a Transformer decoder that predicts a distribution over the vocabulary at each time step given the previous words. Suppose the sequence generated by the decoder has length \(n\) and the input to a self-attention sub-layer is an \(n \times d\) matrix \(\mathbf{S}\). First, \(\mathbf{S}\) is linearly transformed to obtain the queries \(\mathbf{S}^{q} \in \mathbb{R}^{n \times d}\), keys \(\mathbf{S}^{k} \in \mathbb{R}^{n \times d}\), and values \(\mathbf{S}^{v} \in \mathbb{R}^{n \times d}\). To simplify the notation in this subsection, we use \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\) to represent \(\mathbf{S}^{q}\), \(\mathbf{S}^{k}\), and \(\mathbf{S}^{v}\), respectively.

The output of the self-attention sub-layer can then be computed using \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{A} \mathbf{V}, \label{eq:efficient-self-attetention-decoder-as} \end{align}\tag{43}\]

where \(\mathbf{A}\) is an \(n \times n\) attention matrix or attention map \[\begin{align} \mathbf{A} & = & \mathrm{Softmax}(\frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d}} + \mathbf{M}). \label{eq:efficient-self-attetention-decoder-softmax} \end{align}\tag{44}\]

Here, \(\mathbf{M}\) is a masking matrix used to prevent the model from attending to future context at each position. Specifically, for a position \(i\), \(M(i,j) = 0\) if \(j \le i\), and \(M(i,j) = -\infty\) otherwise. Both the time and space complexity of the self-attention sub-layer scale quadratically with \(n\)¹⁰. As a result, standard self-attention becomes computationally prohibitive for large values of \(n\).

The standard implementation of the above model relies on dense matrix computations, such as the multiplications in Eqs. (43 ) and (44 ). One common approach to reducing both memory footprint and floating-point operations is sparsification. By assuming \(\mathbf{A}\) is a sparse matrix, only a fraction \(\varrho \cdot n^2\) of its entries remain non-zero, where \(\varrho\) is the sparsity ratio. Using sparse matrix representations drastically reduces memory requirements. Furthermore, computing \(\frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d}}\) and \(\mathbf{A} \mathbf{V}\) becomes more efficient, as the model only processes a small subset of relevant positions.

Given a position \(i\), we define the attention field \(\pi_i\) to be the set of positions that are considered in computing the representation at this position. We therefore only need to compute the dot-product attention between the given position \(i\) and each position \(j \in \pi_i\). This results in a sparse attention matrix \(\mathbf{A}'\) where \[\begin{align} A'(i,j) & = & \begin{cases} a_{i,j} & j \in \pi_i\;\textrm{and}\;j \le i \\ 0 & \textrm{otherwise} \end{cases}, \end{align}\]

where \(a_{i,j}\) is a non-zero weight. A simple implementation of this model involves a slight modification to \(\mathbf{M}\), leading to a new masking variable \(\mathbf{M}'\) \[\begin{align} M'(i,j) & = & \begin{cases} 0 & j \in \pi_i\;\textrm{and}\;j \le i \\ -\infty & \textrm{otherwise} \end{cases}. \end{align}\]

In practical implementations, a more efficient approach is to employ sparse operations for \(\mathbf{Q} \mathbf{K}^\top\) and \(\mathbf{A}' \mathbf{V}\) by considering \(\mathbf{M}'\) and \(\mathbf{A}'\), respectively. That is, we compute only for pairs with non-zero attention weights and skip the rest.

Several approaches to sparse self-attention can be considered.

• Span-based Attention/Local Attention. As discussed in Section 4.1, the use of context in sequence modeling is local in many cases. The basic idea of local attention is to restrict attention to a local region of the input sequence. We can then write \(\pi_i\) as \[\begin{align} \pi_i & = & [a_i^l,a_i^r], \end{align}\]

  where \(a_i^l\) and \(a_i^r\) are the left and right ends of \(\pi_i\). \(a_i^r - a_i^l + 1\) determines the size of the region, and so we can use it to control the sparsity of the attention model, for example, if \(a_i^r - a_i^l + 1 \ll n\), the model would be very sparse. \(a_i^l\) and \(a_i^r\) can be obtained by using either heuristics or machine learning methods. The reader may refer to related papers for more details [7], [44]–[46]. See Figure 14 (b) for an illustration of local attention.

• Chunked Attention. In this method, we segment a sequence into chunks and run the attention model on each of them [95], [96]. Given a sequence \(\{1,\cdots,n\}\), we define \(\{\mathrm{chunk}_1,\cdots,\mathrm{chunk}_q\}\) to be a segmentation of the sequence. A chunk can be expressed as a span \[\begin{align} \mathrm{chunk}_k & = & [c_k^l,c_k^r]. \end{align}\]

  In the attention step, we treat each chunk as a sequence and perform self-attention on it as usual. In other words, the representation at position \(i\) is computed by using only the context in the chunk that \(i\) belongs to. In this sense, this model can be thought of as a variant of local attention. Figure 14 (c) shows an illustration of this model. A remaining issue is how to segment the sequence. There are several ways to do this. For example, as discussed in Section 3.4, we can do segmentation from a linguistic perspective, and segment the sequence into linguistically motivated units. In practical systems, it is sometimes more convenient to segment the sequence into chunks that are of equal length. Thus, the sparsity of the model is controlled by the size of these chunks, for example, the use of smaller chunks would lead to a more sparse attention model.

• Strided Attention. Since the chunked attention approach enforces a hard segmentation on the input sequence, it may lose the ability to learn representations from inputs in different chunks. An alternative way to achieve chunk-wise attention is to allow overlap between chunks [97]–[99]. This approach is analogous to the family of approaches that are commonly used to apply a local model to 1D or 2D data to generate outputs of the same shape. Like CNNs, we use a context window to represent the field of input of the attention model. The context window slides along the sequence, each time moving forward a step of size \(\mathrm{stride}\). As a special case, if \(\mathrm{stride}\) equals the size of the context window, this model reduces to the chunked attention model mentioned above. If \(\mathrm{stride}\) is smaller than the size of the context window, the attention model will become denser. Figure 14 (d) shows the case of \(\mathrm{stride}=1\) where the chunk overlapping is maximized. A way to achieve relatively sparser attention is to use a dilated context window. Figure 14 (e) shows an example of the dilated strided attention model, where the context window is discontinuous, with gaps of size 1.

• Learning Attention Fields. Because the attention field \(\pi_i\) can be any subset of \(\{1,\cdots,n\}\), we can develop more general sparse attention models by considering attention maps beyond chunk-based patterns. The only question is how to determine which positions the model attends to for a given position. One simple approach is to use a computationally cheaper model to estimate the “importance” of each position. Then, attention weights are computed only for some of the positions which are thought to be most important [100]. A second approach is grouping: positions are grouped, and then the attention weights are computed only for positions in the same group. This can be achieved by clustering keys and queries. For example, we can cluster keys and queries via \(k\)-means clustering. The centroids of the clusters can be treated as additional parameters of the attention model, and so can be learned during optimization [101]. One benefit of learning attention fields is that the model can spread its attention more broadly over the sequence. This is a useful property for many NLP problems because word dependencies are sometimes long-range, not restricted to a local context window. See Figure 14 (f) for an example of the attention map learned through this model. Alternative approaches to learning to attend are to use sorting or hashing functions to group similar key and query vectors [102], [103]. These functions can be either heuristically designed functions or neural networks with learnable parameters. By using these functions, we can reorder the sequence so that the inputs in the same group are adjacent in the reordered sequence. In this way, the resulting attention map follows a chunk-wise pattern, and the model is computationally efficient through the use of the chunked attention approach.

• Hybrid Methods. Above, we have discussed a range of different sparse attention models. It is natural to explore methods that combine multiple models together to make use of their benefits in some way. A simple way to do this is to combine the attention fields of different models. For example, in [104]’s system, the attention map is generated by considering three different sparse models, including local attention (chunked attention), global attention, and random attention¹¹. The resulting model is still a sparse model, but is somewhat more robust as it involves multiple patterns from different perspectives of attention modeling. Another way of combining multiple attention models is to use different models for different heads in multi-head attention [97], [98]. For example, one can use one head as a local attention model, and use another head as a global attention model (see Figure 14 (g-h)).

One disadvantage of sparse models compared to dense models is that they are often less efficient on GPUs/CPUs due to hardware constraints. While sparse models can ideally reduce both the memory and computation requirements, the computational throughput of sparse models is much slower than that of dense models. In practice, it is difficult for sparse models to approach the peak FLOPS of a GPU or CPU¹². Therefore, they are primarily used to improve memory efficiency rather than achieving raw computational speedups on current hardware. On the other hand, sparse models are still of great use to NLP practitioners in the context of memory-efficient Transformer, especially when Transformers are used to deal with extremely long sequences.

5.2 Recurrent and Memory Models↩︎

For sequence generation problems, Transformers can also be thought of as memory systems. Consider again the general autoregressive setting, in which we are given the states of the previous \(i-1\) positions and wish to predict the next state. In self-attention, this is achieved by using the query at position \(i\) (i.e., \(\mathbf{q}_{i}\)) to access the key-value pairs of the previous positions (i.e., \(\{(\mathbf{k}_{1},\mathbf{v}_{1}),\cdots,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}\)). Then, we move to position \(i+1\) and append \((\mathbf{k}_{i},\mathbf{v}_{i})\) to the collection of key-value pairs. This procedure can be intuitively interpreted using memory mechanisms. A Transformer model maintains a memory bank that retains past information. As the model processes the sequence, it repeats the same operation: generating outputs by reading from the memory, and then updating the memory to incorporate new information. This is illustrated in Figure 15.

5.2.1 Cache-based Memory↩︎

The memory here can be viewed as a datastore of vectors. From a machine learning perspective, this is a non-parametric model, and the cost of accessing the model grows as the sequence becomes longer. Clearly, such a variable-length memory will generally be infeasible if the model deals with extremely long sequences. For modeling sequences of arbitrary length, it is common to use a fixed-length memory instead. As in many NLP problems, one of the simplest ways to do this is to maintain a cache of recent information, that is, we restrict the modeling to a context window. Let \(n_c\) be the size of the context window. The model keeps track of the \(n_c - 1\) latest states preceding the current position, so that its closest predecessors can be considered at each step. This means that, for each position, a self-attention sub-layer attends to \(n_c - 1\) previous positions, as follows

If we stack multiple self-attention sub-layers, a larger context window would be considered. For example, a model involving two self-attention sub-layers has a context window of size \(2n_c - 1\), as follows

Therefore, we can capture a sufficiently large context by using a multi-layer Transformer model. Note that the context window model here is essentially the same as the strided attention model presented in the preceding section. Systems of this type are often easy to implement: we slide a window along the sequence, and, in each move, we make predictions at the last position of the window (for inference), or back-propagate errors (for training).

An alternative strategy for training a context window model is chunked attention. The sequence is explicitly partitioned into sub-sequences (or chunks) of a fixed length \(n_c\). These chunks are then treated as independent training samples. This approach, however, completely ignores the relationship between inputs in different chunks. One way to address this issue is to introduce dependence between chunks. For example, the Transformer-XL model allows every chunk to access one or more preceding chunks [51]. In the simplest case, consider an example in which \(\mathrm{chunk}_k\) can see its predecessor \(\mathrm{chunk}_{k-1}\). Each position in \(\mathrm{chunk}_k\) can attend to all its preceding positions in both \(\mathrm{chunk}_k\) and \(\mathrm{chunk}_{k-1}\).

In Transformer-XL, this approach is implemented in a simplified form. First, each position is constrained to attend to \(n_c - 1\) previous positions so that the size of the attention field of a position is the same in the training and inference stages. Such a method turns the problem back to strided attention, making the implementation of the attention model straightforward. However, unlike standard strided attention, Transformer-XL performs training in a chunk-wise manner. Once we finish the training on a chunk, we directly move to the next chunk, rather than sliding the context window a small step forward. Second, while this approach allows for connections between chunks, the parameters of the sub-network corresponding to \(\mathrm{chunk}_{k-1}\) are fixed, and we only update the parameters of the sub-network on \(\mathrm{chunk}_{k}\) in the \(k\)-th step. See Figure 18 for an illustration.

The above model is similar in spirit to recurrent models because both of them require the computation in one step to depend on the states of the preceding steps. However, it is not in the standard form of a recurrent model, in which the output of a recurrent unit in one step is the input in the next step. Instead, the “recurrence” is expressed by involving connections across layers and chunks, that is, the output of one layer in \(\mathrm{chunk}_{k-1}\) is used as the input of a higher-level layer in \(\mathrm{chunk}_{k}\).

5.2.2 Encoding Long-term Memory↩︎

Another idea for representing the states of a sequence is to frame the task as an encoding problem. Instead of storing all key-value vectors during left-to-right generation, we compress the memory of the entire “history” into a fixed number of encoded key-value vectors. These encoded vectors can be either a small subset of the original sequence \(\{(\mathbf{k}_{1},\mathbf{v}_{1}),\cdots,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}\), or a small set of newly generated vectors that encode this entire history.

One way to achieve this encoding is to apply a pooling operation over \(\{(\mathbf{k}_{1},\mathbf{v}_{1}),\cdots,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}\) [105]. For example, by using average pooling, the memory is reduced to a single key-value pair \((\bar{\mathbf{k}},\bar{\mathbf{v}})\): \[\begin{align} \bar{\mathbf{k}} & = & \frac{1}{i-1} \sum_{j=1}^{i-1} \mathbf{k}_{j} \tag{45}, \\ \bar{\mathbf{v}} & = & \frac{1}{i-1} \sum_{j=1}^{i-1} \mathbf{v}_{j} \tag{46}. \end{align}\]

This leads to a very efficient model, as we only need to update the aggregated vectors \((\bar{\mathbf{k}},\bar{\mathbf{v}})\) incrementally [106]. Let \((\bar{\mathbf{k}}[i],\bar{\mathbf{v}}[i])\) represent the state of the memory at position \(i\). A more general formulation is given recursively: \[\begin{align} \bar{\mathbf{k}}[i] & = & \mathrm{KMem}(\bar{\mathbf{k}}[i-1],\mathbf{k}_{i-1}) \tag{47}, \\ \bar{\mathbf{v}}[i] & = & \mathrm{VMem}(\bar{\mathbf{v}}[i-1],\mathbf{v}_{i-1}) \tag{48}. \end{align}\]

where \(\mathrm{KMem}(\cdot)\) and \(\mathrm{VMem}(\cdot)\) are functions that update the memory by taking both the states of the memory at the previous position (i.e., \(\bar{\mathbf{k}}[i-1]\) and \(\bar{\mathbf{v}}[i-1]\)) and the new states (i.e., \(\mathbf{k}_{i-1}\) and \(\mathbf{v}_{i-1}\)). There are many possible choices for functions \(\mathrm{KMem}(\cdot)\) and \(\mathrm{VMem}(\cdot)\). For example, if \(\mathrm{KMem}(\cdot)\) and \(\mathrm{VMem}(\cdot)\) are weighted sum functions, we can derive the same forms as Eqs. (45 ) and (46 ). If \(\mathrm{KMem}(\cdot)\) and \(\mathrm{VMem}(\cdot)\) are recurrent cells (e.g., standard RNN or LSTM cells), we obtain a recurrent model of memory.

Extending the above concept to memories with more than one key-value pair is straightforward. One approach is chunk-level representation. Let \(\{(\bar{\mathbf{k}}_{1},\bar{\mathbf{v}}_{1}),\cdots,(\bar{\mathbf{k}}_{\kappa},\bar{\mathbf{v}}_{\kappa})\}\) be a memory of size \(\kappa\), where each \((\bar{\mathbf{k}}_{j},\bar{\mathbf{v}}_{j})\) is a representation of a chunk of length \(n_c\). This memory can thus encode a sequence of maximum length \(\kappa \cdot n_c\), computing each \((\bar{\mathbf{k}}_{j},\bar{\mathbf{v}}_{j})\) via Eqs. (47 ) and (48 ) over its corresponding chunk. A second approach is to organize the memory into a priority queue. A learned scoring function, which is often an auxiliary neural network, evaluates and assigns a priority score to each generated key-value pair. High-scoring pairs are inserted into the queue via a push operation. In this way, the memory actively maintains a dynamic collection of only the most valuable key-value pairs across the input sequence.

Although representing memory as a set of vectors is an intuitive design choice for Transformers, its capacity is bounded by the number of vectors. An alternative paradigm is continuous memory. This approach leverages function approximation, treating the sets of keys \(\{\mathbf{k}_{1},\cdots,\mathbf{k}_{i-1}\}\) or values \(\{\mathbf{v}_{1},\cdots,\mathbf{v}_{i-1}\}\) as a series of data points to be fitted by a continuous function. Instead of explicitly storing these vectors, the memory is parameterized by the functions themselves. A common method is to use a linear combination of simple basis functions to approximate the complex function defined by the discrete key or value vectors [107].

It is also straightforward to use a short-term memory and a long-term memory simultaneously so that we can combine the merits of both. For example, we use a cache-based memory to capture local context, and use an efficient long-term memory that encodes the entire history to model long-range dependencies. This idea is also similar to that used in combining different sparse attention models as discussed in the previous subsection.

5.2.3 Retrieval-based Methods↩︎

So far in this subsection, we have discussed approaches based on fixed-length models. It is also possible to develop efficient memory models by improving the efficiency of accessing memory, instead of strictly reducing memory capacity. One way to achieve this is to store the past key-value pairs in a database (call it a vector database), and to find the most similar key-value pairs when querying the database. To be more precise, given a query \(\mathbf{q}\), we use the database to find a set of top-\(p\) relevant key-value pairs (denoted by \(\Omega_p\)) by performing a similarity search based on the dot-product between the query and key vectors. The query \(\mathbf{q}\) then attends to \(\Omega_p\) as in standard self-attention models. The idea behind this method is to consider only a small number of elements that contribute most to the attention result. Therefore, the model is essentially a sparse attention model which is computationally efficient. Another advantage of this method is that it allows for fast similarity search over a very large set of vectors because of the highly optimized implementation of vector databases. Framing memory as a retrieval system falls under the broader paradigm of retrieval-augmented generation (RAG) or retrieval-augmented approaches. It provides a scalable framework for incorporating external memory into neural architectures like Transformers [108]–[110].

5.3 Low-dimensional Models↩︎

In many practical applications, Transformers are “high-dimensional” models. This is not only because the input and/or output data is in high-dimensional spaces, but also because some of the intermediate representations of the data in the model are high-dimensional. As discussed in Section 5.1, this high dimensionality arises in part from computing the weighted sum of value vectors, as in Eq. (49 ) (repeated here for convenience) \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{A} \mathbf{V} \label{eq:low-dim-self-attetention-decoder-as} \end{align}\tag{49}\]

and computing the attention matrix itself, as in Eq. (50 ) \[\begin{align} \mathbf{A} & = & \mathrm{Softmax}(\frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d}} + \mathbf{M}), \label{eq:low-dim-self-attetention-decoder-softmax} \end{align}\tag{50}\]

which involves large matrix multiplications \(\mathbf{Q} \mathbf{K}^\top\) and \(\mathbf{A} \mathbf{V}\) when the length \(n\) and the hidden dimensionality \(d\) are large.

The \(\mathbf{A} \mathbf{V}\) and \(\mathbf{Q} \mathbf{K}^\top\) operations scale with a time complexity of \(O(n^2 \cdot d)\) and a space complexity of \(O(n^2 + n \cdot d)\). While previous approaches have mitigated this complexity using sparse models, this subsection focuses on methods that approximate these operations via dense computation. One straightforward idea is to project \(\mathbf{Q}\), \(\mathbf{K}\), and \(\mathbf{V}\) into smaller matrices, thereby reducing the computational burden of matrix multiplication. Since \(\mathbf{Q}\), \(\mathbf{K}\), and \(\mathbf{V}\) all reside in \(\mathbb{R}^{n \times d}\), this can be achieved by reducing either the length dimension \(n\), the hidden dimension \(d\), or both.

5.3.1 Reducing \(n\)↩︎

Note that the output \(\mathrm{Att}_{\mathrm{self}}(\mathbf{S})\) is required to be an \(n \times d\) matrix, and so we cannot reduce the number of queries. We instead consider reducing the number of keys and values. Suppose \(n'\) is an integer with \(n' < n\), and \(\mathbf{K}\) and \(\mathbf{V}\) can be transformed into \(n' \times d\) matrices \(\mathbf{K}'\) and \(\mathbf{V}'\) in some way. We can obtain a “smaller” model simply by replacing \(\mathbf{K}\) and \(\mathbf{V}\) with \(\mathbf{K}'\) and \(\mathbf{V}'\), giving \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{A} {\color{red} \mathbf{V}'} \tag{51} \\ \mathbf{A} & = & \mathrm{Softmax}(\frac{\mathbf{Q} [{\color{red} \mathbf{K}'}]^\top}{\sqrt{d}} + \mathbf{M}). \tag{52} \end{align}\]

This model is in the standard form of self-attention, but has lower time and space complexities, that is, \(O(n' \cdot n \cdot d) < O(n^2 \cdot d)\) and \(O(n' \cdot n + n' \cdot d) < O(n^2 + n \cdot d)\). If \(n' \ll n\), the resulting model will be linear with \(n\).

The key problem here is how to obtain \(\mathbf{K}'\) and \(\mathbf{V}'\) in a way that retains much of the information in \(\mathbf{K}\) and \(\mathbf{V}\). There are several ways to do so. One simple method is to select the keys and values that are thought to be important. The importance of a key (or value) can be computed in terms of some computationally inexpensive measure. For example, we can sample a small number of query-key dot-products and estimate the importance of a key by collecting these dot-product results.

Although straightforward, the above method still relies on sparse operations like sampling and aggregation. Alternatively, dense computation can be used to project \(\mathbf{K}\) and \(\mathbf{V}\) down to \(\mathbf{K}'\) and \(\mathbf{V}'\). A common approach employs CNNs [111]. Let \(\mathrm{Conv}(\cdot)\) denote a 1D convolution operation sliding along the \(n\) dimension. \(\mathbf{K}'\) is then computed as: \[\begin{align} \mathbf{K}' & = & \mathrm{Conv}(\mathbf{K},\mathbf{W}_{c},\mathrm{size}_r,\mathrm{stride}), \end{align}\]

where \(\mathbf{W}_{c}\) is the parameter matrix of the filters, \(\mathrm{size}_r\) is the size of the receptive field, and \(\mathrm{stride}\) is the step size of the convolution. In general, we can achieve a high compression rate by choosing large values for \(\mathrm{size}_r\) and \(\mathrm{stride}\). Likewise, we can compute \(\mathbf{V}'\) using another convolutional function. It is worth noting that, if the parameter \(n'\) is fixed for all samples, compression of \(\mathbf{K}\) and \(\mathbf{V}\) along the length dimension is essentially the same as the fixed-length memory model as described in the preceding subsection. The methods presented here are more general and could be applied to variable-length memories.

We might also be tempted to model the attention function by considering the attention matrix \(\mathbf{A}\) as a high-dimensional representation of data and then applying conventional dimensionality reduction methods. Empirically, in many problems, \(\mathbf{A}\) (or more precisely \(\mathbf{Q} \mathbf{K}^\top\)) is found to be a low-rank matrix. In this case, we can compress \(\mathbf{A}\) while retaining as much information as possible. There are many ways to do so. For example, we might use a product of smaller matrices as an approximation to \(\mathbf{A}\) via the SVD technique. However, this introduces computational overhead in using SVD compared with the standard attention model. A simpler idea is to directly transform \(\mathbf{K}\) and \(\mathbf{V}\) into smaller-sized matrices via linear mappings, given by \[\begin{align} \mathbf{K}' & = & \mathbf{U}^{k} \mathbf{K}, \\ \mathbf{V}' & = & \mathbf{U}^{v} \mathbf{V}, \end{align}\]

where \(\mathbf{U}^{k} \in \mathbb{R}^{n' \times n}\) and \(\mathbf{U}^{v} \in \mathbb{R}^{n' \times n}\) are parameter matrices. Clearly, this leads to a model which is equivalent to that described in Eqs. (51 ) and (52 ). While highly intuitive, this approach has been proven to yield a sufficiently small approximation error \(\epsilon\), provided that \(n'\) scales linearly with \(d/\epsilon^2\) [112].

5.3.2 Reducing \(d\)↩︎

Another approach to reducing computational complexity is to reduce the dimensionality \(d\). One of the simplest methods is to project all queries and keys onto a \(d'\)-dimensional space (\(d' < d\)), and to compute the dot-product between query-key pairs in the new space. For modeling, we only need to replace \(\mathbf{Q} \in \mathbb{R}^{n \times d}\) and \(\mathbf{K} \in \mathbb{R}^{n \times d}\) with new representations \(\mathbf{Q}' \in \mathbb{R}^{n \times d'}\) and \(\mathbf{K}' \in \mathbb{R}^{n \times d'}\). We can easily modify Eq. (50 ) to use \(\mathbf{Q}'\) and \(\mathbf{K}'\) in computing the attention matrix \[\begin{align} \mathbf{A} & = & \mathrm{Softmax}(\frac{{\color{red} \mathbf{Q}'} [{\color{red} \mathbf{K}'}]^\top}{\sqrt{d'}} + \mathbf{M}). \end{align}\]

\(\mathbf{Q}'\) and \(\mathbf{K}'\) are given by \[\begin{align} \mathbf{Q}' & = & \mathbf{Q} \mathbf{U}^{q}, \\ \mathbf{K}' & = & \mathbf{K} \mathbf{U}^{k}, \end{align}\]

where \(\mathbf{U}^{q} \in \mathbb{R}^{d \times d'}\) and \(\mathbf{U}^{k} \in \mathbb{R}^{d \times d'}\) are parameter matrices of linear transformations.

It is also possible to exploit kernel methods to obtain an efficient dot-product attention model. The basic idea is to map data points (represented as vectors) from one space to another, transforming a problem that is difficult in the original space into one that is more tractable in the new space. The “kernel trick” allows us to compute such inner products implicitly via a kernel function, without explicitly constructing the mapping \(\phi(\cdot)\)¹³. Such inner products in feature space are commonly represented by kernel functions, denoted as \(K(\cdot,\cdot)\).

It is interesting to approximate \(\mathbf{A}\) in a fashion analogous to \(K(\cdot,\cdot)\) in kernel methods. To illustrate, note in Eq. (50 ) \(\mathbf{A}\) represents normalized attention weights. The numerator can be written in the form \[\begin{align} \widetilde{\mathbf{A}} & = & \mathrm{Mask}(\exp(\frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d}})). \end{align}\]

Here \(\mathrm{Mask}(\cdot)\) is a function which has the same effect as using the additive masking variable \(\mathbf{M}\). Then, \(\mathbf{A}\) can be expressed as \[\begin{align} \mathbf{A} & = & \mathbf{D}^{-1} \widetilde{\mathbf{A}}, \end{align}\]

where \(\mathbf{D}\) is an \(n \times n\) diagonal matrix. Each entry of the main diagonal is the sum of the entries of the corresponding row in \(\widetilde{\mathbf{A}}\), denoting the normalization factor of softmax. Substituting this equation into Eq. (50 ), we have \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{D}^{-1} \widetilde{\mathbf{A}} \mathbf{V}. \label{eq:attention-das-form} \end{align}\tag{53}\]

In this model, \(\widetilde{A}(i,j)\) can be viewed as a similarity function over all query-key pairs in a \(d\)-dimensional space. Here we assume that this function, which is in the form of the dot-product of vectors, can be approximated by a kernel function \[\begin{align} \widetilde{A}(i,j) & = & K(\mathbf{q}_i,\mathbf{k}_j) \nonumber \\ & = & \langle \phi(\mathbf{q}_i), \phi(\mathbf{k}_j) \rangle \nonumber . \end{align}\]

\(\phi(\cdot)\) is a mapping from \(\mathbb{R}^{d}\) to \(\mathbb{R}^{d'}\). We can represent the queries and keys in the following form \[\begin{align} \mathbf{Q}' & = & \phi(\mathbf{Q}) \nonumber \\ & = & \begin{bmatrix} \phi(\mathbf{q}_1) \\ \vdots \\ \phi(\mathbf{q}_n) \end{bmatrix}, \\ \mathbf{K}' & = & \phi(\mathbf{K}) \nonumber \\ & = & \begin{bmatrix} \phi(\mathbf{k}_1) \\ \vdots \\ \phi(\mathbf{k}_n) \end{bmatrix}. \end{align}\]

Then, we develop a kernelized attention model by approximating the attention weight \(\alpha_{i,j}\) in the form \[\begin{align} \alpha_{i,j} & \approx & \frac{\phi(\mathbf{q}_i) \phi(\mathbf{k}_j)^\top}{\sum_{j'=1}^{n}\phi(\mathbf{q}_i) \phi(\mathbf{k}_{j'})^\top}. \end{align}\]

The key idea behind this kernelized attention model is that we can remove the softmax function if the queries and keys are mapped to a new space. Using this approximation, the \(i\)-th output vector of the attention model (i.e., the \(i\)-th row vector of \(\mathrm{Att}_{\mathrm{self}}(\mathbf{S})\)) is given by \[\begin{align} \mathbf{c}_{i} & = & \sum_{j=1}^{n} \alpha_{i,j} \cdot \mathbf{v}_j \nonumber \\ & \approx & \sum_{j=1}^{n} \Big( \frac{\phi(\mathbf{q}_i) \phi(\mathbf{k}_j)^\top}{\sum_{j'=1}^{n}\phi(\mathbf{q}_i) \phi(\mathbf{k}_{j'})^\top} \cdot \mathbf{v}_j \Big) \nonumber \\ & = & \frac{\sum_{j=1}^{n} \phi(\mathbf{q}_i) \phi(\mathbf{k}_j)^\top \mathbf{v}_j}{\sum_{j'=1}^{n}\phi(\mathbf{q}_i) \phi(\mathbf{k}_{j'})^\top} \nonumber \\ & = & \frac{\phi(\mathbf{q}_i) {\color{red} (} \sum_{j=1}^{n} \phi(\mathbf{k}_j)^\top \mathbf{v}_j {\color{red} )}}{\phi(\mathbf{q}_i) {\color{red} (} \sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^\top {\color{red} )}}. \label{eq:context-vector-kernel-methods} \end{align}\tag{54}\]

Here \(\phi(\mathbf{q}_i) \in \mathbb{R}^{1 \times d'}\), while \(\phi(\mathbf{k}_j)^\top \mathbf{v}_j \in \mathbb{R}^{d' \times d}\). Therefore, the inner term \(\sum_{j=1}^n \phi(\mathbf{k}_j)^\top \mathbf{v}_j\) is a \(d' \times d\) matrix, and the final product yields a vector in \(\mathbb{R}^{1 \times d}\).

Although the equation appears a bit complicated, the idea is simple: instead of attending the query to all keys to obtain the attention weight \(\alpha_{i,j}\), we can compute the sum of the multiplications \(\sum_{j=1}^{n} \phi(\mathbf{k}_j)^\top \mathbf{v}_j \in \mathbb{R}^{d' \times d}\) and then multiply it with the kernelized query \(\phi(\mathbf{q}_i)\). Returning to the notation used in Eq. (53 ), we define the \(i\)-th diagonal entry of \(\mathbf{D}\) to be \(\phi(\mathbf{q}_i) \sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^\top\). Then, the attention model can be re-expressed in the form \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{D}^{-1} \phi(\mathbf{Q}) \phi(\mathbf{K})^\top \mathbf{V} \nonumber \\ & = & \mathbf{D}^{-1} \mathbf{Q}' \mathbf{K}'^\top \mathbf{V} \nonumber \\ & = & \mathbf{D}^{-1} {\color{red} \big(} \mathbf{Q}' {\color{red}(}\mathbf{K}'^\top \mathbf{V} {\color{red})} {\color{red}\big)}. \end{align}\]

Here we change the order of computation from left-to-right to right-to-left using parentheses. Given that \(\mathbf{Q}' \in \mathbb{R}^{n \times d'}\) and \(\mathbf{K}' \in \mathbb{R}^{n \times d'}\), this model has time and space complexities of \(O(n \cdot d \cdot d')\) and \(O(n \cdot d + n \cdot d' + d \cdot d')\), respectively. Therefore, the model is linear with respect to the sequence length \(n\), and is sometimes called the linear attention model. One computational advantage of this model is that we need only compute the multiplication \(\mathbf{K}'^\top \mathbf{V}\) (i.e., \(\sum_{j=1}^{n} \phi(\mathbf{k}_j)^\top \mathbf{v}_j\)) and the corresponding normalization factor (i.e., \(\sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^\top\)) once. The results can then be used for any query [113]. The model needs to maintain \(\sum_{j=1}^{n} \phi(\mathbf{k}_j)^\top \mathbf{v}_j\) and \(\sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^\top\) and update them when new key and value vectors arrive.

Still, there are several limitations regarding this kernelized model, for example, how to develop the feature map \(\phi(\cdot)\) to obtain a good approximation to the standard attention model. Interested readers may refer to [114]’s work for more details.

A second idea for reducing \(d\) is to take sub-space models, in which a problem in a \(d\)-dimensional space is transformed into sub-problems in lower-dimensional spaces, and the solution to the original problem is approximated by some combination of the solutions to these sub-problems. In a general sub-space model, a \(d\)-dimensional key vector \(\mathbf{k}\) can be mapped into a set of \(d'\)-dimensional vectors \(\{ \mathbf{K}'_{1},\cdots,\mathbf{K}'_{\eta} \}\). To simplify modeling, we can do this by vector segmentation, that is, we segment \(\mathbf{k}\) into \(\eta\) sub-vectors, each having \(d' = \frac{d}{\eta}\) dimensions. We can transform all query and value vectors in the same way. Then, the attention model is applied in each of these sub-spaces.

This method, however, does not reduce the total amount of computation. As presented in [115]’s work, we can instead approximate the dot-product attention over a set of key-value pairs by considering top-\(p\) candidates in each sub-space. More precisely, we find \(p\)-best key-value pairs in each sub-space, which is computationally cheaper. The Cartesian product of these \(p\)-best key sets consists of \(p^\eta\) product keys. Likewise, we obtain \(p^\eta\) product values. The remaining work is simple: the \(d\)-dimensional queries attend to these \(d\)-dimensional product keys and values. An interesting difference between this sub-space model and the \(d\)-dimensional space model is that the generated product keys and values may be different from any of the original key-values \(\{(\mathbf{k}_{1},\mathbf{v}_{1}),\cdots,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}\). This provides a way for learning new representations of the past information.

So far we have discussed approaches to dimensionality reduction along either the \(n\) or \(d\) dimension. It is straightforward to combine them to develop a “lower-dimensional” model. As an example, suppose that we have the \(n \to n'\) reduction for keys and values, and the \(d \to d'\) reduction for queries and keys. The model takes the form \[\begin{align} \mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{A} {\color{red} \mathbf{V}'} \nonumber , \\ \mathbf{A} & = & \mathrm{Softmax}(\frac{{\color{red} \mathbf{Q}'} {\color{red} \mathbf{K}'}^\top}{\sqrt{d'}} + \mathbf{M}), \end{align}\]

where \(\mathbf{Q}' \in \mathbb{R}^{n \times d'}\), \(\mathbf{K}' \in \mathbb{R}^{n' \times d'}\), and \(\mathbf{V}' \in \mathbb{R}^{n' \times d'}\) are low-dimensional representations for queries, keys and values. As usual, we can easily obtain these representations through the linear mappings of \(\mathbf{Q}\), \(\mathbf{K}\) and \(\mathbf{V}\). The time and space complexities of this model are \(O(n' \cdot n \cdot d')\) and \(O(n' \cdot n + n' \cdot d')\).

5.4 Parameter and Activation Sharing↩︎

Redundancy is common to most large-scale neural networks. As a result, many of these models are over-parameterized, making training and inference less efficient. One common approach to reducing redundancy is to simplify the model by removing unnecessary components. For example, we can either prune a complex model or share sub-modules across different parts of the model to obtain a more compact model. In this subsection, we discuss methods for sharing parameters and intermediate states in Transformer models. We leave the discussion of model transfer and pruning to Section 5.7.

Shared-parameter architectures are widely used in neural network-based systems. Well-known examples include CNNs and RNNs, where the same set of parameters (or layers) is applied across different regions of the input. This produces a large neural network, parts of which have the same architecture and the same shared parameters. For Transformers as well as other sequence models, the sharing mechanism can be applied to different levels of modeling. A simple example is embedding sharing. In machine translation, a typical strategy for dealing with words in two languages is to develop two separate embedding models. Alternatively, one can use a single embedding model for both languages. The parameters of the model are then learned during the training of both the source-side and target-side networks. Such a strategy is also often adopted in multilingual sequence models, such as language models that are able to deal with text in many different languages.

For multi-layer neural networks, a popular method is layer-wise sharing. Suppose there is a stack of layers, all of which have the same form \[\begin{align} \mathbf{S}^{l} & = & \mathrm{Layer}(\mathbf{S}^{l-1};\theta^{l}). \end{align}\]

We can tie the parameters for some or all of these layers. For example, given a set of layers \(\{l_1, l_2,\cdots,l_n\}\), we enforce the constraint \(\theta^{l_1} = \theta^{l_2} = \cdots = \theta^{l_n}\), so that we can obtain a smaller model and the optimization of the model can be easier. In practice, this shared-layer model is highly advantageous if many layers are involved, because we can repeat the same process many times to construct a very deep neural network [79]. For example, sharing a single FFN sub-layer across all encoder layers is found to be effective in reducing the redundancy in machine translation systems [116].

For Transformers, sharing can also be performed in multi-head attention. An example of this is multi-query attention [117]. Recall from Section 2.3 that the output of a head \(h\) in standard multi-head self-attention can be written as \[\begin{align} \mathbf{C}_{h}^{\mathrm{head}} & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{S}_{h}^{q}, \mathbf{S}_{h}^{k}, \mathbf{S}_{h}^{v}) \nonumber \\ & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{S} \mathbf{W}_{h}^{q}, \mathbf{S} \mathbf{W}_{h}^{k}, \mathbf{S} \mathbf{W}_{h}^{v}). \end{align}\]

Here \(\mathbf{S}_{h}^{q} = \mathbf{S} \mathbf{W}_{h}^{q}\), \(\mathbf{S}_{h}^{k} = \mathbf{S} \mathbf{W}_{h}^{k}\), and \(\mathbf{S}_{h}^{v} = \mathbf{S} \mathbf{W}_{h}^{v}\) are the queries, keys, and values, which are obtained by linearly transforming the input \(\mathbf{S}\) with distinct parameter matrices \(\mathbf{W}_{h}^{q}\), \(\mathbf{W}_{h}^{k}\), and \(\mathbf{W}_{h}^{v}\). In multi-query attention, we share the same keys and values across all the heads, but use different queries for different heads. The form of this model is given by \[\begin{align} \mathbf{C}_{h}^{\mathrm{head}} & = & \mathrm{Att}_{\mathrm{qkv}}(\mathbf{S} \mathbf{W}_{h}^{q}, \mathbf{S} \mathbf{W}_{0}^{k}, \mathbf{S} \mathbf{W}_{0}^{v}). \end{align}\]

Because the keys \(\mathbf{S} \mathbf{W}_{0}^{k}\) and values \(\mathbf{S} \mathbf{W}_{0}^{v}\) are independent of \(h\), they are computed and stored only once, rather than redundantly for each head. This yields significant computational and memory bandwidth savings during inference, especially when the number of attention heads is large. Due to its efficiency, multi-query attention (and its variant, grouped-query attention) has been successfully integrated into modern large language models, including Llama 2 [118] and Falcon¹⁴.

By extending the idea of sharing to more general situations, intermediate states can, in principle, be shared across a neural network. Reusing neuron activations, for instance, allows sub-components to be evaluated multiple times efficiently. In Transformers, this method can be applied directly to the self-attention matrices. Studies have empirically shown that attention maps across adjacent layers are often highly correlated in certain NLP tasks [119]. As a result, a computationally efficient strategy is to compute the dense attention map once and reuse it across subsequent layers.

More broadly, the sharing mechanism can be viewed as a process by which we reuse previously computed results rather than recomputing them on the fly. It is thus possible to reuse the information across different runs of a neural network. A related example is reversible residual networks, in which activations of one layer can be recovered from the activations of the following layer [120]. Hence we only keep the output of the latest layer in the forward pass. Then, in the backward pass of training, we reconstruct the output of each layer from its successor. One advantage of this reversible design is that the information produced in the forward pass is shared implicitly, and the model is memory-efficient [102].

5.5 Alternatives to Self-Attention↩︎

We have seen that the use of self-attention is one of the main sources of high computational and memory costs in Transformers. It is natural to wonder if there are efficient alternatives to self-attention models. Here, we briefly present some Transformer variants in which self-attention sub-layers are replaced with other types of neural architectures.

5.5.1 CNNs as a Replacement for Self-Attention↩︎

CNNs are simple and widely used neural networks, and are considered potential alternatives to self-attention models. To apply CNNs to Transformers, all we need is to construct a convolutional sub-layer to replace the self-attention sub-layer in a Transformer block. While a CNN filter has a restricted receptive field and thus takes inputs from a “local” context window, large contexts can be easily modeled by stacking multiple convolutional sub-layers. One key advantage of CNNs is that their computational complexity scales linearly with the sequence length \(n\), compared to the quadratic scaling of self-attention models. Furthermore, the availability of highly optimized CNN implementations makes them easy to apply to sequence modeling in practical systems. To further improve memory efficiency, we can utilize lightweight CNN variants, such as depthwise CNNs [121] ¹⁵.

5.5.2 Linear Attention↩︎

As with many practical approaches to sequence modeling, there is also considerable interest in developing linear models in order to speed up the processing of long sequences. While there are many ways to define a linear model, one general form that is commonly used in sequence models is \[\begin{align} \mathbf{z}_{i} & = & f(a \cdot \mathbf{z}_{i-1} + b \cdot \mathbf{s}_{i}). \label{eq:linear-base-model-for-transformer} \end{align}\tag{55}\]

Here \(\mathbf{s}_i\) represents some intermediate states of the model at step \(i\), and \(\mathbf{z}_{i}\) represents the summary of the past states up to step \(i\). It is easy to see that this is a recurrent model: the output at step \(i\) depends only on the input at the current step and the output at the previous step. As with the popular design choices in neural network-based systems, the linear part is followed by a transformation \(f(\cdot)\) which can be either an activation function or a feedforward neural network. Note that, Eq. (55 ) defines a standard linear model only if \(f(\cdot)\) is a linear function. The use of \(f(\cdot)\) gives greater flexibility in modeling the problem, although the term linear model may not be applied if \(f(\cdot)\) chooses a non-linear form.

The above formula describes a linearly structured model which can be seen as an instance of a general family of mathematical models. Typically, it can be represented as a chain structure, or an ordered set of nodes. The model repeats the same computation process from the first node to the last, each time taking the information from the current and previous steps and producing an output vector that is used in the following time steps. As a result, the space and time cost of the model scales linearly with the length of the chain.

We can extend Eq. (55 ) to a standard RNN model by simply applying a linear transformation of the current input and the previous state, that is, \(\mathbf{z}_{i} = f(\mathbf{z}_{i-1} \cdot \mathbf{W}_{z} + \mathbf{s}_{i} \cdot \mathbf{W}_{s})\). It is thus straightforward to apply RNNs and their variants to Transformer to obtain a hybrid model. For example, we can use LSTM and GRUs in building some of the Transformer layers to combine the merits of both recurrent models and self-attentive models [122]. As conventional recurrent models are well known, we skip the discussion of them here.

In fact, we may be more interested in developing linear attention models, so that we can obtain an efficient system, while still retaining the benefit of globally attentive sequence modeling. Part of the difficulty in doing this is that the form of self-attention is not linear. Let us take a moment to see how this difficulty arises. Recall that the result of self-attention can be written in the following form \[\begin{align} \mathrm{Att}_{\mathrm{self}} & = & \mathbf{A} \cdot \mathbf{V} \nonumber \\ & = & \psi(\mathbf{Q} \cdot \mathbf{K}^\top) \cdot \mathbf{V}. \end{align}\]

Here \(\psi(\cdot)\) is a function that is composed of scaling, exponentiation, masking, and normalization operations (i.e., \(\psi(\mathbf{a}) = \mathrm{Normalize}(\mathrm{Mask}(\exp(\frac{\mathbf{a}}{\sqrt{d}})))\)). Because \(\psi(\cdot)\) is a complex non-linear function, there is no straightforward way to simplify the computation directly, and we typically need to calculate the two matrix multiplications separately (one inside \(\psi(\cdot)\) and one outside \(\psi(\cdot)\)). As a consequence, we need to store all the key-value pairs explicitly, and visit each of them given a query. Not surprisingly, this leads to a model whose computational cost grows quadratically with the sequence length \(n\).

Although in self-attention keys and values are coupled, they are used in separate steps. An elegant solution decouples the query from the key-value interaction so that the context information can be encoded independently of any specific query. The key insight is to bypass the non-linear softmax by applying a feature map \(\phi(\cdot)\) to the queries and keys. As introduced in Section 5.3, we transform \(\mathbf{Q}\) and \(\mathbf{K}\) into \(\mathbf{Q}' = \phi(\mathbf{Q}) \in \mathbb{R}^{n \times d'}\) and \(\mathbf{K}' = \phi(\mathbf{K}) \in \mathbb{R}^{n \times d'}\). The attention mechanism can then be reformulated as: \[\begin{align} \mathrm{Att}_{\mathrm{self}} & \equiv & \psi'(\mathbf{Q}' \cdot \mathbf{K}'^\top) \cdot \mathbf{V} \nonumber \\ & = & \mathbf{D}^{-1} (\mathbf{Q}' \cdot \mathbf{K}'^\top) \cdot \mathbf{V} \nonumber \\ & = & \mathbf{D}^{-1} \Big( \mathbf{Q}' \cdot {\color{red} \big(} \mathbf{K}'^\top \cdot \mathbf{V} {\color{red} \big)} \Big), \end{align}\]

where \(\mathbf{D}\) is the diagonal normalization matrix, and \(\psi'(\mathbf{Z}) = \mathbf{D}^{-1} \mathbf{Z}\)¹⁶. In this transformed space, the query-key dot product does not require exponential softmax normalization. It corresponds to a normalization by multiplying with \(\mathbf{D}^{-1}\). Because the core operation is now pure matrix multiplication, we can exploit matrix associativity to change the order of computation.

This yields a powerful autoregressive formulation: keys and values are first aggregated via \(\mathbf{K}'^\top \cdot \mathbf{V}\), and then queries attend to this global representation. Given that \(\mathbf{K}'^\top \cdot \mathbf{V} = \sum_{j=1}^{n} \mathbf{k'}_{j}^\top \cdot \mathbf{v}_{j}\), we can write \(\mathbf{K}'^\top \cdot \mathbf{V}\) in the form of Eq. (55 ), as follows \[\begin{align} \boldsymbol{\mu}_j & = & \boldsymbol{\mu}_{j-1} + \mathbf{k'}_{j}^\top \cdot \mathbf{v}_{j}. \label{eq:linear-attention-mu} \end{align}\tag{56}\]

Here \(\boldsymbol{\mu}_j \in \mathbb{R}^{d' \times d}\) is a variable that adds \(\mathbf{k'}_{j}^\top \cdot \mathbf{v}_{j}\) at a time. Likewise, we can define another variable \(\boldsymbol{\nu}_j \in \mathbb{R}^{d'}\) \[\begin{align} \boldsymbol{\nu}_j & = & \boldsymbol{\nu}_{j-1} + \mathbf{k'}_{j}^\top. \label{eq:linear-attention-nu} \end{align}\tag{57}\]

Then, the output of self-attention for the \(j\)-th query can be written as (see also Eq. (54 )) \[\begin{align} \mathrm{Att}_{\mathrm{self},j} & = & \frac{\mathbf{q}'_{j} \cdot \boldsymbol{\mu}_n}{\mathbf{q}'_{j} \cdot \boldsymbol{\nu}_n}. \end{align}\]

Clearly, this yields a linear-time model, because \(\boldsymbol{\mu}_n\) and \(\boldsymbol{\nu}_n\) are linear with respect to \(n\). In simple implementations of this model, only \(\boldsymbol{\mu}_j\) and \(\boldsymbol{\nu}_j\) are kept. Each time a new query is encountered, we update \(\boldsymbol{\mu}_j\) and \(\boldsymbol{\nu}_j\) using Eqs. (56 ) and (57 ), and then compute \(\mathrm{Att}_{\mathrm{self},j} = \frac{\mathbf{q}'_{j} \cdot \boldsymbol{\mu}_j}{\mathbf{q}'_{j} \cdot \boldsymbol{\nu}_j}\)¹⁷.

One straightforward extension to the linear attention model is to allow Eqs. (56 ) and (57 ) to combine different terms with different weights. For example, we can redefine \(\boldsymbol{\mu}_j\) and \(\boldsymbol{\nu}_j\) as \[\begin{align} \boldsymbol{\mu}_j & = & a \cdot \boldsymbol{\mu}_{j-1} + (1 - a) \cdot \mathbf{k'}_{j}^\top \cdot \mathbf{v}_{j}, \\ \boldsymbol{\nu}_j & = & a \cdot \boldsymbol{\nu}_{j-1} + (1 - a) \cdot \mathbf{k'}_{j}^\top \end{align}\]

and train the parameter \(a\) as usual. Also, we can treat \(a\) as a gate and use another neural network to compute \(a\) [123]. Another model design is to add more terms to Eqs. (56 ) and (57 ) in order to provide a more expressive formulation of the linear attention approach [124], [125].

We have seen a general idea of designing linear models for the attention mechanism. The key design choice of such models is to remove the softmax-based normalization, thereby yielding linear forms of representation based on various intermediate states of the models. This motivates several recently developed alternatives to self-attention in which efficient inference systems are developed on the basis of recurrent models of sequence modeling [126], [127]. While these systems have different architectures, the underlying models have a similar form, as described in Eq. (55 ). Note that, by using the general formulation of recurrent models, we are not restricted to standard QKV attention. Instead we may give new meanings and forms to the queries, keys, and values.

The discussion here is also related to the memory models discussed in Section 5.2. From a memory perspective, the keys and values can be treated as encodings of the context. Therefore, in the linear attention model above we have a memory system in which two simple variables \(\boldsymbol{\mu}_j\) and \(\boldsymbol{\nu}_j\) are used to represent all the context information up to position \(j\). This results in a fixed-length memory which is very useful in practice. There are also other linear approaches to encoding long sequences. For example, we can view the moving average model as an instance of Eq. (55 ), and average a series of state vectors of a Transformer model, either weighted or unweighted.

5.5.3 State-Space Models↩︎

In control systems, state-space models (SSMs) are representations of a system whose input and output are related by some state variables (or states for short), and whose dynamics are described by first-order differential equations of these states. As a simple example, we consider a continuous-time linear time-invariant system which is given in the form of the state-space representation. We adopt a row-vector convention throughout this section, where all variables are represented as row vectors and linear transformations are applied on the right: \[\begin{align} \frac{d \mathbf{z}(t)}{d t} & = & \mathbf{z}(t) \cdot \mathbf{A} + \mathbf{s}(t) \cdot \mathbf{B}, \tag{58} \\ \mathbf{o}(t) & = & \mathbf{z}(t) \cdot \mathbf{C} + \mathbf{s}(t) \cdot \mathbf{D}. \tag{59} \end{align}\]

Here, \(\mathbf{s}(t)\), \(\mathbf{o}(t)\), and \(\mathbf{z}(t)\) represent the values of the input variable, output variable, and state variable at time \(t\), respectively¹⁸. In a general setting, these variables may possess different dimensionalities. For simplicity, we assume \(\mathbf{s}(t), \mathbf{o}(t) \in \mathbb{R}^{d}\) and \(\mathbf{z}(t) \in \mathbb{R}^{d_z}\)¹⁹. Eq. (58 ) is called the state equation, where \(\mathbf{A} \in \mathbb{R}^{d_z \times d_z}\) is the state matrix and \(\mathbf{B} \in \mathbb{R}^{d \times d_z}\) is the input matrix. Eq. (59 ) is called the output equation, where \(\mathbf{C} \in \mathbb{R}^{d_z \times d}\) is the output matrix and \(\mathbf{D} \in \mathbb{R}^{d \times d}\) is the feedforward matrix.

These equations describe a continuous mapping from the input \(\mathbf{s}(t)\) to the output \(\mathbf{o}(t)\) over time. They are, therefore, often used to deal with continuous time series data. To apply this model to sequence modeling, we need to modify the above equations to obtain a discrete form of the state-space representation. Suppose that \(\{\mathbf{s}_0,\mathbf{s}_1,\cdots,\mathbf{s}_n\}\) is a sequence of input data points sampled from \(\mathbf{s}(t)\) with time step \(\Delta t\). Similarly, we define \(\{\mathbf{z}_0,\mathbf{z}_1,\cdots,\mathbf{z}_n\}\) and \(\{\mathbf{o}_0,\mathbf{o}_1,\cdots,\mathbf{o}_n\}\) as sequences of the state and output vectors. Given this notation, we now have a discretized version of the SSM, written as \[\begin{align} \mathbf{z}_t & = & \mathbf{z}_{t-1} \cdot \overline{\mathbf{A}} + \mathbf{s}_t \cdot \overline{\mathbf{B}}, \tag{60} \\ \mathbf{o}_t & = & \mathbf{z}_t \cdot \overline{\mathbf{C}} + \mathbf{s}_t \cdot \overline{\mathbf{D}}. \tag{61} \end{align}\]

This formulation of the SSM defines an RNN with a residual connection. To be more precise, Eq. (60 ) describes a recurrent unit that reads the input at step \(t\) and the state at step \(t-1\), without using any activation function. Eq. (61 ) describes an output layer that sums the linear transformations of the state \(\mathbf{z}_t\) and the residual mapping \(\mathbf{s}_t\).

The parameters \(\overline{\mathbf{A}}\), \(\overline{\mathbf{B}}\), \(\overline{\mathbf{C}}\), and \(\overline{\mathbf{D}}\) can be induced from \(\mathbf{A}\), \(\mathbf{B}\), \(\mathbf{C}\) and \(\mathbf{D}\) in several different ways, depending on how Eq. (58 ) is approximated by Eq. (60 )²⁰. One approach to time discretization is the bilinear transform (also known as Tustin’s method). An alternative approach is to use the Zero-Order-Hold (ZOH) discretization. A detailed discussion of these approaches lies beyond the scope of this book, and we refer the interested reader to standard textbooks on control theory for further details [129].

The recurrent form of Eq. (60 ) makes it easy to compute the states and outputs over a sequence of discrete time steps. We can unroll \(\mathbf{z}_t\) and \(\mathbf{o}_t\) in a feedforward fashion

It is easy to write \[\begin{align} \mathbf{z}_t & = & \sum_{i=0}^{t} \mathbf{s}_{i} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-i}, \\ \mathbf{o}_t & = & \sum_{i=0}^{t} \mathbf{s}_{i} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-i} \cdot \overline{\mathbf{C}} + \mathbf{s}_{t} \cdot \overline{\mathbf{D}}. \label{eq:ssm-o-t-unrolled} \end{align}\tag{65}\]

Clearly, the right-hand side of Eq. (65 ) can be interpreted as the sum of a convolutional term and a linear term. Given that \[\begin{align} \sum_{i=0}^{t} \mathbf{s}_{i} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t-i} \cdot \overline{\mathbf{C}} & = & \begin{bmatrix} \mathbf{s}_0 & \mathbf{s}_1 & \cdots & \mathbf{s}_t \end{bmatrix} \cdot \nonumber \\ & & \begin{bmatrix} \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t} \cdot \overline{\mathbf{C}} & \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t - 1} \cdot \overline{\mathbf{C}} & \cdots & \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} \end{bmatrix}^\top, \label{eq:ssm-convolutional-part} \end{align}\tag{66}\]

we define a convolutional filter with parameters \[\begin{align} \mathbf{W}_{\mathrm{ssm}} & = & \begin{bmatrix} \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{n_{\mathrm{max}}} \cdot \overline{\mathbf{C}} & \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{n_{\mathrm{max}} - 1} \cdot \overline{\mathbf{C}} & \cdots & \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} \end{bmatrix}, \end{align}\]

where \(n_{\mathrm{max}}\) is the maximum length of the sequence²¹. Then, the output of the state-space model for a sequence \(\mathbf{S} = \begin{bmatrix} \mathbf{s}_0 \\ \vdots \\ \mathbf{s}_n \end{bmatrix}\) can be expressed as \[\begin{align} \mathbf{O} & = & \mathrm{Conv}(\mathbf{S},\mathbf{W}_{\mathrm{ssm}}) + \mathrm{Linear}(\mathbf{S},\overline{\mathbf{D}}), \end{align}\]

where \(\mathrm{Conv}(\cdot)\) is the convolution operation, and \(\mathrm{Linear}(\cdot)\) is the linear transformation operation. Such a treatment of the state-space model enables the system to be efficiently implemented using fast parallel convolution algorithms.

However, naive implementations of this model often perform poorly. As with many deep neural architectures, careful parameter initialization is crucial. Specifically, restricting the state matrix \(\mathbf{A}\) to particular structures (e.g., HiPPO matrices) has proven essential for memorizing and generalizing over long sequences [130].

Another computational bottleneck is the repeated matrix multiplication required to compute \(\overline{\mathbf{A}}^{n}\). For large \(n\), this operation is computationally expensive and numerically unstable. A common approach in modern SSMs is diagonalization. The idea is to apply a basis transformation to the state space such that \(\mathbf{A}\) (or \(\overline{\mathbf{A}}\)) becomes diagonal. Given an SSM parameterized by \((\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D})\), we can define an equivalent model \((\mathbf{U} \mathbf{A} \mathbf{U}^{-1}, \mathbf{B} \mathbf{U}^{-1}, \mathbf{U} \mathbf{C}, \mathbf{D})\) via an invertible transformation matrix \(\mathbf{U}\). It can be shown that both representations are mathematically equivalent under a change of basis²².

By applying this transformation, and assuming \(\overline{\mathbf{A}}\) can be diagonalized into the canonical form \(\mathbf{P}^{-1} \Lambda \mathbf{P}\) (where \(\Lambda\) is diagonal), we enforce the constraint that the state matrix is diagonal, giving rise to diagonal state-space models. To illustrate, consider the filter used in the convolutional representation (Eq. (66 )). Substituting \(\overline{\mathbf{A}} = \mathbf{P}^{-1} \Lambda \mathbf{P}\), we can rewrite the term \(\overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t} \cdot \overline{\mathbf{C}}\) as: \[\begin{align} \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t} \cdot \overline{\mathbf{C}} & = & \overline{\mathbf{B}} \cdot (\mathbf{P}^{-1} \Lambda \mathbf{P})^{t} \cdot \overline{\mathbf{C}} \nonumber \\ & = & \overline{\mathbf{B}} \cdot (\mathbf{P}^{-1} \Lambda \mathbf{P}) \cdot (\mathbf{P}^{-1} \Lambda \mathbf{P}) \cdots (\mathbf{P}^{-1} \Lambda \mathbf{P}) \cdot \overline{\mathbf{C}} \nonumber \\ & = & (\overline{\mathbf{B}} \cdot \mathbf{P}^{-1}) \cdot \Lambda^t \cdot (\mathbf{P} \cdot \overline{\mathbf{C}}). \end{align}\]

Since \(\Lambda\) is a diagonal matrix, we can efficiently compute \(\Lambda^t\) by simply raising all the entries of \(\Lambda\) to the \(t\)-th power. We then have a computationally cheaper model, in which \[\begin{align} \overline{\mathbf{A}}' & = & \Lambda, \\ \overline{\mathbf{B}}' & = & \overline{\mathbf{B}} \cdot \mathbf{P}^{-1}, \\ \overline{\mathbf{C}}' & = & \mathbf{P} \cdot \overline{\mathbf{C}}, \\ \overline{\mathbf{D}}' & = & \overline{\mathbf{D}}. \end{align}\]

More detailed discussions of diagonal state-space models in sequence modeling can be found in [131]’s work.

The application of state-space models to Transformers is simple. Each self-attention sub-layer is replaced in this case by an SSM sub-layer as described in Eqs. (60 ) and (61 ). As we have seen there is a close relationship between state-space models and both CNNs and RNNs. For sequence modeling, we can deal with a sequence of tokens either sequentially as in RNNs, or in parallel as in CNNs. This leads to a new paradigm that takes both the sequential view and the parallel view of the sequence modeling problem — for training, the system operates like CNNs to make use of fast parallel training algorithms; for prediction, the problem is re-cast as a sequential update problem which can be efficiently solved by using RNN-like models.

While the formalism of state-space models is different from that we discussed in this chapter, it provides a unifying framework for sequence modeling. By viewing the problem through complementary perspectives (parallel and sequential), we can optimize for different hardware and latency constraints. Several recent architectures are motivated by this duality, yielding systems that effectively combine the parallelizability of Transformers with the inference efficiency of RNNs [127], [132].

5.6 Conditional Computation↩︎

So far in our discussion of efficient Transformer models, we have assumed that the model architecture is determined before training and remains fixed throughout. We now turn to the case of learning efficient model architectures. We can write a model in the form \[\begin{align} \mathbf{y} & = & \mathrm{Model}(\mathbf{x},g(\mathbf{x})), \end{align}\]

where \(\mathbf{x}\) and \(\mathbf{y}\) are the input and output of the model. \(g(\mathbf{x})\) is a function that returns the model architecture and corresponding parameters for the given input \(\mathbf{x}\). In general, we adopt the convention prevalent in learning problems of using a fixed model architecture and learning only the parameters, say, \(g(\mathbf{x}) = \theta\). In this case, the goal is to learn optimal parameter values given the architecture and training data. On test data, we make predictions using the same model architecture along with the optimized parameters.

A natural extension of this approach is to consider the learning of both the model architecture and parameters. In architecture learning, we would like to find a function \(\hat{g}(\mathbf{x})\) that produces the optimal model architecture and parameter values given the input \(\mathbf{x}\). However, searching the entire hypothesis space of all possible combinations of architectures and parameters is intractable, and so practical methods are required. Two classes of methods can be applied:

• Neural Architecture Search (NAS). In automated machine learning (AutoML), neural architecture search is the process of exploring a space of neural networks to find one that best fits some criteria [133], [134]. Once the optimal neural network is determined, its parameters are trained as usual, and then applied to new data. In order to make search tractable, several additional techniques, such as search space pruning and fast search algorithms, are typically used. Applying neural architecture search to the development of efficient neural networks is straightforward [135], [136]. We need only incorporate efficiency measures into the performance estimation of neural networks, for example, the search can be guided by a criterion that penalizes neural networks with high latency or excessive memory requirements.

• Dynamic Neural Networks. The key idea of dynamic neural networks is to adapt a neural network to various inputs dynamically [137], [138]. Ideally, we would like to learn \(\hat{g}(\cdot)\), and then, for any input \(\mathbf{x}_{\mathrm{new}}\), we apply the model \(\mathrm{Model}(\mathbf{x}_{\mathrm{new}},\hat{g}(\mathbf{x}_{\mathrm{new}}))\). As a result, at test time we may have different model structures and/or different parameters for different inputs. However, it is infeasible to develop a function \(\hat{g}(\cdot)\) that can model arbitrary neural networks. In practice, \(\hat{g}(\cdot)\) is often considered to represent a family of sub-networks of a super-network. The problem is therefore reframed as a simpler problem to learn to choose which sub-network is used for a given input.

From a machine learning perspective, the approaches to neural architecture search are general and can be applied to any neural network. On the other hand, from a practical perspective, it is still difficult to find an efficient neural network that is sufficiently powerful and generalizes well. While neural architecture search provides interesting ideas for developing efficient Transformer models, we do not discuss it in depth here. Instead, the reader can refer to the above papers to have a general idea of it, and refer to [139], [140], and [141]’s work for its application to Transformers.

In this subsection, we focus on a particular family of approaches to dynamic neural networks, called conditional computation. This concept was originally motivated by the dynamic selection of neurons of a neural network [142], [143]. More recently, it has often been used to refer to the process of dynamically selecting parts of a neural network. A narrow view of conditional computation is to see \(g(\cdot)\) as an adaptive neural network which dynamically reduces or grows the number of computation units (such as neurons and layers). As a result, computation can adapt to changing conditions, and we can seek a good accuracy-latency trade-off by this adaptation mechanism.

A common way to achieve this is to learn how to skip some computation steps so that we can work with a subset of the network [144]. One of the simplest methods, sometimes called early exiting, is to stop the computation at some point during encoding or decoding a sequence. This technique is often used in practical sequence generation applications where a low latency is required. Suppose \(y_1 \cdots y_{n_{\mathrm{max}}}\) is the longest sequence that the system can generate, and \(\mathbf{s}_1\cdots\mathbf{s}_{n_{\mathrm{max}}}\) is the corresponding sequence of the hidden states of the top Transformer layer. Then we develop a model \(f_{\mathrm{stop}}(\cdot)\) that takes one hidden state \(\mathbf{s}_i\) at a time and produces a distribution of a binary variable \(c \in \{\mathrm{stop},\mathrm{nonstop}\}\): \[\begin{align} \mathop{\mathrm{Pr}}(c|\mathbf{s}_i) & = & f_{\mathrm{stop}}(\mathbf{s}_i). \label{eq:early-stop-classification} \end{align}\tag{67}\]

The generation process terminates if \(\mathop{\mathrm{Pr}}(\mathrm{stop}|\mathbf{s}_i)\) is sufficiently large, for example \[\begin{align} \mathop{\mathrm{Pr}}(\mathrm{stop}|\mathbf{s}_i) & \ge & \mathop{\mathrm{Pr}}(\mathrm{nonstop}|\mathbf{s}_i) + \theta_{\mathrm{stop}}, \label{eq:conditional-computation-simple-early-stop} \end{align}\tag{68}\]

where \(\theta_{\mathrm{stop}}\) denotes the minimal margin for distinguishing the two actions²³. This formulation is also related to the stopping criterion problem that is frequently discussed in search algorithms for sequence generation. \(f_{\mathrm{stop}}(\cdot)\) can be designed in several different ways. For example, in many practical applications, the stopping criterion is based on simple heuristics. Alternatively, we can define the function \(f_{\mathrm{stop}}(\cdot)\) as a neural network and train it using labeled data.

The above approach can be easily extended to handle situations in which some of the tokens are skipped. This learning-to-skip approach is typically used in the encoding stage in which all input tokens are given in advance. Let \(\mathbf{h}_1 \cdots\mathbf{h}_m\) be low-level representations of a sequence \(x_1 \cdots x_m\). Like Eq. (67 ), we can develop a model \(\mathop{\mathrm{Pr}}(c|\mathbf{s}_i)\) (\(c \in \{\mathrm{skip},\mathrm{nonskip}\}\)) to determine whether the token \(x_i\) can be skipped. Figure 19 (a) and (b) show illustrations of early exiting and skipping. Note that the learning-to-skip method overlaps with other lines of research on training neural networks. For example, erasing some of the input tokens in training is found to be useful for achieving higher generalization of Transformer models [145], [146]. This method is also related to the downsampling methods which will be discussed in Section 5.8.

A second approach to conditional computation is to resort to sparse expert models, or its popular instance — MoE [91]. In deep learning, a model of this kind is typically built from a number of experts which are neural networks having the same structure but with different parameters. In this way, we can construct a large model by simply increasing the number of experts. When running this model, during either training or prediction, we activate only a small number of the experts by some routing algorithms (see Figure 19 (c)). An MoE model is an adaptive network since each time we have a new input, the model routes it to different experts. In Section 4.4, we presented the basic form of MoE, and showed how Transformer models can be scaled up by this sparse method. For a comprehensive review of the recent advances in MoE, we refer the interested reader to [147]’s work.

A third approach that can be used to adapt a Transformer model to changing input is to dynamically shrink the number of layers. Several methods have been proposed to do this in an attempt to improve inference efficiency. The simplest of these is to exit at some hidden layers by which we can still make accurate predictions for the input (see Figure 19 (d) and (e)). To do this, we can either determine the appropriate depth for the entire sequence (call it a sentence-level depth-adaptive model), or use an adaptive depth for each token (call it a token-level depth-adaptive model). Here we consider token-level depth-adaptive models but the methods can be easily extended to sequence-level depth-adaptive models.

Suppose there are \(L\) stacked layers at position \(i\)²⁴. We would ideally like to find a layer in the stack, which can be used as the last hidden layer for making predictions, and whose depth is as low as possible. However, we cannot simply use the \(L\)-th layer of the stack as the oracle for this problem, because we never know in advance what the last layer generates during inference. Instead, we need to determine whether the network should stop growing at depth \(i\), considering the layers generated so far.

Now suppose we have a Transformer decoder which produces a distribution over a vocabulary \(V\) at each step. As usual, we denote the output of the \(l\)-th layer at step \(i\) by \(\mathbf{s}_i^{l}\). For each \(\mathbf{s}_i^{l}\), we create an output layer that produces a distribution \(\mathbf{p}_i^{l}\) over the vocabulary (call it an early exit classifier), given by \[\begin{align} \mathbf{p}_i^{l} & = & \mathrm{Softmax}(\mathbf{s}_i^{l} \cdot \mathbf{W}_{\mathrm{o}}^{l}), \end{align}\]

where \(\mathbf{W}_{\mathrm{o}}^{l} \in \mathbb{R}^{d \times |V|}\) is the parameter matrix. Hence we have \(L-1\) additional output layers, each corresponding to a hidden layer from depth 1 to \(L-1\). At training time, we consider the cross-entropy losses of \(\{\mathbf{p}_i^{1},\cdots,\mathbf{p}_i^{L-1}\}\), and train these layers together with the Transformer model. At test time, the depth of the network grows as usual, and we use \(\{\mathbf{p}_i^{1},\cdots,\mathbf{p}_i^{l}\}\) and/or \(\{\mathbf{s}_i^{1},\cdots,\mathbf{s}_i^{l}\}\) to determine whether we should exit at the \(l\)-th layer. There are several exit criteria, for example,

• Common criteria are based on measures of the confidence of predictions. A simple method is to compute the entropy of \(\mathbf{p}_i^{l}\), and exit if this entropy is above a pre-defined value.

• Alternatively, one can view the maximum probability of the entries of \(\mathbf{p}_i^{l}\) as the confidence of the prediction.

• Instead of considering the output of a single layer, we can also examine the change in the outputs or hidden states over a number of layers. For example, we can measure the similarity between \(\mathbf{p}_i^{l-1}\) and \(\mathbf{p}_i^{l}\) or between \(\mathbf{s}_i^{l-1}\) and \(\mathbf{s}_i^{l}\). If the similarity is above a given threshold, then we say that the output of the neural network tends to converge and the number of layers can stop growing.

• The above methods can be extended to examine the change in the predictions made by the classifiers associated with the layers. For example, the model can choose to exit if the predictions made by the classifiers remain unchanged for a number of layers.

Discussions of these criteria can be found in the related papers [148]–[150]. There are a variety of ways to improve these early exit methods. One is to explore other forms of the prediction for each layer. For example, we can develop a model that directly predicts how many layers we need to model the input [151]. Another line of research on early exit focuses on better training for these models, for example, we can consider various loss functions for training the classifiers [150], [152]. In addition, there is also interest in learning the combination of the outputs of multiple layers so that we can make predictions by using multiple levels of representation [149], [153].

A problem with token-level adaptive-depth models is that the representations at certain depths may be absent in the previous steps. In this case, standard self-attention is not directly applicable, because we may not attend to the previous tokens in the same level of representation. For training, this can be addressed by using all the \(L\) layers of the full model. For inference, we can either duplicate the layer from which we exit to fill up the layer stack, or modify the self-attention model to enable it to attend to the representations of the previous tokens at different depths.

It is also possible to select any subset of the layers for constructing a shallow network. The adaptive models therefore can be generalized to skipping models (see Figure 19 (f)). As with the early exit problem, the skipping problem can be framed as a learning task, in which a classifier is trained to decide whether a layer should be dropped. The learning-to-skip problem has been studied in the field of computer vision [154], [155]. However, learning a skipping model for large-scale, deep neural networks is difficult. For practical systems, it still seems reasonable to use heuristics or cheap models to obtain a neural network having skipped layers, which has been discussed in recent pre-trained NLP models [156], [157].

5.7 Model Transfer and Pruning↩︎

Many large Transformer models have been successfully developed to address NLP problems. A common question is: can we transform a large, well-trained model into a smaller one that allows for more efficient inference? At a high level, this can be thought of as a transfer learning problem in which the knowledge is transferred from one model to another. However, we will not discuss this general topic, which spans a broad range of issues and models, many of which fall outside the scope of this chapter. Instead, we narrow our discussion to two approaches that are widely used for training smaller neural networks from larger ones.

5.7.1 Knowledge Distillation↩︎

Knowledge distillation is a process of compressing the knowledge in a large neural network (or an ensemble of neural networks) into a smaller neural network [158]. In the supervised learning of neural networks, objective functions are generally designed to quantify the error between the predicted and true answers. Hence, by minimizing this loss, models are trained to output the correct answers. While models are typically optimized on the training data in this manner, what we really want is to generalize to unseen data. This is, however, difficult because training only with ground-truth labels provides no direct supervision signal regarding generalization. In knowledge distillation, instead of forcing a model to match the ground-truth output, we train it to mimic the behavior (or output distributions) of a larger model. To achieve this, we directly transfer the knowledge of a pre-trained model to the model we wish to train.

A frequently used approach to knowledge distillation is teacher-student training. A teacher model is typically a relatively large neural network that has already been trained and can generalize well. A student model is a smaller neural network, such as one with fewer layers, to which we transfer the knowledge. A simple way to distill the knowledge from the teacher model into the student model is to use the output of the teacher model as the “correct” answer for training the student model. Suppose we have a teacher Transformer model that can generate a sequence of distributions \(\{\mathop{\mathrm{Pr}}(\cdot|y_0,\mathbf{x}),\cdots,\mathop{\mathrm{Pr}}(\cdot|y_0 \cdots y_{n-1},\mathbf{x})\}\) for the input \(\mathbf{x}\). To keep the notation simple, we denote the distribution \(\mathop{\mathrm{Pr}}(\cdot|y_0 \cdots y_{i-1},\mathbf{x})\) as \(\widetilde{\mathbf{p}}_i\). Similarly, we denote the output of the student Transformer model for the same input as \(\mathbf{p}_i\). As usual, we consider a loss function \(\mathrm{Loss}(\widetilde{\mathbf{p}}_i,\mathbf{p}_i)\) (such as the cross-entropy function) for computing the distance between \(\widetilde{\mathbf{p}}_i\) and \(\mathbf{p}_i\). Then, we can define the loss over the entire sequence as \[\begin{align} \mathcal{L}(\mathbf{x},\theta) & = & \frac{1}{n} \sum_{i=1}^{n} \mathrm{Loss}(\widetilde{\mathbf{p}}_i,\mathbf{p}_i), \end{align}\]

where \(\theta\) denotes the parameters of the student model²⁵. Using this loss, we can optimize \(\theta\) for any given set of source sequences \(\{\mathbf{x}_1, \cdots, \mathbf{x}_K\}\) by minimizing the objective \(\sum_{k=1}^{K} \mathcal{L}(\mathbf{x}_k,\theta)\).

Several extensions to this basic method have been developed to model the problem of knowledge transfer between two models. A simple way is to use the hidden states instead of the output probabilities as the training targets [159]. In this case, the objective is to minimize the difference between some hidden states of the teacher model and the corresponding states of the student model. Rather than using the outputs of various layers as the targets for training the student model, another technique is to model the relations between samples and train the student model by minimizing the difference between the relation encodings of the teacher and student models [160], [161]. For example, we can develop a relation encoding model based on the Transformer architecture. The goal is then to optimize the student model so that its corresponding relation encoding of a group of samples is as close as possible to that of the teacher model.

For sequence generation problems, a special case of knowledge distillation, which can be viewed as a means of data augmentation, is often used for developing lightweight models [162]. For example, consider the problem of transferring the translation ability of a well-developed machine translation model (i.e., the teacher model) to a new model (i.e., the student model). Given a set of source-side sentences \(\{\mathbf{x}_1,\cdots,\mathbf{x}_K\}\), we can use the teacher model to translate each \(\mathbf{x}_k\) to a target-side sentence \(\widetilde{\mathbf{y}}_k\). Then, by treating \(\mathbf{x}_k\) and \(\widetilde{\mathbf{y}}_k\) as paired sentences, we obtain a bilingual dataset consisting of \(\{(\mathbf{x}_1,\widetilde{\mathbf{y}}_1),\cdots,(\mathbf{x}_K,\widetilde{\mathbf{y}}_K)\}\). We can use this bilingual dataset as the labeled dataset to train the student model as usual. One advantage of this data augmentation method is that it is architecture-agnostic. It does not require access to the internal architectures of either the teacher or the student models. Hence, it can be applied even if the teacher model is used as a black-box model. More detailed discussions on knowledge distillation can be found in the surveys by [163] and [164].

5.7.2 Structured Pruning↩︎

Pruning is among the most popular model compression methods and has been applied to a broad range of systems. One common approach is unstructured pruning, in which we retain only a sparse subset of connections between neurons. However, as with most sparse models, networks pruned in this way typically require specialized implementations and hardware support, which in turn limits their efficiency gains in certain applications. A simpler yet more aggressive alternative is structured pruning. In deep learning, structured pruning is a technique that removes entire groups of neurons or connections simultaneously. For example, we can remove an entire layer of neurons from a neural network to obtain a shallower model. Given their multi-layer and multi-head architecture, Transformers are naturally suited to structured pruning, and we can prune a Transformer network in several different ways. For instance, we can prune specific heads in the multi-head attention mechanism [165], [166], or entire layers within the stack [167], [168].

Formally, we can represent a neural network as a set of parameter groups \(\{\theta_1, \cdots, \theta_{R}\}\), each corresponding to a specific component or sub-module of the network. Our goal is to find a subset of \(\{\theta_1, \cdots, \theta_{R}\}\) that allows us to build a model yielding strong performance while maintaining significantly lower computational complexity. However, an exhaustive search for such an optimal subset is infeasible because there is a combinatorial number of possible candidates, and evaluating all of them is computationally prohibitive.

One approach to structured pruning is to randomly discard components of a model. We can repeat this random pruning process multiple times to generate a pool of model candidates, and select the best-performing one. Another approach relies on heuristics to determine which components are redundant and can be safely removed. Common measures for evaluating the importance of a parameter group \(\theta_r\) include various metrics based on the norms of its weights or gradients [169]. We can then prune \(\theta_r\) if these metrics fall below or exceed predefined thresholds. A third approach frames the pruning problem as a continuous optimization task by introducing trainable gates that indicate the presence or absence of different components [170]–[172]. The final pruned model can be directly derived using these trained gates. It is worth noting that, in many cases, pruning is no longer merely a post-processing step for a fully trained model, but rather an integral part of the training process itself.

5.8 Sequence Compression↩︎

In sequence modeling and generation problems, time and space complexities are strongly influenced by the length of the input or output sequences. Therefore, shorter sequences are generally preferred. This is particularly important for Transformer models, as their time and space complexities scale quadratically with sequence length, making memory footprint and latency significant computational bottlenecks for very long sequences. While previous subsections discussed modifications to the Transformer architecture to handle long sequences, in this section, we instead explore methods for compressing sequences to more manageable lengths.

One simple approach is to map the input sequence to a fixed-size representation. For example, using the recurrent models discussed in Section 5.2, we can encode a sequence of vectors into a single vector. This method can be easily extended to generate a larger representation so that this representation can retain more information from the original input. For example, we can select a fixed number of the hidden states over the sequence to form a new sequence of fixed length. Another way to represent a variable-length sequence as a fixed-length sequence is to perform cross-attention between the input vectors and a set of learnable hidden representations. In the work of [173], this is done by introducing \(r\) hidden representations \(\{\mathbf{u}_{1},\cdots,\mathbf{u}_{r}\}\), and then applying cross-attention between the input vectors \({\mathbf{x}_1,\cdots,\mathbf{x}_m}\) and these hidden representations. The attention model can be a standard QKV attention model in which we view \(\{\mathbf{u}_{1},\cdots,\mathbf{u}_{r}\}\) as queries and \(\{\mathbf{x}_1,\cdots,\mathbf{x}_m\}\) as keys and values. The output of this model is a sequence of \(r\) vectors, which can be used as fixed-length input to downstream systems.

A second approach is to use downsampling to compress the sequence into a shorter one. A typical method of downsampling is strided convolution, which has been widely used in computer vision and speech processing. For example, suppose there is a sequence of \(m\) vectors in \(\mathbb{R}^d\). We can define a filter with a width of \(2\) and a stride of \(2\). By taking the sequence as input, the filter produces a sequence of \(\frac{m}{2}\) new vectors in \(\mathbb{R}^d\), and so we have a reduction factor of \(2\). Also, we can stack multiple convolutional layers or pooling layers to achieve a desired level of length reduction, called progressive downsampling. However, it seems inevitable that downsampling will lead to information loss [174], [175]. We need to consider a trade-off between the compressed sequence length and the performance of downstream systems [176].

In NLP, the problem of sequence compression is also closely related to the problem of tokenizing input strings. Therefore, tokenization is a practical approach that can be taken to address the length issue. Segmenting a string into small tokens (such as characters) generally reduces the sparsity of the data, which makes it easier to learn the embeddings of these tokens, but such approaches often lead to a long sequence. By contrast, we will have a shorter sequence if we segment the input string into larger units, but this may suffer from data sparsity. In deterministic tokenization methods, which produce tokenization results using statistics collected from the entire dataset, the sequence length can be controlled to some extent by adjusting specific hyperparameters, for example, in byte pair encoding [177], increasing the size of the vocabulary generally reduces the number of resulting tokens. Another way to obtain an appropriate sequence of tokens is to use a model for choosing among tokenization candidates [178], [179]. As with many probabilistic models for text generation, in this case, we can incorporate priors to the criterion for tokenization selection so that we can express a preference for shorter sequences over longer sequences.

A fourth approach to sequence compression is to drop some of the tokens in the sequence. For instance, in many practical applications, sequences are simply truncated when their lengths exceed a predefined threshold. We can relate this to the early exiting and skipping approaches in conditional computation. Thus the methods discussed in Section 5.6 can be directly applied. The token dropping methods can also be viewed as pruning methods, called token pruning. By discarding tokens that are less important for representing the entire sequence, token pruning can significantly reduce the sequence length while maintaining the performance of NLP systems on downstream tasks [180].

5.9 High Performance Computing Methods↩︎

So far in this section, we have discussed efficient Transformer models from the perspective of deep learning. However, we have not considered their hardware efficiency. As modern hardware provides a variety of execution paradigms, the practical time and memory footprint savings generally depend on the specifications of the hardware systems. One line of research on the efficient use of computing resources explores parallel computing methods. There have been many attempts to develop large-scale Transformer models by using a cluster of machines. Typically, scaling Transformers to models with billions or even tens of billions of parameters requires a careful design of parallelism strategies for sharding the big networks. More efficient implementations of such systems also require consideration of communication overhead in the cluster, as well as the utilization of sparse models that activate only a small subset of the parameters for each sample, thereby enabling the use of very large models. Most of these methods have been extensively studied in the literature [86], [87], [181]. The results of these studies are foundational to many follow-on works investigating the scaling laws for large language models [3], [94]. Since large-scale distributed models are generic and not specialized to the case of Transformers, we skip their discussion here. Interested readers can refer to the above papers for more detailed discussions.

In this subsection, we consider hardware-aware methods to seek greater computational efficiency for Transformer models. We first consider a simple but widely used method that aims to represent and execute neural networks using lower- or mixed-precision formats [182]. Conventional neural networks are typically based on single-precision and/or double-precision floating-point representations of data. While single-precision floating-point data types provide a sufficiently precise way to represent parameters and intermediate states in most cases, in some applications, they are not essential. As an alternative, one can use half-precision (or even lower-precision) formats for representing floating-point numbers in neural networks. The size of the resulting model is thus half the size of the original model. One advantage of using half-precision floating-point representations is that, although processing such data types requires new APIs for linear algebra operations and hardware support, it does not change the model architecture, requiring only slight modifications to the system. For example, half-precision floating-point representations can be applied to either the training or inference of Transformers, or both.

Recently, the deployment of large Transformer models has been further improved by quantizing these models. In signal processing, quantization is a process of mapping continuous values (i.e., floating-point representations) to a set of discrete values (i.e., fixed-point representations). This process is generally implemented using a system called a quantizer. In the context of neural networks, a quantizer involves two functions — the quantization function and the de-quantization function. The quantization function maps a floating-point number to a (lower-bit) integer. A simple quantization function is given by \[\begin{align} Q(x) & = & \lfloor \frac{x}{s} \rceil, \end{align}\]

where \(\lfloor \cdot \rceil\) is a rounding function²⁶, \(x\) is the real-valued input, and \(s\) is the quantization step size that controls the level of quantization. The quantization function is coupled with a de-quantization function \[\begin{align} D(r) & = & s \cdot r. \end{align}\]

With this notation, the quantizer can be expressed as \(D(Q(x)) = s \cdot \lfloor \frac{x}{s} \rceil\). The difference between \(D(Q(x))\) and \(x\) is called the quantization error. A smaller value of \(s\) typically reduces the quantization error. In practice, however, we wish to choose an appropriate value of \(s\) in order to spread the possible values of \(Q(x)\) evenly across the integer representation range. For example, \(s = \frac{\max\{D(r)\}}{2^p-1}\), where \(p\) is the number of bits used to represent an integer and \(\max\{D(r)\}\) is the maximum representable value after de-quantization. The equations above show one of the simplest cases of quantization. More general discussions of quantization can be found in books on digital signal processing and related surveys [183]–[185].

Applying quantization to Transformers is relatively straightforward. The idea is that we quantize the inputs and model parameters using \(Q(x)\) and feed them to a quantized Transformer model in which all layers operate on integer-valued tensors. In other words, we implement the model using integer-only arithmetic. However, the price to be paid for this compressed model, as with many approximation approaches to deep learning, is that its predictions are less accurate than those of the standard Transformer. Using integer operations to approximate continuous-valued operations generally leads to approximation errors. These errors will accumulate if the quantized neural network is deep. Furthermore, Transformer models involve components (such as self-attention sub-layers) that require relatively complex linear algebra operations. Simply applying quantization to these modules will lead to a significant loss in accuracy. One solution is to simplify the model architecture and develop new sub-models that are more suited for quantization. Alternatively, a more common paradigm in quantized neural networks is to add de-quantization functions to the neural networks so that the outputs of a layer are floating-point tensors and can be used as usual in the following steps. Consider a simple example where we multiply a real-valued input matrix \(\mathbf{a}\) with a real-valued parameter matrix \(\mathbf{A}\). We first quantize \(\mathbf{a}\) and \(\mathbf{A}\), and multiply them using integer-based matrix multiplication. The result is then de-quantized to a real-valued matrix. In this way, we obtain an approximation \(D(Q(\mathbf{a}) \cdot Q(\mathbf{A}))\) to \(\mathbf{a} \cdot \mathbf{A}\) at a very low computational cost.

However, sandwiching each layer between \(Q(\cdot)\) and \(D(\cdot)\) will lead to the additional cost of running \(Q(\cdot)\) and \(D(\cdot)\). In some practical applications, the computational overhead introduced by \(Q(\cdot)\) and \(D(\cdot)\) is even larger than the time savings of performing integer-based operations. In general, the benefits of quantizing neural networks are greater than the costs only if the neural networks are sufficiently large. Therefore, in practice, it is common to perform quantized computation only for operations whose computational costs are high. For example, in recent large language models, quantization is primarily applied to the multiplication of large matrices, yielding significant time and memory savings.

While quantization approaches can be used in both training and inference, a widely used approach is to quantize Transformer models after training (called post-training quantization). In this approach, quantization is performed on well-trained floating-point-based neural networks, resulting in fewer quantization-related errors. However, these errors cannot be easily compensated for because they are introduced after training is complete. A more promising approach is to incorporate quantization during training so that the model can learn to compensate for quantization-related errors [186], [187]. There have been several attempts to apply quantization-aware training to Transformers [188]–[190]. In addition to computational efficiency, another important consideration for high-performance systems is the constraints of the memory hierarchy. In general, better system design requires considering the speeds and capacities of different memory levels. The problem is even more complicated when training large Transformer models on modern hardware where both GPUs and CPUs are utilized. A general principle of system design is that memory transfers between different memory levels should be minimized. While we would ideally like to have a large, high-level memory on which we can store all the data we need to process, in many practical situations, the size of fast, on-chip memory is orders of magnitude smaller than the total size of the data. In this case, we can re-order the memory access in the algorithms so that data used in adjacent computation steps can be loaded into high-speed memory all at once. This idea motivates the development of many fast linear algebra libraries. For example, there are matrix multiplication algorithms that are highly optimized for different shapes of input matrices.

It is relatively straightforward to use these optimized linear algebra algorithms to build a Transformer-based system, but the modules of this system are not jointly optimized for efficiency. For example, a self-attention sub-layer involves a series of scaling, normalization, and matrix multiplication operations. Although each of these operations is implemented in supported and efficient linear algebra libraries, successive calls to them still require multiple memory transfers when switching from one operation to another. In practice, a better approach would be to keep some of the intermediate states in the on-chip memory and reuse them in subsequent computation steps instead of fetching them again from slower memory. For example, on modern GPUs, a simple way to achieve this is to merge multiple operations into a single operation, a technique known as kernel fusion. For Transformer models, a general idea is to design data partitioning and layout strategies that maximize the computation on each data block loaded into high-performance memory while simultaneously minimizing memory transfers. There have been several attempts to use these strategies to improve the attention models in Transformers [191], [192]. Some of these methods, such as FlashAttention and PagedAttention, have been successfully incorporated into recent large language models [193], [194].

6 Applications↩︎

Transformers have a wide range of applications across NLP tasks. While the Transformer model introduced by [1] is based on a standard encoder-decoder architecture, it is primarily used in three different configurations:

• Decoder-only Models. If the cross-attention sub-layers are removed, the decoder becomes a standard language model. Hence, this decoder-only model can be applied to text generation problems. For example, given a sequence of left-context tokens, we use the model to predict subsequent tokens.

• Encoder-only Models. Transformer encoders can be treated as sequence models that take a sequence of tokens as input and produce a sequence of representations, each corresponding to an input token. These representations can be viewed as a form of encoding of the input sequence, and are often taken as input to a prediction model. This encoder-predictor architecture forms the basis of many NLP systems, such as systems for sentence classification and sequence labeling. Pre-trained Transformer encoders can also be used to map texts into the same vector space so that we can compute the distance or similarity between any two texts.

• Encoder-Decoder Models. Encoder-decoder models are typically used to model sequence-to-sequence problems. Applications include many tasks in NLP and related fields, such as machine translation and image captioning.

Note that while most Transformer-based systems can be categorized into the above three types, the same NLP problem can generally be addressed using different types of models. For example, recent decoder-only models have demonstrated strong performance on a broad range of problems by framing them as text generation tasks, even though many of these problems were historically addressed using encoder-decoder or encoder-only models. To illustrate how the above models are applied, this section considers a few applications where Transformers are chosen as backbone models.

6.1 Language Modeling↩︎

Language modeling is an NLP task in which we predict the next token given its preceding tokens. This is generally formulated as a problem of estimating the distribution of tokens at position \(i+1\) given tokens at positions \(0 \sim i\). Here we denote this distribution by \(\mathop{\mathrm{Pr}}(\cdot \mid x_0, \ldots, x_i)\), where \(x_0, \ldots, x_i\) are the observed tokens. The predicted token is obtained by maximizing the probability: \[\begin{align} \hat{x}_{i+1} & = & \mathop{\mathrm{arg\,max}}_{x_{i+1} \in V} \mathop{\mathrm{Pr}}(x_{i+1}|x_0,\cdots,x_{i}), \end{align}\]

where \(V\) is the vocabulary. The prediction can be extended to the tokens following \(\hat{x}_{i+1}\)

\[\begin{align} \hat{x}_{k+1} & = & \mathop{\mathrm{arg\,max}}_{x_{k+1} \in V} \mathop{\mathrm{Pr}}(x_{k+1}|x_0,\cdots,x_{i},\hat{x}_{i+1},\cdots,\hat{x}_{k}) \label{eq:transformer-app-lm-long}. \end{align}\tag{69}\]

This model forms the basis of many systems for text generation: given the context tokens \(x_0 \cdots x_{i}\), we generate the remaining tokens \(\hat{x}_{i+1} \cdots \hat{x}_{k+1}\) to make the sequence complete and coherent.

As discussed in Section 2.1, Transformer decoders are essentially language models. A key difference between the problem of decoding in an encoder-decoder Transformer and the problem of language modeling is that the Transformer decoder makes predictions conditioned on the context tokens on both the encoder and decoder sides, rather than being conditioned only on preceding tokens. To modify the Transformer decoder to implement a standard language model, the cross-attention sub-layers are simply removed and a Transformer decoding block can be expressed as \[\begin{align} \mathbf{S}^{l} & = & \mathrm{Layer}_{\mathrm{ffn}}(\mathbf{S}_{\mathrm{self}}^{l}), \\ \mathbf{S}_{\mathrm{self}}^{l} & = & \mathrm{Layer}_{\mathrm{self}}(\mathbf{S}^{l-1}). \end{align}\]

Here, \(\mathbf{S}^{l}\) denotes the output of the block at depth \(l\). \(\mathrm{Layer}_{\mathrm{self}}(\cdot)\) denotes the self-attention sub-layer, and \(\mathrm{Layer}_{\mathrm{ffn}}(\cdot)\) denotes the FFN sub-layer. We see that this decoding block has the same form as an encoding block. The difference between the decoding and encoding blocks arises from the masking strategies adopted during training: the former masks the attention from a position \(i\) to any right-context position \(k > i\), whereas the latter has no such restriction. A softmax layer is stacked on top of the last block and is used to produce the distribution over the vocabulary at each position. For inference, the Transformer decoder works in an autoregressive manner, as described in Eq. (69 ).

Training these models follows a standard paradigm. We learn the model by repeatedly updating the parameters, based on the gradients of the loss on the training samples. This paradigm can be extended to the training of large Transformer-based language models, which have been widely applied in generative AI. However, training Transformer models at scale, including decoder-only, encoder-only, and encoder-decoder models, may lead to new difficulties, such as training instabilities.

6.2 Text Encoding↩︎

For many NLP problems, a widely used paradigm is to first learn a representation of an input sequence, and then make predictions for downstream tasks based on this representation. As a result, we separate sequence representation from specific NLP tasks. One of the advantages of this paradigm is that we can train a sequence model that is not specialized to particular problems, thereby enabling it to generalize well.

Clearly, Transformer encoders are a type of sequence model and can be used as text encoders. Consider a Transformer encoder with \(L\) encoding layers. The output of the last encoding layer can be viewed as the sequence representation. Here, we prepend a special token \(x_0\) to each sequence, indicating the beginning of the sequence (often written as \(\langle \mathrm{SOS} \rangle\) or \([\mathrm{CLS}]\)). Given a sequence of \(m+1\) input tokens \(x_0, x_1,\cdots, x_m\), the output of the encoder will be a sequence of \(m+1\) vectors \(\mathbf{h}_0^{L}, \mathbf{h}_1^{L},\cdots, \mathbf{h}_m^{L}\). Since \(x_0\) does not correspond to an actual word and has a fixed positional embedding, it serves as a global representation token for collecting information from other positions via the self-attention mechanism. Hence, \(\mathbf{h}_0^{L}\) acts as a representation of the entire sequence, without being biased toward any specific tokens or positions. In many cases, we require a single representation of a sequence to serve as input for downstream components of the system. For example, we can construct a sentence classification system based on a single vector generated from \(\{\mathbf{h}_0^{L}, \dots, \mathbf{h}_m^{L}\}\). In this scenario, we can simply use \(\mathbf{h}_0^{L}\) as the sequence representation. A more general approach is to add a pooling layer to the encoder, which allows us to explore various pooling methods to generate the sequence embedding from \(\{\mathbf{h}_0^{L}, \cdots, \mathbf{h}_m^{L}\}\).

In text encoding, token sequences are represented by real-valued vectors, often referred to as sentence representations or sentence embeddings, which can be interpreted as points in a multi-dimensional space [195]. Consequently, another use of text encoding is to compute the semantic or syntactic similarity of token sequences based on their relative proximity in this space. A straightforward method for this is to compute the Euclidean distance or cosine similarity between sequence embeddings. The shorter the distance between two sequences, the more similar they are considered to be. There are many distance metrics available, and it is possible to combine them to obtain a more robust measure of sequence similarity. Such similarity computations are applied in areas such as textual entailment, information retrieval, and translation evaluation [196], [197]. Additionally, they are frequently used to assess the inherent quality of text encoding models.

Text encoding is also a crucial component of sequence-to-sequence models. In sequence-to-sequence tasks, we can develop a separate Transformer encoder for source-side sequence modeling within an encoder-decoder system (see Figure 20). For instance, we can pre-train a Transformer encoder on a large-scale corpus of source-side texts and subsequently use it as the encoder in a downstream encoder-decoder model. It is worth noting that while the encoder is based on the Transformer architecture, the decoder is not restricted to this architecture. This flexibility enables us to incorporate pre-trained Transformer encoders into hybrid sequence-to-sequence architectures, such as systems that combine a Transformer encoder with an LSTM decoder.

In supervised learning scenarios, training a Transformer encoder is straightforward. We can treat it as a regular component of the target model and train this model on task-specific labeled data. However, such a method requires the encoder to be optimized on each task, and the resulting encoder might not always generalize well to other tasks, especially given that labeled data is scarce in most cases. A more prevalent approach is to frame the training of text encoders as an independent task in which supervision signals are derived solely from raw text. This led researchers to develop self-supervised Transformer encoders, such as BERT, which make use of large-scale unlabeled text, and these encoders were found to generalize well across many downstream tasks. Further discussions of pre-trained Transformer encoders can be found in the literature on pre-trained models.

6.3 Speech Translation↩︎

As illustrated in Section 2, the standard encoder-decoder Transformer model was proposed to model sequence-to-sequence problems. Here we consider the problem of translating speech in one language to text in another language — a problem that is conventionally addressed using both automatic speech recognition (ASR) and machine translation techniques. Instead of cascading an automatic speech recognition system and a machine translation system, we can use Transformer models to build an end-to-end speech-to-text (S2T) translation system to directly translate the input speech to the output text.

To simplify the discussion, we assume that the input to an S2T translation system is a sequence of source-side acoustic feature vectors, denoted by \(\mathbf{a}_1 \cdots \mathbf{a}_m\), and the output of the system is a sequence of target-side tokens, denoted by \(y_1 \cdots y_n\).²⁷ Mapping \(\mathbf{a}_1 \cdots \mathbf{a}_m\) to \(y_1 \cdots y_n\) is a sequence-to-sequence problem. Thus it is straightforward to model the problem using an encoder-decoder Transformer model, and the training and inference of this model are standard, as in neural machine translation.

In S2T translation, however, we have to deal with sequence mappings across both modalities and languages. This poses new challenges compared to conventional machine translation problems and influences the design of S2T translation models. There have been several improvements to Transformers to adapt them better for S2T translation tasks. Some improvements focus on the design of Transformer layers [198]. For example, in [48]’s system, a CNN sub-layer and relative positional embeddings are integrated into each Transformer layer, enabling the model to efficiently capture both local and global features.

Another line of research on S2T translation focuses on improving the encoder-decoder architecture. This involves modifying the encoder, the decoder, or both. To illustrate, Figure 21 shows the architectures of three S2T translation models. All of them are based on Transformers but have different encoder architectures. As shown in the figure, the standard encoder-decoder architecture has one Transformer encoder for reading the source-side input \(\mathbf{a}_1, \cdots, \mathbf{a}_m\) and one Transformer decoder for producing the target-side output \(y_1, \cdots, y_n\). By contrast, the decoupled encoder model decomposes the encoder into two stacked encoders — one for acoustic modeling (referred to as the speech encoder) and one for textual modeling (referred to as the text encoder) [199], [200]. This design reflects a modeling hierarchy in which representations at different levels of the network are concerned with different aspects of the problem. For example, the speech encoder models low-level features when mapping acoustic embeddings into larger language units, and the text encoder models the semantic or syntactic features when representing the entire input sequence. An advantage of separating out the text encoder is that the encoding process follows our prior knowledge: we need to first transcribe the speech input and then translate the transcript into the target language. Therefore, we can train the speech encoder in a manner similar to how we train an ASR system. This enables us to pre-train the speech encoder and the text encoder on unlabeled data and subsequently incorporate the pre-trained encoders into S2T translation systems.

An alternative encoding architecture is the two-stream architecture, as shown in Figure 21 (c). Like the decoupled encoder architecture, this architecture has a speech encoder and a text encoder, but the two encoders work in parallel rather than in sequence [201]. The speech encoder takes acoustic features as input, and the text encoder takes tokens (or their embeddings) as input. A third encoder, called the shared encoder, integrates the outputs from both the speech and text encoders, merging the representations from the two modalities. This two-stream architecture is flexible because it provides multiple ways to train S2T translation models. A common approach is to train each branch individually. For example, if we mask the speech encoder, the model will transform into a machine translation model which can be trained using bilingual texts. Conversely, if we mask the text encoder, we can train the model as a standard S2T translation model. For inference, the text encoder can be dropped, and the speech input is modeled using the speech encoder and the shared encoder.

In deep learning, training is often intertwined with architecture design. Because we are dealing with data in two modalities and two languages, we can develop multiple supervision signals for the multi-task learning of S2T translation models. A widely used method is to introduce an ASR-related loss into the training of speech encoders. For example, in the decoupled encoder model, a classifier can be constructed based on the output from the speech encoder. By minimizing the connectionist temporal classification (CTC) loss for this classifier, the speech encoder can be optimized in a manner similar to ASR. In general, training S2T translation models is challenging because speech-to-text aligned data is scarce. Typical responses to this challenge include data augmentation, pre-training, knowledge distillation with machine translation, and so on. However, an in-depth discussion of these methods goes beyond the scope of this section on Transformers. Interested readers can refer to a recent survey on speech translation for more detailed information [202].

6.4 Vision Models↩︎

While Transformers were first used in NLP, their application to other domains has been a prominent research topic. In computer vision, for instance, there is a notable trend of shifting from CNNs to Transformers as backbone models. In this subsection, we consider the Vision Transformer (ViT) — an interesting application of Transformers to image classification [203]. The Vision Transformer is a milestone model that opened the door to purely Transformer-based vision models. Here, we consider the basic structure of the Vision Transformer to keep this section focused and coherent, although extensive literature exists on the Vision Transformer and its variants. More detailed discussions of vision models can be found in recent surveys [204], [205].

The core idea behind the Vision Transformer is to transform an image into a sequence of visual tokens, and input this sequence into a Transformer encoder to generate a representation of the image. The Transformer encoder is standard, and so we will not discuss it here, given the introduction to Transformers we have presented so far in this chapter. Mapping a 2D image into a sequence of tokens requires additional processing. Suppose we have an image represented as an \(H \times W \times C\) feature map, where \(H\) is the height, \(W\) is the width, and \(C\) is the number of channels. The first step is to partition this image into a set of patches. Suppose all patches are squares of side length \(P\). Then, the resulting patches can be represented by feature maps of shape \(P \times P \times C\). By ordering these patches in a specific manner, we obtain a sequence of \(\frac{HW}{P^2}\) patches, with each patch being treated as a “token.”

Given this patch sequence, the subsequent steps are straightforward. For the patch at each position, we obtain a \(d\)-dimensional embedding via a linear transformation of the flattened patch. The input to the Transformer encoder is a sequence of \(d\)-dimensional vectors, each of which is the sum of the corresponding patch embedding and positional embedding. Figure 22 illustrates the patching and embedding steps in the Vision Transformer.

Once we have a sequence of vectors for representing the image, we can employ the Transformer encoder to encode the sequence. The encoding process is exactly the same as that in text encoding, as discussed in Section 6.2. For classification problems, we need only a single representation of the input. It is convenient to take the output of the encoder at position 0 (denoted by \(\mathbf{h}_0^{L}\)) and feed it into a classifier. Given that the first token (\([\mathrm{CLS}]\)) serves as a special global token that is attended to by all other tokens, \(\mathbf{h}_0^{L}\) can be viewed as an aggregated representation of the entire sequence.

A standard way to train the Vision Transformer is to minimize a task-specific loss on labeled data, such as ImageNet. More recently, inspired by self-supervised learning in BERT-like models, there have been successful attempts to train Transformer-based image encoders on large-scale unlabeled data [206]–[208]. Note that one of the most significant contributions of the Vision Transformer is that it unifies the representation models for different modalities. This suggests that if an object, whether an image or text, is represented as a sequence of embeddings, it can be easily modeled using the Transformer architecture.

6.5 Multimodal Models↩︎

The above discussion of the Vision Transformer offers the possibility of unifying the representations from multiple modalities using the same Transformer architecture. In fact, many recent multimodal systems draw heavy inspiration from Transformers [209]. Such systems convert objects from different modalities into vector sequences and feed these vectors into a single Transformer model. The output is a fused representation of all inputs, which can then be used in downstream systems.

As a simple example, consider the task of encoding a pair consisting of text and its corresponding image. First, we represent both the text and the image as sequences of embeddings that have the same dimensionality. This is a common step in sequence modeling, which we have encountered many times so far. We can do this by using either a simple embedding model (e.g., a word or patch embedding model) or a pre-trained model (e.g., a vision model). Then, these two sequences are concatenated into a long sequence comprising both textual and visual embeddings. The subsequent step is standard: a Transformer encoder takes the concatenated sequence of embeddings as input and produces representations of the text and image as output. Note that concatenating textual and visual sequences is one of the simplest methods for vision-text modeling. There are several alternative ways to merge information from different modalities. For example, we can feed visual representations into the attention layers of a text encoder or decoder [210], [211].

The above multimodal encoder can be used in both encoder-only and encoder-decoder systems. For encoder-only systems, consider an example where, given an image and its description, we predict the class of the image using a classifier built on top of the encoder [212]. For encoder-decoder systems, we pair the encoder with a decoder, as in sequence-to-sequence modeling [213]. For example, we might employ a Transformer decoder to generate text based on the output of the encoder. A common application of this architecture is visual question answering (VQA), where an image and a question about the image are provided, and the system is tasked with generating an answer [214]. The architectures of these models are illustrated in Figure 23 (a-b).

Large language models have also been extended to handle both textual and other modalities, such as images, video, and audio, leading to new advances in multimodal processing [215], [216]. By representing all inputs as a sequence of token embeddings, the problem becomes straightforward: we predict the next token given its context. This can be done using decoder-only systems, as shown in Figure 23 (c).

7 Summary↩︎

Transformer models have been widely adopted since the architecture was first proposed by [1]. This adoption has accelerated the development of a wide range of algorithms, systems, and concepts. A thorough discussion of Transformers requires broad coverage, and so it is impossible to address all topics or provide an exhaustive list of references. Figure 24 provides a high-level overview of the Transformer landscape. Note that these models and related techniques can be categorized in various ways, and we present one such perspective. To summarize, we highlight the following points:

• Foundations of Transformers. Although the impact of Transformers has been revolutionary, they are not completely “new” models. From a deep learning perspective, Transformers are composed of common building blocks, including word and positional embeddings [4], [5], [217], attention mechanisms [6], [7], residual connections [59], layer normalization [9], and so on. Many of these components were presented in earlier systems, for example, similar ideas with QKV attention can be found in memory networks [218] and hierarchical attention networks [219]. Transformers offer a novel approach to integrating these components, resulting in a unique architecture. For example, in Transformers, the combination of multi-head attention and dot-product QKV attention, along with the incorporation of layer normalization and residual connections, gives rise to a distinctive neural network block, specifically a self-attention sub-layer. This design has since become a de facto standard in many follow-on sequence modeling systems.

• Attention Models. The success of Transformers on NLP tasks has largely been attributed to the use of multi-head self-attention for sequence modeling. This has led to a surge of interest in enhancing the attention mechanisms within Transformers. While it is impossible to detail every attention model, there are several notable research directions. One prominent direction involves modifying the forms of QKV attention and multi-head attention for improved performance. The scope of this direction is vast, as there are numerous aspects to consider when enhancing Transformers [220]. For example, one may add new components to self-attention sub-layers to adapt them to specific tasks, resulting in various Transformer variants. A second direction is to incorporate prior knowledge into the design of attention models. This makes sense, because much of the emphasis in traditional NLP has been on using linguistic insights to guide system design, and we generally want NLP systems to be linguistically explainable. For example, many Transformer-based systems take syntactic parses as input in various forms and make use of syntax in sequence modeling. A third direction is to develop efficient attention models [221]. Self-attention has long been criticized for its quadratic time complexity and dependency on all previous tokens for each new token. In response, many researchers have focused on simplifying the structure of self-attention, or on approximating it using sparse or recurrent models. This concern for efficiency also motivates the development of alternatives to self-attention, such as attention models with linear time complexity. In addition to exploring stronger and more efficient attention models, it is natural to examine what knowledge is learned by such models. Interestingly, researchers have found that the underlying structure of languages can be learned by multi-head self-attention models, although these models are not trained to represent such knowledge [34].

• Word and Positional Embeddings. Transformers represent each input word (or token) through a combination of word and positional embeddings. Learning these word embeddings is not unique to Transformers; one can either utilize pre-trained embeddings, such as Word2Vec or GloVe, or treat them as learnable parameters within the model. A closely related issue is the tokenization of input sequences. Generally, tokenization determines the number of resulting tokens and affects the difficulty of learning their corresponding embeddings. Therefore, selecting an appropriate tokenization method is crucial for many applications. Furthermore, positional embeddings play a vital role in Transformers, as the underlying attention mechanisms are order-insensitive by design [222]. While positional encoding is a general research topic, much effort has been dedicated to tailoring it for Transformers, leading to various architectural modifications [49], [50]. Additionally, studies show that for sequences significantly longer than those encountered during training, better extrapolation can be achieved by replacing sinusoidal embeddings with rotary positional embeddings (RoPE) or by scaling attention weights with positional scalars [54], [223], [224].

• Training and Model Scaling. In the era of deep learning, high-performance systems are typically built using large-scale neural networks. A straightforward approach to increasing Transformer capacity is to stack more layers or enlarge the hidden representation size. Indeed, deep and wide Transformer models consistently outperform their smaller counterparts. However, training extremely large models poses significant challenges, particularly when applying gradient descent over vast datasets, which demands substantial computational resources. Engineering solutions involve distributing training across computer clusters [86], [94]. Although distributed training is a general paradigm not restricted to Transformers, it heavily influences architectural design; for instance, sparse mixture-of-experts (MoE) models can facilitate training with massive parameter sets and serve as the foundation for many expansive systems. Scaling also enables the study of scaling laws — quantifying how performance scales with model size, data volume, and compute budget [88], [225]. This scaling is occasionally accompanied by “emergence” [226], where models manifest unpredictable new abilities. In recent research, the acquisition of such emergent abilities has been viewed as a hallmark of powerful large language models.

• Efficient Models. There are different goals for efficiency. A system may prioritize memory efficiency for memory-bound problems or low latency when inference speed is critical. Balancing these requirements leads to diverse optimization strategies. In Transformers, many optimizations involve modifying attention mechanisms; for instance, several variants reduce the memory footprint when processing long sequences [221], while others minimize computation to decrease latency. Furthermore, Transformers can benefit from architecture-independent optimizations typical of neural networks, such as conditional computation, knowledge distillation, structured pruning, and sequence compression. Efficiency can also be addressed from a computer architecture perspective [180]. In sequence-to-sequence tasks, for example, the encoding process is typically compute-intensive, whereas the autoregressive decoding process is often IO-intensive. Therefore, we can employ different optimization methods for different components of Transformers.

• Inference. The inference problem is commonly discussed in sequence generation. In NLP, we often need to find the “best” hypothesis in a space involving sequences of hundreds or even thousands of tokens over a vocabulary. As a classic search problem in artificial intelligence, various algorithms are applicable, including breadth-first, depth-first, and A* search. In many practical applications of NLP, the efficiency of the search systems is an important consideration. As a result, optimized search algorithms are required. Most of these algorithms have been explored in machine translation and ASR, and are directly applicable to neural text generation models like Transformers. Beyond traditional search, recent optimizations such as speculative decoding are specifically tailored to the Transformer architecture [227]. Moreover, the efficient attention models mentioned previously are now widely deployed to accelerate large language models and neural machine translation systems [228], [229].

• Applications. Applications of Transformers cover a wide variety of NLP problems. During the development of Transformers, they were at first used to build supervised models that perform particular tasks. Later, a greater success was achieved by using them as backbone networks for large scale self-supervised learning of foundation models [230]. This shift markedly changed the NLP paradigm: a model is first pre-trained on vast amounts of text to acquire general linguistic knowledge, and then adapted to downstream tasks through efficient methods such as fine-tuning or prompting. Beyond NLP, we have witnessed an explosion of Transformer applications in fields like computer vision, speech processing, and bioinformatics. The core idea is to represent diverse data types as sequences of tokens, extending Transformers into general-purpose representation models capable of handling multi-modal data.

• Large Language Models as Foundation Models. Transformers form the basis of recent large language models, such as the GPT series, which show surprising breakthroughs in NLP [8], [231]. Much of the research in large language models is more or less related to Transformers. As discussed in Section 6.1, training these language models is often synonymous with training large-scale Transformer decoders, and architectural refinements to the latter directly benefit the former. On the other hand, the rapid development of large language models has spurred further innovation in Transformer-related techniques, particularly in developing efficient and low-cost adaptation methods for massive models.

• Theoretical Analysis. While Transformers dominate empirically, their theoretical foundations have historically received less focus compared to engineering improvements — an empirical-first trend common across the broader machine learning community. In response, researchers have begun to analyze Transformers more deeply. One approach is to view Transformers as deep neural networks and interpret them via mathematical tools. For instance, the residual connections within Transformers can be viewed as Euler discretizations of ODEs. This equivalence suggests that we can leverage insights from numerical ODE methods to inform model design. Another promising avenue aims to develop a theoretical understanding of the self-attention mechanism, which distinguishes Transformers from earlier deep learning models. Studies have interpreted self-attention through various theoretical lenses, including data compression [232], optimization [233], and function approximation [234]. Moreover, Transformers can be mapped to formal systems, such as Turing machines [235], counter machines [236], regular and context-free languages [237], Boolean circuits [238], [239], programming languages [240], and first-order logic [241]. These connections provide theoretical tools to rigorously study the expressivity of Transformers. It is worth noting, however, that while we can analyze Transformers through these diverse perspectives, no unified theory yet explains the fundamental nature of these models. Addressing this theoretical gap remains a central challenge in the field, yet it is essential for developing complex neural systems with explainable and predictable behaviors.

Acknowledgements↩︎

We would like to thank Yongyu Mu, Chenglong Wang, Bei Li, Weiqiao Shan, Yuchun Fan, Kaiyan Chang, Tong Zheng, and Huiwen Bao for their suggestions on improving the early version of this work.

8 Sinusoidal Positional Encoding↩︎

In Transformers, positions are represented as vectors. Although vectorizing the representations of positions sounds complicated, a simple idea is to use a carrying system which describes how a natural number is expressed by a polynomial with respect to a base [242]. For example, \(i\) can be written as \[\begin{align} i & = & \sum_{k=0}^{k_{\mathrm{max}}} a(i,k) b^k \label{eq:pos-sin-encoding-carry-math} \end{align}\tag{70}\]

where \(a(i,k)\) is the \(k\)-th digit, \(k_{\mathrm{max}} + 1\) is the maximum number of digits, and \(b\) is the base of the system. The carrying occurs when \(a(i,k)\) reaches \(b\): we increase \(a(i,k+1)\) by 1 and roll back \(a(i,k)\) to 0. In this way we can change \(a(i,k)\) with a period of \(b^k\), that is, \(a(i,0)\) changes with a period of \(b^0\), \(a(i,1)\) changes with a period of \(b^1\), \(a(i,2)\) changes with a period of \(b^2\), and so on.

Using this system, \(i\) can be represented as a vector \[\begin{align} \mathrm{PE}(i) & = & \begin{bmatrix} a(i,0) & a(i,1) & ... & a(i,k_{\mathrm{max}})\end{bmatrix} \label{eq:pos-sin-encoding-carrying-vector} \end{align}\tag{71}\]

For example, when \(b=2\), \(\mathrm{PE}(11) = \begin{bmatrix} 1 & 1 & 0 & 1\end{bmatrix}\). However, in Eq. (71 ), \(\mathrm{PE}(i)\) is still a discrete function. We may want a continuous vector representation that can describe intermediate states between discrete events. Considering \(a(i,k)\) as a periodic function, a common choice is the sine function. Thus \(a(i,k)\) can be re-defined, as follows \[\begin{align} a(i,k) & = & \mathrm{sin}(i \cdot \omega_k) \label{eq:eq:pos-sin-encoding-sin} \end{align}\tag{72}\]

This function has an amplitude of 1 and a period of \(\frac{2 \pi}{\omega_k}\). Using an analogous form of periods to that used in Eq. (70 ), we define \(\omega_k\) as \[\begin{align} \omega_k & = & \frac{1}{(b_{\textrm{model}})^{k/d_{\textrm{model}}}} \end{align}\]

where \(b_{\textrm{model}} > 0\) and \(d_{\textrm{model}}>0\) are hyper-parameters of the model. Obviously, we have \(\frac{2 \pi}{\omega_0} < \frac{2 \pi}{\omega_1} < ... < \frac{2 \pi}{\omega_{k_{\mathrm{max}}}}\).

Similarly, we can define \(a(i,k)\) via the cosine function \[\begin{align} a(i,k) & = & \mathrm{cos}(i \cdot \omega_k) \label{eq:eq:pos-sin-encoding-cos} \end{align}\tag{73}\]

Taking both Eqs. (72 ) and (73 ), we create a new representation of \(i\), as follows \[\begin{align} \mathrm{PE}(i) & = & \begin{bmatrix} \mathrm{sin}(i \cdot \omega_0) & \mathrm{cos}(i \cdot \omega_0) & ... & \mathrm{sin}(i \cdot \omega_{k_{\mathrm{max}}}) & \mathrm{cos}(i \cdot \omega_{k_{\mathrm{max}}}) \end{bmatrix} \end{align}\]

[1] instantiated the above form by setting \(b_{\textrm{model}} = 10,000\). Let \(\mathrm{PE}(i,k)\) be the \(k\)-th dimension of \(\mathrm{PE}(i)\). [1]’s version of positional encoding is written as \[\begin{align} \mathrm{PE}(i,2k) & = & \mathrm{sin}(i \cdot \frac{1}{10000^{2k/d_{\mathrm{model}}}}) \tag{74} \\ \mathrm{PE}(i,2k+1) & = & \mathrm{cos}(i \cdot \frac{1}{10000^{2k/d_{\mathrm{model}}}}) \tag{75} \end{align}\]

Choosing \(b_{\textrm{model}} = 10,000\) is empirical. One can adjust it for specific tasks. Figure 25 plots the positional encoding for different positions. We see that, when \(k\) becomes larger, the change of the color follows a larger period.

Note that Eqs. (74 ) and (75 ) have a useful property that \(\mathrm{PE}(i+\mu)\) can be easily expressed by a linear function of \(\mathrm{PE}(i)\) for a given offset \(\mu\)²⁸ \[\begin{align} \mathrm{PE}(i+\mu,2k) & = & \mathrm{PE}(i,2k) \cdot \mathrm{PE}(\mu,2k+1) + \nonumber \\ & & \mathrm{PE}(i,2k+1) \cdot \mathrm{PE}(\mu,2k) \\ \mathrm{PE}(i+\mu,2k+1) & = & \mathrm{PE}(i,2k+1) \cdot \mathrm{PE}(\mu,2k+1) - \nonumber \\ & & \mathrm{PE}(i,2k) \cdot \mathrm{PE}(\mu,2k) \end{align}\]

The resulting benefit is that the encoding can somewhat model relative positions. That is, the state at position \(i + \mu\) can be described by starting with \(i\) and then appending it with the offset \(\mu\).

When applying the sinusoidal positional encoding, one way is to concatenate \(\mathbf{x}_i\) and \(\mathrm{PE}(i)\). In [1]’s work, they instead assume \(\mathrm{PE}(i)\) to be a vector of the same size as \(\mathbf{x}_i\) (i.e., \(|\mathrm{PE}(i)|= |\mathbf{x}_i| = d_e\)), and add \(\mathrm{PE}(i)\) to \(\mathbf{x}_i\), like this \[\begin{align} \mathbf{xp}_i & = & \mathbf{x}_i + \mathrm{PE}(i) \end{align}\]

This sinusoidal addictive model has been the basis of many positional encoding approaches [79], [223], [243].