A Survey of Transformers

Introduction

Transformer (Vaswani et al., 2017) is a prominent deep learning model that has been widely adopted in various fields, such as natural language processing (NLP), computer vision (CV) and speech processing. Transformer was originally proposed as a sequence-to-sequence model (Sutskever et al., 2014) for machine translation. Later works show that Transformer-based pre-trained models (PTMs) (Qiu et al., 2020) can achieve state-of-the-art performances on various tasks. As a consequence, Transformer has become the go-to architecture in NLP, especially for PTMs. In addition to language related applications, Transformer has also been adopted in CV (Parmar et al., 2018; Carion et al., 2020; Dosovitskiy et al., 2020), audio processing (Dong et al., 2018; Gulati et al., 2020; Chen et al., 2021) and even other disciplines, such as chemistry (Schwaller et al., 2019) and life sciences (Rives et al., 2021).

Due to the success, a variety of Transformer variants (a.k.a. X-formers) have been proposed over the past few years. These X-formers improve the vanilla Transformer from different perspectives.

Model Efficiency. A key challenge of applying Transformer is its inefficiency at processing long sequences mainly due to the computation and memory complexity of the self-attention module. The improvement methods include lightweight attention (e.g. sparse attention variants) and Divide-and-conquer methods (e.g., recurrent and hierarchical mechanism).

Model Generalization. Since the transformer is a flexible architecture and makes few assumptions on the structural bias of input data, it is hard to train on small-scale data. The improvement methods include introducing structural bias or regularization, pre-training on large-scale unlabeled data, etc.

Model Adaptation. This line of work aims to adapt the Transformer to specific downstream tasks and applications.

In this survey, we aim to provide a comprehensive review of the Transformer and its variants. Although we can organize X-formers on the basis of the perspectives mentioned above, many existing X-formers may address one or several issues. For example, sparse attention variants not only reduce the computational complexity but also introduce structural prior on input data to alleviate the overfitting problem on small datasets. Therefore, it is more methodical to categorize the various existing X-formers and propose a new taxonomy mainly according to their ways to improve the vanilla Transformer: architecture modification, pre-training, and applications. Considering the audience of this survey may be from different domains, we mainly focus on the general architecture variants and just briefly discuss the specific variants on pre-training and applications.

The rest of the survey is organized as follows. Sec. 2 introduces the architecture and the key components of Transformer. Sec. 3 clarifies the categorization of Transformer variants. Sec. 4$\sim$5 review the module-level modifications, including attention module, position encoding, layer normalization and feed-forward layer. Sec. 6 reviews the architecture-level variants. Sec. 7 introduces some of the representative Transformer-based PTMs. Sec. 8 introduces the application of Transformer to various different fields. Sec. 9 discusses some aspects of Transformer that researchers might find intriguing and summarizes the paper.

Background

The vanilla Transformer (Vaswani et al., 2017) is a sequence-to-sequence model and consists of an encoder and a decoder, each of which is a stack of $L$ identical blocks. Each encoder block is mainly composed of a multi-head self-attention module and a position-wise feed-forward network (FFN). For building a deeper model, a residual connection (He et al., 2016) is employed around each module, followed by Layer Normalization (Ba et al., 2016) module. Compared to the encoder blocks, decoder blocks additionally insert cross-attention modules between the multi-head self-attention modules and the position-wise FFNs. Furthermore, the self-attention modules in the decoder are adapted to prevent each position from attending to subsequent positions. The overall architecture of the vanilla Transformer is shown in Fig. 1.

In the following subsection, we shall introduce the key modules of the vanilla Transformer.

Instead of simply applying a single attention function, Transformer uses multi-head attention, where the $D_{m}$-dimensional original queries, keys and values are projected into $D_{k}$, $D_{k}$ and $D_{v}$ dimensions, respectively, with $H$ different sets of learned projections. For each of the projected queries, keys and values, and output is computed with attention according to Eq. (1). The model then concatenates all the outputs and projects them back to a $D_{m}$-dimensional representation.

In Transformer, there are three types of attention in terms of the source of queries and key-value pairs:

Self-attention. In Transformer encoder, we set $\mathbf{Q}=\mathbf{K}=\mathbf{V}=\mathbf{X}$ in Eq. (2), where $\mathbf{X}$ is the outputs of the previous layer.

Masked Self-attention. In the Transformer decoder, the self-attention is restricted such that queries at each position can only attend to all key-value pairs up to and including that position. To enable parallel training, this is typically done by applying a mask function to the unnormalized attention matrix $\hat{\mathbf{A}}=\exp(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{D_{k}}})$, where the illegal positions are masked out by setting $\hat{A}_{ij}=-\infty\textrm{ if }i<j$. This kind of self-attention is often referred to as autoregressive or causal attentionThis term seems to be borrowed from the causal system, where the output depends on past and current inputs but not future inputs..

Cross-attention. The queries are projected from the outputs of the previous (decoder) layer, whereas the keys and values are projected using the outputs of the encoder.

1.2. Position-wise FFN

The position-wise FFNThe parameters are shared across different positions, thus the position-wise FFN can also be understood as two convolution layers with kernel size of 1. is a fully connected feed-forward module that operates separately and identically on each position

1.3. Residual Connection and Normalization

In order to build a deep model, Transformer employs a residual connection (He et al., 2016) around each module, followed by Layer Normalization (Ba et al., 2016). For instance, each Transformer encoder block may be written as

1.4. Position Encodings

Since Transformer doesn’t introduce recurrence or convolution, it is ignorant of positional information (especially for the encoder). Thus additional positional representation (Detailed discussion in Sec. 5.1) is needed to model the ordering of tokens.

2. Model Usage

Generally, the Transformer architecture can be used in three different ways:

Encoder-Decoder. The full Transformer architecture as introduced in Sec. 2.1 is used. This is typically used in sequence-to-sequence modeling (e.g., neural machine translation).

Encoder only. Only the encoder is used and the outputs of the encoder are utilized as a representation for the input sequence. This is usually used for classification or sequence labeling problems.

Decoder only. Only the decoder is used, where the encoder-decoder cross-attention module is also removed. This is typically used for sequence generation, such as language modeling.

3. Model Analysis

To illustrate the computation time and parameter requirements of the Transformer, we analyze the two core components of the Transformer (i.e., the self-attention module and the position-wise FFN) in Table 1. We assume that the hidden dimension $D_{m}$ of the model is $D$, and that the input sequence length is $T$. The intermediate dimension of FFN is set to $4D$ and the dimension of keys and values are set to $D/H$ as in Vaswani et al. (2017).

When the input sequences are short, the hidden dimension $D$ dominates the complexity of self-attention and position-wise FFN. The bottleneck of Transformer thus lies in FFN. However, as the input sequences grow longer, the sequence length $T$ gradually dominates the complexity of these modules, in which case self-attention becomes the bottleneck of Transformer. Furthermore, the computation of self-attention requires that a $T\times T$ attention distribution matrix is stored, which makes the computation of Transformer infeasible for long-sequence scenarios (e.g., long text documents and pixel-level modeling of high-resolution images). One shall see that the goal of increasing the efficiency of Transformer generally leads to the long-sequence compatibility of self-attention, as well as the computation and parameter efficiency of position-wise FFN for ordinary settings.

4. Comparing Transformer to Other Network Types

As a central piece of Transformer, self-attention comes with a flexible mechanism to deal with variable-length inputs. It can be understood as a fully connected layer where the weights are dynamically generated from pairwise relations from inputs. Table 2 compares the complexity, sequential operations, and maximum path lengthThe maximum length of the paths forward and backward signals have to traverse to get from any input position to arbitrary output position. Shorter length implies a better potential for learning long-range dependencies. of self-attention with three commonly used layer types. We summarize the advantages of self-attention as follows:

It has the same maximum path length as fully connected layers, making it suitable for long-range dependencies modeling. Compared to fully connected layers, it is more parameter-efficient and more flexible in handling variable-length inputs.

Due to the limited receptive field of convolutional layers, one typically needs to stack a deep network to have a global receptive field. On the other hand, the constant maximum path length enables self-attention to model long-range dependencies with a constant number of layers.

The constant sequential operations and maximum path length make self-attention more parallelizable and better at long-range modeling than recurrent layers.

4.2. In Terms of Inductive Bias

Transformer is often compared against convolutional and recurrent networks. Convolutional networks are known to impose the inductive biases of translation invariance and locality with shared local kernel functions. Similarly, recurrent networks carry the inductive biases of temporal invariance and locality via their Markovian structure (Battaglia et al., 2018). On the other hand, the Transformer architecture makes few assumptions about structural information of data. This makes Transformer a universal and flexible architecture. As a side effect, the lack of structural bias makes Transformer prone to overfitting for small-scale data.

