{
  "source_file": "0930",
  "source_platform": "个人博客",
  "author": "Lilian Weng / Lil'Log",
  "title_listed": "The Transformer Family Version 2.0",
  "keywords": "Transformer,注意力,架构",
  "domain": "大模型算法",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/",
  "publish_date": "2023-01",
  "corpus_stage": "预训练/SFT",
  "notes": "Transformer架构家族全景综述（2.0版）",
  "id": "bd2a1c61491f",
  "collected_at": "2026-10-02T14:12:43",
  "title": "The Transformer Family Version 2.0 | Lil'Log",
  "body_markdown": "Many new Transformer architecture improvements have been proposed since my last post on\n“The Transformer Family”\nabout three years ago. Here I did a big refactoring and enrichment of that 2020 post — restructure the hierarchy of sections and improve many sections with more recent papers. Version 2.0 is a superset of the old version, about twice the length.\nNotations\n#\nSymbol\nMeaning\n$d$\nThe model size / hidden state dimension / positional encoding size.\n$h$\nThe number of heads in multi-head attention layer.\n$L$\nThe segment length of input sequence.\n$N$\nThe total number of attention layers in the model; not considering MoE.\n$\\mathbf{X} \\in \\mathbb{R}^{L \\times d}$\nThe input sequence where each element has been mapped into an embedding vector of shape $d$, same as the model size.\n$\\mathbf{W}^k \\in \\mathbb{R}^{d \\times d_k}$\nThe key weight matrix.\n$\\mathbf{W}^q \\in \\mathbb{R}^{d \\times d_k}$\nThe query weight matrix.\n$\\mathbf{W}^v \\in \\mathbb{R}^{d \\times d_v}$\nThe value weight matrix. Often we have $d_k = d_v = d$.\n$\\mathbf{W}^k_i, \\mathbf{W}^q_i \\in \\mathbb{R}^{d \\times d_k/h}; \\mathbf{W}^v_i \\in \\mathbb{R}^{d \\times d_v/h}$\nThe weight matrices per head.\n$\\mathbf{W}^o \\in \\mathbb{R}^{d_v \\times d}$\nThe output weight matrix.\n$\\mathbf{Q} = \\mathbf{X}\\mathbf{W}^q \\in \\mathbb{R}^{L \\times d_k}$\nThe query embedding inputs.\n$\\mathbf{K} = \\mathbf{X}\\mathbf{W}^k \\in \\mathbb{R}^{L \\times d_k}$\nThe key embedding inputs.\n$\\mathbf{V} = \\mathbf{X}\\mathbf{W}^v \\in \\mathbb{R}^{L \\times d_v}$\nThe value embedding inputs.\n$\\mathbf{q}_i, \\mathbf{k}_i \\in \\mathbb{R}^{d_k}, \\mathbf{v}_i \\in \\mathbb{R}^{d_v}$\nRow vectors in query, key, value matrices, $\\mathbf{Q}$, $\\mathbf{K}$ and $\\mathbf{V}$.\n$S_i$\nA collection of key positions for the $i$-th query $\\mathbf{q}_i$ to attend to.\n$\\mathbf{A} \\in \\mathbb{R}^{L \\times L}$\nThe self-attention matrix between a input sequence of lenght $L$ and itself. $\\mathbf{A} = \\text{softmax}(\\mathbf{Q}\\mathbf{K}^\\top / \\sqrt{d_k})$.\n$a_{ij} \\in \\mathbf{A}$\nThe scalar attention score between query $\\mathbf{q}_i$ and key $\\mathbf{k}_j$.\n$\\mathbf{P} \\in \\mathbb{R}^{L \\times d}$\nposition encoding matrix, where the $i$-th row $\\mathbf{p}_i$ is the positional encoding for input $\\mathbf{x}_i$.\nTransformer Basics\n#\nThe\nTransformer\n(which will be referred to as “vanilla Transformer” to distinguish it from other enhanced versions;\nVaswani, et al., 2017\n) model has an encoder-decoder architecture, as commonly used in many\nNMT\nmodels. Later simplified Transformer was shown to achieve great performance in language modeling tasks, like in encoder-only\nBERT\nor decoder-only\nGPT\n.\nAttention and Self-Attention\n#\nAttention\nis a mechanism in neural network that a model can learn to make predictions by selectively attending to a given set of data. The amount of attention is quantified by learned weights and thus the output is usually formed as a weighted average.\nSelf-attention\nis a type of attention mechanism where the model makes prediction for one part of a data sample using other parts of the observation about the same sample. Conceptually, it feels quite similar to\nnon-local means\n. Also note that self-attention is permutation-invariant; in other words, it is an operation on sets.\nThere are various forms of attention / self-attention, Transformer (\nVaswani et al., 2017\n) relies on the\nscaled dot-product attention\n: given a query matrix $\\mathbf{Q}$, a key matrix $\\mathbf{K}$ and a value matrix $\\mathbf{V}$, the output is a weighted sum of the value vectors, where the weight assigned to each value slot is determined by the dot-product of the query with the corresponding key:\n$$\n\\text{attn}(\\mathbf{Q}, \\mathbf{K}, \\mathbf{V}) = \\text{softmax}(\\frac{\\mathbf{Q} {\\mathbf{K}}^\\top}{\\sqrt{d_k}})\\mathbf{V}\n$$\nAnd for a query and a key vector $\\mathbf{q}_i, \\mathbf{k}_j \\in \\mathbb{R}^d$ (row vectors in query and key matrices), we have a scalar score:\n$$\na_{ij} = \\text{softmax}(\\frac{\\mathbf{q}_i {\\mathbf{k}_j}^\\top}{\\sqrt{d_k}})\n= \\frac{\\exp(\\frac{\\mathbf{q}_i {\\mathbf{k}_j}^\\top}{\\sqrt{d_k}})}{ \\sum_{r \\in \\mathcal{S}_i} \\exp(\\frac{\\mathbf{q}_i {\\mathbf{k}_r}^\\top}{\\sqrt{d_k}}) }\n$$\nwhere $\\mathcal{S}_i$ is a collection of key positions for the $i$-th query to attend to.\nSee my old\npost for other types of attention\nif interested.\nMulti-Head Self-Attention\n#\nThe\nmulti-head self-attention\nmodule is a key component in Transformer. Rather than only computing the attention once, the multi-head mechanism splits the inputs into smaller chunks and then computes the scaled dot-product attention over each subspace in parallel. The independent attention outputs are simply concatenated and linearly transformed into expected dimensions.\n$$\n\\begin{aligned}\n\\text{MultiHeadAttn}(\\mathbf{X}_q, \\mathbf{X}_k, \\mathbf{X}_v) &= [\\text{head}_1; \\dots; \\text{head}_h] \\mathbf{W}^o \\\\ \n\\text{where head}_i &= \\text{Attention}(\\mathbf{X}_q\\mathbf{W}^q_i, \\mathbf{X}_k\\mathbf{W}^k_i, \\mathbf{X}_v\\mathbf{W}^v_i)\n\\end{aligned}\n$$\nwhere $[.;.]$ is a concatenation operation. $\\mathbf{W}^q_i, \\mathbf{W}^k_i \\in \\mathbb{R}^{d \\times d_k/h}, \\mathbf{W}^v_i \\in \\mathbb{R}^{d \\times d_v/h}$ are weight matrices to map input embeddings of size $L \\times d$ into query, key and value matrices. And $\\mathbf{W}^o \\in \\mathbb{R}^{d_v \\times d}$ is the output linear transformation. All the weights should be learned during training.\nIllustration of the multi-head scaled dot-product attention mechanism. (Image source: Figure 2 in\nVaswani, et al., 2017\n)\nEncoder-Decoder Architecture\n#\nThe\nencoder\ngenerates an attention-based representation with capability to locate a specific piece of information from a large context. It consists of a stack of 6 identity modules, each containing two submodules, a\nmulti-head self-attention\nlayer and a\npoint-wise\nfully connected feed-forward network. By point-wise, it means that it applies the same linear transformation (with same weights) to each element in the sequence. This can also be viewed as a convolutional layer with filter size 1. Each submodule has a residual connection and layer normalization. All the submodules output data of the same dimension $d$.\nThe function of Transformer\ndecoder\nis to retrieve information from the encoded representation. The architecture is quite similar to the encoder, except that the decoder contains two multi-head attention submodules instead of one in each identical repeating module. The first multi-head attention submodule is\nmasked\nto prevent positions from attending to the future.\nThe architecture of the vanilla Transformer model. (Image source:\nFigure 17\n)\nPositional Encoding\n#\nBecause self-attention operation is permutation invariant, it is important to use proper\npositional encoding\nto provide\norder information\nto the model. The positional encoding $\\mathbf{P} \\in \\mathbb{R}^{L \\times d}$ has the same dimension as the input embedding, so it can be added on the input directly. The vanilla Transformer considered two types of encodings:\nSinusoidal Positional Encoding\n#\nSinusoidal positional encoding is defined as follows, given the token position $i=1,\\dots,L$ and the dimension $\\delta=1,\\dots,d$:\n$$\n\\text{PE}(i,\\delta) = \n\\begin{cases}\n\\sin(\\frac{i}{10000^{2\\delta'/d}}) & \\text{if } \\delta = 2\\delta'\\\\\n\\cos(\\frac{i}{10000^{2\\delta'/d}}) & \\text{if } \\delta = 2\\delta' + 1\\\\\n\\end{cases}\n$$\nIn this way each dimension of the positional encoding corresponds to a sinusoid of different wavelengths in different dimensions, from $2\\pi$ to $10000 \\cdot 2\\pi$.\nSinusoidal positional encoding with $L=32$ and $d=128$. The value is between -1 (black) and 1 (white) and the value 0 is in gray.\nLearned Positional Encoding\n#\nLearned positional encoding assigns each element with a\nlearned\ncolumn vector which encodes its absolute position (\nGehring, et al. 2017\n) and furthermroe this encoding can be learned differently per layer (\nAl-Rfou et al. 2018\n).\nRelative Position Encoding\n#\nShaw et al. (2018)\n) incorporated relative positional information into $\\mathbf{W}^k$ and $\\mathbf{W}^v$. Maximum relative position is clipped to a maximum absolute value of $k$ and this clipping operation enables the model to generalize to unseen sequence lengths. Therefore, $2k + 1$ unique edge labels are considered and let us denote $\\mathbf{P}^k, \\mathbf{P}^v \\in \\mathbb{R}^{2k+1}$ as learnable relative position representations.\n$$\nA_{ij}^k = P^k_{\\text{clip}(j - i, k)} \\quad\nA_{ij}^v = P^v_{\\text{clip}(j - i, k)} \\quad\n\\text{where }\\text{clip}(x, k) = \\text{clip}(x, -k, k)\n$$\nTransformer-XL\n(\nDai et al., 2019\n) proposed a type of relative positional encoding based on reparametrization of dot-product of keys and queries. To keep the positional information flow coherently across segments, Transformer-XL encodes the\nrelative\nposition instead, as it could be sufficient enough to know the position offset for making good predictions, i.e. $i-j$, between one key vector $\\mathbf{k}_{\\tau, j}$ and its query $\\mathbf{q}_{\\tau, i}$.\nIf omitting the scalar $1/\\sqrt{d_k}$ and the normalizing term in softmax but including positional encodings, we can write the attention score between query at position $i$ and key at position $j$ as:\n$$\n\\begin{aligned}\na_{ij} \n&= \\mathbf{q}_i {\\mathbf{k}_j}^\\top = (\\mathbf{x}_i + \\mathbf{p}_i)\\mathbf{W}^q ((\\mathbf{x}_j + \\mathbf{p}_j)\\mathbf{W}^k)^\\top \\\\\n&= \\mathbf{x}_i\\mathbf{W}^q {\\mathbf{W}^k}^\\top\\mathbf{x}_j^\\top + \\mathbf{x}_i\\mathbf{W}^q {\\mathbf{W}^k}^\\top\\mathbf{p}_j^\\top + \\mathbf{p}_i\\mathbf{W}^q {\\mathbf{W}^k}^\\top\\mathbf{x}_j^\\top + \\mathbf{p}_i\\mathbf{W}^q {\\mathbf{W}^k}^\\top\\mathbf{p}_j^\\top\n\\end{aligned}\n$$\nTransformer-XL reparameterizes the above four terms as follows:\n$$\na_{ij}^\\text{rel} = \n\\underbrace{ \\mathbf{x}_i\\mathbf{W}^q \\color{blue}{ {\\mathbf{W}_E^k}^\\top } \\mathbf{x}_j^\\top }_\\text{content-based addressing} + \n\\underbrace{ \\mathbf{x}_i\\mathbf{W}^q \\color{blue}{ {\\mathbf{W}_R^k}^\\top } \\color{green}{\\mathbf{r}_{i-j}^\\top} }_\\text{content-dependent positional bias} + \n\\underbrace{ \\color{red}{\\mathbf{u}} \\color{blue}{ {\\mathbf{W}_E^k}^\\top } \\mathbf{x}_j^\\top }_\\text{global content bias} + \n\\underbrace{ \\color{red}{\\mathbf{v}} \\color{blue}{ {\\mathbf{W}_R^k}^\\top } \\color{green}{\\mathbf{r}_{i-j}^\\top} }_\\text{global positional bias}\n$$\nReplace $\\mathbf{p}_j$ with relative positional encoding $\\mathbf{r}_{i-j} \\in \\mathbf{R}^{d}$;\nReplace $\\mathbf{p}_i\\mathbf{W}^q$ with two trainable parameters $\\mathbf{u}$ (for content) and $\\mathbf{v}$ (for location) in two different terms;\nSplit $\\mathbf{W}^k$ into two matrices, $\\mathbf{W}^k_E$ for content information and $\\mathbf{W}^k_R$ for location information.\nRotary Position Embedding\n#\nRotary position embedding (\nRoPE\n;\nSu et al. 2021\n) encodes the absolution position with a\nrotation matrix\nand multiplies key and value matrices of every attention layer with it to inject relative positional information at every layer.\nWhen encoding relative positional information into the inner product of the $i$-th key and the $j$-th query, we would like to formulate the function in a way that the inner product is only about the relative position $i-j$. Rotary Position Embedding (RoPE) makes use of the rotation operation in Euclidean space and frames the relative position embedding as simply rotating feature matrix by an angle proportional to its position index.\nGiven a vector $\\mathbf{z}$, if we want to rotate it counterclockwise by $\\theta$, we can multiply it by a rotation matrix to get $R\\mathbf{z}$ where the rotation matrix $R$ is defined as:\n$$\nR = \\begin{bmatrix}\n\\cos\\theta & -\\sin\\theta \\\\\n\\sin\\theta & \\cos\\theta\n\\end{bmatrix}\n$$\nWhen generalizing to higher dimensional space, RoPE divide the $d$-dimensional space into $d/2$ subspaces and constructs a rotation matrix $R$ of size $d \\times d$ for token at position $i$:\n$$\nR^d_{\\Theta, i} = \\begin{bmatrix}\n\\cos i\\theta_1 & -\\sin i\\theta_1 & 0 & 0 & \\dots & 0 & 0 \\\\\n\\sin i\\theta_1 & \\cos i\\theta_1 & 0 & 0 & \\dots & 0 & 0 \\\\\n0 & 0 & \\cos i\\theta_2 & -\\sin i\\theta_2 & \\dots & 0 & 0 \\\\\n0 & 0 & \\sin i\\theta_2 & \\cos i\\theta_2 & \\dots & 0 & 0 \\\\\n\\vdots & \\vdots & \\vdots & \\vdots & \\ddots & \\vdots & \\vdots \\\\\n0 & 0 & 0 & 0 & \\dots & \\cos i\\theta_{d/2} & -\\sin i\\theta_{d/2} \\\\\n0 & 0 & 0 & 0 & \\dots & \\sin i\\theta_{d/2} & \\cos i\\theta_{d/2} \\\\\n\\end{bmatrix}\n$$\nwhere in the paper we have $\\Theta = {\\theta_i = 10000^{-2(i−1)/d}, i \\in [1, 2, …, d/2]}$. Note that this is essentially equivalent to sinusoidal positional encoding but formulated as a rotation matrix.\nThen both key and query matrices incorporates the positional information by multiplying with this rotation matrix:\n$$\n\\begin{aligned}\n& \\mathbf{q}_i^\\top \\mathbf{k}_j = (R^d_{\\Theta, i} \\mathbf{W}^q\\mathbf{x}_i)^\\top (R^d_{\\Theta, j} \\mathbf{W}^k\\mathbf{x}_j) = \\mathbf{x}_i^\\top\\mathbf{W}^q R^d_{\\Theta, j-i}\\mathbf{W}^k\\mathbf{x}_j \\\\\n& \\text{ where } R^d_{\\Theta, j-i} = (R^d_{\\Theta, i})^\\top R^d_{\\Theta, j}\n\\end{aligned}\n$$\nVisual illustration of how rotary position embedding is implemented.(Image source:\nSu et al., 2021\n) Note: I used $i$ instead of $m$ to represent the position index compared to the original figure in the paper.\nLonger Context\n#\nThe length of an input sequence for transformer models at inference time is upper-bounded by the context length used for training. Naively increasing context length leads to high consumption in both time ($\\mathcal{O}(L^2d)$) and memory ($\\mathcal{O}(L^2)$) and may not be supported due to hardware constraints.\nThis section introduces several improvements in transformer architecture to better support long context at inference; E.g. using additional memory, design for better context extrapolation, or recurrency mechanism.\nContext Memory\n#\nThe vanilla Transformer has a fixed and limited attention span. The model can only attend to other elements in the same segments during each update step and no information can flow across separated fixed-length segments. This\ncontext segmentation\ncauses several issues:\nThe model cannot capture very long term dependencies.\nIt is hard to predict the first few tokens in each segment given no or thin context.\nThe evaluation is expensive. Whenever the segment is shifted  to the right by one, the new segment is re-processed from scratch, although there are a lot of overlapped tokens.\nTransformer-XL\n(\nDai et al., 2019\n; “XL” means “extra long”) modifies the architecture to reuse hidden states between segments with an additional memory. The recurrent connection between segments is introduced into the model by continuously using the hidden states from the previous segments.\nA comparison between the training phrase of vanilla Transformer & Transformer-XL with a segment length 4. (Image source: left part of Figure 2 in\nDai et al., 2019\n).\nLet’s label the hidden state of the $n$-th layer for the $(\\tau + 1)$-th segment in the model as $\\mathbf{h}_{\\tau+1}^{(n)} \\in \\mathbb{R}^{L \\times d}$. In addition to the hidden state of the last layer for the same segment $\\mathbf{h}_{\\tau+1}^{(n-1)}$, it also depends on the hidden state of the same layer for the previous segment $\\mathbf{h}_{\\tau}^{(n)}$.  By incorporating information from the previous hidden states, the model extends the attention span much longer in the past, over multiple segments.\n$$\n\\begin{aligned}\n\\color{red}{\\widetilde{\\mathbf{h}}_{\\tau+1}^{(n-1)}} &= [\\text{stop-gradient}(\\mathbf{h}_{\\tau}^{(n-1)}) \\circ \\mathbf{h}_{\\tau+1}^{(n-1)}] \\\\\n\\mathbf{Q}_{\\tau+1}^{(n)} &= \\mathbf{h}_{\\tau+1}^{(n-1)}\\mathbf{W}^q \\\\\n\\mathbf{K}_{\\tau+1}^{(n)} &= \\color{red}{\\widetilde{\\mathbf{h}}_{\\tau+1}^{(n-1)}} \\mathbf{W}^k \\\\\n\\mathbf{V}_{\\tau+1}^{(n)} &= \\color{red}{\\widetilde{\\mathbf{h}}_{\\tau+1}^{(n-1)}} \\mathbf{W}^v \\\\\n\\mathbf{h}_{\\tau+1}^{(n)} &= \\text{transformer-layer}(\\mathbf{Q}_{\\tau+1}^{(n)}, \\mathbf{K}_{\\tau+1}^{(n)}, \\mathbf{V}_{\\tau+1}^{(n)})\n\\end{aligned}\n$$\nNote that both keys and values rely on extended hidden states, while queries only consume hidden states at the current step. The concatenation operation $[. \\circ .]$ is along the sequence length dimension. And Transformer-XL needs to use\nrelative positional encoding\nbecause previous and current segments would be assigned with the same encoding if we encode absolute positions, which is undesired.\nCompressive Transformer\n(\nRae et al. 2019\n) extends Transformer-XL by compressing past memories to support longer sequences. It explicitly adds\nmemory\nslots of size $m_m$ per layer for storing past activations of this layer to preserve long context. When some past activations become old enough, they are compressed and saved in an additional\ncompressed memory\nof size $m_{cm}$ per layer.\nCompressive transformer maintains two types of memory slots, memory and compressed memory, to support long context. (Image source:\nRae et al. 2019\n).\nBoth memory and compressed memory are FIFO queues. Given the model context length $L$, the compression function of compression rate $c$ is defined as $f_c: \\mathbb{R}^{L \\times d} \\to \\mathbb{R}^{[\\frac{L}{c}] \\times d}$, mapping $L$ oldest activations to $[\\frac{L}{c}]$ compressed memory elements. There are several choices of compression functions:\nMax/mean pooling of kernel and stride size $c$;\n1D convolution with kernel and stride size $c$ (need to learn additional parameters);\nDilated convolution (need to learn additional parameters). In their experiments, convolution compression works out the best on\nEnWik8\ndataset;\nMost used memories.\nCompressive transformer has two additional training losses:\nAuto-encoding loss\n(lossless compression objective) measures how well we can reconstruct the original memories from compressed memories\n$$\n \\mathcal{L}_{ac} = \\| \\textbf{old_mem}^{(i)} - g(\\textbf{new_cm}^{(i)}) \\|_2\n $$\nwhere $g: \\mathbb{R}^{[\\frac{L}{c}] \\times d} \\to \\mathbb{R}^{L \\times d}$ reverses the compression function $f$.\nAttention-reconstruction loss\n(lossy objective) reconstructs content-based attention over memory vs compressed memory and minimize the difference:\n$$\n \\mathcal{L}_{ar} = \\|\\text{attn}(\\mathbf{h}^{(i)}, \\textbf{old_mem}^{(i)}) − \\text{attn}(\\mathbf{h}^{(i)}, \\textbf{new_cm}^{(i)})\\|_2\n $$\nTransformer-XL with a memory of size $m$ has a maximum temporal range of $m \\times N$, where $N$ is the number of layers in the model, and attention cost $\\mathcal{O}(L^2 + Lm)$. In comparison, compressed transformer has a temporal range of $(m_m + c \\cdot m_{cm}) \\times N$ and attention cost $\\mathcal{O}(L^2 + L(m_m + m_{cm}))$. A larger compression rate $c$ gives better tradeoff between temporal range length and attention cost.\nAttention weights, from oldest to newest, are stored in three locations: compressed memory → memory → causally masked sequence. In the experiments, they observed an increase in attention weights from oldest activations stored in the regular memory, to activations stored in the compressed memory, implying that the network is learning to preserve salient information.\nAttention weights with one standard deviation as error bars versus memory positions, from oldest (left) to newest (right). (Image source:\nRae et al. 2019\n).\nNon-Differentiable External Memory\n#\n$k$NN-LM\n(\nKhandelwal et al. 2020\n) enhances a pretrained LM with a separate $k$NN model by linearly interpolating the next token probabilities predicted by both models. The $k$NN model is built upon an external key-value store  which can store any large pre-training dataset or OOD new dataset. This datastore is preprocessed to save a\nlarge\nnumber of pairs, (LM embedding representation of context, next token) and the nearest neighbor retrieval happens in the LM embedding space. Because the datastore can be gigantic, we need to rely on libraries for fast dense vector search such as\nFAISS\nor\nScaNN\n. The indexing process only happens once and parallelism is easy to implement at inference time.\nAt inference time, the next token probability is a weighted sum of two predictions:\n$$\n\\begin{aligned}\np(y \\vert \\mathbf{x}) &= \\lambda \\; p_\\text{kNN}(y \\vert \\mathbf{x}) + (1- \\lambda) \\; p_\\text{LM}(y \\vert \\mathbf{x}) \\\\\np_\\text{kNN}(y \\vert \\mathbf{x}) &\\propto \\sum_{(k_i, w_i) \\in \\mathcal{N}} \\mathbb{1}[y = w_i] \\exp(-d(k_i, f(\\mathbf{x})))\n\\end{aligned}\n$$\nwhere $\\mathcal{N}$ contains a set of nearest neighbor data points retrieved by $k$NN; $d(., .)$ is a distance function such as L2 distance.\nAccording to the experiments, larger datastore size or larger $k$ is correlated with better perplexity. The weighting scalar $\\lambda$ should be tuned, but in general it is expected to be larger for out-of-domain data compared to in-domain data and larger datastore can afford a larger $\\lambda$.\nSPALM\n(\nAdaptive semiparametric language models\n;\nYogatama et al. 2021\n) incorporates both (1) Transformer-XL style memory for hidden states from external context as short-term memory and (2) $k$NN-LM style key-value store as long memory.\nIllustration of how SPALM combines context memory of past hidden states (short term memory) with an external key-value datastore (long term memory) to support longer context. (Image source:\nYogatama et al. 2021\n).\nSPALM runs $k$NN search to fetch $k$ tokens with most relevant context. For each token we can get the same embedding representation provided by a pretrained LM, denoted as $\\{\\mathbf{y}_i\\}_{i=1}^k$. The gating mechanism first aggregates the retrieved token embeddings with a simple attention layer using $\\mathbf{h}^R_t$ (the hidden state for token $x_t$ at layer $R$) as a query and then learns a gating parameter $\\mathbf{g}_t$ to balance between local information $\\mathbf{h}^R_t$ and long-term information $\\mathbf{m}_t$.\n$$\n\\begin{aligned}\n\\mathbf{m}_t &= \\sum_{i=1}^k \\frac{\\exp(\\mathbf{y}_i^\\top \\mathbf{h}^R_t)}{\\sum_{j=1}^k \\exp(\\mathbf{y}_j^\\top \\mathbf{h}^R_t)} \\cdot \\mathbf{y}_i \\\\\n\\mathbf{g}_t &= \\sigma(\\mathbf{w}_g^\\top \\mathbf{h}_t^R) \\\\\n\\mathbf{z}_t &= (1 - \\mathbf{g}_t) \\odot \\mathbf{m}_t + \\mathbf{g}_t \\odot \\mathbf{h}^R_t \\\\\np(x_{t+1}\\mid \\mathbf{x}_{\\leq t}) &= \\text{softmax}(\\mathbf{z}_t; \\mathbf{W})\n\\end{aligned}\n$$\nwhere $\\mathbf{w}_g$ is a parameter vector to learn; $\\sigma(.)$ is sigmoid; $\\mathbf{W}$ is the word embedding matrix shared between both input and output tokens. Different from $k$NN-LM, they didn’t find the nearest neighbor distance to be helpful in the aggregation of retrieved tokens.\nDuring training, the key representations in the long-term memory stay constant, produced by a pretrained LM, but the value encoder, aka the word embedding matrix, gets updated.\nMemorizing Transformer\n(\nWu et al. 2022\n) adds a $k$NN-augmented attention layer near the top stack of a decoder-only Transformer. This special layer maintains a Transformer-XL style FIFO cache of past key-value pairs.\nThe same QKV values are used for both local attention and $k$NN mechanisms. The $k$NN lookup returns top-$k$ (key, value) pairs for each query in the input sequence and then they are processed through the self-attention stack to compute a weighted average of retrieved values. Two types of attention are combined with a learnable per-head gating parameter. To prevent large distributional shifts in value magnitude, both keys and values in the cache are normalized.\nWhat they found during experiments with Memorizing Transformer:\nIt is observed in some experiments that training models with a small memory and then finetuned with a larger memory works better than training with a large memory from scratch.\nThe smaller Memorizing Transformer with just 8k tokens in memory can match the perplexity of a larger vanilla Transformer with 5X more trainable parameters.\nIncreasing the size of external memory provided consistent gains up to a size of 262K.\nA  non-memory transformer can be finetuned to use memory.\nFine-tuning a vanilla Transformer with a key-value memory can achieve similar performance as training a memorizing transformer from scratch. (Image source:\nWu et al. 2022\n).\nDistance-Enhanced Attention Scores\n#\nDistance Aware Transformer\n(\nDA-Transformer\n;\nWu, et al. 2021\n) and\nAttention with Linear Biases\n(\nALiBi\n;\nPress et al. 2022\n) are motivated by similar ideas — in order to encourage the model to extrapolate over longer context than what the model is trained on, we can explicitly attach the positional information to every pair of attention score based on the distance between key and query tokens.\nNote that the default positional encoding in vanilla Transformer only adds positional information to the input sequence, while later improved encoding mechanisms alter attention scores of every layer, such as\nrotary position embedding\n, and they take on form very similar to distance enhanced attention scores.\nDA-Transformer\n(\nWu, et al. 2021\n) multiplies attention scores at each layer by a learnable bias that is formulated as a function of the distance between key and query. Different attention heads use different parameters to distinguish diverse preferences to short-term vs long-term context. Given two positions, $i, j$, DA-Transformer uses the following weighting function to alter the self-attention score:\n$$\n\\begin{aligned}\n\\mathbf{R}^{(i)} &= \\alpha_i \\mathbf{R} \\quad \\text{where }R_{ij} = \\vert i-j \\vert\\\\\nf(\\mathbf{R}^{(i)}; \\beta_i) &= \\frac{1 + \\exp(\\beta_i)}{1 + \\exp(\\beta_i - \\mathbf{R}^{(i)})} \\\\\n\\text{attn}(\\mathbf{Q}^{(i)}, \\mathbf{K}^{(i)}, \\mathbf{V}^{(i)}) &= \\text{row-softmax}\\Big(\\frac{\\text{ReLU}(\\mathbf{Q}^{(i)}\\mathbf{K}^{(i)\\top})f(\\mathbf{R}^{(i)})}{\\sqrt{d}}\\Big) \\mathbf{V}^{(i)}\n\\end{aligned}\n$$\nwhere $\\alpha_i$ is a learnable parameters to weight relative distance differently per head where the head is indexed by superscript $^{(i)}$; $\\beta_i$ is a learnable parameter to control the upper bound and ascending slope wrt the distance for the $i$-th attention head. The weighting function $f(.)$ is designed in a way that: (1) $f(0)=1$; (2) $f(\\mathbf{R}^{(i)}) = 0$ when $\\mathbf{R}^{(i)} \\to -\\infty$; (3) $f(\\mathbf{R}^{(i)})$ is bounded when $\\mathbf{R}^{(i)} \\to +\\infty$; (4) the scale is tunable; (5) and the function is monotonic. The extra time complexity brought by $f(\\mathbf{R}^{(i)})$ is $\\mathcal{O}(L^2)$ and it is small relative to the self attention time complexity $\\mathcal{O}(L^2 d)$. The extra memory consumption is minimal, ~$\\mathcal{O}(2h)$.\nInstead of multipliers,\nALiBi\n(\nPress et al. 2022\n) adds a constant bias term on query-key attention scores, proportional to pairwise distances. The bias introduces a strong recency preference and penalizes keys that are too far away. The penalties are increased at different rates within different heads.\n$$\n\\text{softmax}(\\mathbf{q}_i \\mathbf{K}^\\top + \\alpha_i \\cdot [0, -1, -2, \\dots, -(i-1)])\n$$\nwhere $\\alpha_i$ is a head-specific weighting scalar. Different from DA-transformer, $\\alpha_i$ is not learned but fixed as a geometric sequence; for example, for 8 heads, ${\\alpha_i} = {\\frac{1}{2}, \\frac{1}{2^2}, \\dots, \\frac{1}{2^8}}$. The overall idea is very much similar to what relative positional encoding aims to solve.\nIllustration of how ALiBi enhances attention scores with a positional bias term. (Image source:\nPress et al. 2021\n).\nWith ALiBi,\nPress et al. (2022)\ntrained a 1.3B model on context length 1024 during training and extrapolated to 2046 at inference time.\nExtrapolation experiments for running inference with Transformers of different configs, including sinusoidal positional encoding, rotary positional encoding, simplified relative positional encoding in T5 and ALiBi. All models were trained with small context length but inference ran for much longer context. (Image source:\nPress et al. 2021\n).\nMake it Recurrent\n#\nUniversal Transformer\n(\nDehghani, et al. 2019\n) combines self-attention in Transformer with the recurrent mechanism in RNN, aiming to benefit from both a long-term global receptive field of Transformer and learned inductive biases of RNN. Rather than going through a fixed number of layers, Universal Transformer dynamically adjusts the number of steps using\nadaptive computation time\n. If we fix the number of steps, an Universal Transformer is equivalent to a multi-layer Transformer with shared parameters across layers.\nOn a high level, the universal transformer can be viewed as a recurrent function for learning the hidden state representation per token. The recurrent function evolves in parallel across token positions and the information between positions is shared through self-attention.\nHow the Universal Transformer refines a set of hidden state representations repeatedly for every position in parallel. (Image source: Figure 1 in\nDehghani, et al. 2019\n).\nGiven an input sequence of length $L$, Universal Transformer iteratively updates the representation $\\mathbf{h}^t \\in \\mathbb{R}^{L \\times d}$ at step $t$ for an adjustable number of steps. At step 0, $\\mathbf{h}^0$ is initialized to be same as the input embedding matrix. All the positions are processed in parallel in the multi-head self-attention mechanism and then go through a recurrent transition function.\n$$\n\\begin{aligned}\n\\mathbf{A}^t &= \\text{LayerNorm}(\\mathbf{h}^{t-1} + \\text{MultiHeadAttention}(\\mathbf{h}^{t-1} + \\mathbf{P}^t) \\\\\n\\mathbf{h}^t &= \\text{LayerNorm}(\\mathbf{A}^{t-1} + \\text{Transition}(\\mathbf{A}^t))\n\\end{aligned}\n$$\nwhere $\\text{Transition}(.)$ is either a\nseparable convolution\nor a fully-connected neural network that consists of two position-wise (i.e. applied to each row of $\\mathbf{A}^t$ individually) affine transformation + one ReLU.\nThe positional encoding $\\mathbf{P}^t$ uses\nsinusoidal position signal\nbut with an additional time dimension:\n$$\n\\text{PE}(i, t, \\delta) = \n\\begin{cases}\n\\sin(\\frac{i}{10000^{2\\delta'/d}}) \\oplus \\sin(\\frac{t}{10000^{2\\delta'/d}}) & \\text{if } \\delta = 2\\delta'\\\\\n\\cos(\\frac{i}{10000^{2\\delta'/d}}) \\oplus \\cos(\\frac{t}{10000^{2\\delta'/d}}) & \\text{if } \\delta = 2\\delta' + 1\\\\\n\\end{cases}\n$$\nA simplified illustration of Universal Transformer. The encoder and decoder share the same basic recurrent structure. But the decoder also attends to final encoder representation $\\mathbf{h}^T$. (Image source: Figure 2 in\nDehghani, et al. 2019\n)\nIn the adaptive version of Universal Transformer, the number of recurrent steps $T$ is dynamically determined by\nACT\n. Each position is equipped with a dynamic ACT halting mechanism. Once a per-token recurrent block halts, it stops taking more recurrent updates but simply copies the current value to the next step until all the blocks halt or until the model reaches a maximum step limit.\nAdaptive Modeling\n#\nAdaptive modeling refers to a mechanism that can adjust the amount of computation according to different inputs. For example, some tokens may only need local information and thus demand a shorter attention span; Or some tokens are relatively easier to predict and do not need to be processed through the entire attention stack.\nAdaptive Attention Span\n#\nOne key advantage of Transformer is the capability of capturing long-term dependencies. Depending on the context, the model may prefer to attend further sometime than others; or one attention head may had different attention pattern from the other. If the attention span could adapt its length flexibly and only attend further back when needed, it would help reduce both computation and memory cost to support longer maximum context size in the model.\nThis is the motivation for\nAdaptive Attention Span\n.\nSukhbaatar et al (2019)\nproposed a self-attention mechanism that seeks an optimal attention span. They hypothesized that different attention heads might assign scores differently within the same context window (See Fig. 14) and thus the optimal span would be trained separately per head.\nTwo attention heads in the same model, A & B, assign attention differently within the same context window. Head A attends more to the recent tokens, while head B look further back into the past uniformly. (Image source:\nSukhbaatar, et al. 2019\n)\nGiven the $i$-th token, we need to compute the attention weights between this token and other keys within its attention span of size $s$:\n$$\n\\begin{aligned}\ne_{ij} &= \\mathbf{q}_i {\\mathbf{k}_j}^\\top \\\\ \na_{ij} &= \\text{softmax}(e_{ij}) = \\frac{\\exp(e_{ij})}{\\sum_{r=i-s}^{i-1} \\exp(e_{ir})} \\\\\n\\mathbf{y}_i &= \\sum_{r=i-s}^{i-1}a_{ir}\\mathbf{v}_r = \\sum_{r=i-s}^{i-1}a_{ir}\\mathbf{x}_r\\mathbf{W}^v\n\\end{aligned}\n$$\nA\nsoft mask function\n$m_z$ is added to control for an effective adjustable attention span, which maps the distance between query and key into a [0, 1] value. $m_z$ is parameterized by $z \\in [0, s]$ and $z$ is to be learned:\n$$\nm_z(x) = \\text{clip}(\\frac{1}{R}(R+z-x), 0, 1)\n$$\nwhere $R$ is a hyper-parameter which defines the softness of $m_z$.\nThe soft masking function used in the adaptive attention span. (Image source:\nSukhbaatar, et al. 2019\n.)\nThe soft mask function is applied to the softmax elements in the attention weights:\n$$\na_{ij} = \\frac{m_z(i-j)\\exp(s_{ij})}{\\sum_{r=i-s}^{i-1}m_z(i-r) \\exp(s_{ir})}\n$$\nIn the above equation, $z$ is differentiable so it is trained jointly with other parts of the model. Parameters $z^{(i)}, i=1, \\dots, h$ are learned\nseparately per head\n. Moreover, the loss function has an extra L1 penalty on $\\sum_{i=1}^h z^{(i)}$.\nUsing\nAdaptive Computation Time\n, the approach can be further enhanced to have flexible attention span length, adaptive to the current input dynamically. The span parameter $z_t$ of an attention head at time $t$ is a sigmoidal function, $z_t = S \\sigma(\\mathbf{v} \\cdot \\mathbf{x}_t +b)$, where the vector $\\mathbf{v}$ and the bias scalar $b$ are learned jointly with other parameters.\nIn the experiments of Transformer with adaptive attention span,\nSukhbaatar, et al. (2019)\nfound a general tendency that lower layers do not require very long attention spans, while a few attention heads in higher layers may use exceptionally long spans. Adaptive attention span also helps greatly reduce the number of FLOPS, especially in a big model with many attention layers and a large context length.\nDepth-Adaptive Transformer\n#\nAt inference time, it is natural to assume that some tokens are easier to predict and thus do not require as much computation as others. Therefore we may only process its prediction through a limited number of layers to achieve a good balance between speed and performance.\nBoth\nDepth-Adaptive Transformer\n(\nElabyad et al. 2020\n) and\nConfident Adaptive Language Model\n(\nCALM\n;\nSchuster et al. 2022\n) are motivated by this idea and learn to predict optimal numbers of layers needed for different input tokens.\nDepth-adaptive transformer\n(\nElabyad et al. 2020\n) attaches an output classifier to every layer to produce exit predictions based on activations of that layer. The classifier weight matrices can be different per layer or shared across layers. During training, the model sample different sequences of exits such that the model is optimized with hidden states of different layers. The learning objective incorporates likelihood probabilities predicted at different layers, $n=1, \\dots, N$:\n$$\n\\text{LL}^n_t = \\log p(y_t \\vert \\mathbf{h}^n_{t-1}) \\quad\n\\text{LL}^n = \\sum_{t=1}^{\\vert\\mathbf{y}\\vert} LL^n_t\n$$\nAdaptive depth classifiers outputs a parametric distribution $q_t$. It is trained with cross entropy loss against an oracle distribution $q^*_t$. The paper explored three confiurations for how to learn such a classifier $q_t$.\nIllustration of three types of adaptive depth classifiers.\n(Image source:\nElabyad et al. 2020\n).\nSequence-specific depth classifier\n: All tokens of the same sequence share the same exit block. It depends on the average of the encoder representation of the sequence. Given an input sequence $\\mathbf{x}$ of length $L$, the classifier takes $\\bar{\\mathbf{x}} = \\frac{1}{L} \\sum_{t=1}^L \\mathbf{x}_t$ as input and outputs a multinomial distribution of $N$ dimensions, corresponding to $N$ layers.\n$$\n \\begin{aligned}\n q(n \\vert \\mathbf{x}) &=\\text{softmax}(\\mathbf{W}_n \\bar{\\mathbf{x}} + b_n) \\in \\mathbb{R}^N \\\\\n q_\\text{lik}^*(\\mathbf{x}, \\mathbf{y}) &= \\delta(\\arg\\max_n \\text{LL}^n - \\lambda n) \\\\\n \\text{or }q_\\text{corr}^*(\\mathbf{x}, \\mathbf{y}) &= \\delta(\\arg\\max_n C^n - \\lambda n) \\text{ where }C^n = \\vert\\{t \\vert y_t = \\arg\\max_y p(y \\vert \\mathbf{h}^n_{t-1})\\}\\vert \\\\\n \\end{aligned}\n $$\nwhere $\\delta$ is\ndirac delta\n(unit impulse) function and $-\\lambda n$ is a regularization term to encourage lower layer exits. The ground truth $q^*$ can be prepared in two way, based on maximum likelihood $q_\\text{lik}^*$ or correctness $q_\\text{corr}^*$.\nToken-specific depth classifier (multinomial)\n: Each token is decoded with different exit block, predicted conditioned on the first decoder hidden state $\\mathbf{h}^1_t$:\n$$\n q_t(n \\vert \\mathbf{x}, \\mathbf{y}_{< t}) = \\text{softmax}(\\mathbf{W}_n \\mathbf{h}^1_t + b_n)\n $$\nToken-specific depth classifier (geometric-like)\n:  A binary exit prediction distribution is made per layer per token, $\\mathcal{X}^n_t$. The RBF kernel $\\kappa(t, t’) = \\exp(\\frac{\\vert t - t’ \\vert^2}{\\sigma})$ is used to smooth the predictions to incorporate the impact of current decision on future time steps.\n$$\n \\begin{aligned}\n \\mathcal{X}^n_t &= \\text{sigmoid}(\\mathbf{w}_n^\\top \\mathbf{h}^n_t + b_n)\\quad \\forall n \\in [1, \\dots, N-1] \\\\\n q_t(n \\vert \\mathbf{x}, \\mathbf{y}_{< t}) &= \\begin{cases}\n \\mathcal{X}^n_t \\prod_{n' < n} (1 - \\mathcal{X}^{n'}_t) & \\text{if } n < N\\\\\n \\prod_{n' < N} (1 - \\mathcal{X}^{n'}_t) & \\text{otherwise}\n \\end{cases} \\\\\n q_\\text{lik}^*(\\mathbf{x}, \\mathbf{y}) &= \\delta(\\arg\\max_n \\widetilde{\\text{LL}}^n_t - \\lambda n) \\text{ where } \\widetilde{\\text{LL}}^n_t = \\sum_{t'=1}^{\\vert\\mathbf{y}\\vert}\\kappa(t, t') LL^n_{t'} \\\\\n \\text{or }q_\\text{cor}^*(\\mathbf{x}, \\mathbf{y}) &= \\delta(\\arg\\max_n \\tilde{C}_t^n - \\lambda n) \\text{ where }C_t^n = \\mathbb{1}[y_t = \\arg\\max_y p(y \\vert \\mathbf{h}^n_{t-1})],\\; \\tilde{C}^n_t = \\sum_{t'=1}^{\\vert\\mathbf{y}\\vert}\\kappa(t, t') C^n_{t'} \\\\\n \\end{aligned}\n $$\nAt inference time, the confidence threshold for making an exit decision needs to be calibrated. Depth-adaptive transformer finds such a threshold on a validation set via grid search.\nCALM\n(\nSchuster et al. 2022\n) applied the Learn then Test (LTT) framework (\nAngelopoulos et al. 2021\n) to identify a subset of valid thresholds and chose the minimum value as the threshold for inference. Except for training per-layer exit classifier, CALM also explored other methods for adaptive depth prediction, including the softmax responses (i.e. difference between top two softmax outputs) and hidden state saturation (i.e. $\\cos(\\mathbf{h}^n_t, \\mathbf{h}^{n+1}_t)$) as confidence scores for exit decisions. They found softmax responses result in best inference speedup.\nEfficient Attention\n#\nThe computation and memory cost of the vanilla Transformer grows quadratically with sequence length and hence it is hard to be applied on very long sequences. Many efficiency improvements for Transformer architecture have something to do with the self-attention module - making it cheaper, smaller or faster to run. See the survey paper on\nEfficient Transformers\n(\nTay et al. 2020\n).\nSparse Attention Patterns\n#\nFixed Local Context\n#\nA simple alternation to make self-attention less expensive is to restrict the attention span of each token to\nlocal\ncontext only, so that self-attention grows linearly with the sequence length.\nThe idea was introduced by\nImage Transformer\n(\nParmer, et al 2018\n), which formulates image generation as sequence modeling using an encoder-decoder transformer architecture:\nThe encoder generates a contextualized, per-pixel-channel representation of the source image;\nThen the decoder autoregressively generates an output image, one channel per pixel at each time step.\nLet’s label the representation of the current pixel to be generated as the query $\\mathbf{q}$. Other positions whose representations will be used for computing $\\mathbf{q}$ are key vector $\\mathbf{k}_1, \\mathbf{k}_2, \\dots$ and they together form a memory matrix $\\mathbf{M}$. The scope of $\\mathbf{M}$ defines the context window for pixel query $\\mathbf{q}$.\nImage Transformer introduced two types of localized $\\mathbf{M}$, as illustrated below.\nIllustration of 1D and 2D attention span for visual inputs in Image Transformer. The black line marks a query block and the cyan outlines the actual attention span for pixel q. (Image source: Figure 2 in\nParmer et al, 2018\n)\n1D Local Attention\n: The input image is flattened in the\nraster scanning\norder, that is, from left to right and top to bottom. The linearized image is then partitioned into non-overlapping query blocks. The context window consists of pixels in the same query block as $\\mathbf{q}$ and a fixed number of additional pixels generated before this query block.\n2D Local Attention\n: The image is partitioned into multiple non-overlapping rectangular query blocks. The query pixel can attend to all others in the same memory blocks. To make sure the pixel at the top-left corner can also have a valid context window, the memory block is extended to the top, left and right by a fixed amount, respectively.\nStrided Context\n#\nSparse Transformer\n(\nChild et al., 2019\n) introduced\nfactorized self-attention\n, through sparse matrix factorization, making it possible to train dense attention networks with hundreds of layers on sequence length up to 16,384, which would be infeasible on modern hardware otherwise.\nGiven a set of attention connectivity pattern $\\mathcal{S} = \\{S_1, \\dots, S_n\\}$, where each $S_i$ records a set of key positions that the $i$-th query vector attends to.\n$$\n\\begin{aligned}\n\\text{Attend}(\\mathbf{X}, \\mathcal{S}) &= \\Big( a(\\mathbf{x}_i, S_i) \\Big)_{i \\in \\{1, \\dots, L\\}} \\\\\n\\text{ where } a(\\mathbf{x}_i, S_i) &= \\text{softmax}\\Big(\\frac{(\\mathbf{x}_i \\mathbf{W}^q)(\\mathbf{x}_j \\mathbf{W}^k)_{j \\in S_i}^\\top}{\\sqrt{d_k}}\\Big) (\\mathbf{x}_j \\mathbf{W}^v)_{j \\in S_i}\n\\end{aligned}\n$$\nNote that although the size of $S_i$ is not fixed, $a(\\mathbf{x}_i, S_i)$ is always of size $d_v$ and thus $\\text{Attend}(\\mathbf{X}, \\mathcal{S}) \\in \\mathbb{R}^{L \\times d_v}$.\nIn auto-regressive models, one attention span is defined as $S_i = \\{j: j \\leq i\\}$ as it allows each token to attend to all the positions in the past.\nIn factorized self-attention, the set $S_i$ is decomposed into a\ntree\nof dependencies, such that for every pair of $(i, j)$ where $j \\leq i$, there is a path connecting $i$ back to $j$ and $i$ can attend to $j$ either directly or indirectly.\nPrecisely, the set $S_i$ is divided into $p$\nnon-overlapping\nsubsets, where the $m$-th subset is denoted as $A^{(m)}_i \\subset S_i, m = 1,\\dots, p$. Therefore the path between the output position $i$ and any $j$ has a maximum length $p + 1$. For example, if $(j, a, b, c, \\dots, i)$ is a path of indices between $i$ and $j$, we would have $j \\in A_a^{(1)}, a \\in A_b^{(2)}, b \\in A_c^{(3)}, \\dots$, so on and so forth.\nSparse Factorized Attention\nSparse Transformer proposed two types of fractorized attention. It is easier to understand the concepts as illustrated in Fig. 10 with 2D image inputs as examples.\nThe top row illustrates the attention connectivity patterns in (a) Transformer, (b) Sparse Transformer with strided attention, and (c) Sparse Transformer with fixed attention. The bottom row contains corresponding self-attention connectivity matrices. Note that the top and bottom rows are not in the same scale. (Image source:\nChild et al., 2019\n+ a few of extra annotations.)\nStrided\nattention with stride $\\ell \\sim \\sqrt{n}$. This works well with image data as the structure is aligned with strides. In the image case, each pixel would attend to all the previous $\\ell$ pixels in the raster scanning order (naturally cover the entire width of the image) and then those pixels attend to others in the same column (defined by another attention connectivity subset).\n$$\n \\begin{aligned}\n A_i^{(1)} &= \\{ t, t+1, \\dots, i\\} \\text{, where } t = \\max(0, i - \\ell) \\\\\n A_i^{(2)} &= \\{j: (i-j) \\mod \\ell = 0\\}\n \\end{aligned}\n $$\nFixed\nattention. A small set of tokens summarize previous locations and propagate that information to all future locations.\n$$\n \\begin{aligned}\n A_i^{(1)} &= \\{j: \\lfloor \\frac{j}{\\ell} \\rfloor = \\lfloor \\frac{i}{\\ell} \\rfloor \\} \\\\\n A_i^{(2)} &= \\{j: j \\mod \\ell \\in \\{\\ell-c, \\dots, \\ell-1\\} \\}\n \\end{aligned}\n $$\nwhere $c$ is a hyperparameter. If $c=1$, it restricts the representation whereas many depend on a few positions. The paper chose $c\\in \\{ 8, 16, 32 \\}$ for $\\ell \\in \\{ 128, 256 \\}$.\nUse Factorized Self-Attention in Transformer\nThere are three ways to use sparse factorized attention patterns in Transformer architecture:\nOne attention type per residual block and then interleave them,\n$\\text{attn}(\\mathbf{X}) = \\text{Attend}(\\mathbf{X}, A^{(n \\mod p)}) \\mathbf{W}^o$, where $n$ is the index of the current residual block.\nSet up a single head which attends to locations that all the factorized heads attend to,\n$\\text{attn}(\\mathbf{X}) = \\text{Attend}(\\mathbf{X}, \\cup_{m=1}^p A^{(m)}) \\mathbf{W}^o $.\nUse a multi-head attention mechanism, but different from vanilla Transformer, each head might adopt a pattern presented above, 1 or 2. $\\rightarrow$ This option often performs the best.\nSparse Transformer also proposed a set of changes so as to train the Transformer up to hundreds of layers, including gradient checkpointing, recomputing attention & FF layers during the backward pass, mixed precision training, efficient block-sparse implementation, etc. Please check the\npaper\nfor more details or my previous post on\ntechniques for scaling up model training\n.\nBlockwise Attention\n(\nQiu et al. 2019\n) introduces a\nsparse block matrix\nto only allow each token to attend to a small set of other tokens. Each attention matrix of size $L \\times L$ is partitioned into $n \\times n$ smaller blocks of size $\\frac{L}{n}\\times\\frac{L}{n}$ and a sparse block matrix $\\mathbf{M} \\in \\{0, 1\\}^{L \\times L}$ is defined by a permutation $\\pi$ of ${1, \\dots, n}$, which records the column index per row in the block matrix.\n$$\n\\begin{aligned}\n\\text{attn}(\\mathbf{Q}, \\mathbf{K}, \\mathbf{V}, \\mathbf{M}) &= \\text{softmax}\\Big(\\frac{\\mathbf{Q}\\mathbf{K}^\\top}{\\sqrt{d}} \\odot \\mathbf{M}\\Big)\\mathbf{V} \\\\\n(\\mathbf{A} \\odot \\mathbf{M})_{ij} &= \\begin{cases}\nA_{ij} & \\text{if }M_{ij} = 1 \\\\\n-\\infty & \\text{if }M_{ij} = 0 \\\\\n\\end{cases} \\\\\n\\text{where } M_{ij} &= \\begin{cases}\n1 & \\text{if }\\pi\\big(\\lfloor\\frac{(i-1)n}{L} + 1\\rfloor\\big) = \\lfloor\\frac{(j-1)n}{L} + 1\\rfloor \\\\\n0 & \\text{otherwise}\n\\end{cases}\n\\end{aligned}\n$$\nThe actual implementation of Blockwise Attention only stores QKV as block matrices, each of size $n\\times n$:\n$$\n\\text{Blockwise-attn}(\\mathbf{Q}, \\mathbf{K}, \\mathbf{V}, \\mathbf{M}) = \\begin{bmatrix}\n\\text{softmax}\\big(\\frac{\\hat{\\mathbf{q}}_1\\hat{\\mathbf{k}}_{\\pi(1)}^\\top}{\\sqrt{d}} \\Big)\\hat{\\mathbf{v}}_{\\pi(1)} \\\\\n\\vdots \\\\\n\\text{softmax}\\big(\\frac{\\hat{\\mathbf{q}}_n\\hat{\\mathbf{k}}_{\\pi(n)}^\\top}{\\sqrt{d}} \\odot \\Big)\\hat{\\mathbf{v}}_{\\pi(n)} \\\\\n\\end{bmatrix}\n$$\nwhere $\\hat{\\mathbf{q}}_i$, $\\hat{\\mathbf{k}}_i$ and $\\hat{\\mathbf{v}}_i$ are the $i$-the row in the QKV block matrix respectively. Each $\\mathbf{q}_i\\mathbf{k}_{\\pi(i)}^\\top, \\forall i = 1, \\dots, n$ is of size $\\frac{N}{n}\\times\\frac{N}{n}$ and therefore Blockwise Attention is able to reduce the memory complexity of attention matrix from $\\mathcal{O}(L^2)$ to $\\mathcal{O}(\\frac{L}{n}\\times\\frac{L}{n} \\times n) = \\mathcal{O}(L^2/n)$.\nCombination of Local and Global Context\n#\nETC\n(\nExtended Transformer Construction\n;\nAinslie et al. 2019\n),\nLongformer\n(\nBeltagy et al. 2020\n) and\nBig Bird\n(\nZaheer et al. 2020\n) models combine both local and global context when building an attention matrix. All these models can be initialized from existing pretrained models.\nGlobal-Local Attention\nof\nETC\n(\nAinslie et al. 2019\n) takes two inputs, (1) the long input $\\mathbf{x}^l$ of size $n_l$ which is the regular input sequence and (2) the global input $\\mathbf{x}^g$ of size $n_g$ which contains a smaller number of auxiliary tokens, $n_g \\ll n_l$. Attention is thus split into four components based on directional attention across these two inputs: g2g, g2l, l2g and l2l. Because the l2l attention piece can be very large, it is restricted to a fixed size attention span of radius $w$ (i.e. local attention span) and the l2l matrix can be reshaped to $n_l \\times (2w+1)$.\nETC utilizes four binary matrices to handle structured inputs, $\\mathbf{M}^{g2g}$, $\\mathbf{M}^{g2l}$, $\\mathbf{M}^{l2g}$ and $\\mathbf{M}^{l2l}$. For example, each element $z^g_i \\in \\mathbb{R}^d$ in the attention output $z^g = (z^g_1, \\dots, z^g_{n_g})$ for g2g attention piece is formatted as:\n$$\n\\begin{aligned}\na^{g2g}_{ij} = \\frac{1}{\\sqrt{d}} x^g_i \\mathbf{W}^Q (x^g_j \\mathbf{W}^K + P^K_{ij})^\\top - (1- M^{g2g}_{ij})C \\\\\nA^{g2g}_{ij} = \\frac{\\exp(a^{g2g}_{ij})}{\\sum_{k=1}^{n_g} \\exp(a^{g2g}_{ik})} \\quad\nz^g_i = \\sum^{n_g}_{j=1} A^{g2g}_{ij} x^g_j \\mathbf{W}^V\n\\end{aligned}\n$$\nwhere $P^K_{ij}$ is a learnable vector for relative position encoding and $C$ is a very large constant ($C=10000$ in the paper) to offset any attention weights when mask is off.\nAttention patterns of ETC, Longformer and Big Bird.\nOne more update in ETC is to incorporate a CPC (contrastive predictive coding) task using\nNCE loss\ninto the pretraining stage, besides the\nMLM\ntask: The representation of one sentence should be similar to the representation of context around it when this sentence is masked.\nThe global input $\\mathbf{x}^g$ for ETC is constructed as follows: Assuming there are some segments within the long inputs (e.g. by sentence), each segment is attached with one auxiliary token to learn global inputs.\nRelative position encoding\nis used to mark the global segment tokens with the token position. Hard masking in one direction (i.e., tokens before vs after are labeled differently) is found to bring performance gains in some datasets.\nAttention pattern in Longformer contains three components:\nLocal attention\n: Similar to ETC, local attention is controlled by a sliding window of fixed size $w$;\nGlobal attention of preselected tokens\n: Longformer has a few pre-selected tokens (e.g.\n[CLS]\ntoken) assigned with global attention span, that is, attending to all other tokens in the input sequence.\nDilated attention\n: Dilated sliding window of fixed size $r$ and gaps of dilation size $d$, similar to Sparse Transformer;\nBig Bird\nis quite similar to Longformer, equipped with both local attention and a few preselected tokens with global attention span, but Big Bird replaces dilated attention with a new mechanism where all tokens attend to a set of random tokens. The design is motivated by the fact that attention pattern can be viewed as a\ndirected graph\nand a\nrandom graph\nhas the property that information is able to rapidly flow between any pair of nodes.\nLongformer\nuses smaller window size at lower layers and larger window sizes at higher layers. Ablation studies showed that this setup works better than reversed or fixed size config. Lower layers do not have dilated sliding windows to better learn to use immediate local context. Longformer also has a staged training procedure where initially the model is trained with small window size to learn from local context and then subsequent stages of training have window sizes increased and learning rate decreased.\nContent-based Attention\n#\nThe improvements proposed by\nReformer\n(\nKitaev, et al. 2020\n) aim to solve the following pain points in vanilla Transformer:\nQuadratic time and memory complexity within self-attention module.\nMemory in a model with $N$ layers is $N$-times larger than in a single-layer model because we need to store activations for back-propagation.\nThe intermediate FF layers are often quite large.\nReformer proposed two main changes:\nReplace the dot-product attention with\nlocality-sensitive hashing (LSH) attention\n, reducing the complexity from $\\mathcal{O}(L^2)$ to $\\mathcal{O}(L\\log L)$.\nReplace the standard residual blocks with\nreversible residual layers\n, which allows storing activations only once during training instead of $N$ times (i.e. proportional to the number of layers).\nLocality-Sensitive Hashing Attention\nIn $\\mathbf{Q} \\mathbf{K}^\\top$ part of the\nattention formula\n, we are only interested in the largest elements as only large elements contribute a lot after softmax. For each query $\\mathbf{q}_i \\in \\mathbf{Q}$, we are looking for row vectors in $\\mathbf{K}$ closest to $\\mathbf{q}_i$. In order to find nearest neighbors quickly in high-dimensional space, Reformer incorporates\nLocality-Sensitive Hashing (LSH)\ninto its attention mechanism.\nA hashing scheme $x \\mapsto h(x)$ is\nlocality-sensitive\nif it preserves the distancing information between data points, such that close vectors obtain similar hashes while distant vectors have very different ones. The Reformer adopts a hashing scheme as such, given a fixed random matrix $\\mathbf{R} \\in \\mathbb{R}^{d \\times b/2}$ (where $b$ is a hyperparam), the hash function is $h(x) = \\arg\\max([xR; −xR])$.\nIllustration of Locality-Sensitive Hashing (LSH) attention. (Image source: right part of Figure 1 in\nKitaev, et al. 2020\n).\nIn LSH attention, a query can only attend to positions in the same hashing bucket, $S_i = \\{j: h(\\mathbf{q}_i) = h(\\mathbf{k}_j)\\}$. It is carried out in the following process, as illustrated in Fig. 20:\n(a) The attention matrix for full attention is often sparse.\n(b) Using LSH, we can sort the keys and queries to be aligned according to their hash buckets.\n(c) Set $\\mathbf{Q} = \\mathbf{K}$ (precisely $\\mathbf{k}_j = \\mathbf{q}_j / |\\mathbf{q}_j|$), so that there are equal numbers of keys and queries in one bucket, easier for batching. Interestingly, this “shared-QK” config does not affect the performance of the Transformer.\n(d) Apply batching where chunks of $m$ consecutive queries are grouped together.\nThe LSH attention consists of 4 steps: bucketing, sorting, chunking, and attention computation. (Image source: left part of Figure 1 in\nKitaev, et al. 2020\n).\nReversible Residual Network\nAnother improvement by Reformer is to use\nreversible residual layers\n(\nGomez et al. 2017\n). The motivation for reversible residual network is to design the architecture in a way that activations at any given layer can be recovered from the activations at the following layer, using only the model parameters. Hence, we can save memory by recomputing the activation during backprop rather than storing all the activations.\nGiven a layer $x \\mapsto y$, the normal residual layer does $y = x + F(x)$, but the reversible layer splits both input and output into pairs $(x_1, x_2) \\mapsto (y_1, y_2)$ and then executes the following:\n$$\ny_1 = x_1 + F(x_2),\\; y_2 = x_2 + G(y_1) \n$$\nand reversing is easy:\n$$\nx_2 = y_2 - G(y_1), \\; x_1 = y_1 − F(x_2)\n$$\nReformer applies the same idea to Transformer by combination attention ($F$) and feed-forward layers ($G$) within a reversible net block:\n$$\nY_1 = X_1 + \\text{Attention}(X_2), \\; Y_2 = X_2 + \\text{FeedForward}(Y_1)\n$$\nThe memory can be further reduced by chunking the feed-forward computation:\n$$\nY_2 = [Y_2^{(1)}; \\dots; Y_2^{(c)}] = [X_2^{(1)} + \\text{FeedForward}(Y_1^{(1)}); \\dots; X_2^{(c)} + \\text{FeedForward}(Y_1^{(c)})]\n$$\nThe resulting reversible Transformer does not need to store activation in every layer.\nRouting Transformer\n(\nRoy et al. 2021\n) is also built on content-based clustering of keys and queries. Instead of using a static hashing function like LSH, it utilizes online $k$-means clustering and combines it with local, temporal sparse attention to reduce the attention complexity from $O(L^2)$ to $O(L^{1.5})$.\nWithin routing attention, both keys and queries are clustered with $k$-means clustering method and the same set of centroids $\\boldsymbol{\\mu} = (\\mu_1, \\dots, \\mu_k) \\in \\mathbb{R}^{k \\times d}$. Queries are routed to keys that get assigned to the same centroid. The total complexity is $O(Lkd + L^2d/k)$, where $O(Lkd)$ is for running clustering assignments and $O(L^2d/k)$ is for attention computation. The cluster centroids are updated by EMA (exponential moving average) using all associated keys and queries.\nIn the experiments for Routing Transformer, some best config only has routing attention enabled in the last two layers of the model and half of the attention heads, while the other half utilizing local attention. They also observed that local attention is a pretty strong baseline and larger attention window always leads to better results.\nLow-Rank Attention\n#\nLinformer\n(\nWang et al. 2020\n) approximates the full attention matrix with a\nlow rank\nmatrix, reducing the time & space complexity to be\nlinear\n. Instead of using expensive SVD to identify low rank decomposition, Linformer adds two linear projections $\\mathbf{E}_i, \\mathbf{F}_i \\in \\mathbb{R}^{L \\times k}$ for key and value matrices, respectively, reducing their dimensions from $L \\times d$ to $k \\times d$. As long as $k \\ll L$, the attention memory can be greatly reduced.\n$$\n\\begin{aligned}\n\\overline{\\text{head}}_i \n&= \\text{attn}(\\mathbf{X}_q\\mathbf{W}^q_i, \\mathbf{E}_i\\mathbf{X}_k\\mathbf{W}^k_i, \\mathbf{F}_i\\mathbf{X}_v\\mathbf{W}^v_i) \\\\\n&= \\underbrace{\\text{softmax}\\Big( \\frac{\\mathbf{X}_q\\mathbf{W}^q_i (\\mathbf{E}_i \\mathbf{X}_k\\mathbf{W}^k_i)^\\top}{\\sqrt{d}} \\Big)}_{\\text{low rank attention matrix }\\bar{A} \\in \\mathbb{R}^{k \\times d}} \\mathbf{F}_i \\mathbf{X}_v\\mathbf{W}^v_i\n\\end{aligned}\n$$\nAdditional techniques can be applied to further improve efficiency of Linformer:\nParameter sharing between projection layers, such as head-wise, key-value  and layer-wise (across all layers) sharing.\nUse different $k$ at different layers, as heads in higher layers tend to have a more skewed distribution (lower rank) and thus we can use smaller $k$ at higher layers.\nUse different types of projections; e.g. mean/max pooling, convolution layer with kernel and stride $L/k$.\n(Left) Informer has two projection layers added for keys and values. (Right) Plot of inference time as a function of sequence length. (Image source:\nWang et al. 2020\n).\nRandom Feature Attention\n(\nRFA\n;\nPeng et al. 2021\n) relies on\nrandom feature methods\n(\nRahimi & Recht, 2007\n) to approximate softmax operation in self-attention with low rank feature maps in order to achieve linear time and space complexity.\nPerformers\n(\nChoromanski et al. 2021\n) also adopts random feature attention with improvements on the kernel construction to further reduce the kernel approximation error.\nThe main theorem behind RFA is from\nRahimi & Recht, 2007\n:\nLet $\\phi: \\mathbb{R}^d \\to \\mathbb{R}^{2D}$ be a nonlinear transformation:\n$$\n\\phi(\\mathbf{x}) = \\frac{1}{\\sqrt{D}}[\\sin(\\mathbf{w}_1^\\top \\mathbf{x}), \\dots, \\sin(\\mathbf{w}_D^\\top \\mathbf{x}), \\cos(\\mathbf{w}_1^\\top \\mathbf{x}), \\dots, \\cos(\\mathbf{w}_D^\\top \\mathbf{x})]^\\top\n$$\nWhen $d$-dimensional random vectors $\\mathbf{w}_i$ are i.i.d. from $\\mathcal{N}(\\mathbf{0}, \\sigma^2\\mathbf{I}_d)$,\n$$\n\\mathbb{E}_{\\mathbf{w}_i} [\\phi(\\mathbf{x}) \\cdot \\phi(\\mathbf{y})] = \\exp(-\\frac{\\| \\mathbf{x} - \\mathbf{y} \\|^2}{2\\sigma^2})\n$$\nAn unbiased estimation of $\\exp(\\mathbf{x} \\cdot \\mathbf{y})$ is:\n$$\n\\begin{aligned}\n\\exp(\\mathbf{x} \\cdot \\mathbf{y} / \\sigma^2) \n&= \\exp(\\frac{1}{2\\sigma^2}(\\|\\mathbf{x}\\|^2 + \\|\\mathbf{y}\\|^2 - \\|\\mathbf{x} - \\mathbf{y}\\|^2) \\\\\n&= \\exp(\\frac{\\|\\mathbf{x}\\|^2}{2\\sigma^2}) \\exp(\\frac{\\|\\mathbf{y}\\|^2}{2\\sigma^2}) ( - \\frac{\\|\\mathbf{x} - \\mathbf{y}\\|^2}{2\\sigma^2}) \\\\\n&\\approx \\exp(\\frac{\\|\\mathbf{x}\\|^2}{2\\sigma^2}) \\exp(\\frac{\\|\\mathbf{y}\\|^2}{2\\sigma^2})\\;\\phi(\\mathbf{x})\\cdot\\phi(\\mathbf{y}) \\\\\n&= \\exp(\\frac{1}{\\sigma^2})\\;\\phi(\\mathbf{x})\\cdot\\phi(\\mathbf{y}) & \\text{; unit vectors}\n\\end{aligned}\n$$\nThen we can write the attention function as follows, where $\\otimes$ is outer product operation and $\\sigma^2$ is the temperature:\n$$\n\\begin{aligned}\n\\text{attn}(\\mathbf{q}_t, \\{\\mathbf{k}_i\\}, \\{\\mathbf{v}_i\\}) \n&= \\sum_i \\frac{\\exp(\\mathbf{q}_t\\cdot\\mathbf{k}_i/\\sigma^2)}{\\sum_j \\exp(\\mathbf{q}_t\\cdot\\mathbf{k}_j/\\sigma^2)}\\mathbf{v}_i^\\top\n\\approx \\sum_i \\frac{\\phi(\\mathbf{q}_t)\\phi(\\mathbf{k}_i)\\mathbf{v}_i^\\top}{\\sum_j \\phi(\\mathbf{q}_t)\\phi(\\mathbf{k}_j)} \\\\\n&= \\color{green}{\\frac{\\phi(\\mathbf{q}_t)^\\top \\sum_i \\phi(\\mathbf{k}_i)\\otimes\\mathbf{v}_i}{\\phi(\\mathbf{q}_t)^\\top \\sum_j \\phi(\\mathbf{k}_j)}\n= \\text{RFA}(\\mathbf{q}_t, \\{\\mathbf{k}_i\\}, \\{\\mathbf{v}_i\\})}\n\\end{aligned}\n$$\n(Left) The order of computation for default softmax operation. (Right) The order of computation when using random feature attention, a lot cheaper than default softmax. (Image source:\nPeng et al. 2021\n).\nCausal Attention RFA\nhas token at time step $t$ only attend to earlier keys and values $\\{\\mathbf{k}_i\\}_{i \\leq t}, \\{\\mathbf{v}_i\\}_{i \\leq t}$. Let us use a tuple of variables, $(\\mathbf{S}_t \\in \\mathbb{R}^{2D \\times d}, \\mathbf{z} \\in \\mathbb{R}^{2D})$, to track the hidden state history at time step $t$, similar to RNNs:\n$$\n\\begin{aligned}\n&\\text{causal-RFA}(\\mathbf{q}_t, \\{\\mathbf{k}_i\\}_{i \\leq t}, \\{\\mathbf{v}_i\\}_{i \\leq t}) = \\frac{\\phi(\\mathbf{q}_t)^\\top \\mathbf{S}_t}{\\phi(\\mathbf{q}_t) \\cdot \\mathbf{z}_t} \\\\\n&\\text{where } \n\\mathbf{S}_t = \\mathbf{S}_{t-1} + \\phi(\\mathbf{k}_t)\\otimes\\mathbf{v}_t,\n\\quad \n\\mathbf{z}_t = \\mathbf{z}_{t-1} + \\phi(\\mathbf{k}_t)\n\\end{aligned}\n$$\nwhere $2D$ is the size of $\\phi(.)$ and $D$ should be no less than the model size $d$ for reasonable approximation.\nRFA leads to significant speedup in autoregressive decoding and the memory complexity mainly depends on the choice of $D$ when constructing the kernel $\\phi(.)$.\nPerformer modifies the random feature attention with positive random feature maps to reduce the estimation error. It also keeps the randomly sampled $\\mathbf{w}_1, \\dots, \\mathbf{w}_D$ to be orthogonal to further reduce the variance of the estimator.\nComparison of approximation error when using (Left) i.i.d vs orthogonal features and (Right) sin/cos vs positive random features. (Image source:\nChoromanski et al. 2021\n).\nTransformers for Reinforcement Learning\n#\nThe self-attention mechanism avoids compressing the whole past into a fixed-size hidden state and does not suffer from vanishing or exploding gradients as much as RNNs. Reinforcement Learning tasks can for sure benefit from these traits.\nHowever\n, it is quite difficult to train Transformer even in supervised learning, let alone in the RL context. It could be quite challenging to stabilize and train a LSTM agent by itself, after all.\nThe\nGated Transformer-XL\n(\nGTrXL\n;\nParisotto, et al. 2019\n) is one attempt to use Transformer for RL. GTrXL succeeded in stabilizing training with two changes on top of\nTransformer-XL\n:\nThe layer normalization is only applied on the input stream in a residual module, but NOT on the shortcut stream. A key benefit to this reordering is to allow the original input to flow from the first to last layer.\nThe residual connection is replaced with a GRU-style (Gated Recurrent Unit;\nChung et al., 2014\n)\ngating\nmechanism.\n$$\n\\begin{aligned}\nr &= \\sigma(W_r^{(l)} y + U_r^{(l)} x) \\\\\nz &= \\sigma(W_z^{(l)} y + U_z^{(l)} x - b_g^{(l)}) \\\\\n\\hat{h} &= \\tanh(W_g^{(l)} y + U_g^{(l)} (r \\odot x)) \\\\\ng^{(l)}(x, y) &= (1-z)\\odot x + z\\odot \\hat{h}\n\\end{aligned}\n$$\nThe gating function parameters are explicitly initialized to be close to an identity map - this is why there is a $b_g$ term. A $b_g > 0$ greatly helps with the learning speedup.\nComparison of the model architecture of Transformer-XL, Transformer-XL with the layer norm reordered, and Gated Transformer-XL. (Image source: Figure 1 in\nParisotto, et al. 2019\n)\nDecision Transformer\n(\nDT\n;\nChen et al 2021\n) formulates Reinforcement Learning problems as a process of\nconditional sequence modeling\n, outputting the optimal actions conditioned on the desired return, past states and actions. It therefore becomes straightforward to use Transformer architecture. Decision Transformer is for\noff-policy RL\n, where the model only has access to a fixed collection of trajectories collected by other policies.\nTo encourage the model to learn how to act in order to achieve a desired return, it feeds the model with desired future return $\\hat{R} = \\sum_{t’=t}^T r_{t’}$ instead of the current reward. The trajectory consists of a list of triplets, (return-to-go $\\hat{R}_t, state $s_t$, action $a_t$), and it is used as an input sequence for Transformer:\n$$\n\\tau = (\\hat{R}_1, s_1, a_1, \\hat{R}_2, s_2, a_2, \\dots, \\hat{R}_T, s_T, a_T)\n$$\nThree linear layers are added and trained for return-to-go, state and action respectively to extract token embeddings. The prediction head learns to predict $a_t$ corresponding to the input token $s_t$. The training uses cross-entropy loss for discrete actions or MSE for continuous actions. Predicting the states or return-to-go was not found to help improve the performance in their experiments.\nThe experiments compared DT with several model-free RL algorithm baselines and showed that:\nDT is more efficient than behavior cloning in low data regime;\nDT can model the distribution of returns very well;\nHaving a long context is crucial for obtaining good results;\nDT can work with sparse rewards.\nCitation\n#\nCited as:\nWeng, Lilian. (Jan 2023). The transformer family version 2.0. Lil’Log. https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/.\nOr\n@article{weng2023transformer,\n  title   = \"The Transformer Family Version 2.0\",\n  author  = \"Weng, Lilian\",\n  journal = \"lilianweng.github.io\",\n  year    = \"2023\",\n  month   = \"Jan\",\n  url     = \"https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/\"\n}\nReferences\n#\n[1] Ashish Vaswani, et al.\n“Attention is all you need.”\nNIPS 2017.\n[2] Rami Al-Rfou, et al.\n“Character-level language modeling with deeper self-attention.”\nAAAI 2019.\n[3] Olah & Carter,\n“Attention and Augmented Recurrent Neural Networks”\n, Distill, 2016.\n[4] Sainbayar Sukhbaatar, et al.\n“Adaptive Attention Span in Transformers”\n. ACL 2019.\n[5] Rewon Child, et al.\n“Generating Long Sequences with Sparse Transformers”\narXiv:1904.10509 (2019).\n[6] Nikita Kitaev, et al.\n“Reformer: The Efficient Transformer”\nICLR 2020.\n[7] Alex Graves. (“Adaptive Computation Time for Recurrent Neural Networks”)[https://arxiv.org/abs/1603.08983]\n[8] Niki Parmar, et al.\n“Image Transformer”\nICML 2018.\n[9] Zihang Dai, et al.\n“Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context.”\nACL 2019.\n[10] Aidan N. Gomez, et al.\n“The Reversible Residual Network: Backpropagation Without Storing Activations”\nNIPS 2017.\n[11] Mostafa Dehghani, et al.\n“Universal Transformers”\nICLR 2019.\n[12] Emilio Parisotto, et al.\n“Stabilizing Transformers for Reinforcement Learning”\narXiv:1910.06764 (2019).\n[13] Rae et al.\n“Compressive Transformers for Long-Range Sequence Modelling.”\n2019.\n[14] Press et al.\n“Train Short, Test Long: Attention With Linear Biases Enables Input Length Extrapolation.”\nICLR 2022.\n[15] Wu, et al.\n“DA-Transformer: Distance Aware Transformer”\n2021.\n[16] Elabyad et al.\n“Depth-Adaptive Transformer.”\nICLR 2020.\n[17] Schuster et al.\n“Confident Adaptive Language Modeling”\n2022.\n[18] Qiu et al.\n“Blockwise self-attention for long document understanding”\n2019\n[19] Roy et al.\n“Efficient Content-Based Sparse Attention with Routing Transformers.”\n2021.\n[20] Ainslie et al.\n“ETC: Encoding Long and Structured Inputs in Transformers.”\nEMNLP 2019.\n[21] Beltagy et al.\n“Longformer: The long-document transformer.”\n2020.\n[22] Zaheer et al.\n“Big Bird: Transformers for Longer Sequences.”\n2020.\n[23] Wang et al.\n“Linformer: Self-Attention with Linear Complexity.”\narXiv preprint arXiv:2006.04768 (2020).\n[24] Tay et al. 2020\n“Sparse Sinkhorn Attention.”\nICML 2020.\n[25] Peng et al.\n“Random Feature Attention.”\nICLR 2021.\n[26] Choromanski et al.\n“Rethinking Attention with Performers.”\nICLR 2021.\n[27] Khandelwal et al.\n“Generalization through memorization: Nearest neighbor language models.”\nICLR 2020.\n[28] Yogatama et al.\n“Adaptive semiparametric language models.”\nACL 2021.\n[29] Wu et al.\n“Memorizing Transformers.”\nICLR 2022.\n[30] Su et al.\n“Roformer: Enhanced transformer with rotary position embedding.”\narXiv preprint arXiv:2104.09864 (2021).\n[31] Shaw et al.\n“Self-attention with relative position representations.”\narXiv preprint arXiv:1803.02155 (2018).\n[32] Tay et al.\n“Efficient Transformers: A Survey.”\nACM Computing Surveys 55.6 (2022): 1-28.\n[33] Chen et al.,\n“Decision Transformer: Reinforcement Learning via Sequence Modeling”\narXiv preprint arXiv:2106.01345 (2021).",
  "metadata": {
    "title": "The Transformer Family Version 2.0 | Lil'Log",
    "description": "Many new Transformer architecture improvements have been proposed since my last post on “The Transformer Family” about three years ago. Here I did a big refactoring and enrichment of that 2020 post — restructure the hierarchy of sections and improve many sections with more recent papers. Version 2.0 is a superset of the old version, about twice the length.\nNotations\n\n\t\n\t\t\t\n\t\t\t\t\tSymbol\n\t\t\t\t\tMeaning\n\t\t\t\n\t\n\t\n\t\t\t\n\t\t\t\t\t$d$\n\t\t\t\t\tThe model size / hidden state dimension / positional encoding size.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$h$\n\t\t\t\t\tThe number of heads in multi-head attention layer.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$L$\n\t\t\t\t\tThe segment length of input sequence.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$N$\n\t\t\t\t\tThe total number of attention layers in the model; not considering MoE.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{X} \\in \\mathbb{R}^{L \\times d}$\n\t\t\t\t\tThe input sequence where each element has been mapped into an embedding vector of shape $d$, same as the model size.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^k \\in \\mathbb{R}^{d \\times d_k}$\n\t\t\t\t\tThe key weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^q \\in \\mathbb{R}^{d \\times d_k}$\n\t\t\t\t\tThe query weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^v \\in \\mathbb{R}^{d \\times d_v}$\n\t\t\t\t\tThe value weight matrix. Often we have $d_k = d_v = d$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^k_i, \\mathbf{W}^q_i \\in \\mathbb{R}^{d \\times d_k/h}; \\mathbf{W}^v_i \\in \\mathbb{R}^{d \\times d_v/h}$\n\t\t\t\t\tThe weight matrices per head.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^o \\in \\mathbb{R}^{d_v \\times d}$\n\t\t\t\t\tThe output weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{Q} = \\mathbf{X}\\mathbf{W}^q \\in \\mathbb{R}^{L \\times d_k}$\n\t\t\t\t\tThe query embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{K} = \\mathbf{X}\\mathbf{W}^k \\in \\mathbb{R}^{L \\times d_k}$\n\t\t\t\t\tThe key embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{V} = \\mathbf{X}\\mathbf{W}^v \\in \\mathbb{R}^{L \\times d_v}$\n\t\t\t\t\tThe value embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{q}_i, \\mathbf{k}_i \\in \\mathbb{R}^{d_k}, \\mathbf{v}_i \\in \\mathbb{R}^{d_v}$\n\t\t\t\t\tRow vectors in query, key, value matrices, $\\mathbf{Q}$, $\\mathbf{K}$ and $\\mathbf{V}$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$S_i$\n\t\t\t\t\tA collection of key positions for the $i$-th query $\\mathbf{q}_i$ to attend to.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{A} \\in \\mathbb{R}^{L \\times L}$\n\t\t\t\t\tThe self-attention matrix between a input sequence of lenght $L$ and itself. $\\mathbf{A} = \\text{softmax}(\\mathbf{Q}\\mathbf{K}^\\top / \\sqrt{d_k})$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$a_{ij} \\in \\mathbf{A}$\n\t\t\t\t\tThe scalar attention score between query $\\mathbf{q}_i$ and key $\\mathbf{k}_j$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{P} \\in \\mathbb{R}^{L \\times d}$\n\t\t\t\t\tposition encoding matrix, where the $i$-th row $\\mathbf{p}_i$ is the positional encoding for input $\\mathbf{x}_i$.\n\t\t\t\n\t\n\nTransformer Basics\nThe Transformer (which will be referred to as “vanilla Transformer” to distinguish it from other enhanced versions; Vaswani, et al., 2017) model has an encoder-decoder architecture, as commonly used in many NMT models. Later simplified Transformer was shown to achieve great performance in language modeling tasks, like in encoder-only BERT or decoder-only GPT.",
    "keywords": "architecture, attention, transformer, foundation, long-read, reinforcement-learning",
    "author": "Lilian Weng",
    "og:title": "The Transformer Family Version 2.0",
    "og:description": "Many new Transformer architecture improvements have been proposed since my last post on “The Transformer Family” about three years ago. Here I did a big refactoring and enrichment of that 2020 post — restructure the hierarchy of sections and improve many sections with more recent papers. Version 2.0 is a superset of the old version, about twice the length.\nNotations\n\n\t\n\t\t\t\n\t\t\t\t\tSymbol\n\t\t\t\t\tMeaning\n\t\t\t\n\t\n\t\n\t\t\t\n\t\t\t\t\t$d$\n\t\t\t\t\tThe model size / hidden state dimension / positional encoding size.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$h$\n\t\t\t\t\tThe number of heads in multi-head attention layer.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$L$\n\t\t\t\t\tThe segment length of input sequence.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$N$\n\t\t\t\t\tThe total number of attention layers in the model; not considering MoE.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{X} \\in \\mathbb{R}^{L \\times d}$\n\t\t\t\t\tThe input sequence where each element has been mapped into an embedding vector of shape $d$, same as the model size.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^k \\in \\mathbb{R}^{d \\times d_k}$\n\t\t\t\t\tThe key weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^q \\in \\mathbb{R}^{d \\times d_k}$\n\t\t\t\t\tThe query weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^v \\in \\mathbb{R}^{d \\times d_v}$\n\t\t\t\t\tThe value weight matrix. Often we have $d_k = d_v = d$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^k_i, \\mathbf{W}^q_i \\in \\mathbb{R}^{d \\times d_k/h}; \\mathbf{W}^v_i \\in \\mathbb{R}^{d \\times d_v/h}$\n\t\t\t\t\tThe weight matrices per head.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^o \\in \\mathbb{R}^{d_v \\times d}$\n\t\t\t\t\tThe output weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{Q} = \\mathbf{X}\\mathbf{W}^q \\in \\mathbb{R}^{L \\times d_k}$\n\t\t\t\t\tThe query embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{K} = \\mathbf{X}\\mathbf{W}^k \\in \\mathbb{R}^{L \\times d_k}$\n\t\t\t\t\tThe key embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{V} = \\mathbf{X}\\mathbf{W}^v \\in \\mathbb{R}^{L \\times d_v}$\n\t\t\t\t\tThe value embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{q}_i, \\mathbf{k}_i \\in \\mathbb{R}^{d_k}, \\mathbf{v}_i \\in \\mathbb{R}^{d_v}$\n\t\t\t\t\tRow vectors in query, key, value matrices, $\\mathbf{Q}$, $\\mathbf{K}$ and $\\mathbf{V}$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$S_i$\n\t\t\t\t\tA collection of key positions for the $i$-th query $\\mathbf{q}_i$ to attend to.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{A} \\in \\mathbb{R}^{L \\times L}$\n\t\t\t\t\tThe self-attention matrix between a input sequence of lenght $L$ and itself. $\\mathbf{A} = \\text{softmax}(\\mathbf{Q}\\mathbf{K}^\\top / \\sqrt{d_k})$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$a_{ij} \\in \\mathbf{A}$\n\t\t\t\t\tThe scalar attention score between query $\\mathbf{q}_i$ and key $\\mathbf{k}_j$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{P} \\in \\mathbb{R}^{L \\times d}$\n\t\t\t\t\tposition encoding matrix, where the $i$-th row $\\mathbf{p}_i$ is the positional encoding for input $\\mathbf{x}_i$.\n\t\t\t\n\t\n\nTransformer Basics\nThe Transformer (which will be referred to as “vanilla Transformer” to distinguish it from other enhanced versions; Vaswani, et al., 2017) model has an encoder-decoder architecture, as commonly used in many NMT models. Later simplified Transformer was shown to achieve great performance in language modeling tasks, like in encoder-only BERT or decoder-only GPT.",
    "og:type": "article",
    "og:url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/",
    "twitter:card": "summary",
    "twitter:title": "The Transformer Family Version 2.0",
    "twitter:description": "Many new Transformer architecture improvements have been proposed since my last post on “The Transformer Family” about three years ago. Here I did a big refactoring and enrichment of that 2020 post — restructure the hierarchy of sections and improve many sections with more recent papers. Version 2.0 is a superset of the old version, about twice the length.\nNotations\n\n\t\n\t\t\t\n\t\t\t\t\tSymbol\n\t\t\t\t\tMeaning\n\t\t\t\n\t\n\t\n\t\t\t\n\t\t\t\t\t$d$\n\t\t\t\t\tThe model size / hidden state dimension / positional encoding size.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$h$\n\t\t\t\t\tThe number of heads in multi-head attention layer.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$L$\n\t\t\t\t\tThe segment length of input sequence.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$N$\n\t\t\t\t\tThe total number of attention layers in the model; not considering MoE.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{X} \\in \\mathbb{R}^{L \\times d}$\n\t\t\t\t\tThe input sequence where each element has been mapped into an embedding vector of shape $d$, same as the model size.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^k \\in \\mathbb{R}^{d \\times d_k}$\n\t\t\t\t\tThe key weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^q \\in \\mathbb{R}^{d \\times d_k}$\n\t\t\t\t\tThe query weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^v \\in \\mathbb{R}^{d \\times d_v}$\n\t\t\t\t\tThe value weight matrix. Often we have $d_k = d_v = d$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^k_i, \\mathbf{W}^q_i \\in \\mathbb{R}^{d \\times d_k/h}; \\mathbf{W}^v_i \\in \\mathbb{R}^{d \\times d_v/h}$\n\t\t\t\t\tThe weight matrices per head.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{W}^o \\in \\mathbb{R}^{d_v \\times d}$\n\t\t\t\t\tThe output weight matrix.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{Q} = \\mathbf{X}\\mathbf{W}^q \\in \\mathbb{R}^{L \\times d_k}$\n\t\t\t\t\tThe query embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{K} = \\mathbf{X}\\mathbf{W}^k \\in \\mathbb{R}^{L \\times d_k}$\n\t\t\t\t\tThe key embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{V} = \\mathbf{X}\\mathbf{W}^v \\in \\mathbb{R}^{L \\times d_v}$\n\t\t\t\t\tThe value embedding inputs.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{q}_i, \\mathbf{k}_i \\in \\mathbb{R}^{d_k}, \\mathbf{v}_i \\in \\mathbb{R}^{d_v}$\n\t\t\t\t\tRow vectors in query, key, value matrices, $\\mathbf{Q}$, $\\mathbf{K}$ and $\\mathbf{V}$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$S_i$\n\t\t\t\t\tA collection of key positions for the $i$-th query $\\mathbf{q}_i$ to attend to.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{A} \\in \\mathbb{R}^{L \\times L}$\n\t\t\t\t\tThe self-attention matrix between a input sequence of lenght $L$ and itself. $\\mathbf{A} = \\text{softmax}(\\mathbf{Q}\\mathbf{K}^\\top / \\sqrt{d_k})$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$a_{ij} \\in \\mathbf{A}$\n\t\t\t\t\tThe scalar attention score between query $\\mathbf{q}_i$ and key $\\mathbf{k}_j$.\n\t\t\t\n\t\t\t\n\t\t\t\t\t$\\mathbf{P} \\in \\mathbb{R}^{L \\times d}$\n\t\t\t\t\tposition encoding matrix, where the $i$-th row $\\mathbf{p}_i$ is the positional encoding for input $\\mathbf{x}_i$.\n\t\t\t\n\t\n\nTransformer Basics\nThe Transformer (which will be referred to as “vanilla Transformer” to distinguish it from other enhanced versions; Vaswani, et al., 2017) model has an encoder-decoder architecture, as commonly used in many NMT models. Later simplified Transformer was shown to achieve great performance in language modeling tasks, like in encoder-only BERT or decoder-only GPT.",
    "article:section": "posts",
    "article:published_time": "2023-01-27T00:00:00+00:00",
    "article:modified_time": "2023-01-27T00:00:00+00:00"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/bd2a1c61491f/code/body-000.txt",
      "chars": 272,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/multi-head-attention.png",
      "bytes": 106503,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-cfd86d1d041d4e587dd2.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/multi-head-attention.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/transformer.png",
      "bytes": 376109,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-57e8505d379090b6a5c2.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/transformer.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/sinoidual-positional-encoding.png",
      "bytes": 66540,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-eb72333120a9bf61db22.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/sinoidual-positional-encoding.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/RoPE.png",
      "bytes": 517373,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-94806f3485ee4c7d4137.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/RoPE.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/transformer-XL-training.png",
      "bytes": 420355,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-10c3f7a4432b939f41e9.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/transformer-XL-training.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/compressive-transformer.png",
      "bytes": 110695,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-9af813630ff27b91e9f8.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/compressive-transformer.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/compressive-transformer-memory.png",
      "bytes": 104844,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-2d52fb35835b8ffb9f5c.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/compressive-transformer-memory.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/SPALM2.png",
      "bytes": 441748,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-b512db56c30433d4c252.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/SPALM2.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/memorizing-transformer.png",
      "bytes": 81702,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-701853184c4d260333df.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/memorizing-transformer.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/ALiBi-bias.png",
      "bytes": 75150,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-440b77c01199a8fa47a8.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/ALiBi-bias.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/ALiBi-exp.png",
      "bytes": 308982,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-321a0b363e9a6d1aa98b.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/ALiBi-exp.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/universal-transformer-loop.png",
      "bytes": 136227,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-65f1b4ee283f16066c80.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/universal-transformer-loop.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/universal-transformer.png",
      "bytes": 118602,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-f6d1acbfbe9a2ebc5720.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/universal-transformer.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/attention-per-head.png",
      "bytes": 128314,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-6801941ec960782dbcdd.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/attention-per-head.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/soft-masking-function.png",
      "bytes": 49638,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-aa7a4901bbb57ffba3a8.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/soft-masking-function.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/depth-adaptive-classifier.png",
      "bytes": 144593,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-61300c84161edec245ef.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/depth-adaptive-classifier.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/image-transformer-attention.png",
      "bytes": 92635,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-e33cbf677e130ee48f58.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/image-transformer-attention.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/sparse-attention.png",
      "bytes": 244624,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-f5d3b8b9ebc8beb78c77.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/sparse-attention.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/combined-attention.png",
      "bytes": 135708,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-30a7b1c7f2d502623782.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/combined-attention.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/LSH-attention-matrix.png",
      "bytes": 181556,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-528e7d446bed9aeb5917.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/LSH-attention-matrix.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/LSH-attention.png",
      "bytes": 306995,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-08fa8454100a632db05c.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/LSH-attention.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/linformer.png",
      "bytes": 296726,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-5fade868b98da3fc740e.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/linformer.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/RFA.png",
      "bytes": 150541,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-1864a0421abec41db939.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/RFA.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/performer.png",
      "bytes": 300912,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-9bbffc5733b6530c7240.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/performer.png",
      "alt": "",
      "downloaded": true
    },
    {
      "status": "ok",
      "url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/gated-transformer-XL.png",
      "bytes": 278596,
      "reused": false,
      "local_path": "attachments/bd2a1c61491f/images/body-8ed5854240a34aa93f46.png",
      "src": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/gated-transformer-XL.png",
      "alt": "",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://lilianweng.github.io/posts/2020-04-07-the-transformer-family/",
      "label": "“The Transformer Family”",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1706.03762",
      "label": "Vaswani, et al., 2017",
      "type": "a"
    },
    {
      "href": "https://lilianweng.github.io/posts/2018-06-24-attention/",
      "label": "NMT",
      "type": "a"
    },
    {
      "href": "https://lilianweng.github.io/posts/2019-01-31-lm/",
      "label": "BERT",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Non-local_means",
      "label": "non-local means",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1705.03122",
      "label": "Gehring, et al. 2017",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1808.04444",
      "label": "Al-Rfou et al. 2018",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1803.02155",
      "label": "Shaw et al. (2018)",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1901.02860",
      "label": "Dai et al., 2019",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2104.09864",
      "label": "Su et al. 2021",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Rotation_matrix",
      "label": "rotation matrix",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1911.05507",
      "label": "Rae et al. 2019",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1911.00172",
      "label": "Khandelwal et al. 2020",
      "type": "a"
    },
    {
      "href": "https://github.com/facebookresearch/faiss",
      "label": "FAISS",
      "type": "a"
    },
    {
      "href": "https://github.com/google-research/google-research/tree/master/scann",
      "label": "ScaNN",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2102.02557",
      "label": "Yogatama et al. 2021",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2203.08913",
      "label": "Wu et al. 2022",
      "type": "a"
    },
    {
      "href": "https://aclanthology.org/2021.naacl-main.166",
      "label": "Wu, et al. 2021",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2108.12409",
      "label": "Press et al. 2022",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1807.03819",
      "label": "Dehghani, et al. 2019",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1610.02357",
      "label": "separable convolution",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1905.07799",
      "label": "Sukhbaatar et al (2019)",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1910.10073",
      "label": "Elabyad et al. 2020",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2207.07061",
      "label": "Schuster et al. 2022",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Dirac_delta_function",
      "label": "dirac delta",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2110.01052",
      "label": "Angelopoulos et al. 2021",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2009.06732",
      "label": "Tay et al. 2020",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1802.05751",
      "label": "Parmer, et al 2018",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Raster_scan",
      "label": "raster scanning",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1904.10509",
      "label": "Child et al., 2019",
      "type": "a"
    },
    {
      "href": "https://lilianweng.github.io/posts/2021-09-25-train-large/",
      "label": "techniques for scaling up model training",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1911.02972",
      "label": "Qiu et al. 2019",
      "type": "a"
    },
    {
      "href": "https://aclanthology.org/2020.emnlp-main.19/",
      "label": "Ainslie et al. 2019",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2004/05150",
      "label": "Beltagy et al. 2020",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2007.14062",
      "label": "Zaheer et al. 2020",
      "type": "a"
    },
    {
      "href": "https://lilianweng.github.io/posts/2021-05-31-contrastive/",
      "label": "NCE loss",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Directed_graph",
      "label": "directed graph",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Random_graph",
      "label": "random graph",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2001.04451",
      "label": "Kitaev, et al. 2020",
      "type": "a"
    },
    {
      "href": "https://en.wikipedia.org/wiki/Locality-sensitive_hashing",
      "label": "Locality-Sensitive Hashing (LSH)",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1707.04585",
      "label": "Gomez et al. 2017",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2003.05997",
      "label": "Roy et al. 2021",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2006.04768",
      "label": "Wang et al. 2020",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2103.02143",
      "label": "Peng et al. 2021",
      "type": "a"
    },
    {
      "href": "https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf",
      "label": "Rahimi & Recht, 2007",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2009.14794",
      "label": "Choromanski et al. 2021",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1910.06764",
      "label": "Parisotto, et al. 2019",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/1412.3555",
      "label": "Chung et al., 2014",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2106.01345",
      "label": "Chen et al 2021",
      "type": "a"
    },
    {
      "href": "https://lilianweng.github.io/posts/2018-02-19-rl-overview/",
      "label": "off-policy RL",
      "type": "a"
    },
    {
      "href": "http://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf",
      "label": "“Attention is all you need.”",
      "type": "a"
    },
    {
      "href": "http://doi.org/10.23915/disti",
      "label": "“Attention and Augmented Recurrent Neural Networks”",
      "type": "a"
    },
    {
      "href": "https://arxiv.org/abs/2002.11296",
      "label": "“Sparse Sinkhorn Attention.”",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:51:12Z",
  "source_url": "https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/",
  "papers": [
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1706.03762",
      "bytes": 2215244,
      "reused": true,
      "local_path": "papers/repair/c3ab01e47e614fe3cf33.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1705.03122",
      "bytes": 573311,
      "reused": false,
      "local_path": "papers/repair/758d2db10bbacb1a4049.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1808.04444",
      "bytes": 726375,
      "reused": false,
      "local_path": "papers/repair/f7a4eccd97bbfdf62ee3.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1803.02155",
      "bytes": 235373,
      "reused": false,
      "local_path": "papers/repair/f32df5aeac8714d373be.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1901.02860",
      "bytes": 4570910,
      "reused": false,
      "local_path": "papers/repair/8514ec65c2953ee68ccf.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2104.09864",
      "bytes": 599289,
      "reused": true,
      "local_path": "papers/repair/b64aa8c5e278403b64ef.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1911.05507",
      "bytes": 584093,
      "reused": false,
      "local_path": "papers/repair/244ccad8664efd18e77c.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1911.00172",
      "bytes": 451453,
      "reused": false,
      "local_path": "papers/repair/43f160fbf0fa81d77393.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2102.02557",
      "bytes": 1018474,
      "reused": false,
      "local_path": "papers/repair/b5c9e91bd79c1b9ef250.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2203.08913",
      "bytes": 1047591,
      "reused": false,
      "local_path": "papers/repair/691204beb9e3b5b17a09.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2108.12409",
      "bytes": 759030,
      "reused": true,
      "local_path": "papers/repair/238b286ab4f338940590.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1807.03819",
      "bytes": 2371799,
      "reused": false,
      "local_path": "papers/repair/933b1c22c5c54fb670e9.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1610.02357",
      "bytes": 804466,
      "reused": false,
      "local_path": "papers/repair/9aca5987caa4ba692507.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1905.07799",
      "bytes": 381620,
      "reused": false,
      "local_path": "papers/repair/71b8a5bd387ec43c94cc.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1910.10073",
      "bytes": 480824,
      "reused": false,
      "local_path": "papers/repair/929832372275e4ca6acd.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2207.07061",
      "bytes": 3006455,
      "reused": false,
      "local_path": "papers/repair/3e493b8dcc10d3f89436.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2110.01052",
      "bytes": 2934353,
      "reused": false,
      "local_path": "papers/repair/f48952de2c7830ad0fc7.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2009.06732",
      "bytes": 847824,
      "reused": true,
      "local_path": "papers/repair/5442c18a5d3ec7c8eca5.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1802.05751",
      "bytes": 1108088,
      "reused": false,
      "local_path": "papers/repair/e71a6184638a873330d8.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1904.10509",
      "bytes": 4436935,
      "reused": false,
      "local_path": "papers/repair/b20b3e8f92bfa6d99c9d.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1911.02972",
      "bytes": 716586,
      "reused": false,
      "local_path": "papers/repair/e8d64283562a841b64f9.pdf"
    },
    {
      "status": "http_404",
      "url": "https://arxiv.org/pdf/2004/05150",
      "retried_at": "2026-10-20T18:03:23Z"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2007.14062",
      "bytes": 1153823,
      "reused": false,
      "local_path": "papers/repair/f56d7d4a1317893ab3ad.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2001.04451",
      "bytes": 658515,
      "reused": false,
      "local_path": "papers/repair/362c024f67418a537ddf.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1707.04585",
      "bytes": 543235,
      "reused": false,
      "local_path": "papers/repair/5416fd20cbe5674f95d1.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2003.05997",
      "bytes": 727287,
      "reused": false,
      "local_path": "papers/repair/d2f0607b565e49aa10fe.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2006.04768",
      "bytes": 922879,
      "reused": false,
      "local_path": "papers/repair/f20cd0da3a7534f96d1a.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2103.02143",
      "bytes": 1241988,
      "reused": false,
      "local_path": "papers/repair/e438f2e0a286bcd79311.pdf"
    },
    {
      "status": "ok",
      "url": "https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf",
      "bytes": 488453,
      "reused": false,
      "local_path": "papers/repair/38e94b7d3685dc5b22d9.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2009.14794",
      "bytes": 5880177,
      "reused": true,
      "local_path": "papers/repair/26948568ccf004004970.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1910.06764",
      "bytes": 2179274,
      "reused": false,
      "local_path": "papers/repair/51f85523d9bf32a25ab8.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/1412.3555",
      "bytes": 683896,
      "reused": false,
      "local_path": "papers/repair/50678b2017a71df661ae.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2106.01345",
      "bytes": 704784,
      "reused": false,
      "local_path": "papers/repair/9e056437cb3e5b4f695d.pdf"
    },
    {
      "status": "ok",
      "url": "http://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf",
      "bytes": 569417,
      "reused": false,
      "local_path": "papers/repair/95fa654499d86d006948.pdf"
    },
    {
      "status": "ok",
      "url": "https://arxiv.org/pdf/2002.11296",
      "bytes": 450138,
      "reused": true,
      "local_path": "papers/repair/0008d87e3417535c824a.pdf"
    }
  ],
  "actual_video_files": [],
  "body_chars": 68430,
  "expected_body_images": 25,
  "discovery_status": "ok",
  "delivery_status": "partial",
  "missing_reason": "http_404",
  "raw_html": "attachments/bd2a1c61491f/pages/source.html",
  "local_html": "attachments/bd2a1c61491f/pages/article.html",
  "body_base_path": "attachments/bd2a1c61491f/pages/",
  "body_markdown_path": "attachments/bd2a1c61491f/pages/article.md"
}