$\mathbf{M}$ is a masking matrix that is used to prevent the model from seeing the right context words at each position, that is, for a position $i$, $M(i,j) = 0$ for $j \le i$, and $M(i,j) = -\infty$ otherwise. Both the time and space complexities of the self-attention sub-layer are quadratic functions of $n$[More precisely, the amount of memory used by the self-attention function is $n^2 + n \cdot d$, and so it will be dominated by the quadratic term $n^2$ if $n >> d$.]. Therefore, if $n$ is large, the model would be computationally expensive.

The usual implementation of the above model depends on dense matrix computation, for example, the dense matrix multiplications in Eqs. (<ref>-<ref>). One approach to reducing the amount of memory and the number of floating-point calculations in a dense computation system is to sparsify the problem. To do this, we assume that $\mathbf{A}$ is a sparse matrix, for example, only $\varrho \cdot n^2$ entries of $\mathbf{M}$ have non-zero values, where $\varrho$ indicates how sparse the matrix is, also called sparsity ratio. Since we only need to store these non-zero entries, the memory requirement of $\mathbf{A}$ can be reduced by using sparse matrix representations. Another advantage of using a sparse attention matrix is that the models of $\frac{\mathbf{Q} \mathbf{K}^{\mathrm{T}}}{\sqrt{d}}$ and $\mathbf{A} \mathbf{V}$ can be simplified, as we consider only a “small” number of related positions when learning a representation.

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{eqnarray}
A'(i,j) & = &
\begin{cases}
\textrm{some weight} & j \in \pi_i\ \textrm{and}\ j \le i \\
0 & \textrm{otherwise}
\end{cases}
\end{eqnarray}

A simple implementation of this model involves a slight modification to $\mathbf{M}$, leading to a new masking variable $\mathbf{M}'$

\begin{eqnarray}
M'(i,j) & = &
\begin{cases}
0 & j \in \pi_i\ \textrm{and}\ j \le i \\
-\infty & \textrm{otherwise}
\end{cases}
\end{eqnarray}

In practical implementation, a more efficient approach is to employ sparse operations for $\mathbf{Q} \mathbf{K}^{\mathrm{T}}$ and $\mathbf{A}' \mathbf{V}$ by considering $\mathbf{M}'$ and $\mathbf{A}'$, respectively. That is, we save on computation for pairs of positions whose attention weights are non-zero, and skip the rest.

There are several approaches that we can take to the sparse modeling of self-attention. We describe briefly some of them as follows

Span-based Attention/Local Attention. As discussed in Section <ref>, the use of context in sequence modeling is local in many cases. The basic idea of local attention is to span the attention weights to a restricted region of the input sequence. We can then write $\pi_i$ as

\begin{eqnarray}
\pi_i & = & [a_i^l,a_i^r]
\end{eqnarray}

where $a_i^l$ and $a_i^r$ and the left and right ends of $\pi_i$. $a_i^r - a_i^l + 1$ determines how small the region is, and so we can use it to control the sparsity of the attention model, for example, if $a_i^r - a_i^l + 1 << 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 [Luong et al., 2015, Sperber et al., 2018, Yang et al., 2018, Sukhbaatar et al., 2019]. See Figure <ref> (b) for an illustration of local attention.

Chunked Attention. When a problem is too difficult to solve, one can transform it into easier problems and solve each of them separately, as is often the case in practice. This motivates the chunked attention approach in which we segment a sequence into chunks and run the attention model on each of them [Parmar et al., 2018, Qiu et al., 2020]. Given a sequence $\{1,...,n\}$, we define $\{\mathrm{chunk}_1,...,\mathrm{chunk}_q\}$ to be a segmentation of the sequence. A chunk can be expressed as a span

\begin{eqnarray}
\mathrm{chunk}_k & = & [c_k^l,c_k^r]
\end{eqnarray}

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 some sort of local attention model. Figure <ref> (c) shows an illustration of this model. There remains the issue of how to segment the sequence. There are several ways to do this. For example, as discussed in Section <ref>, 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 [Child et al., 2019, Beltagy et al., 2020, Ainslie et al., 2020]. 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 $stride$. As a special case, if $stride$ equals the size of the context window, this model is the same as the chunked attention model mentioned above. If $stride$ chooses a value smaller than the size of the context window, the attention model will become denser. Figure <ref> (d) shows the case of $strdie=1$ where the chunk overlapping is maximized. A way to achieve relatively sparser attention is to use a dilated context window. Figure <ref> (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 sub-set of $\{1,...,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 [Zhou et al., 2021]. A second approach is grouping: positions are grouped, and then the attention weights are computed only for positions in the same group. It is often relatively easy to achieve this by running clustering algorithms on 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 [Roy et al., 2021]. One benefit of learning attention fields is that the model can spread its attention broader 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 <ref> (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 [Kitaev et al., 2020, Tay et al., 2020]. 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 Zaheer et al., 2020's system, the attention map is generated by considering three different sparse models, including local attention (chunked attention), global attention, and random attention[Here the global attention model attends each word only to a special word which accounts for the entire sequence and is often placed at the beginning of the sequence. The random attention model attends each word to a random set of the words of the sequence.]. 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 [Child et al., 2019, Beltagy et al., 2020]. For example, one can use one head as a local attention model, and use another head as a global attention model (see Figure <ref> (g-h)).

Illustration of the attention maps of different models (self-attention on the decoder side). Dark cells mean $A'(i,j) \ne 0$ (i.e., $i$ attends to $j$), and light cells mean $A'(i,j) = 0$ (i.e., $i$ does not attend to $j$). In all these attention maps, we assume that every position attends to itself by default (see diagonals).

One disadvantage of sparse models compared to dense models is that they are not computationally efficient on GPUs or CPUs. While sparse models can ideally reduce both the memory and computation requirements, the actual rate at which work can be done by sparse models is much slower than by dense models. In practice, it is difficult for sparse models to approach the peak FLOPS of a GPU or CPU[FLOPS = floating point operations per second.]. Therefore, they are often used for the purpose of high memory efficiency, not really for the purpose of efficient computation. On the other hand, sparse models are still of great use to NLP practitioners in the context of memory-efficient Transformer, especially when Transformer systems are used to deal with extremely long sequences.

§.§ Recurrent and Memory Models

For sequence generation problems, Transformer can also be thought of as a memory system. Consider again the general setting, in which we are given the states of previous $i-1$ positions, and we wish to predict the next state. In self-attention, this is done 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}),...,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}$). Then, we move to position $i+1$, and add $(\mathbf{k}_{i},\mathbf{v}_{i})$ to the collection of key-value pairs. This procedure can be interpreted in terms of memory mechanisms. The Transformer model maintains a memory that retains the information of the past. When moving along the sequence, we repeat the same operation, each time generating some output by reading the memory, and then updating the memory so that new information could be stored in some way. This is illustrated in Figure <ref>.

Transformer as a memory system. At position $i$, the collection of the key-value pairs of positions $\{1,...,i-1\}$ is used as a memory of the past information. The Transformer model accesses this memory to generate some output, and then adds the key-value pair of position $i$ to the memory. Moving to the next position, we repeat the same procedure of memory access and update.

§.§.§ 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 a longer sub-sequence is observed. Clearly, such a variable-length memory will generally be infeasible if the model deals with a very, very long sequence. For the modeling problem of arbitrary length sequences, it is common to use a fix-length memory instead. As in many NLP problems,  one of the simplest ways to do this is to consider a cache saving 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 to the current position, so that its closest successors can be considered at each step. This means that, for each position, a self-attention sub-layer attends to $n_c - 1$ positions ahead, like this

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 take 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 way to train this context window model is by chunked attention. We divide the sequence into chunks (or sub-sequences) which are of the same length $n_c$. Then, we treat these chunks as individual training samples, and run the training program on each of them as usual. 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 [Dai et al., 2019]. In the simplest case, consider an example in which $\mathrm{chunk}_k$ can see its successor $\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. On the other hand, the difference between the standard strided attention model and the Transformer-XL model is that in Transformer-XL, we perform 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 on $\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 <ref> for an illustration.

Illustration of chunk-wise training [Dai et al., 2019]. The input sequence is divided into chunks of the same length $n_c$. Training is performed on these chunks, each time dealing with a chunk. In $\mathrm{chunk}_k$, the attention field for every position in this chunk is a left context window of size $n_c$. Hence this model allows for attention across chunks, for example, position $i - 2$ in $\mathrm{chunk}_k$ can attend to positions $i - 3$ and $i - 4$ in $\mathrm{chunk}_{k-1}$ (see sub-figure (a)). For training, errors are back-propagated only in the sub-network for $\mathrm{chunk}_k$, leaving other parts of the model unchanged. Here we use dashed lines to denote information flow that we consider in the forward pass but not in the backward pass. Once we finish the training on $\mathrm{chunk}_k$, we move to the next chunk, and repeat the same training procedure.

The above model is similar in spirit to recurrent models because all 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 between two different layers, 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}$.