Another closely related network type is Graph Neural Networks (GNNs) with message passing (Wu et al., 2021a). Transformer can be viewed as a GNN defined over a complete directed graph (with self-loop) where each input is a node in the graph. The key difference between Transformer and GNNs is that Transformer introduces no prior knowledge over how input data are structured — the message passing process in Transformer solely depends on similarity measures over the content.

Taxonomy of Transformers

A wide variety of models have been proposed so far based on the vanilla Transformer from three perspectives: types of architecture modification, pre-training methods, and applications. Fig. 2 gives an illustrations of our categorization of Transformer variants.

Fig. 3 illustrates our taxonomy and some representative models.

In this survey, we focus on reviewing the works on architecture modifications. Since the attention module is the key component of Transformer, we solely describe the attention-related variants in Sec. 4 and introduce the other module-level variants in Sec. 5. Then Sec. 6 describes the other architecture-level variants. Finally, we briefly review the works on pre-training in Sec. 7 and applications in Sec. 8. There are some comprehensive surveys on the latter two categories of work, such as pre-trained models (PTMs) (Qiu et al., 2020) and visual Transformers(Han et al., 2021a; Khan et al., 2021).

Attention

Self-attention plays an important role in Transformer, but there are two challenges in practical applications.

Complexity. As discussion in Sec. 2.3, the complexity of self-attention is $\mathcal{O}(T^{2}\cdot D)$. Therefore, the attention module becomes a bottleneck when dealing with long sequences.

Structural prior. Self-attention does no assume any structural bias over inputs. Even the order information is also needed to be learned from training data. Therefore, Transformer (w/o pre-training) is usually easy to overfit on small or moderate-size data.

The improvements on attention mechanism can be divided into several directions:

Sparse Attention. This line of work introduces sparsity bias into the attention mechanism, leading to reduced complexity.

Linearized Attention. This line of work disentangles the attention matrix with kernel feature maps. The attention is then computed in reversed order to achieve linear complexity.

Prototype and Memory Compression. This class of methods reduces the number of queries or key-value memory pairs to reduce the size of the attention matrix.

Low-rank Self-Attention. This line of work capture the low-rank property of self-attention.

Attention with Prior. The line of research explores supplementing or substituting standard attention with prior attention distributions.

Improved Multi-Head Mechanism. The line of studies explores different alternative multi-head mechanisms.

We will describe these attention variants at length in the rest of this section.

In the standard self-attention mechanism, every token needs to attend to all other tokens. However, it is observed that for the trained Transformers the learned attention matrix $\mathbf{A}$ is often very sparse across most data points (Child et al., 2019). Therefore, it is possible to reduce computation complexity by incorporating structural bias to limit the number of query-key pairs that each query attends to. Under this limitation, we just compute the similarity score of the query-key pairs according to pre-defined patterns

From another perspective, the standard attention can be regarded as a complete bipartite graph where each query receives information from all memory nodes and updates its representation. The sparse attention can be considered as a sparse graph where some of the connections between nodes are removed.

Based on the metrics of determining the sparse connection, we categorize these approaches into two classes: position-based and content-based sparse attention.

In position-based sparse attention, the attention matrix is limited according to some pre-defined patterns. Although these sparse patterns vary in different forms, we find that some of them can be decomposed into some atomic sparse patterns.

We first identify some atomic sparse patterns and then describe how these patterns are composed in some existing work. Finally, we introduce some extended sparse patterns for specific data types.

There are mainly five types of atomic sparse attention patterns, as shown in Fig. 4.

Global Attention. To alleviate the degradation of the ability to model the long-range dependencies in sparse attention, one can add some global nodesIn practice, these global nodes can be selected from the sequence (internal global nodes) or virtual nodes with trainable parameters (external global nodes). as the hub for information propagation between nodes. These global nodes can attend all nodes in the sequence and the whole sequence attend to these global nodes, as illustrated in Fig. 4(a).

Band Attention(a.k.a sliding window attention or local attention). Since most data come with a strong property of locality, it is natural to restrict each query to attend to its neighbor nodes. A widely adopted class of such sparse pattern is band attention, in which the attention matrix is a band matrix as illustrated in Fig. 4(b).

Dilated Attention. Analogous to dilated CNNs (van den Oord et al., 2016), one can potentially increase the receptive field of the band attention without increasing computation complexity by using a dilated window with gaps of dilation $w_{d}\geq 1$, as depicted in Fig. 4(c). This can be easily extended to strided attention, where the window size is not limited but the dilation $w_{d}$ is set to a large value.

Random Attention. To increase the ability of non-local interactions, a few edges are randomly sampled for each query, as illustrated in Fig. 4(d). This is based on the observation that random graphs (e.g., Erdős–Rényi random graph) can have similar spectral properties with complete graphs that leads to a fast mixing time for random walking on graphs.

Block Local Attention. This class of attention segments input sequence into several non-overlapping query blocks, each of which is associated with a local memory block. All the queries in a query block attend to only the keys in the corresponding memory block. Fig. 4(e) depicts a commonly used case where the memory blocks are identical to their corresponding query blocks.

Existing sparse attentions are often composed of more than one of the above atomic patterns. Fig. 5 illustrates some representative compound sparse attention patterns.

Star-Transformer (Guo et al., 2019a) uses a combination of band attention and global attention. Specifically, Star-Transformer just includes only a global node and a band attention with the width of 3, in which any pair of non-adjacent nodes are connected through a shared global node and adjacent nodes are connected directly with each other. This kind of sparse pattern forms a star-shaped graph among nodes. Longformer (Beltagy et al., 2020) uses a combination of band attention and internal global-node attention. The global nodes are chosen to be [CLS] token for classification and all question tokens for Question Answering tasks. They also replace some of the band attention heads in upper layers with dilated window attention to increase the receptive field without increasing computation. As a concurrent work to Longformer (Beltagy et al., 2020), Extended Transformer Construction (ETC) (Ainslie et al., 2020) utilizes combination of band attention and external global-node attention. ETC also includes a masking mechanism to handle structured inputs and adapt Contrastive Predictive Coding (CPC) (van den Oord et al., 2018) for pre-training. In addition to the band and global attention, BigBird (Zaheer et al., 2020) uses additional random attention to approximate full attention. Their theoretical analysis also reveals that the usage of a sparse encoder and sparse decoder can simulate any Turing Machine, which explains the success of those sparse attention models.

Sparse Transformer (Child et al., 2019) uses a factorized attention where different sparse patterns are designed for different types of data. For data with a periodic structure (e.g., images), it uses a composition of band attention and strided attention. Whereas for data without a periodic structure (e.g., text), it uses a composition of block local attention combined with global attention, where global nodes are from fixed positions in the input sequence.

Apart from the above patterns, some existing studies have explored extended sparse patterns for specific data types.

For text data, BP-Transformer (Ye et al., 2019) constructs a binary tree where all tokens are leaf nodes and the internal nodes are span nodes containing many tokens. The edges in this graph are constructed so that each leaf node is connected to its neighbor leaf nodes and higher-level span nodes containing tokens from a longer distance. This approach can be seen as an extension of global attention, where global nodes are hierarchically organized and any pair of tokens are connected with paths in the binary tree. An abstract view of this method is illustrated in Fig. 6(a).

There are also some extensions for vision data. Image Transformer (Parmar et al., 2018) explores two types of attention: (1) flattening image pixels in raster-scan order and then applying block local sparse attention. (2) 2D block local attention, where query blocks and memory blocks are arranged directly in 2D plate, as depicted in Fig. 6(b). As another example of sparse pattern on vision data, Axial Transformer (Ho et al., 2019) applies independent attention modules over each axis of the image. Each attention module mixes information along one axis while keeping information along the other axis independent, as illustrated in Fig. 6(c). This can be understood as horizontally and vertically flattening image pixels in raster-scan order and then applying strided attention with gaps of image width and height, respectively.

1.2. Content-based Sparse Attention

Another line of work creates a sparse graph based on input content, i.e., the sparse connections are conditioned on inputs.

A straightforward way of constructing a content-based sparse graph is to select those keys that are likely to have large similarity scores with the given query. To efficiently construct the sparse graph, we can recur to Maximum Inner Product Search (MIPS) problem, where one tries to find the keys with maximum dot product with a query without computing all dot product terms. Routing Transformer (Roy et al., 2020) uses k-means clustering to cluster both queries $\{\mathbf{q}_{i}\}_{i=1}^{T}$ and keys $\{\mathbf{k}_{i}\}_{i=}^{T}$ on the same set of centroid vectors $\{\mathbf{\mu}_{i}\}_{i=1}^{k}$. Each query only attends to the keys that belong to the same cluster. During training, the cluster centroid vectors are updated using the exponentially moving average of vectors assigned to it, divided by the exponentially moving average of cluster counts:

where $|\mu|$ denotes the number of vectors currently in cluster $\mu$ and $\lambda\in(0,1)$ is a hyperparameter.

Let $\mathcal{P}_{i}$ denote the set of indices of keys that the $i$-th query attend to. $\mathcal{P}_{i}$ in Routing Transformer is defined as

Reformer (Kitaev et al., 2020) uses locality-sensitive hashing (LSH) to select key-value pairs for each query. The proposed LSH attention allows each token to attend only to the tokens within the same hashing bucket. The basic idea is to use an LSH function to hash queries and keys into several buckets, with similar items fall in the same bucket with high probability. Specifically, they use the random matrix method for the LSH function. Let $b$ be the number of buckets, given a random matrix $R$ of size $[D_{k},b/2]$, the LSH function is computed by :

The LSH attention allows the $i$-th query to attend only to key-value pairs with indices

Sparse Adaptive Connection (SAC) (Li et al., 2020b) views the input sequence as a graph and learns to construct attention edges to improve task-specific performances using an adaptive sparse connection. SAC uses an LSTM edge predictor to construct edges between tokens. With no ground truth for edges, the edge predictor is trained with reinforcement learning.

Sparse Sinkhorn Attention (Tay et al., 2020b) first splits queries and keys into several blocks and assigns a key block to each query block. Each query is only allowed to attend to the keys in the key block that is assigned to its corresponding query block. The assignment of key blocks is controlled by a sorting network, which uses Sinkhorn normalization to produce a doubly stochastic matrix as the permutation matrix representing the assignment. They use this content-based block sparse attention along with block local attention introduced in Sec. 4.1.1 to enhance the ability of the model to model locality.

2. Linearized Attention

Linearized attention is a class of methods that approximate or replace the unnormalized attention matrix $\exp(\mathbf{Q}\mathbf{K}^{\top})$ with $\phi(\mathbf{Q})\phi(\mathbf{K})^{\top}$, where $\phi$ is a feature map that is applied in row-wise manner. Hence the computation of unnormalized attention matrix can be linearized by computing \phi(\mathbf{Q})\left({\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi(\mathbf{K})^{\top}\mathbf{V}}\right)Similarly, the partition term $\mathbf{D}$ can be computed with \phi(\mathbf{Q})\left({\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\phi(\mathbf{K})^{\top}\mathbf{1}_{T}^{\top}}\right) in linear time., as illustrated in Fig. 7(b).

To gain further insights into linearized attention, we derive the formulation in vector form. We consider a general form of attention

where $\otimes$ denotes outer product of vectors. Based on this formulation, attention can be linearized by first computing the highlighted terms $\sum_{j}\phi(\mathbf{k}_{j})\otimes\mathbf{v}_{j}$ and$\sum_{j^{\prime}}\phi(\mathbf{k}_{j^{\prime}})^{\top}$. This could be especially beneficial for autoregressive attention, as the cumulative sums $\mathbf{S}_{i}=\sum_{j=1}^{i}\phi(\mathbf{k}_{j})\otimes\mathbf{v}_{j}$ and $\mathbf{u}_{i}=\sum_{j=1}^{i}\phi(\mathbf{k}_{j})$ can be computed from $\mathbf{S}_{i-1}$ and $\mathbf{u}_{i-1}$ in constant time. The effectively enables Transformer decoders to run like RNNs.

An interpretation of Eq. (16) is that the model maintains a memory matrix by aggregating associations represented by outer products of (feature mapped) keys and values, and then retrieve a value by multiplying the memory matrix with feature mapped query with proper normalization. There are two key components in this approach: (1) feature map $\phi(\cdot)$, and (2) aggregation rule.

Although the trigonometric random feature map leads to an unbiased approximation, it does not guarantee non-negative attention scores and thus could lead to unstable behaviors and abnormal behaviors. To mitigate this issue, the second version of Performer (Choromanski et al., 2020b) proposes positive random feature maps, which uses $h(\mathbf{x})=\exp(-\frac{\|\mathbf{x}\|^{2}}{2}),l=1,f_{1}=\exp$ and thus guarantees unbiased and non-negative approximation of dot-product attention. This approach is more stable than Choromanski et al. (2020a) and reports better approximation results.

2.2. Aggregation Rule

In Eq. (16) the associations $\{\phi(\mathbf{k})_{j}\otimes\mathbf{v}_{j}\}$ are aggregated into the memory matrix by simple summation. This is adopted by several studies (Katharopoulos et al., 2020; Choromanski et al., 2020a; Choromanski et al., 2020b). However, it could be more beneficial for the network to selectively drop associations as new associations are added to the memory matrix.

RFA (Peng et al., 2021) introduces a gating mechanism to the summation to model local dependency in sequence data. Specifically, when adding a new association to the memory matrix $\mathbf{S}$, at a particular time step, they weigh $\mathbf{S}$ by a learnable, input-dependent scalar $g$, and the new association by $(1-g)$ (and a similar mechanism to $\mathbf{u}$). With this modification, history associations are exponentially decayed and recent context is favored in each timestep.

Schlag et al. (2021) argue that simple summation limits the capacity of the memory matrix and thus propose to enlarge the capacity in a write-and-remove fashion. Specifically, given a new input key-value pair $(\mathbf{k}_{i},\mathbf{v}_{i})$, the model first retrieve the value $\bar{\mathbf{v}}_{i}$ currently associated with $\mathbf{k}_{i}$ using matrix multiplication. It then writes to the memory matrix a convex combination of $\bar{\mathbf{v}}_{i}$ and $\mathbf{v}_{i}$, using a input-dependent gating scalar $g$, and removes the association $\bar{\mathbf{v}}_{i}$. They also propose sum normalization (normalizing $\phi(\mathbf{q}_{i}),\phi(\mathbf{k}_{i})$ by the sum of their components before updating the memory matrix) instead of normalizing with the denominator in Eq. (16) for this aggregation rule.

3. Query Prototyping and Memory Compression

Apart from using sparse attention or kernel-based linearized attention, one could also reduce the complexity of attention by reducing the number of queries or key-value pairs, which leads to query prototyping and memory compressionThe key-value pairs are often referred to as a key-value memory (hence the name memory compression). methods, respectively.

In query prototyping, several prototypes of queries serve as the main source to compute attention distributions. The model either copies the distributions to the positions of represented queries or filling those positions with discrete uniform distributions. Fig. 8(a) illustrates the computing flow of query prototyping.

Clustered Attention (Vyas et al., 2020) groups queries into several clusters and then computes attention distributions for cluster centroids. All queries in a cluster share the attention distribution calculated with the corresponding centroid.

Informer (Zhou et al., 2021) selects prototypes from queries using explicit query sparsity measurement, which is derived from an approximation of the Kullback-Leibler divergence between the query’s attention distribution and the discrete uniform distribution. Attention distributions are then only calculated for the top-$u$ queries under query sparsity measurement. The rest of the queries are assigned with discrete uniform distributions.

3.2. Attention with Compressed Key-Value Memory

Apart from decreasing the number of queries with query prototyping, one can also reduce the complexity by reducing the number of the key-value pairs before applying the attention mechanism, as depicted in Fig. 8(b).

Liu et al. (2018) propose Memory Compressed Attention (MCA) that reduces the number of keys and values using a strided convolution. This modification is used as a complement to local attention proposed in the same work (as discussed in Sec. 4.1), in that it can capture global context. The mechanism reduces the number of keys and values by a factor of kernel size $k$ and thus allowing to process significantly longer sequences than vanilla Transformer given the same computation resources.

Set Transformer (Lee et al., 2019) and Luna (Ma et al., 2021) use a number of external trainable global nodes to summarize information from inputs and then the summarized representations serve as a compressed memory that the inputs attend to. This reduces the quadratic complexity of self-attention to linear complexity w.r.t. sequence length.

Linformer (Wang et al., 2020a) utilizes linear projections to project keys and values from length $n$ to a smaller length $n_{k}$. This also reduces the complexity of self-attention to linear. The drawback of this approach is that an input sequence length has to be assumed and hence it cannot be used in autoregressive attention.

Poolingformer (Zhang et al., 2021) adopts two-level attention that combines a sliding window attention and a compressed memory attention. The compressed memory module is used after the sliding window attention to increase the receptive field. They explore a few different pooling operations as the compression operation to compress the number of keys and values, including max pooling and pooling with Dynamic Convolution (Wu et al., 2019).

4. Low-rank Self-Attention

The fact that the rank of the attention matrix is less than sequence length implies that, for scenarios where the inputs are typically short, setting $D_{k}>T$ would be more than an over-parameterization and lead to overfitting. It is thus reasonable to limit the dimension of $D_{k}$ to explicitly model the low-rank property as an inductive bias. Guo et al. (2019b) decompose self-attention matrix into a low-rank attention module with small $D_{k}$ that captures long-range non-local interactions, and a band attention module that captures local dependencies.