§.§.§ 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 the key-value vectors during left-to-right generation, we construct the memory of the entire “history” as a fixed number of encoded key-value vectors. These encoded key-value vectors can be either a small sub-set of $\{(\mathbf{k}_{1},\mathbf{v}_{1}),...,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}$ or a small set of newly-generated vectors that encodes $\{(\mathbf{k}_{1},\mathbf{v}_{1}),...,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}$.

One way to do the encoding is to apply a pooling operation to $\{(\mathbf{k}_{1},\mathbf{v}_{1}),...,(\mathbf{k}_{i-1},\mathbf{v}_{i-1})\}$ [Rae et al., 2019]. For example, by using average pooling, the memory contains only one key-value pair $(\bar{\mathbf{k}},\bar{\mathbf{v}})$

\begin{eqnarray}
\bar{\mathbf{k}} & = & \frac{1}{i-1} \sum_{j=1}^{i-1} \mathbf{k}_{j} \label{eq:avearge-attention-k} \\
\bar{\mathbf{v}} & = & \frac{1}{i-1} \sum_{j=1}^{i-1} \mathbf{v}_{j} \label{eq:avearge-attention-v}
\end{eqnarray}

This leads to a very efficient model, and we only need to update the vectors $(\bar{\mathbf{k}},\bar{\mathbf{v}})$ at a time [Zhang et al., 2018]. Let $(\bar{\mathbf{k}}[i],\bar{\mathbf{v}}[i])$ be the state of the memory at position $i$. A more general definition of $(\bar{\mathbf{k}}[i],\bar{\mathbf{v}}[i])$ is given in a recursive form

\begin{eqnarray}
\bar{\mathbf{k}}[i] & = & \mathrm{KMem}(\bar{\mathbf{k}}[i-1],\mathbf{k}_{i-1}) \label{eq:memory-attention-k} \\
\bar{\mathbf{v}}[i] & = & \mathrm{VMem}(\bar{\mathbf{v}}[i-1],\mathbf{v}_{i-1}) \label{eq:memory-attention-v}
\end{eqnarray}

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 forms of the functions like $\mathrm{KMem}(\cdot)$ and $\mathrm{VMem}(\cdot)$ in common use. For example, if $\mathrm{KMem}(\cdot)$ and $\mathrm{VMem}(\cdot)$ are weighted sum functions, we can derive the same forms as Eqs. (<ref>) and (<ref>). If $\mathrm{KMem}(\cdot)$ and $\mathrm{VMem}(\cdot)$ are recurrent cells in RNNs or LSTM, we obtain a recurrent model of memory.

Extension of the above model to memories having more than one key-value pair is straightforward. One approach is to use the memory to represent sub-sequences. Let $\{(\bar{\mathbf{k}}_{1},\bar{\mathbf{v}}_{1}),...,(\bar{\mathbf{k}}_{\kappa},\bar{\mathbf{v}}_{\kappa})\}$ be a memory of size $\kappa$. Each $(\bar{\mathbf{k}}_{j},\bar{\mathbf{v}}_{j})$ is a snapshot of a chunk of length $n_c$. Thus, this memory can encode a sequence with maximum length $\kappa \cdot n_c$. Then, we can compute $(\bar{\mathbf{k}}_{j},\bar{\mathbf{v}}_{j})$ on the corresponding chunk using Eqs. (<ref>) and (<ref>). A second approach is to organize $\{(\bar{\mathbf{k}}_{1},\bar{\mathbf{v}}_{1}),...,(\bar{\mathbf{k}}_{\kappa},\bar{\mathbf{v}}_{\kappa})\}$ into a priority queue. We design some function to assign a score to any given key-value pair. The key-value pair can be inserted into the priority queue through the push operation. Ideally, we wish to develop a scoring function to estimate the value of a key-value pair, for example, we use another neural network to evaluate the key-value pair. In this way, the memory is a collection of the most valuable key-value pairs over the input sequence.

Although representing the memory as a set of vectors is an obvious choice for the model design in Transformer, the memory is discrete and its capacity is determined by the number of the vectors. An alternative form of memory is continuous memory. This type of model typically builds on the idea of function approximation, in which $\{\mathbf{k}_{1},...,\mathbf{k}_{i-1}\}$ or $\{\mathbf{v}_{1},...,\mathbf{v}_{i-1}\}$ is viewed as a series of data points, and a continuous function is developed to fit these data points. Then, we no longer need to store $\{\mathbf{k}_{1},...,\mathbf{k}_{i-1}\}$ and $\{\mathbf{v}_{1},...,\mathbf{v}_{i-1}\}$. Instead, the memory is represented by the functions fitting these vectors. A simple method is to combine simple functions to fit complex curves of data points. For example, we can develop a set of basis functions and use a linear combination of them to approximate the key or value vectors [Martins et al., 2022]. The resulting model is parameterized by these basis functions and the corresponding weights in the combination.

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 dependency. This idea is also similar to that used in combining different sparse attention models as discussed in the previous sub-section.

§.§.§ Retrieval-based Methods

So far in this sub-section, we have discussed approaches based on fixed-length models. It is also possible to develop efficient memory models by improving the efficiency of accessing the memories, instead of just reducing the memory capacities. 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 ones 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 similarity search based on the dot-product similarity measure between query and key vectors. Then, we attend $\mathbf{q}$ 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. Building a memory as a retrieval system can fall under the general framework called the retrieval-augmented approach. It provides a simple way to incorporate external memories into neural models like Transformer [Guu et al., 2020, Lewis et al., 2020, Wu et al., 2021].

§.§ Low-dimensional Models

In many practical applications, Transformer models 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 <ref>, this high dimensionality arises in part from the steps of computing the attention matrix as in Eq. (<ref>) (for ease of presentation, we repeat the equation here)

\begin{eqnarray}
\mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{A} \mathbf{V} \label{eq:low-dim-self-attetention-decoder-as}
\end{eqnarray}

and the weighted sum of value vectors as in Eq. (<ref>)

\begin{eqnarray}
\mathbf{A} & = & \mathrm{Softmax}(\frac{\mathbf{Q} \mathbf{K}^{\mathrm{T}}}{\sqrt{d}} + \mathbf{M}) \label{eq:low-dim-self-attetention-decoder-softmax}
\end{eqnarray}

which involves large matrix multiplications $\mathbf{Q} \mathbf{K}^{\mathrm{T}}$ and $\mathbf{A} \mathbf{V}$ when the length $n$ and the hidden dimensionality $d$ have large values.

The $\mathbf{A} \mathbf{V}$ and $\mathbf{Q} \mathbf{K}^{\mathrm{T}}$ operations have a time complexity of $O(n^2 \cdot d)$ and a space complexity of $O(n^2 + n \cdot d)$. Several previously described approaches have reduced this complexity by using sparse models. In this sub-section, we focus on methods that approximate these operations via dense computation. One simple idea is to transform $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ into smaller matrices, and thus to reduce the computational burden of matrix multiplication. Since $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$ are all in $\mathbb{R}^{n \times d}$, we can achieve this by reducing either the $n$ dimension or the $d$ dimension, or both.