4.2. Low-rank Approximation

Another implication of the low-rank property of the attention matrix is that one can use a low-rank matrix approximation to reduce the complexity of self-attention. A closely related methodology is the low-rank approximation of kernel matrices. We believe some existing works are inspired by kernel approximation.

Some of the aforementioned linearized attention methods in Sec. 4.2 are inspired from kernel approximation with random feature maps. For example, Performer (Choromanski et al., 2020a) follows the Random Fourier feature map originally proposed to approximate Gaussian kernels. The method first decomposes the attention distribution matrix $\mathbf{A}$ into $\mathbf{C}_{Q}\mathbf{G}\mathbf{C}_{K}$ where $\mathbf{G}$ is a Gaussian kernel matrix and the random feature map is used to approximate $\mathbf{G}$.

5. Attention with Prior

Some types of data (e.g., text) can exhibit a strong preference for the locality. This property can be explicitly encoded as a prior attention. A simple method would be to use a Gaussian distribution over positions. Specifically, one could multiply the generated attention distribution with some Gaussian density and then renormalize, which is equivalent to adding to the generated attention scores $\mathbf{A}$ a bias term $\mathbf{G}$, where higher $G_{ij}$ indicates a higher prior probability that the $i$-th input attend to the $j$-th input.

Yang et al. (2018) proposes to first predict a central position $p_{i}$ for each $\mathbf{q}_{i}$ using a simple feed-forward network. The Gaussian bias is then defined to be

where $\sigma$ denotes standard deviation for the Gaussian and can be determined as a hyperparameter or predicted from inputs.

Gaussian Transformer (Guo et al., 2019c) assumes the central position to be $i$ for each $\mathbf{q}_{i}$ and defines the bias to bes

where $w\geq 0,b\leq 0$ are scalar parameters that controls the deviation and reduce the weight for central position, respectively.

5.2. Prior from Lower Modules

In Transformer architecture, it is often observed the attention distributions are similar in adjacent layers. It is thus natural to provide attention distribution from previous layer as a prior for attention computation. The final attention scores can be defined as

Predictive Attention Transformer (Wang et al., 2021) proposes to apply a 2D-convolutional layer to previous attention scores and compute the final attention scores as a convex combination of the generated attention scores and the convolved scores. This is equivalent to setting $w_{1}=\alpha,w_{2}=1-\alpha$ and $g(\cdot)$ to be a convolutional layer in Eq. (22). They experiment training such a model from scratch and finetune after adapting the pre-trained BERT model, and both sets of experiments show improvements over baseline models.

Realformer (He et al., 2020b) uses adds the previous attention scores directly to the generated attention scores, thus resembles a residual skip connection on attention maps. It’s equivalent to setting $w_{1}=w_{2}=1$ and $g(\cdot)$ to be identity map in Eq. (22). They conduct pre-training experiments on this model. The results show that this model outperforms the baseline BERT model in multiple datasets and surpasses the baseline model even when pre-training budgets are significantly lower.

As an extreme case, Lazyformer (Ying et al., 2021) proposes to share attention maps between a number of adjacent layers. This is equivalent to setting $g(\cdot)$ to identity and switch the settings of $w_{1}=0,w_{2}=1$ and $w_{1}=1,w_{2}=0$ alternatingly. The benefit of this approach is that the attention maps are computed only once and reused several times in the succeeding layers, thus reducing the computation cost. Their pre-training experiments show that the resulting model remains effective while being much more efficient to compute.

5.3. Prior as Multi-task Adapters

5.4. Attention with Only Prior

Some works have explored using an attention distribution that is independent of pair-wise interaction between inputs. In other words, their models exploit only a prior attention distribution.

Zhang et al. (2018) design an efficient Transformer decoder variant called average attention network that uses a discrete uniform distribution as the sole source of attention distribution. The values are thus aggregated as a cumulative-average of all values. To improve the expressiveness of the network, they further adds a feed-forward gating layer on top of the average attention module. The advantage of this approach is that the adapted Transformer decoder can train in a parallel manner as usual Transformers do and decode like an RNN, thus avoiding the $\mathcal{O}(T^{2})$ complexity in decoding.

You et al. (2020) utilize a Gaussian distribution as the hardcoded attention distribution for attention calculation. The intuition is very similar to Yang et al. (2018) and Guo et al. (2019c) in that attention distribution should be focused on a certain local window. Distinctively, they drop the generated attention completely and use only the Gaussian distribution for attention computation. In this approach, the mean (central position) and variance are designed to be hyperparameters. The experiments show that the hardcoded attention, when applied only to self-attention, can achieve comparable performance to the baseline model in machine translation tasks.

Synthesizer (Tay et al., 2020a) proposes to replace generated attention scores with: (1) a learnable, randomly initialized attention scores, and (2) attention scores output by a feed-forward network that is only conditioned on the querying input itself. The experiments on machine translation and language modeling show that these variants can achieve competitive performance with vanilla Transformer. It is not explained why these variants work but the empirical results are intriguing.

6. Improved Multi-Head Mechanism

Multi-head attention is appealing for the ability to jointly attend to information from different representation subspaces at different positions. However, there is no mechanism to guarantee that different attention heads indeed capture distinct features.

A basic motivation for using multi-head attention is to allow the model to jointly attend to information from different representation subspaces at different positions (Vaswani et al., 2017). However, in vanilla Transformer there is no explicit mechanism to guarantee different behavior across attention heads, nor is there any mechanism for heads to interact with each other. A line of work is dedicated to improving multi-head mechanism by introducing incorporating more sophisticated mechanisms that guide the behavior of different attention heads or allow interaction across attention heads.

Li et al. (2018) introduce an auxiliary disagreement regularization term into loss function to encourage diversity among different attention heads. Two regularization terms are respectively to maximize cosine distances of the input subspaces and output representations, while the last one is to disperse the positions attended by multiple heads with element-wise multiplication of the corresponding attention matrices.

Several probing works have revealed that pre-trained Transformer models exhibit certain patterns of self-attention that are of little linguistic backing. As a representative work, Kovaleva et al. (2019) identify several simple attention patterns in BERT. For instance, many of the attention heads simply pay attention to special BERT tokens [CLS] and [SEP]. As a result, some constraints can be introduced to boost the training of Transformer models. To this end, Deshpande and Narasimhan (2020) propose to use an auxiliary loss, which is defined to be the Frobenius norm between attention distribution maps and predefined attention patterns.

Talking-head Attention (Shazeer et al., 2020) uses a talking head mechanism that linearly projects the generated attention scores from $h_{k}$ to $h$ heads, applies softmax in that space, and then projects to $h_{v}$ heads for value aggregation. The motivation is to encourage the model to move information between attention heads in a learnable fashion.

Collaborative Multi-head Attention (Cordonnier et al., 2020) uses shared query and key projection $\mathbf{W}^{Q}$ and $\mathbf{W}^{K}$ and a mixing vector $\mathbf{m}_{i}$ for the $i$-th head to filter from the projection parameters such that Eq. (3) is adapted to

where $\mathbf{W}^{Q}$ and $\mathbf{W}^{K}$ are shared by all the attention heads.

6.2. Multi-head with Restricted Spans

Vanilla attention adopts full attention spans assume, where a query can attend to all of the key-value pairs. However, it is often observed that some heads focus their attention distribution mainly in a local context while some other heads attend to broader contexts. It could thus be beneficial to restrict the attention spans:

Locality. Restricting attention spans induce explicit local constraints. This is advantageous in cases where locality is an important prior.

Efficiency. If implemented appropriately, such a model can scale to very long sequences without introducing additional memory footprint and computational time.

Restricting attention spans can be expressed as multiplying each attention distribution value with a mask value and then re-normalize, where the mask can be expressed as a non-increasing function that maps a distance to a value in $$. A vanilla attention assigns a mask value of$1$ for all distances, as depicted in Fig. 10(a).

Sukhbaatar et al. (2019a) propose to use a learnable attention span, as depicted in Fig. 10(b) . The mask is parameterized by a learnable scalar $z$ and a hyperparameter $R$. The experiments on character-level language modeling show that the adaptive-span models outperform baseline models while having significantly fewer FLOPS. It is also observed that lower layers generally have smaller learned spans and higher layers otherwise. This indicates that the model can learn a hierarchical composition of features.

Multi-Scale Transformer (Guo et al., 2020) proposes to use a fixed attention span, with different heads in different layers using a different max span. The fixed attention span is depicted in Fig. 10(c). The attention is restricted within a fixed window which is controlled by a scale value $w$. They design the scales from an intuitive linguistic perspective and empirical observation from BERT such that higher layers tend to have more large scales (e.g., large span size), and lower layers should be confined with a smaller scale. Their experiments on several tasks show that the model can outperform baseline models while accelerating inference on long sequences.