§.§.§ 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 a number less than $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{eqnarray}
\mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{A} {\color{red} \mathbf{V}^{'}} \label{eq:low-dim-self-attetention-decoder-smaller-matrix-as} \\
\mathbf{A} & = & \mathrm{Softmax}(\frac{\mathbf{Q} [{\color{red} \mathbf{K'}}]^{\mathrm{T}}}{\sqrt{d}} + \mathbf{M}) \label{eq:low-dim-self-attetention-decoder-smaller-matrix-softmax}
\end{eqnarray}

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' << 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 cheap 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.

The above method is straightforward but still requires sparse operations, such as sampling and collection. As an alternative, we can use dense computation to transform $\mathbf{K}$ and $\mathbf{V}$ to $\mathbf{K'}$ and $\mathbf{V}^{'}$. A typical choice is to use CNNs [Liu et al., 2018]. Let $\mathrm{Conv}(\cdot)$ be a function describing a set of filters that slide along the $n$ dimension. $\mathbf{K'}$ is then given by

\begin{eqnarray}
\mathbf{K'} & = & \mathrm{Conv}(\mathbf{K},\mathbf{W}_{c},size_r,stride)
\end{eqnarray}

where $\mathbf{W}_{c}$ is the parameter matrix of the filters, $size_r$ is the size of the receptive field, and $stride$ is the number of units the filters are translated at a time. In general, we can achieve a high compression rate by choosing large values for $size_r$ and $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 sub-section. 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. For many problems, it is found that $\mathbf{A}$ (or more precisely $\mathbf{Q} \mathbf{K}^{\mathrm{T}}$) is 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 to directly transform $\mathbf{K}$ and $\mathbf{V}$ into smaller-sized matrices via linear mappings, given by

\begin{eqnarray}
\mathbf{K'} & = & \mathbf{U}^{k} \mathbf{K} \\
\mathbf{V}^{'} & = & \mathbf{U}^{v} \mathbf{V}
\end{eqnarray}

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. (<ref>) and (<ref>). While such a method is intuitive and simple, it is proven to obtain a sufficiently small approximation error $\epsilon$ if $n'$ is a linear function of $d/\epsilon^2$ [Wang et al., 2020].

§.§.§ Reducing $d$

Another approach to working in a low-dimensional space is to reduce the $d$ dimension. 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 of any key-value pair 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}$ by new representations $\mathbf{Q'} \in \mathbb{R}^{n \times d'}$ and $\mathbf{K'} \in \mathbb{R}^{n \times d'}$. We can easily modify Eq. (<ref>) to use $\mathbf{Q'}$ and $\mathbf{K'}$ in computing the attention matrix

\begin{eqnarray}
\mathbf{A} & = & \mathrm{Softmax}(\frac{{\color{red} \mathbf{Q'}} [{\color{red} \mathbf{K'}}]^{\mathrm{T}}}{\sqrt{d}} + \mathbf{M})
\end{eqnarray}

$\mathbf{Q'}$ and $\mathbf{K'}$ are given by

\begin{eqnarray}
\mathbf{Q'} & = & \mathbf{Q} \mathbf{U}^{q} \\
\mathbf{K'} & = & \mathbf{K} \mathbf{U}^{k}
\end{eqnarray}

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 all data points (represented as vectors) from one space to another space, so that the problem, which might be difficult to solve in the original space, is easier to solve in the new space. The “trick” of kernel methods is that we actually do not need to know the mapping function, but only need to know how to compute the inner product of vectors in the new space in one operation[In mathematical analysis, the inner product is a generalized notion of the dot-product. It is typically denoted by $\langle \cdot,\cdot \rangle$. A formal definition of the inner product requires that $\langle \cdot,\cdot \rangle$ satisfies several properties in a vector space. Although the inner product has different forms in different contexts, in the Euclidean space $\mathbb{R}^{d}$, it is the same thing as the dot-product, that is, given two vectors $\mathbf{a} \in \mathbb{R}^{d}$ and $\mathbf{b} \in \mathbb{R}^{d}$, we have

\begin{eqnarray}
\langle \mathbf{a}, \mathbf{b}\rangle & = & \mathbf{a} \cdot \mathbf{b} \nonumber \\
& = & \sum_{i=1}^{d} a_i \cdot b_i
\end{eqnarray}

]. This operation of the inner product is usually called the kernel and denoted by $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. (<ref>) $\mathbf{A}$ is a fraction denoting the normalized attention weights. The numerator can be written in the form

\begin{eqnarray}
\widetilde{\mathbf{A}} & = & \mathrm{Mask}(\exp(\frac{\mathbf{Q} \mathbf{K}^{\mathrm{T}}}{\sqrt{d}})) \\
\end{eqnarray}

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{eqnarray}
\mathbf{A} & = & \mathbf{D}^{-1} \widetilde{\mathbf{A}}
\end{eqnarray}

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. (<ref>), we have

\begin{eqnarray}
\mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{D}^{-1} \widetilde{\mathbf{A}} \mathbf{V} \label{eq:attention-das-form}
\end{eqnarray}

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{eqnarray}
\widetilde{A}(i,j) & = & K(\mathbf{q}_i,\mathbf{k}_j) \nonumber \\
& = & \langle \phi(\mathbf{q}_i), \phi(\mathbf{k}_j) \rangle \nonumber
\end{eqnarray}

$\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{eqnarray}
\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{eqnarray}

Then, we develop a kernelized attention model by approximating the attention weight $\alpha_{i,j}$ in the form

\begin{eqnarray}
\alpha_{i,j} & \approx & \frac{\phi(\mathbf{q}_i) \phi(\mathbf{k}_j)^{\mathrm{T}}}{\sum_{j'=1}^{n}\phi(\mathbf{q}_i) \phi(\mathbf{k}_{j'})^{\mathrm{T}}}
\end{eqnarray}

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{eqnarray}
\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)^{\mathrm{T}}}{\sum_{j'=1}^{n}\phi(\mathbf{q}_i) \phi(\mathbf{k}_{j'})^{\mathrm{T}}} \cdot \mathbf{v}_j \Big) \nonumber \\
& = & \frac{\sum_{j=1}^{n} \phi(\mathbf{q}_i) \phi(\mathbf{k}_j)^{\mathrm{T}} \mathbf{v}_j}{\sum_{j'=1}^{n}\phi(\mathbf{q}_i) \phi(\mathbf{k}_{j'})^{\mathrm{T}}} \nonumber \\
& = & \frac{\phi(\mathbf{q}_i) {\color{red} (} \sum_{j=1}^{n} \phi(\mathbf{k}_j)^{\mathrm{T}} \mathbf{v}_j {\color{red} )}}{\phi(\mathbf{q}_i) {\color{red} (} \sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^{\mathrm{T}} {\color{red} )}} \label{eq:context-vector-kernel-methods}
\end{eqnarray}

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)^{\mathrm{T}} \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. (<ref>), we define the $i$-th entry of $\mathbf{D}$ to be $\phi(\mathbf{q}_i) \sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^{\mathrm{T}}$. Then, the attention model can be re-expressed in the form

\begin{eqnarray}
\mathrm{Att}_{\mathrm{self}}(\mathbf{S}) & = & \mathbf{D}^{-1} \phi(\mathbf{Q}) \phi(\mathbf{K})^{\mathrm{T}} \mathbf{V} \nonumber \\
& = & \mathbf{D}^{-1} \mathbf{Q'} \mathbf{K'}^{\mathrm{T}} \mathbf{V} \nonumber \\
& = & \mathbf{D}^{-1} {\color{red} \big(} \mathbf{Q'} {\color{red}(}\mathbf{K'}^{\mathrm{T}} \mathbf{V} {\color{red})} {\color{red}\big)}
\end{eqnarray}

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'}^{\mathrm{T}} \mathbf{V}$ (i.e., $\sum_{j=1}^{n} \phi(\mathbf{k}_j)^{\mathrm{T}} \mathbf{v}_j$) and the corresponding normalization factor (i.e., $\sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^{\mathrm{T}}$) once. The results can then be used for any query [Katharopoulos et al., 2020]. The memory needs to maintain $\sum_{j=1}^{n} \phi(\mathbf{k}_j)^{\mathrm{T}} \mathbf{v}_j$ and $\sum_{j'=1}^{n} \phi(\mathbf{k}_{j'})^{\mathrm{T}}$ and update them when new key and value vectors come.

Still, there are several problems 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 Choromanski et al., 2020'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},...,\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 Lample et al., 2019'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}),...,(\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{eqnarray}
\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'}}^{\mathrm{T}}}{\sqrt{d'}} + \mathbf{M})
\end{eqnarray}

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')$.

§.§ Parameter and Activation Sharing

Redundancy is common to most large-scale neural networks. As a result, many of these models are over-parameterized, making the training and inference less efficient. One common approach to redundancy reduction is to simplify the modeling by removing useless components of the models, for example, we can either prune a complex model or share sub-models among different components of it to obtain a reasonably small model. In this sub-section, we discuss methods of parameter and intermediate state sharing in Transformer models. We leave the discussion of model transfer and pruning to Section <ref>.

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 “big” 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, which might be not related to architecture design, is shared embedding. 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 multi-lingual sequence models, such as language models that are able to deal with texts 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{eqnarray}
\mathbf{S}^{l} & = & \mathrm{Layer}(\mathbf{S}^{l-1};\theta^{l})
\end{eqnarray}

We can tie the parameters for some or all of these layers. For example, given a set of layers $\{l_1, l_2,...,l_n\}$, we enforce the constraint $\theta^{l_1} = \theta^{l_2} = ... = \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 [Dehghani et al., 2018]. For example, sharing a single FFN sub-layer across the Transformer encoder is found to be effective in reducing the redundancy in machine translation systems [Pires et al., 2023].

For Transformers, sharing can also be performed in multi-head attention. An example of this is multi-query attention [Shazeer, 2019]. Recall from Section <ref> that the output of a head $h$ in standard multi-head self-attention can be written as

\begin{eqnarray}
\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{eqnarray}

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}_{v}^{v}$ are the query, key, and value, 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 key and value across all the heads, but use different queries for different heads. The form of this model is given by

\begin{eqnarray}
\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{eqnarray}

Here the key $\mathbf{S} \mathbf{W}_{0}^{k}$ and value $\mathbf{S} \mathbf{W}_{0}^{v}$ are irrelevant to $h$. Hence we need only compute them once rather than computing them several times. As a result, we can make a significant saving in computational cost, especially if the number of heads is large. Multi-query attention has been successfully incorporated into recent large language models, such as Llama 2 [Touvron et al., 2023] and Falcon[https://falconllm.tii.ae/index.html].

By extending the idea of sharing to more general situations, any intermediate states can be shared across a neural network. For example, reusing neuron activations allows a sub-model to be applied multiple times. For Transformers, sharing can be considered inside the process of self-attention. It is found that the attention maps of different layers are similar in some NLP tasks [Xiao et al., 2019]. Therefore, it is reasonable to compute the attention map only once and then use it in the following layers.

If we make a further generalization of the sharing mechanism, we can view it as a process by which we use the result produced previously rather than computing it 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 [Gomez et al., 2017]. 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 treatment is that the information produced in the forward pass is shared implicitly, and the model is memory-efficient [Kitaev et al., 2020].

§.§ Alternatives to Self-Attention

We have seen that the use of self-attention is a primary source of the large computation and memory requirements for Transformer systems. It is natural to wonder if there are efficient alternatives to self-attention models. Here we present briefly some of the Transformer variants in which self-attention sub-layers are not required and we instead replace them with other types of neural networks.

§.§.§ CNN as A Replacement of Self-Attention

CNNs are simple and widely used neural networks, and are considered as 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 filter of CNNs 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 the number of elementary operations required to run CNNs is a linear function of the sequence length $n$, compared with the quadratic function for self-attention networks. In practical systems, there have been many highly-optimized implementations for CNNs, making it easier to apply them to sequence modeling. For further improvements to memory efficiency, we can use lightweight CNN variants, for example, depth-wise CNNs [Wu et al., 2018].
In CNNs, a filter (or a set of filters) combines the input variables in the receptive field into an output variable (or a set of output variables) via linear mapping. Suppose that the input and output of a problem are represented as sequences of feature vectors. Given a filter having a $d \times k$ receptive field, we slide it along the sequence. At each step, the filter takes $d \times k$ input features and produces an output feature. This procedure is typically expressed by

\begin{eqnarray}
y & = & \mathrm{ReduceSum}(\mathbf{x} \odot \mathbf{W})
\end{eqnarray}

where $\mathbf{x} \in \mathbb{R}^{k \times d}$ is the vector representation of the input, $y \in \mathbb{R}$ is the output feature, and $\mathbf{W} \in \mathbb{R}^{k \times d}$ is the weight matrix. The function $\mathrm{ReduceSum}(\cdot)$ computes the sum of all element-wise products between $\mathbf{x}$ and $\mathbf{W}$. If we want the input and output to have the same number of features, we can design $d$ filters and the number of parameters will be $d^2 \cdot k$.

In depth-wise CNNs, we tie the weights across different feature dimensions. More precisely, all the column vectors of $\mathbf{W}$ are the same. Thus, the number of the unique parameters of the model is reduced to $d \cdot k$ (each $\mathbf{W}$ corresponding to a filter having $k$ unique parameters).

§.§.§ 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{eqnarray}
\mathbf{z}_{i} & = & f(a \cdot \mathbf{z}_{i-1} + b \cdot \mathbf{s}_{i}) \label{eq:linear-base-model-for-transformer}
\end{eqnarray}

Here $\mathbf{s}_i$ represents some intermediate states of the model at step $i$, and $\mathbf{z}_{i}$ represents the summary of the history 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. (<ref>) 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. (<ref>)  to a standard RNN model by simply making 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 RNN and its 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 [Chen et al., 2018].

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{eqnarray}
\mathrm{Att}_{\mathrm{self}} & = & \mathbf{A} \cdot \mathbf{V} \nonumber \\
& = & \psi(\mathbf{Q} \cdot \mathbf{K}^{\mathrm{T}}) \cdot \mathbf{V}
\end{eqnarray}

Here $\psi(\cdot)$ is a function that is composed by taking the scaling, exponentiating, 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 obvious equivalent that simplifies the computation, and we have 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 form of this model might be that allows for a direct interaction between the keys and queries, so that we can encode the context information in a way that is irrelevant to the queries. A trick here is that we can remove the non-linearity from $\psi(\cdot)$ by using a feature space mapping $\phi(\cdot)$ on the queries and keys, and reformulate $\psi(\mathbf{Q} \cdot \mathbf{K}^{\mathrm{T}})$ (i.e., $\mathbf{A}$) in a form of matrix products. For example, recall from Section <ref> that we can transform $\mathbf{Q}$ and $\mathbf{K}$ to $\mathbf{Q}' = \phi(\mathbf{Q}) \in \mathbb{R}^{n \times d'}$ and $\mathbf{K}' = \phi(\mathbf{K}) \in \mathbb{R}^{n \times d'}$ through the mapping $\phi(\cdot)$. Then, we define the form of the attention model to be

\begin{eqnarray}
\mathrm{Att}_{\mathrm{self}} & \equiv & \psi'(\mathbf{Q}' \cdot \mathbf{K'}^{\mathrm{T}}) \cdot \mathbf{V} \nonumber \\
& =      & \frac{\mathbf{Q}' \cdot \mathbf{K'}^{\mathrm{T}}}{\mathbf{D}} \cdot \mathbf{V} \nonumber \\
& =      & \frac{\mathbf{Q}' \cdot {\color{red} \big(} \mathbf{K'}^{\mathrm{T}} \cdot \mathbf{V} {\color{red} \big)}}{\mathbf{D}}
\end{eqnarray}

where $\psi'(\mathbf{a}) = \frac{\mathbf{a}}{\mathbf{D}}$. From this definition, we see that, in the case of transformed queries and keys, the query-key product needs not be normalized via Softmax, but needs only be normalized via a simple factor $\mathbf{D}$. Hence the model has a very simple form involving only matrix multiplication and division, allowing us to change the order of the operations using the associativity of matrix multiplication.

This leads to an interesting procedure: keys and values are first encoded via $\mathbf{K'}^{\mathrm{T}} \cdot \mathbf{V}$, and then each query attends to this encoding result. Given that $\mathbf{K'}^{\mathrm{T}} \cdot \mathbf{V} = \sum_{j=1}^{n} \mathbf{k'}_{j}^{\mathrm{T}} \cdot \mathbf{v}_{j}$, we can write $\mathbf{K'}^{\mathrm{T}} \cdot \mathbf{V}$ in the form of Eq. (<ref>), as follows