6.3. Multi-head with Refined Aggregation

where each $\mathbf{W}_{i}^{O}$ is of dimension $D_{v}\times D_{m}$. It’s thus easy to see that multi-head attention can be reformulated as

One might argue that this simple aggregate-by-summation paradigm does not fully exploit the expressiveness of multi-head attention and that it is more desirable to use a more complex aggregation.

Li et al. (2019b); Gu and Feng (2019) propose to use routing methods, originally proposed for capsule networks (Sabour et al., 2017), to further aggregate information produced by different attention heads. The outputs of attention heads are first transformed into input capsules, then output capsules are obtained after the iterative routing process. The output capsules are then concatenated as a final output of multi-head attention. These two works both utilizes two routing mechanisms, namely dynamic routing(Sabour et al., 2017) and EM routing(Hinton et al., 2018). One would notice that iterative routing introduces additional parameters and computational overhead. Li et al. (2019b) empirically show that applying the routing mechanism only to the lower layers can best balance the translation performance and computational efficiency.

6.4. Other Modifications

Several other modifications to the multi-head mechanism have been proposed to improve multi-head attention.

Shazeer (2019) propose multi-query attention, where key-value pairs are shared among attention heads (i.e., to use only one key projection and one value projection for all attention heads). The advantage of this method is that it reduces the memory bandwidth requirements for decoding and results in a model that is faster to decode, while incurring only minor quality degradation from the baseline.

Bhojanapalli et al. (2020) establish that small attention key size can affect its ability to represent arbitrary distribution. They thus propose to disentangle head size from the number of heads $h$, as opposed to the common practice that sets the head size to be $D_{m}/h$. It is observed empirically that setting attention head size to be input sequence length is beneficial.

Other Module-level Modifications

Let $\Pi_{n}$ be the set of all permutations of indices $\{1,2,\cdots,T\}$. A function $f:\mathcal{X}^{T}\rightarrow\mathcal{Y}^{T}$ is said to be permutation equivariant if and only if for any $\pi\in\Pi_{T}$

It is easy to verify that Convolution and Recurrence networks are not permutation equivariant. However, both self-attention modules and position-wise feed-forward layers in Transformer are permutation equivariant, which could be a problem when it comes to modeling problems other than set-input problems where the structure of inputs is needed. For example, when modeling sequences of text, the ordering of words matters and it’s thus crucial to properly encode the positions of words in Transformer architecture. Therefore, additional mechanisms are required to inject positional information into Transformers. A common design is to first represent positional information using vectors and then infuse the vectors to the model as an additional input.

where $\omega_{i}$ is the hand-crafted frequency for each dimension. The position encoding of each position in the sequence is then added to the token embeddings and fed to Transformer.

Another way of representing absolute positions is to learn a set of positional embeddings for each position (Gehring et al., 2017; Devlin et al., 2019). Compared to hand-crafted position representation, learned embeddings are more flexible in that position representation can adapt to tasks through back-propagation. But the number of embeddings is limited up to a maximum sequence length determined before training, which makes this approach no longer inductive, i.e., not able to handle sequences longer than sequences seen in the training time(Liu et al., 2020b; Chu et al., 2021).

Wang et al. ([n.d.]) propose to use sinusoidal position representation, but with each frequency $\omega_{i}$ (in Eq. (28)) learned from data. This approach retains inductiveness but is more flexible than hand-crafted sinusoidal encoding. FLOATER (Liu et al., 2020b) frames positional representation as a continuous dynamical system and adopts Neural ODE to enable end-to-end training with backpropagation. This method is inductive and flexible while being parameter efficient compared to a fully learnable approach.

The Vanilla approach to incorporating absolute position representations is to add position encodings/embeddings to token embeddings. However, as the input signals propagate through the layers, the positional information might get lost in the upper layers. Later works find it beneficial to add position representations to inputs to each Transformer layer (Al-Rfou et al., 2019; Dehghani et al., 2019; Liu et al., 2020b; Guo et al., 2019b).

1.2. Relative Position Representations

Another line of works focuses on representing positional relationships between tokens instead of positions of individual tokens. The intuition is that in self-attention, pairwise positional relationships between input elements (direction and distance) could be more beneficial than positions of elements. Methods following this principles are called relative positional representation. Shaw et al. (2018) propose to add a learnable relative position embedding to keys of attention mechanism

Transformer-XL (Dai et al., 2019) use a sinusoidal encoding to represent positional relationships but fuses contents and position information by redesign the computation of attention scoresthe scaling factor is omitted without loss of generality.

DeBERTa (He et al., 2020a) utilizes position embeddings like Shaw et al. (2018) and applies the embeddings to the model in a disentangled style similar to Transformer-XL (Dai et al., 2019)

1.3. Other Representations

Some research studies have explored using hybrid positional representations that contains both absolute and relative positional information. Transformer with Untied Position Encoding (TUPE) (Ke et al., 2020) re-designs the computation of attention scores as a combination of a content-to-content term, an absolute position-to-position term and a bias term representing relative positional relationships

One can also design a single set of positional representations that express both absolute and relative information. Roformer (Su et al., 2021) uses Rotary Position Embedding (RoPE) to represent the position of a token by multiplying the affine-transformed embedding of the $t$-th input $x_{t}$ by a rotatory matrix $\mathbf{R}_{\Theta,t}$

where $\bigoplus$ denotes direct sum of matrices. Each $\mathbf{M}(t,\theta_{j})$ is a 2-D clockwise rotatory matrix of angle $t\cdot\theta_{j}$

The key advantage of this formulation is that the induced representation is translation invariant, i.e., the attention score of $(\mathbf{q}_{i},\mathbf{k}_{j})$ is only related to their relative position offset

In practice, the embedding matrix multiplication can be implemented by two element-wise multiplication for lower memory footprint. The RoPE uses the form of absolute embedding but can capture relative positional relations. This approach is compatible with linearized attention in Sec. 4.2.

1.4. Position Representations without Explicit Encoding

Instead of explicitly introducing additional positional encodings, Wang et al. (2020b) propose to encode positional information in word embeddings, by generalizing embedding to continuous (complex-valued) functions over positions.

R-Transformer (Wang et al., 2019b) model locality of sequential data with a local RNN. Specifically, inputs to each block of R-Transformer are first fed to a local RNN and then to multi-Head self-attention module. The RNN structure introduces ordering information and captures local dependencies as a complement to self-attention.

Conditional positional encoding (CPE) (Chu et al., 2021) generate conditional position encodings at each layer for ViT with a 2-D convolution with zero-paddings. The intuition behind this approach is that convolution networks can implicitly encode absolute positional information with zero-paddings (Islam et al., 2020).

1.5. Position Representation on Transformer Decoders

It is worth noticing that masked self-attention is not permutation equivariant (Tsai et al., 2019). Thus a model that exploits only the decoder of Transformer has the potential of sensing positional information without incorporating explicit positional representation. This is confirmed by some empirical results on language modeling tasks (Irie et al., 2019; Schlag et al., 2021), where the authors find that removing position encodings even improves performance.

2. Layer Normalization

Layer Normalization (LN), along with residual connection, is considered as a mechanism to stabilizing training of deep networks (e.g., alleviating ill-posed gradients and model degeneration). There are some studies that are dedicated to analyzing and improving LN module.

In vanilla Transformer, the LN layer lies between the residual blocks, called post-LN (Wang et al., 2019a). Later Transformer implementations (Vaswani et al., 2018; Klein et al., 2017) place the LN layer inside the residual connection before the attention or FFN, with an additional LN after the final layer to control the magnitude of final outputs, which is referred to as pre-LNTo the best of our knowledge, this approach is adopted since v1.1.7 in the Tensor2Tensor implementation (Vaswani et al., 2018).. The pre-LN has been adopted by numerous following research studies and implementations, e.g., (Baevski and Auli, 2019; Child et al., 2019; Wang et al., 2019a). The difference between pre-LN and post-LN is shown in Fig. 11.

Xiong et al. (2020) theoretically investigate the gradients of Transformers and find that the gradients near the output layer are large at initialization in post-LN Transformers, which could be the reason why post-LN Transformers without learning rate warm-up (Vaswani et al., 2017)Learning rate warm-up refers to starting optimization with an extremely small learning rate and then gradually increasing it to a pre-defined maximum value in a certain number of iterations. leads to unstable training, whereas pre-LN Transformers do not suffer from the same problem. They thus deduce and empirically verify that warm-up stage can be safely removed for pre-LN Transformers.