\begin{eqnarray}
\mu_j & = & \mu_{j-1} + \mathbf{k'}_{j}^{\mathrm{T}} \cdot \mathbf{v}_{j} \label{eq:linear-attention-mu}
\end{eqnarray}

Here $\mu_j \in \mathbb{R}^{d' \times d}$ is a variable that adds $\mathbf{k'}_{j}^{\mathrm{T}} \cdot \mathbf{v}_{j}$ at a time. Likewise, we can define another variable $\nu_j \in \mathbb{R}^{d'}$

\begin{eqnarray}
\nu_j & = & \nu_{j-1} + \mathbf{k'}_{j}^{\mathrm{T}} \label{eq:linear-attention-nu}
\end{eqnarray}

Then, the output of self-attention for the $j$-th query can be written as (see also Eq. (<ref>))

\begin{eqnarray}
\mathrm{Att}_{\mathrm{self},j} & = & \frac{\mathbf{q}'_{j} \cdot \mu_n}{\mathbf{q}'_{j} \cdot \nu_n}
\end{eqnarray}

Clearly, this is a linear model, because $\mu_n$ and $\nu_n$ are linear with respect to $n$. In simple implementations of this model, only $\mu_j$ and $\nu_j$ are kept. Each time a new query is encountered, we update $\mu_j$ and $\nu_j$ using Eqs. (<ref>) and (<ref>), and then compute $\mathrm{Att}_{\mathrm{self},j} = \frac{\mathbf{q}'_{j} \cdot \mu_j}{\mathbf{q}'_{j} \cdot \nu_j}$[In autoregressive generation, we generate a sequence from left to right. In this case, we need not consider the keys and values for positions $>j$.].

One straightforward extension to the linear attention model is to allow Eqs. (<ref>) and (<ref>) to combine different terms with different weights. For example, we can redefine $\mu_j$ and $\nu_j$ as

\begin{eqnarray}
\mu_j & = & a \cdot \mu_{j-1} + (1 - a) \cdot \mathbf{k'}_{j}^{\mathrm{T}} \cdot \mathbf{v}_{j} \\
\nu_j & = & a \cdot \nu_{j-1} + (1 - a) \cdot \mathbf{k'}_{j}^{\mathrm{T}}
\end{eqnarray}

and train the parameter $a$ as usual. Also, we can treat $a$ as a gate and use another neural network to compute $a$ [Peng et al., 2021]. Another model design is to add more terms to Eqs. (<ref>) and (<ref>) in order to give a more powerful treatment of the linear attention approach [Bello, 2020, Schlag et al., 2021].

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 taking linear forms of representations 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 [Peng et al., 2023, Sun et al., 2023]. While these systems have different architectures, the underlying models have a similar form, as described in Eq. (<ref>). Note that, by using the general formulation of recurrent models, we need not restrict the modeling to the 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 <ref>. From the memory viewpoint, 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 $\mu_j$ and $\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. (<ref>), and average a series of state vectors of a Transformer system, either weighted or unweighted.

§.§.§ 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 is described by first-order differential equations of these states. As a simple example, we consider a continuous time-invariant linear system which is given in the form of the state-space representation

\begin{eqnarray}
\frac{d \mathbf{z}(t)}{d t} & = & \mathbf{z}(t) \cdot \mathbf{A} + \mathbf{s}(t) \cdot \mathbf{B} \label{eq:ssm-state-equation} \\
\mathbf{o}(t) & = & \mathbf{z}(t) \cdot \mathbf{C} + \mathbf{s}(t) \cdot \mathbf{D} \label{eq:ssm-output-equation}
\end{eqnarray}

Here $\mathbf{s}(t)$, $\mathbf{o}(t)$, and $\mathbf{z}(t)$ are the values of the input variable, output variable and state variable at time $t$[We use boldface letters to emphasize that the variables are vectors.]. In a general setting, $\mathbf{s}(t)$, $\mathbf{o}(t)$, and $\mathbf{z}(t)$ may have different numbers of dimensions. To simplify the discussion here, we assume that $\mathbf{s}(t), \mathbf{o}(t) \in \mathbb{R}^{d}$ and $\mathbf{z}(t) \in \mathbb{R}^{d_z}$[In a general state-space model, all these variables are represented as vectors of complex numbers. Because the models defined on the field of complex numbers is applicable to case of real number-based state-spaces, we restrict our discussion to variables in the multi-dimensional real number field.]. Eq. (<ref>) 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. (<ref>) 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 variable $\mathbf{s}(t)$ to the variable $\mathbf{o}(t)$ over time. They are, therefore, often used to deal with continuous time series data. To apply this model to the sequence modeling problem discussed in this chapter, we need to modify the above equations to give a discrete form of the state-space representation. Suppose that $\{\mathbf{s}_0,\mathbf{s}_1,...,\mathbf{s}_n\}$ is a sequence of input data points sampled from $s(t)$ with time step $\Delta t$. Similarly, we define $\{\mathbf{z}_0,\mathbf{z}_1,...,\mathbf{z}_n\}$ and $\{\mathbf{o}_0,\mathbf{o}_1,...,\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{eqnarray}
\mathbf{z}_t & = & \mathbf{z}_{t-1} \cdot \overline{\mathbf{A}} + \mathbf{s}_t \cdot \overline{\mathbf{B}} \label{eq:ssm-state-equation-discrete} \\
\mathbf{o}_t & = & \mathbf{z}_t \cdot \overline{\mathbf{C}} + \mathbf{s}_t \cdot \overline{\mathbf{D}} \label{eq:ssm-output-equation-discrete}
\end{eqnarray}

This formulation of the SSM defines an RNN with a residual connection. To be more precise, Eq. (<ref>) describes a recurrent unit that reads the input at step $t$ and the state at step $t-1$, without using any activation function. Eq. (<ref>) describes an output layer that sums both the linear transformations of the state $\mathbf{z}_t$ and the identity 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. (<ref>) is approximated by Eq. (<ref>)[The discretization process can be interpreted as a numerical method of solving the differential equation. Note that Eq. (<ref>) is an ODE

\begin{eqnarray}
\frac{d z(t)}{d t} & = & g(z(t),t) \label{eq:ssm-ode}
\end{eqnarray}

\begin{eqnarray}
g(z(t),t) & = & z(t) \cdot \mathbf{A} + s(t) \cdot \mathbf{B} \label{eq:ssm-ode-g}
\end{eqnarray}

There are many numerical approximations to the solutions to the ODE. For example, the Euler method of solving the ODE can be expressed in the form (see in Section <ref>)

\begin{eqnarray}
\mathbf{z}_t & = & \mathbf{z}_{t - 1} + \Delta t \cdot g(z_{t-1},t) \label{eq:ssm-ode-euler-method}
\end{eqnarray}

Substituting Eq. (<ref>) into Eq. (<ref>) yields

\begin{eqnarray}
\mathbf{z}_t & = & \mathbf{z}_{t - 1} + \Delta t (\mathbf{z}_{t - 1} \cdot \mathbf{A} + \mathbf{s}_{t} \cdot \mathbf{B}) \nonumber \\
& = & \mathbf{z}_{t - 1} \cdot (\mathbf{I} + \Delta t \cdot \mathbf{A}) + \mathbf{s}_{t} \cdot (\Delta t \cdot \mathbf{B})
\end{eqnarray}

This gives one of the simplest forms of the discretized state equations [Gu et al., 2022], that is,

\begin{eqnarray}
\overline{\mathbf{A}} & = & \mathbf{I} + \Delta t \cdot \mathbf{A} \\
\overline{\mathbf{B}} & = & \Delta t \cdot \mathbf{B}
\end{eqnarray}

One approach to time discretization, called bilinear transform or Tustin's method, gives a model in which the parameters take the form

\begin{eqnarray}
\overline{\mathbf{A}} & = & (\mathbf{I} - \frac{\Delta t}{2} \cdot \mathbf{A}) \cdot (\mathbf{I} - \frac{\Delta t}{2} \cdot \mathbf{A})^{-1} \\
\overline{\mathbf{B}} & = & \Delta t \cdot \mathbf{B} \cdot (\mathbf{I} - \frac{\Delta t}{2} \cdot \mathbf{A})^{-1} \\
\overline{\mathbf{C}} & = & \mathbf{C} \\
\overline{\mathbf{D}} & = & \mathbf{D}
\end{eqnarray}

An alternative approach is to use the Zero-Order-Hold (ZOH) discretization which has the form

\begin{eqnarray}
\overline{\mathbf{A}} & = & \exp(\Delta t \cdot \mathbf{A}) \\
\overline{\mathbf{B}} & = & \Delta t \cdot \mathbf{B} \cdot (\exp(\Delta t \cdot \mathbf{A}) - \mathbf{I}) \cdot (\Delta t \cdot \mathbf{A})^{-1} \\
\overline{\mathbf{C}} & = & \mathbf{C} \\
\overline{\mathbf{D}} & = & \mathbf{D}
\end{eqnarray}

A detailed discussion of these approaches lies beyond the scope of this paper, and we refer the interested reader to standard textbooks on control theory for further details [Åström and Wittenmark, 2013].

The recurrent form of Eq. (<ref>) 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

$\mathbf{z}_0 = \mathbf{s}_{0} \cdot \overline{\mathbf{B}}$

$\mathbf{o}_0 = \mathbf{s}_{0} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} + \mathbf{s}_0 \cdot \overline{\mathbf{D}}$

$\mathbf{z}_1 = \mathbf{s}_{0} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}} + \mathbf{s}_1 \cdot \overline{\mathbf{B}}$

$\mathbf{o}_1 = \mathbf{s}_{0} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}} \cdot \overline{\mathbf{C}} + \mathbf{s}_{1} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} + \mathbf{s}_1 \cdot \overline{\mathbf{D}}$

$\mathbf{z}_2 = \mathbf{s}_{0} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^2 + \mathbf{s}_1 \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}} + \mathbf{s}_2 \cdot \overline{\mathbf{B}}$

$\mathbf{o}_2 = \mathbf{s}_{0} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^2 \cdot \overline{\mathbf{C}} + \mathbf{s}_{1} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{A}} \cdot \overline{\mathbf{C}} + $

$\mathbf{s}_{2} \cdot \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} + \mathbf{s}_2 \cdot \overline{\mathbf{D}}$

It is easy to write

\begin{eqnarray}
\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{eqnarray}

Clearly, the right-hand side of Eq. (<ref>) can be interpreted as a merged output of a convolutional layer and a linear layer. Given that

\begin{eqnarray}
\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 & ... & \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}} & ... & \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} \end{bmatrix} \label{eq:ssm-convolutional-part}
\end{eqnarray}

we define a filter having the parameters

\begin{eqnarray}
\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}} & ... & \overline{\mathbf{B}} \cdot \overline{\mathbf{C}} \end{bmatrix}
\end{eqnarray}

where $n_{\mathrm{max}}$ is the maximum length of the sequence[Here $\mathbf{W}_{\mathrm{ssm}}$ can be represented as an $n_{\mathrm{max}} \times d \times d$ tensor.]. 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{eqnarray}
\mathbf{O} & = & \mathrm{Conv}(\mathbf{S},\mathbf{W}_{\mathrm{ssm}}) + \mathrm{Linear}(\mathbf{S},\overline{\mathbf{D}})
\end{eqnarray}

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.

Unfortunately, the above model performs poorly in many cases. As with many deep neural networks, careful initialization of the model parameters plays an important role in such models. For example, restricting the state matrix to particular types of matrices is found to be useful for learning and generalizing on long sequences [Gu et al., 2022].

Another problem with the basic state-space model is that it involves multiplication of multiple matrices. If the sequence is long (i.e., $n$ is a large number), computing $\overline{\mathbf{A}}^{n}$ will be computationally expensive and numerically unstable. One of the most popular approaches to developing practical state-space models for sequence modeling is diagonalization. The basic idea is that we can transform a state-space model into a new state-space where $\mathbf{A}$ (or $\overline{\mathbf{A}}$) is diagonalized. Given a state-space model parameterized by $(\mathbf{A},\mathbf{B},\mathbf{C},\mathbf{D})$, we can define a new state-space model $(\mathbf{U} \mathbf{A} \mathbf{U}^{-1},\mathbf{B} \mathbf{U}^{-1},\mathbf{U} \mathbf{C},\mathbf{D})$ by introducing an invertible matrix $\mathbf{U}$. It is easy to prove that the two models are equivalent under the state-space transformation $\mathbf{U}$[A state space transformation can be seen as a process of mapping all states from the old space to the new space, by

\begin{eqnarray}
\mathbf{s}'(t) & = & \mathbf{s}(t) \cdot \mathbf{U}
\end{eqnarray}

By using this state-space transformation, and by noting that $\mathbf{A}$ (or $\overline{\mathbf{A}}$) can be written as a canonical form $\mathbf{P}^{-1} \Lambda \mathbf{P}$[$\Lambda$ denotes a diagonal matrix.], we can enforce the constraint that $\mathbf{A}$ (or $\overline{\mathbf{A}}$) is a diagonal matrix, giving rise to diagonal state-space models. To illustrate, consider the filter used in the convolutional representation of the state-space model (see Eq. (<ref>)). Assuming that $\overline{\mathbf{A}} = \mathbf{P}^{-1} \Lambda \mathbf{P}$, we can write $\overline{\mathbf{B}} \cdot \overline{\mathbf{A}}^{t} \cdot \overline{\mathbf{C}}$ as

\begin{eqnarray}
\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{eqnarray}

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{eqnarray}
\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{eqnarray}

More detailed discussions of diagonal state-space models in sequence modeling can be found in Gu et al., 2021's work.

The application of state-space models to Transformer is simple. Each self-attention sub-layer is replaced in this case by an SSM sub-layer as described in Eqs. (<ref>) and (<ref>). 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. It should be noted, however, that state-space models are found to underperform Transformer models for NLP problems, such as language modeling, although they have achieved promising results in several other fields. Further refinements are often needed to make them competitive with other widely used sequence models [Fu et al., 2022].

While the formalism of state-space models is different from those we discussed in this chapter, it provides a general framework of sequence modeling in which the problem can be viewed from either of two different perspectives and we choose different ones for different purposes. Several recent sequence models were motivated by this idea, leading to systems exhibiting properties of both parallel training and RNN-style inference [Orvieto et al., 2023, Sun et al., 2023].

§.§ Conditional Computation

So far in our discussion of efficient Transformer models, we have assumed that the model architecture is given before beginning the training of a model and is then fixed throughout. We now turn to the case of learning efficient model architectures. Without loss of generality, we can write a model in the form

\begin{eqnarray}
\mathbf{y} & = & \mathrm{Model}(\mathbf{x},g(\mathbf{x}))
\end{eqnarray}

where $\mathbf{x}$ and $\mathbf{y}$ are the input and output of the model. $g(\mathbf{x})$ is a model 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 of learning is to find the optimal values of the parameters given the model 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 model function $\hat{g}(\mathbf{x})$ that produces the optimal model architecture and parameter values given the input $\mathbf{x}$. However, searching a hypothesis space of all possible combinations of architectures and parameter choices is extremely difficult, and so we need practical methods to achieve the goal. 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 criterion [Zoph and Le, 2016, Elsken et al., 2019]. Once the optimal neural network is determined, its parameters will be trained as usual, and then be 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 [Howard et al., 2019, Tan and Le, 2019]. 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 dynamically to various inputs [Gupta et al., 2004, Han et al., 2021]. 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 make no attempt to discuss it here. Instead, the reader can refer to the above papers to have a general idea of it, and refer to So et al., 2019, Wang et al., 2020, and Hu et al., 2021's work for its application to Transformers.

In this sub-section 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 [Bengio et al., 2013, Bengio et al., 2015]. More recently, it has often been used to refer to as a 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 necessary sub-set of the network [Xu and Mcauley, 2023]. One of the simplest methods, sometimes called early stopping, is to stop the computation at some point during reading or generating a sequence. This technique is often used in practical sequence generation applications where a low latency is required. Suppose $y_1...y_{n_{\mathrm{max}}}$ is the longest sequence that the system can generate, and $\mathbf{s}_1...\mathbf{s}_{n_{\mathrm{max}}}$ is the corresponding sequence of the states of the top-most 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{eqnarray}
\Pr(c|\mathbf{s}_i) & = & f_{\mathrm{stop}}(\mathbf{s}_i) \label{eq:early-stop-classification}
\end{eqnarray}

The generation process terminates if $\Pr(\mathrm{stop}|\mathbf{s}_i)$ is sufficiently large, for example

\begin{eqnarray}
\Pr(\mathrm{stop}|\mathbf{s}_i) & \ge & \Pr(\mathrm{nonstop}|\mathbf{s}_i) + \theta_{\mathrm{stop}} \label{eq:conditional-computation-simple-early-stop}
\end{eqnarray}

where $\theta_{\mathrm{stop}}$ denotes the minimal margin for distinguishing the two actions[An equivalent form of Eq. (<ref>) is $\Pr(\mathrm{stop}|\mathbf{s}_i) \ge \frac{1 + \theta_{\mathrm{stop}}}{2}$.]. 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...\mathbf{h}_m$ be low-level representations of a sequence $x_1...x_m$. Like Eq. (<ref>), we can develop a model $\Pr(c|\mathbf{s}_i)$ ($c \in \{\mathrm{skip},\mathrm{nonskip}\}$) to determine whether the token $x_i$ can be skipped. Figure <ref> (a) and (b) show illustrations of early stopping and skipping. Note that the learning-to-skip method has overlap 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 [Shen et al., 2020, Kim and Cho, 2021]. This method is also related to the downsampling methods which will be discussed in Section <ref>.

Methods of conditional computation, including early stopping, token skipping, MoE, sentence-level depth adaptation, token-level depth adaptation, and layer skipping. While these methods are illustrated using either the encoding or decoding process, most of them can be applied to both Transformer encoders and decoders.

A second approach to conditional computation is to resort to sparse expert models, or its popular instance — MoE [Yuksel et al., 2012]. 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 big 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 algorithm (see Figure <ref> (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 <ref>, 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 Fedus et al., 2022'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 layer by which we can still make accurate predictions for the sample (see Figure <ref> (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$[A layer is a standard Transformer block consisting of a few sub-layers.]. 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{eqnarray}
\mathbf{p}_i^{l} & = & \mathrm{Softmax}(\mathbf{s}_i^{l} \cdot \mathbf{W}_{\mathrm{o}}^{l})
\end{eqnarray}

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},...,\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},...,\mathbf{p}_i^{l}\}$ and/or $\{\mathbf{s}_i^{1},...,\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 tends to converge and the number of layers can stop growing.

The above method 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 [Xin et al., 2020, Zhou et al., 2020, Schuster et al., 2022]. 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 [Elbayad et al., 2020]. 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 [Schwartz et al., 2020, Schuster et al., 2022]. 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 [Zhou et al., 2020, Liao et al., 2021].

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 sub-set of the layers for constructing a shallow network. The adaptive models therefore can be generalized to skipping models (see Figure <ref> (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 [Wang et al., 2018, Wu et al., 2018]. 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 [Wang et al., 2022, Del Corro et al., 2023].

§.§ 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. But we will not discuss this general topic, which spans a broad range of issues and models, many outside the scope of this chapter. Instead, we narrow our discussion to two kinds of approaches that are widely used in learning small neural networks from large neural networks.

§.§.§ Knowledge Distillation

Knowledge distillation is a process of compressing the knowledge in a large neural network (or an ensemble of neural networks) into a small neural network [Hinton et al., 2015]. In supervised learning of neural networks, the objective functions are generally designed to represent some loss of replacing the true answer with the predicted answer. Hence we can minimize this loss so that the models are trained to output the true answer. While models are typically optimized on the training data in this manner, what we really want is to generalize them to new data. This is, however, difficult because we have no information about generalization in training with the ground-truth. In knowledge distillation, instead of forcing a model to stay close to the ground-truth output, we train this model to generalize. To do this, we directly transfer the knowledge (i.e., the generalization ability) of a pre-trained model to the model that we want 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 relatively small neural network, such as a neural network 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 $\{\Pr(\cdot|y_0,\mathbf{x}),...,\Pr(\cdot|y_0...y_{n-1},\mathbf{x})\}$ for the input $\mathbf{x}$. To keep the notation simple, we denote the distribution $\Pr(\cdot|y_0...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 $Loss(\widetilde{\mathbf{p}}_i,\mathbf{p}_i)$ (such as the cross-entropy function) for computing some distance between $\widetilde{\mathbf{p}}_i$ and $\mathbf{p}_i$. Then, we can define the loss over the entire sequence as

\begin{eqnarray}
L(\mathbf{x},\theta) & = & \frac{1}{n} \sum_{i=1}^{n} Loss(\widetilde{\mathbf{p}}_i,\mathbf{p}_i)
\end{eqnarray}

where $\theta$ denotes the parameters of the student model[We omit the parameters of the teacher model because they are fixed throughout the training process.]. Using this loss, we can optimize $\theta$, for any given set of source sequences $\{\mathbf{x}_1,...,\mathbf{x}_K\}$, in such a way as to minimize the quality $\sum_{k=1}^{N} L(\mathbf{x}_k,\theta)$.

Several different 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 [Romero et al., 2014]. 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 some difference between the relation encodings of the teacher and student models [Park et al., 2019, Peng et al., 2019]. 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 [Kim and Rush, 2016]. 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,...,\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),...,(\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 argumentation method is that it is architecture free, and we do not even need to understand the internal architectures of the teacher and student models. Hence we can apply this method if we have a black-box teacher model. More detailed discussions of knowledge distillation can be found in Gou et al., 2021 and Wang and Yoon, 2021's surveys.

§.§.§ Structured Pruning

Pruning is among the most popular of the model compression methods and has been applied to a broad range of systems. One common approach to pruning is unstructured pruning, by which we activate only some of the connections between neurons. However, as with most sparse models, models pruned in this way typically require special implementations and hardware support, which in turn reduces their efficiency in some applications. A simple but more aggressive way to do pruning is to use structured pruning. In deep learning, structured pruning is a technique that removes a group of neurons or connections together. For example, we can remove an entire layer of neuron from a neural network to obtain a shallower model. As multi-layer, multi-head neural networks, Transformers are naturally suited to structured pruning, and we can prune a Transformer network in several different ways. For example, we can prune some of the heads in multi-head attention [Voita et al., 2019, Michel et al., 2019], or some of the layers in the layer stack [Hou et al., 2020, Kim and Awadalla, 2020].

Formally, we can represent a neural network as a set of parameter groups $\{\theta_1,...,\theta_{R}\}$, each corresponding to a component or sub-model of the model. Our goal is to find a sub-set of $\{\theta_1,...,\theta_{R}\}$ by which we can build a model that yields good performance, while having a lower model complexity. However, a simple search of such a model is infeasible because there are a combinatorially large number of possible model candidates and evaluating all of these models is computationally expensive.

One approach to structured pruning is to randomly prune components of a model. One can run the random pruning process a number of times to generate a pool of model candidates and select the best one from the pool. Another approach is to use heuristics to decide which components are not important and can be removed. Common measures of the importance of a parameter group $\theta_r$ include various qualities based on norms of the weights or gradients of $\theta_r$ [Santacroce et al., 2023]. We can prune $\theta_r$ if the values of these measures are below (or above) given thresholds. A third approach is to frame the pruning problem as an optimization task by introducing trainable gates indicating the presence of different components [McCarley et al., 2019, Wang et al., 2020, Lagunas et al., 2021]. The pruned model can be induced by using the trained gates. Note that, in many cases, pruning is not a post-processing step for a given trained model, but part of the training.

§.§ Sequence Compression

In sequence modeling and generation problems, the time and space complexities are strongly influenced by the length of the input or output sequence, and we prefer the sequence to be short. This is particularly important for Transformer models, as their time and space complexities are quadratic with the sequence length, and the memory footprint and latency can be heavy burdens if the sequence is very long. In the previous subsections, we have discussed modifications to the Transformer architecture for dealing with long sequences. Here we instead consider methods for compressing the sequences into ones with acceptable lengths.

One simple approach is to map the input sequence to a fixed-size representation. For example, using the recurrent models discussed in Section <ref>, 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 of 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 attend the input vectors to some hidden states, usually a fixed number of learnable hidden representations. In Jaegle et al., 2021's work, this is done by introducing $r$ hidden representations $\{\mathbf{u}_{1},...,\mathbf{u}_{r}\}$, and then attending the input vectors $\{\mathbf{x}_1,...,\mathbf{x}_m\}$ to these hidden representations. The attention model can be a standard QKV attention model in which we view $\{\mathbf{u}_{1},...,\mathbf{u}_{r}\}$ as queries and $\{\mathbf{x}_1,...,\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 develop 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 rate 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 [Han et al., 2020, Burchi and Vielzeuf, 2021]. We need to consider a trade-off between the compressed sequence length and the performance of downstream systems [Xu et al., 2023].

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 will suffer from sparse data. In deterministic tokenization methods, which produce tokenization results using statistics collected from the entire dataset, the sequence length can be somehow controlled by adjusting some hyper-parameter, for example, in byte pair encoding [Sennrich et al., 2016], increasing the size of the vocabulary generally reduces the number of the resulting tokens. Another way to obtain an appropriate sequence of tokens is to use a model for choosing among tokenization candidates [Kudo, 2018, Provilkov et al., 2020]. As with many probabilistic models for text generation, in this case, we can add 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 example, in many practical applications, we chop the sequence when its length exceeds a threshold. We can relate this to the early stopping and skipping approaches in conditional computation. Thus the methods discussed in Section <ref> are 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 [Kim et al., 2023].

§.§ High Performance Computing Methods

So far in this section, we have discussed efficient Transformer models from the perspectives of deep learning and NLP. However, we have not considered their efficiency on hardware. As modern hardware provides a variety of modes for running a program, the practical time and memory footprint savings generally depend on the specifications of hardware systems. One line of research on efficient use of computing resources explores methods of parallel computing. 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 need considerations of networking and communication in the cluster, as well as the utilization of sparse models that activate only a small sub-set of the parameters for each sample, enabling the use of very large models. Most of these methods have been studied in an extensive literature on how to scale up the training of deep neural networks like Transformers efficiently [Lepikhin et al., 2021, Barham et al., 2022, Fedus et al., 2022]. The results of these studies were foundational to many follow-on works on investigating the scaling laws for large language models [Brown et al., 2020, Chowdhery et al., 2022]. Since large-scale distributed models are generic and not specialized to the case of Transformers, we skip the discussion of them here. The interested readers can refer to the above papers for more detailed discussions.

In this sub-section 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 store and execute neural networks using lower or mixed-precision number representations [Gholami et al., 2022]. 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 in storing floating-point numbers for 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 of linear algebra operations and hardware support, it does not change the model architecture, and so we need only a slight modification to the systems. For example, half-precision floating-point representations  can be applied to either 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., fix-point representations) . This process is in general implemented using a system called 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{eqnarray}
Q(x) & = & \lfloor \frac{x}{s} \rceil
\end{eqnarray}

where $\lfloor \cdot \rceil$ is a rounding function[$\lfloor a \rceil$ returns the integer closest to $a$.], $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{eqnarray}
D(r) & = & s \cdot r
\end{eqnarray}

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 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 possible values of $Q(r)$  evenly across values of an integer, 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 value for $D(r)$. The above equations 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 [Oppenheim and Schafer, 1975, Rabiner and Gold, 1975, Gray, 1998].

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 the 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 prediction is not as accurate as that of the standard Transformer model. Using integer operations to approximate continuous-valued operations generally leads to approximation errors. These errors will be accumulated 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 sub-models will lead to high accuracy loss. One solution is to simplify the model architecture and develop new sub-models that is more feasible for quantization. Alternatively, a more common paradigm in quantized neural networks is to add de-quantization functions to the neural networks so that the output of a layer is 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}$ in a very cheap way.

However, sandwiching each layer between $Q(\cdot)$ and $D(\cdot)$ will lead to 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 bigger than the time saving of performing integer-based operations. In general, the benefit of quantizing neural networks would be larger than its cost if the neural networks are 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 the quantization approaches can be used in both training and inference, a widely-used approach is to get Transformer models quantized after training (call it post-training quantization). In this approach, quantization is performed on well-trained floating-point-based neural networks and there will be fewer quantization-related errors. However, these errors cannot be compensated for because they exist after training. A more promising idea is to involve quantization in training so that the model can learn to compensate for quantization-related errors [Jacob et al., 2018, Nagel et al., 2021]. There have been several attempts to apply quantization-aware training to Transformers [Bondarenko et al., 2021, Stock et al., 2021, Yang et al., 2023]. In addition to computational efficiency, another important consideration for high-performance systems is the restrictions of the memory hierarchy. In general, better system design requires considering the speeds and sizes of different levels of memory. The problem is even more complicated when we train large Transformer models on modern hardware where both GPUs and CPUs are used. A general principle of system design is that memory transfer 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 that we need to process, in many practical situations the size of the fast, on-chip memory is orders of magnitude smaller than the size of data. In this case, we can re-order the memory access in the algorithms so that the data used in nearby computation steps can be loaded into the high-speed memory at one time. 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 system. But the modules of this system are not optimized as a whole for efficiency improvement. For example, a self-attention sub-layer involves a series of operations of scaling, normalization, and matrix multiplication. Although each of these operations has been implemented in several supported and efficient libraries of linear algebra, successive calls to them still require multiple times of memory transfer when we switch from one operation to another. In practice, a better approach would be that we keep some of the intermediate states in the on-chip memory, and reuse them in the following computation steps instead of fetching them again from the slow memory. For example, on modern GPUs, a simple way to achieve this is to merge multiple operations into a single operation, known as kernel fusion. For Transformer models, a general idea is to design data partitioning and layout strategies by which we maximize the computation on each data block loaded into the high-performance memory, while at the same time minimizing the memory transfer. There have been several attempts to use these strategies to improve the attention models in Transformers [Ivanov et al., 2021, Pope et al., 2023]. Some of these methods, such as flash attention and paged attention, have been successfully incorporated into recent large language models [Dao et al., 2022, Kwon et al., 2023].

§ APPLICATIONS

Transformers have a wide range of applications, covering numerous NLP problems. While the Transformer model introduced by Vaswani et al., 2017 is based on a standard encoder-decoder architecture, it is mainly used in three different ways.

Decoder-only Models. By removing the cross-attention sub-layers from a Transformer decoder, 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 the next and following tokens.

Encoder-only Models. Transformer encoders can be treated as sequence models that take a sequence of tokens at once and produce a sequence of representations, each of which corresponds to an input token. These representations can be seen as some sort 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, for example, systems of sentence classification, sequence labeling, and so on. 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 fall into the above three categories, the same NLP problem can generally be addressed using different types of models. For example, recent decoder-only models have demonstrated good performance on a broad range of problems by framing them as text generation tasks, though some of these problems were often addressed by using encoder-decoder or encoder-only models. To illustrate how the above models are applied, this section considers a few applications where Transformers as chosen as the backbone models.

§.§ 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$ (denoted by $\Pr(\cdot|x_0,...,x_{i})$ where $\{x_0,...,x_{i}\}$ denote the tokens up to position $i$). The best predicted token is the one which maximizes the probability, given by

\begin{eqnarray}
\hat{x}_{i+1} & = & \underset{x_{i+1} \in V}{\arg \max } \operatorname{Pr}\left(x_{i} \mid x_{0}, \ldots, x_{i}\right)
\end{eqnarray}

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

\begin{eqnarray}
\hat{x}_{k+1} & = & \underset{x_{k+1} \in V}{\arg \max } \operatorname{Pr}\left(x_{k} \mid x_{0}, \ldots, x_{i}, \hat{x}_{i+1}, \ldots, \hat{x}_{k}\right) \label{eq:transformer-app-lm-long}
\end{eqnarray}

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

As discussed in Section <ref>, Transformer decoders are essentially language models. The only 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 on preceding tokens solely on one side. 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{eqnarray}
\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{eqnarray}

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 in training, because 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 the 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 auto-regressive manner, as described in Eq. (<ref>).

The training of this model is standard. 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. We will discuss these issues further in the following chapters, where large-scale pre-training is the primary focus.

§.§ Text Encoding

For many NLP problems, a widely used paradigm is to first represent an input sequence in some form, and then make predictions for downstream tasks based on this representation. As a result, we separate sequence modeling or sequence representation from NLP tasks. One of the advantages of this paradigm is that we can train a sequence model that is not specialized to particular tasks, thereby generalizing well.

Clearly, Transformer encoders are a type of sequence model, and can be used as text encoders. Consider a Transformer encoder with $L$ encoding blocks. The output of the last encoding block can be seen as the encoding result. Here add a special token $x_0$ to any sequence, indicating the beginning of a sequence (written as $\langle \mathrm{SOS} \rangle$ or $[\mathrm{CLS}]$). If there is a sequence of $m+1$ input tokens $x_0x_1...x_m$, the output of the encoder will be a sequence of $m+1$ vectors $\mathbf{h}_0^{L}\mathbf{h}_1^{L}...\mathbf{h}_m^{L}$. Since $x_0$ is not a real token and has a fixed positional embedding, it serves as a tag for collecting information from other positions using the self-attention mechanism. Hence $\mathbf{h}_0^{L}$ is a representation of the entire sequence, with no biases for any specific tokens or positions. In many cases, we need a single representation of a sequence and take it as input to 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},...,\mathbf{h}_m^{L}\}$. In this case, we can simply use $\mathbf{h}_0^{L}$ as the representation of the sequence. A more general approach is to add a pooling layer to the encoder. This allows us to explore various pooling methods to generate the sequence embedding from $\{\mathbf{h}_0^{L},...,\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 seen as points in a multi-dimensional space [Hill et al., 2016]. Another way to make use of text encoding, therefore, is to obtain semantic or syntactic similarities of token sequences based on their relative positions or proximity in this space. A straightforward method for this is to compute the Euclidean distances between sequence embeddings. The shorter the distance between two sequences, the more similar they are considered to be. There are many distance metrics we can choose, and it is possible to combine them to obtain a better measure of sequence similarity. Such similarity computations are applied in areas such as text entailment, information retrieval, translation evaluation, among others [Cer et al., 2018, Reimers and Gurevych, 2019]. Additionally, they are often used to assess the quality of text encoding models.

Text encoding is also a crucial component of sequence-to-sequence models. Given this, we can develop a separate Transformer encoder for source-side sequence modeling in an encoder-decoder system (see Figure <ref>). For example, we can pre-train a Transformer encoder on large-scale source-side texts, and use it as the encoder in a downstream encoder-decoder model. It is worth noting that while the encoder is designed based on the Transformer architecture, the decoder is not confined to just Transformers. Such 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.

Integrating Transformer encoders as components of different systems. A common approach is to feed the output of the encoder (with pooling) into a classifier to obtain a sequence classification system. Another way to utilize Transformer encoders is to compute the similarity between two sequences. We use the same encoder to represent the two sequences, and then construct a neural network on top of the two representations for producing a similarity score between them. As usual, Transformer encoders can also be used in encoder-decoder systems to model sequence-to-sequence problems.

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.

§.§ Speech Translation

As illustrated in Section <ref>, 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 of an S2T translation system is a sequence of source-side acoustic feature vectors, denoted by $\mathbf{a}_1...\mathbf{a}_m$, and the output of the system is a sequence of target-side tokens, denoted by $y_1...y_n$.[In order to obtain the input sequence to the system, we need to discretize continuous speech into signals represented by feature vectors. This process is typically nontrivial, requiring either a feature extractor based on a variety of signal processing operations or a neural network that learns feature mappings in an end-to-end manner. But we will not dive into the details of these methods and simply treat the input feature extractor as an upstream system.] Mapping $\mathbf{a}_1...\mathbf{a}_m$ to $y_1...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, like in neural machine translation.

In S2T translation, however, we have to deal with sequence mappings between modalities and between languages simultaneously. This poses new challenges compared with conventional machine translation problems and influences the design of S2T translation models. There have been several improvements to Transformer models for adapting them better to S2T translation tasks. Some of the improvements concern the design of Transformer blocks [Di Gangi et al., 2019]. For example, in Gulati et al., 2020's system, a CNN sub-layer and relative positional embeddings are integrated into each Transformer block, 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 modifications to either encoders or decoders, or both. To illustrate, Figure <ref> 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...\mathbf{a}_m$ and one Transformer decoder for producing the target-side output $y_1...y_n$. By contrast, the decoupled encoder model separates the encoder into two stacked encoders — one for acoustic modeling (call it the speech encoder), and one for textual modeling (call it the text encoder) [Liu et al., 2020, Xu et al., 2021]. This design reflects a modeling hierarchy in which representations in different levels of the network are concerned with different aspects of the problem, for example, the speech encoder models low-level features in mapping acoustic embeddings into larger language units, and the text encoder models the semantic or syntactic features in representing the entire input sequence. An advantage of separating out the text encoder is that the encoding process follows our prior knowledge that 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 some way we train an ASR system. This enables us to pre-train the speech encoder and the text encoder on unlabeled data, and incorporate the pre-trained encoders into S2T translation systems.

Architectures of speech-to-text translation models based on Transformers. In addition to the standard encoder-decoder architecture, we can explicitly model the acoustic and textual (semantic) information using two separate encoders, called the speech encoder and the text encoder. In the decoupled encoder architecture, the two encoders are stacked, that is, text encoding is a subsequent process after speech encoding. In the two-stream encoder architecture, the two encoders work in parallel, and their outputs are merged using an additional encoder, called the shared encoder. The dotted line indicates the potential for interaction between the two encoders. For example, we could define a loss function to minimize the difference between their outputs, thereby guiding the model towards more aligned representations.

An alternative encoding architecture is the two-stream architecture, as shown in Figure <ref> (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 [Ye et al., 2021]. The speech encoder takes acoustic features as input and the text encoder takes tokens (or their embeddings) as input. A third encoder, called 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, then the model will transform into a machine translation model which can be trained using bilingual texts. Conversely, if we mask the text encoder, then 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 related to architecture design. Here, we have data in two modalities and two languages, and so we can develop multiple supervision signals for multi-task learning of S2T translation models. A widely used method is to introduce 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. Among typical responses to this challenge are 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 discussion on Transformers. The interested reader can refer to a recent survey on speech translation for more information [Xu et al., 2023].

§.§ 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 the backbone models. In this sub-section, we consider Vision Transformer (ViT) - an interesting application of Transformers to image classification [Dosovitskiy et al., 2021]. Vision Transformer is a milestone model which opens the door to purely Transformer-based vision models. Here we consider the basic structure of Vision transformer to make this section concentrated and coherent, although there has been an extensive literature on Vision transformer and its variants. More detailed discussions of vision transformer can be found in recent surveys [Han et al., 2022, Liu et al., 2023].