Although Post-LN often results in unstable training and divergence, it usually outperforms pre-LN variants after convergence (Liu et al., 2020a). Similar to Xiong et al. (2020), Liu et al. (2020a) conduct theoretical and empirical analysis and find that post-LN encoders do not suffer from gradient imbalance. They thus conjecture that the gradient issue is not the direct cause of unstable post-LN Transformer training and further identify the amplification effect in post-LN Transformers — at initialization, the heavier dependency on residual branch leads to a larger output shift in post-LN Transformers, thus resulting in unstable training. In light of this finding, they introduce additional parameters to post-LN Transformers to control residual dependencies of Post-LN. These parameters are initialized according to activation variations of sample data so that the output shift of post-LN Transformers is not amplified. This approach ensures and boosts convergence of post-LN Transformers and reaches better performance than pre-LN Transformers.

2.2. Substitutes of Layer Normalization

Xu et al. (2019) empirically observe that the learnable parameters in the LN module do not work in most experiments, and even increase the risk of overfitting. They further conclude from controlled experiments that the forward normalization is not the reason why LN works for Transformer. From analysis and experiments, it is concluded that the derivatives of the mean and variance re-center and re-scale the gradients and play a significant role in LN. They thus propose AdaNorm, a normalization technique without learnable parameters

where $C,k$ are hyperparameters and $\odot$ denotes element-wise multiplication. $\mu$ and $\sigma$ are the mean and standard deviation of input $\mathbf{x}$, respectively.

where $g$ is a learnable scalar. It is more parameter efficient compared to normal LN and is shown to be effective in machine translation datasets, especially in low-resource settings.

Shen et al. (2020) discuss why Batch Normalization (BN) (Ioffe and Szegedy, 2015) performs poorly in Transformer for text data and conclude that BN’s significant performance degradation stems from the instabilities associated with its batch statistics. They thus propose PowerNorm (PN) that has three modifications over BN: (1) it relaxes the zero-mean normalization; (2) it uses the quadratic mean of the signal, instead of the variance; (3) it uses running statistics for the quadratic mean, instead of using per-batch statistics. Specifically, for the $t$-th iteration, the PN computes the outputs as

where $0<\alpha<1$ is the moving average coefficient and $\gamma,\beta$ are the learnable parameters as in BN formulation.

2.3. Normalization-free Transformer

Besides LN, there is another mechanism to construct deeper neural network. ReZero (Bachlechner et al., 2020) replace LN module with a learnable residual connection. For each module $F(\cdot)$, ReZero re-scales $F(\cdot)$ in the residual formulation:

where $\alpha$ is a learnable parameter with zero-initialization.

Replacing LN in Transformer with ReZero mechanism is verified to induce better dynamic isometry for input signals and leads to faster convergence.

3. Position-wise FFN

Despite its simplicity, the position-wise feed-forward network (FFN) layers are important for a Transformer to achieve good performance. Dong et al. (2021) observe that simply stacking self-attention modules causes a rank collapse problem, leading to token-uniformity inductive bias, and that the feed-forward layer is one of the important building blocks that mitigate this issue. Various works have explored modifications on the FFN module.

The vanilla Transformer (Vaswani et al., 2017) adopts the Rectified Linear Units (ReLU) activation for non-linearity in between the two FFN layers. Over time, several studies have explored different activation other than ReLU.

GPT (Radford et al., 2018) replace ReLU with Gaussian Error Linear Unit (GELU) (Hendrycks and Gimpel, 2020) on language pre-training. It becomes the default practice for many pre-trained language models (Devlin et al., 2019; He et al., 2020a).

Shazeer (2020) explore using Gated Linear Units (GLU) (Dauphin et al., 2017) and its variants as a drop-in replacement for ReLU in FFN. Their pre-training experiments show that the GLU variants consistently improve vanilla Transformer with ReLU activation. Note that GLU introduces extra parameters and the experiments are conducted with the intermediate dimension of FFN reduced to match the parameter count with baseline.

3.2. Adapting FFN for Larger Capacity

Several works have focused on expanding FFNs in order for a larger model capacity. The basic idea is to replace FFNs with similar structures with much more parameters.

Lample et al. (2019) replace some of the FFNs with the product-key memory layers. A product-key memory is composed of three components: a query network, a key selection module containing two sets of sub-keys, and a value lookup table. The model first projects an input to a latent space using the query network, and then compares the generated query to keys that are Cartesian product of the two sets of sub-keys from key selection module to get $k$ nearest neighbors, and finally finds the corresponding values in a value lookup table using the $k$ nearest keys and aggregates them to produce the final output. This process resembles the attention mechanism, in that the generated query attends to a large number of global key-value pairs. They thus propose a multi-head mechanism for the key-product memory to further enlarge the capacity of this module. The experiments on large-scale language modeling suggest that this mechanism significantly improves performance with negligible computational overhead.

Several studies exploits the idea of Mixture-of-Experts (MoE)(Shazeer et al., 2017) to increase the capacity of FFNs. Gshard(Lepikhin et al., 2020) uses sparsely-gated MoE layers to replace FFNs in Transformer. Each MoE layer consists of several FFNs (each called an expert) that are the same structure as position-wise FFNs in vanilla Transformer. The output of the layer is a weighted sum of the outputs of the FFNs, using gate values computed by a routing function $g(\cdot)$. They design a learnable routing function that assigns tokens to experts, with auxiliary loss to satisfy balanced loads between experts and efficiency at the scale of length such that the experts can be distributed across multiple devices. For each forward pass of the MoE layer, only the experts with top-$k$ gate values are activated.

Instead of using $k$ experts for each forward pass, Switch Transformer (Fedus et al., 2021) proposes to route using only a single expert with the largest gate value, leading to a much smaller computational footprint. The authors also design an auxiliary loss to encourage load balance between experts. It is reported to speed up pre-training by a large margin compared to the non-MoE counterpart while having a similar number of FLOPS.

Yang et al. (2021) propose to replace top-$k$ routing with expert prototyping strategy. Specifically, the proposed strategy splits experts into $k$ different groups and applies top-1 routing within each group. The outputs of prototype groups are combined linearly to form the final output of the MoE layer. This strategy is proved to improve the model quality while maintaining constant computational costs.

As opposed to using a learnable routing function for expert assignment, Roller et al. (2021) design hash layers where tokens are hashed into a fixed number of buckets, each bucket corresponding to an expert. This approach requires no routing parameters or any auxiliary loss function, while showing competitive results with existing methods such as Switch Transformer (Fedus et al., 2021).

3.3. Dropping FFN Layers

Notably, one might argue that under some circumstances, FFN layers can be dropped completely, resulting in a simplified network.

Yang et al. (2020) empirically show that FFNs in the decoder of Transformer, despite its large number of parameters, is not efficient and can be removed safely with only slight or no loss of performance. This approach significantly boosts the training and inference speed.

Architecture-level Variants

In this section, we introduce the X-formers that modify the vanilla Transformer beyond modules.

Apart from the efforts made at the module level to alleviate computation overheads, there are several attempts to adapt Transformer to be lightweight by modifications at a higher level.

Similar to low-rank self-attention (Guo et al., 2019b) that decomposes attention into a locality-constrained attention and a low-rank global attention, Lite Transformer (Wu et al., 2020b) proposes to replace each attention module in Transformer with a two-branch structure, where one branch uses attention to capture long-range contexts while the other branch uses depth-wise convolution and linear layers to capture local dependencies. The architecture is lightweight both in terms of model size and computation, and is thus more suitable for mobile devices.

Funnel Transformer (Dai et al., 2020) utilizes a funnel-like encoder architecture where the length of the hidden sequence is gradually reduced using pooling along the sequence dimension, and then recovered using up-sampling. The architecture effectively reduces the FLOPs and memory compared to the vanilla Transformer encoder. Naturally, one can use this architecture to build a deeper or wider model using the same computation resources.

DeLighT (Mehta et al., 2020) replaces the standard Transformer block with DeLighT block, which consists of three sub-modules: (1) a “expand-and-reduce” DeLighT transformation module to learn wider representations with low computation requirements; (2) a single-head self-attention to learn pair-wise interaction; (3) a lightweight “reduce-and-expand” FFN (as opposed to vanilla Transformer that first expands the dimension of hidden representations and then reduces them back to $D_{m}$). They also propose a block-wise scaling strategy that allows for shallower and narrower blocks near the input and wider and deeper blocks near the output. The induced network is much deeper than the vanilla Transformer but with fewer parameters and operations.

2. Strengthening Cross-Block Connectivity

In vanilla Transformer, each block takes outputs from the previous block as inputs and outputs a sequence of hidden representations. One might be interested in creating more paths along which input signals can run through the networks. In Sec. 4.5.2, we introduced Realformer (He et al., 2020b) and Predictive Attention Transformer (Wang et al., 2021) that reuses attention distributions from previous block to guide attention of current block. This can be seen as creating a forward path between adjacent Transformer blocks.

In a deep Transformer encoder-decoder model, the cross-attention modules in the decoder only utilize the final outputs of the encoder, therefore the error signal will have to traverse along the depth of the encoder. This makes Transformer more susceptible to optimization issues (e.g., vanishing gradients). Transparent Attention (Bapna et al., 2018) uses a weighted sum of encoder representations at all encoder layers (including the embedding layer) in each cross-attention module. For the $j$-th decoder block, the cross-attention module is modified to attend to

where each $w_{ij}$ is a trainable parameter. This effectively shortens the path from each layer in the encoder to the error signal and thus eases the optimization of deeper Transformer models.

Another issue associated with vanilla Transformer is that each position can only attend to history representations from lower layers. Feedback Transformer (Fan et al., 2021b) proposes to add a feedback mechanism to Transformer decoder, where each position attends to a weighted sum of history representations from all layers

3. Adaptive Computation Time

Vanilla Transformer, like most neural models, utilizes a fixed (learned) computation procedure to process each input. An intriguing and promising modification is to make computation time conditioned on the inputs, i.e., to introduce Adaptive Computation Time (ACT) (Graves, 2016) into Transformer models. Such modifications potentially give rise to the following advantages:

Feature refinement for hard examples. For data that are hard to process, a shallow representation might not be adequate to fulfill the task at hand. It would be more ideal to apply more computations to acquire a deeper and more refined representation.

Efficiency for easy examples. When processing easy examples, a shallow representation might be enough for the task. In this case, it would be beneficial if the network can learn to extract features using reduced computation time.

Universal Transformer (UT) (Dehghani et al., 2019) incorporates a recurrence-over-depth mechanism that iteratively refines representations for all symbols using a module that is shared over depth, as illustrated in Fig. 12(a). It also adds a per-position dynamic halting mechanism that calculates a halting probability for each symbol at every time step. If a symbol’s halting probability is greater than a predefined threshold, then the symbol’s representation will remain unchanged for subsequent timesteps. The recurrence is stopped when all symbols halt or when a predefined maximum step is reached.

Conditional Computation Transformer (CCT) (Bapna et al., 2020) adds a gating module at each self-attention and feed-forward layer to decide whether to skip the current layer, as illustrated in Fig. 12(b). The authors also introduce an auxiliary loss that encourages the model to adjust the gating modules to match the practical computation cost to the available computation budget.

Similar to the dynamic halting mechanism used in UT, there is a line of work dedicated to adapting the number of layers to each input in order to achieve a good speed-accuracy trade-off, which is called early exit mechanism, as illustrated in Fig. 12(c). A commonly used technique is to add an internal classifier at each layer and jointly train all classifiers. The core of these methods is the criteria used to decide whether to exit at each layer. DeeBERT (Xin et al., 2020) uses the entropy of the output probability distribution of the current layer to determine whether to exit. PABEE (Zhou et al., 2020) counts the number of times that the predictions remain unchanged to decide whether to exit. Li et al. (2021) design a window-based uncertainty criterion to achieve token-level partial exiting for sequence labeling tasks. Sun et al. (2021) introduces a voting-based exiting strategy that considers at each layer predictions of all the past internal classifiers to infer the correct label and to decide whether to exit.

4. Transformers with Divide-and-Conquer Strategies

The quadratic complexity of self-attention on sequences length can significantly limit the performance of some downstream tasks. For example, language modeling usually needs long-range context. Apart from the techniques introduced in Sec. 4, another effective way of dealing with long sequences is to use divide-and-conquer strategy, i.e., to decompose an input sequence into finer segments that can be efficiently processed by Transformer or Transformer modules. We identify two representative class of methods, recurrent and hierarchical Transformers, as illustrated in Fig. 13. These techniques can be understood as a wrapper for the Transformer model in which Transformer acts as an elementary component that is reused to process different input segments.

In recurrent Transformers, a cache memory is maintained to incorporate the history information. While processing a segment of text, the network reads from the cache as an additional input. After the processing is done, the network writes to the memory by simply copying hidden states or using more complex mechanisms. The abstract process is illustrated in Fig. 13(a).

Transformer-XL (Dai et al., 2019) address the limitation of a fixed length context by caching representations from the previous segment and reuse it as an extended context when the model processes the current segment. For the $l$-th layer and the $(\tau+1)$-th segment, the input representation $\mathbf{H}_{\tau+1}^{(l-1)}$ is concatenated with the representation $\mathbf{H}_{\tau}^{(l-1)}$ from previous segment to produce the keys and values

Compressive Transformer (Rae et al., 2020) extends this idea further by extending the cache with two levels of memory. In Transformer-XL, the activations from the previous segment are cached as a memory that is used to augment the current segment, and activations from older segments are discarded. Compressive Transformer, on the other hand, applies a compression operation (e.g., Convolution, Pooling, etc.) on older activations and stores them in the compressed memory. In order to avoid the expensive backpropagating-through-time (BPTT) from training compression sub-network with gradients from the loss, they propose to use local loss functions where original memories are constructed from the compressed memories. This approach further extends the theoretical maximum history context length from $L\times N_{\text{mem}}$ of Transformer-XL to $L\times(N_{\text{mem}}+c\times N_{\text{cm}})$, where $c$ is the compression rate and $N_{\text{cm}}$ is the length of compressed memory.

Memformer (Wu et al., 2020a) extends the recurrence mechanism from decoder-only architecture to an encoder-decoder architecture. They introduce to the encoder a memory cross attention similar to the cross attention in vanilla Transformer to allow the Transformer encoder to attend to the memory. They also introduce a memory slot attention on top of the encoder output to explicitly write the memory for the next segment. To avoid BPTT over a long range of timesteps, they propose Memory Replay Back-Propagation (MRBP) algorithm, which replays the memory at each timestep to accomplish gradient back-propagation over long unrolls.

Yoshida et al. (2020) propose a simple fine-tuning mechanism to add recurrence to a pre-trained language model (e.g., GPT-2 (Radford et al., 2019)). They first compress the representations produced by the $\tau$-th segment into one single vector representation, using a weighted average of pooled representations from each layer $l\in\{1,\cdots,L\}$

ERNIE-Doc (Ding et al., 2020) proposes an enhanced recurrence mechanism based on the recurrence mechanism used in Transformer-XL, by replacing the memory with the history representations from the $l$-th layer.

as opposed to using representations from the $(l-1)$-th layer in Eq. (48). This modification essentially leads to a larger effective context length.

4.2. Hierarchical Transformers

Hierarchical Transformer decomposes inputs hierarchically into elements of finer granularity. Low-level features are first fed to a Transformer encoder, producing output representations that are then aggregated (using pooling or other operations) to form a high-level feature, which is then processed by a high-level Transformer. This class of methods can be understood as a process of hierarchical abstraction. The overview of this approach is depicted in Fig. 13(b). The advantages of this approach are twofold: (1) Hierarchical modeling allows the model to handle long inputs with limited resources; (2) It has the potential to generate richer representations that are beneficial to tasks.

For tasks with inherently long input length, one can use hierarchical Transformers for effective modeling of long-range dependencies. For document-level machine translation tasks, Miculicich et al. (2018) introduce dependencies on the previous sentences from both the source and target sides when translating a sentence. They use an attention mechanism as the aggregation operation to summarize low-level information. For document summarization, HIBERT (Zhang et al., 2019) encodes a document of text by first learn sentence representations for all sentences and then use these sentence representations to encode document-level representations that are then used to generate the summary. The model uses the last hidden representation (corresponding to the EOS token) as the representation for each sentence. Liu and Lapata (2019) propose a similar hierarchical Transformer for multi-document summarization where the extracted low-level representations are aggregated using an attention layer with a global trainable query node and low-level representations as the source of key-value pairs. Hi-Transformer (Wu et al., 2021b) first utilizes a sentence Transformer and a document Transformer to hierarchically learn document context-aware sentence representations. The document context-aware sentence representations are then fed to another sentence Transformer to further improve the sentence context modeling.

One might also be interested in using hierarchical models to acquire richer representations that are beneficial to the tasks at hand. For example, TENER (Yan et al., 2019) uses a low-level Transformer encoder to encode character features, which is then concatenated with word embeddings as the inputs to the high-level Transformer encoder. This incorporates more features and alleviates the problems of data sparsity and out-of-vocabulary (OOV). Recently emerging Vision Transformer (Dosovitskiy et al., 2020) divides an input image into several patches that serve as the basic input elements of Transformer, which potentially loses intrinsic pixel-level information within patches. To address this issue, Transformer in Transformer (TNT) (Han et al., 2021c) uses at each layer an inner Transformer block that transforms pixel representations and an outer Transformer block that takes fused vectors of patch representations and pixel representations as input.

5. Exploring Alternative Architecture

Despite the success of Transformer architecture, one might question whether the current Transformer architecture is optimal. Interestingly, several studies have explored alternative architectures for Transformer.

Lu et al. (2020) interpret Transformer as a numerical Ordinary Differential Equation (ODE) solver for a convection-diffusion equation in a multi-particle dynamic system and design Macaron Transformer, which replaces each Transformer block with a FFN-attention-FFN variant.

Sandwich Transformer (Press et al., 2020) explores reorganizing attention modules and FFN modules such that attention modules are mainly located in lower layers and FFN modules in upper layers. The induced model improves perplexity on multiple language modeling benchmarks, without increasing parameters, memory or training time.

Mask Attention Network (MAN) (Fan et al., 2021a) prepends a dynamic mask attention module to the self-attention module in each Transformer block. The mask is conditioned on token representations, the relative distance between tokens and head indices. The proposed dynamic mask attention is shown to effectively model locality in text data and the induced model consistently outperforms the baseline model in machine translation and abstractive summarization.

Notably, there’s a line of work that uses Neural Architecture Search (NAS) to search for alternative Transformer architectures. The Evolved Transformer (ET) (So et al., 2019) employs evolution-based architecture search with the standard Transformer architecture seeding the initial population. The searched model demonstrates consistent improvement over Transformer on several language tasks. As another representative work, DARTSformer(Zhao et al., 2021) applies differentiable architecture search (DARTS) (Liu et al., 2019b), combined with a multi-split reversible network and a backpropagation-with-reconstruction algorithm for memory efficiency. The resulting model consistently outperforms standard Transformer and compares favorably to larger ET models, with a significantly reduced search cost.

Pre-trained Transformers

As a key difference from convolutional networks and recurrent networks that inherently incorporates the inductive bias of locality, Transformer does not make any assumption about how the data is structured. On the one hand, this effectively makes Transformer a very universal architecture that has the potential of capturing dependencies of different ranges. On the other hand, this makes Transformer prone to overfitting when the data is limited. One way to alleviate this issue is to introduce inductive bias into the model.

Recent studies suggest that Transformer models that are pre-trained on large corpora can learn universal language representations that are beneficial for downstream tasks (Qiu et al., 2020). The models are pre-trained using various self-supervised objectives, e.g., predicting a masked word given its context. After pre-training a model, one can simply fine-tune it on downstream datasets, instead of training a model from scratch. To illustrate typical ways of using Transformers in pre-training, we identify some of the pre-trained Transformers and categorize them as follows.

Encoder only. A line of work uses the Transformer encoder as its backbone architecture. BERT (Devlin et al., 2019) is a representative PTM that is typically used for natural language understanding tasks. It utilizes Masked Language Modeling (MLM) and Next Sentence Prediction (NSP) as the self-supervised training objective. RoBERTa (Liu et al., 2019a) further adapts the training of BERT and removes the NSP objective as it is found to hurt performance on downstream tasks.

Decoder only. Several studies focus on pre-training Transformer decoders on language modeling. For example, the Generative Pre-trained Transformer (GPT) series (i.e., GPT (Radford et al., 2018), GPT-2 (Radford et al., 2019), and GPT-3 (Brown et al., 2020)) is dedicated to scaling pre-trained Transformer decoders and has recently illustrated that a large-scale PTM can achieve impressive few-shot performance with the task and examples fed to the model as constructed prompts (Brown et al., 2020).

Encoder-Decoder. There are also PTMs that adopt Transformer encoder-decoder as the overall architecture. BART (Lewis et al., 2020) extends the denoising objective of BERT to encoder-decoder architecture. The benefit of using an encoder-decoder architecture is that the inducing model is equipped with the ability to perform both natural language understanding and generation. T5 (Raffel et al., 2020) adopts similar architecture and was one of the earliest studies that use task-specific text prefix in downstream tasks.

Some of the Transformer architecture variants can also be applied to Transformer-based PTMs. For instance, BigBird (Zaheer et al., 2020) introduced in Sec. 4.1 is a encoder-based PTM that uses compound position-based sparse attention to enable long sequence inputs. GPT-3 (Brown et al., 2020) uses alternating dense and locally banded sparse attention (which was also introduced in Sec. 4.1) in self-attention modules. Switch Transformer (Fedus et al., 2021) is an encoder-based PTM that replaces FFN layers with mixture-of-experts layers and can increase parameter count while keeping the FLOPs per example constant.

Applications of Transformer

Transformer was originally designed for machine translation but has been widely adopted in various fields besides NLP, including CV and audio processing, due to its flexible architecture.

(1) Natural Language Processing. Transformer and its variants have been extensively explored and applied in NLP tasks, e.g., machine translation (Vaswani et al., 2017; Mehta et al., 2020; Raffel et al., 2020; So et al., 2019; Fan et al., 2021a), language modeling (Dai et al., 2019; Shoeybi et al., 2020; Roy et al., 2020; Rae et al., 2020) and named entity recognition (Yan et al., 2019; Li et al., 2020c). Massive effort has been dedicated to pre-training Transformer models on large-scale text corpora, which we believe is one of the major reasons of Transformer’s wide application in NLP.

(2) Computer Vision. Transformer have also been adapted for various vision tasks, e.g., image classification (Chen et al., 2020b; Dosovitskiy et al., 2020; Liu et al., 2021), object detection (Carion et al., 2020; Zhu et al., 2020; Zheng et al., 2020a; Liu et al., 2021), image generation (Parmar et al., 2018; Jiang et al., 2021) and video processing (Shao et al., 2021; Arnab et al., 2021). Han et al. (2021a) and Khan et al. (2021) provide reviews on existing work of visual Transformers. We encourage readers to refer to these surveys for further understand the current research progress on Transformers in CV.

(3) Audio Applications. Transformer can also be extended for audio-related applications, e.g., speech recognition (Dong et al., 2018; Pham et al., 2019; Chen et al., 2021; Gulati et al., 2020), speech synthesis (Li et al., 2019a; Zheng et al., 2020b; Ihm et al., 2020), speech enhancement (Kim et al., 2020; Yu et al., 2021) and music generation (Huang et al., 2019).

(4) Multimodal Applications. Owing to its flexible architecture, Transformer has also been applied in various multimodal scenarios, e.g., visual question answering (Li et al., 2019c; Hu et al., 2020; Su et al., 2020; Li et al., 2020a), visual commonsense reasoning (Li et al., 2019c; Su et al., 2020), caption generation (Sun et al., 2019; Cornia et al., 2020; Lin et al., 2021), speech-to-text translation (Han et al., 2021b) and text-to-image generation (Ramesh et al., 2021; Lin et al., 2021; Ding et al., 2021).

Conclusion and Future Directions

In this survey, we conduct a comprehensive overview of X-formers and propose a new taxonomy. Most of the existing works improve Transformer from different perspectives, such as efficiency, generalization, and applications. The improvements include incorporating structural prior, designing lightweight architecture, pre-training, and so on.

Although X-formers have proven their power for various tasks, challenges still exist. Besides the current concerns (e.g. efficiency and generalization), the further improvements of Transformer may lie in the following directions:

(1) Theoretical Analysis. The architecture of Transformer has been demonstrated to be capable of supporting large-scale training datasets with enough parameters. Many works show that Transformer has a larger capacity than CNNs and RNNs and hence has the ability to handle a huge amount of training data. When Transformer is trained on sufficient data, it usually has better performances than CNNs or RNNs. An intuitive explanation is that Transformer has few prior assumptions on the data structure and therefore is more flexible than CNNs and RNNs. However, the theoretical reason is unclear and we need some theoretical analysis of Transformer ability.

(2) Better Global Interaction Mechanism beyond Attention. A main advantage of Transformer is the use of the attention mechanism to model the global dependencies among nodes within input data. However, many studies have shown that full attention is unnecessary for most nodes. It is, to some degree, inefficient to indistinguishably calculate attention for all nodes. Therefore, there is still plenty of room for improvements in efficiently modeling global interactions. On the one hand, the self-attention module can be regarded as a fully-connected neural network with dynamical connection weights, which aggregates non-local information with dynamic routing. Therefore, other dynamic routing mechanisms are alternative approaches worth exploring. On the other hand, the global interaction can also be modeled by other types of neural networks, such as memory-enhanced models.

(3) Unified Framework for Multimodal Data. In many application scenarios, integrating multimodal data is useful and necessary to boost the task performance. Moreover, the general AI also needs the ability to capture the semantic relations across different modalities. Since Transformer achieves great success on text, image, video, and audio, we have a chance to build a unified framework and better capture the inherent connections among multimodal data. However, the design of the intra-modal and cross-modal attention still remains to be improved.

Finally, we wish this survey to be a hands-on reference for better understanding the current research progress on Transformers and help readers to further improve Transformers for various applications.

References