Transformer (machine learning model)

The basic building blocks of the Transformer are scaled dot-product attention units. When a sentence is passed into a Transformer model, attention weights are calculated between every token simultaneously. The attention unit produces embeddings for every token in context that contain information not only about the token itself, but also a weighted combination of other relevant tokens weighted by the attention weights.

Concretely, for each attention unit the Transformer model learns three weight matrices; the query weights ${\displaystyle W_{Q}}$, the key weights ${\displaystyle W_{K}}$, and the value weights ${\displaystyle W_{V}}$. For each token ${\displaystyle i}$, the input word embedding ${\displaystyle x_{i}}$ is multiplied with each of the three weight matrices to produce a query vector ${\displaystyle q_{i}=x_{i}W_{Q}}$, a key vector ${\displaystyle k_{i}=x_{i}W_{K}}$, and a value vector ${\displaystyle v_{i}=x_{i}W_{V}}$. Attention weights are calculated using the query and key vectors: the attention weight ${\displaystyle a_{ij}}$ from token ${\displaystyle i}$ to token ${\displaystyle j}$ is the dot product between ${\displaystyle q_{i}}$ and ${\displaystyle k_{j}}$. The attention weights are divided by the square root of the dimension of the key vectors, ${\displaystyle {\sqrt {d_{k}}}}$, which stabilizes gradients during training, and passed through a softmax which normalizes the weights to sum to ${\displaystyle 1}$. The fact that ${\displaystyle W_{Q}}$ and ${\displaystyle W_{K}}$ are different matrices allows attention to be non-symmetric: if token ${\displaystyle i}$ attends to token ${\displaystyle j}$ (i.e. ${\displaystyle q_{i}\cdot k_{j}}$ is large), this does not necessarily mean that token ${\displaystyle j}$ will attend to token ${\displaystyle i}$ (i.e. ${\displaystyle q_{j}\cdot k_{i}}$ is large). The output of the attention unit for token ${\displaystyle i}$ is the weighted sum of the value vectors of all tokens, weighted by ${\displaystyle a_{ij}}$, the attention from ${\displaystyle i}$ to each token.

The attention calculation for all tokens can be expressed as one large matrix calculation, which is useful for training due to computational matrix operation optimizations which make matrix operations fast to compute. The matrices ${\displaystyle Q}$, ${\displaystyle K}$ and ${\displaystyle V}$ are defined as the matrices where the ${\displaystyle i}$th rows are vectors ${\displaystyle q_{i}}$, ${\displaystyle k_{i}}$, and ${\displaystyle v_{i}}$ respectively.

${\displaystyle {\begin{aligned}{\text{Attention}}(Q,K,V)={\text{softmax}}\left({\frac {QK^{\mathrm {T} }}{\sqrt {d_{k}}}}\right)V\end{aligned}}}$

One set of ${\displaystyle \left(W_{Q},W_{K},W_{V}\right)}$ matrices is called an attention head, and each layer in a Transformer model has multiple attention heads. While one attention head attends to the tokens that are relevant to each token, with multiple attention heads the model can learn to do this for different definitions of "relevance". Research has shown that many attention heads in Transformers encode relevance relations that are interpretable by humans. For example there are attention heads that, for every token, attend mostly to the next word, or attention heads that mainly attend from verbs to their direct objects.[7] Since Transformer models have multiple attention heads, they have the possibility of capturing many levels and types of relevance relations, from surface-level to semantic. The multiple outputs for the multi-head attention layer are concatenated to pass into the feed-forward neural network layers.

Each encoder consists of two major components: a self-attention mechanism and a feed-forward neural network. The self-attention mechanism takes in a set of input encodings from the previous encoder and weighs their relevance to each other to generate a set of output encodings. The feed-forward neural network then further processes each output encoding individually. These output encodings are finally passed to the next encoder as its input, as well as the decoders.

The first encoder takes positional information and embeddings of the input sequence as its input, rather than encodings. The positional information is necessary for the Transformer to make use of the order of the sequence, because no other part of the Transformer makes use of this.[1]

Each decoder consists of three major components: a self-attention mechanism, an attention mechanism over the encodings, and a feed-forward neural network. The decoder functions in a similar fashion to the encoder, but an additional attention mechanism is inserted which instead draws relevant information from the encodings generated by the encoders.[1][6]

Like the first encoder, the first decoder takes positional information and embeddings of the output sequence as its input, rather than encodings. Since the transformer should not use the current or future output to predict an output though, the output sequence must be partially masked to prevent this reverse information flow.[1] The last decoder is followed by a final linear transformation and softmax layer, to produce the output probabilities over the vocabulary.

Below is pseudo code for an implementation of the Transformer variant known as the "vanilla" transformer:

def vanilla_transformer(enc_inp, dec_inp):
    """Transformer variant known as the "vanilla" transformer."""
    x = embedding(enc_inp) * sqrt(d_m)
    x = x + pos_encoding(x)
    x = dropout(x)
    for _ in range(n_enc_layers):
        attn = multi_head_attention(x, x, x, None)
        attn = dropout(attn)
        attn = layer_normalization(x + attn)

        x = pointwise_ff(attn) # GELU activation(affine map(attn))
        x = layer_normalization(x + attn)

    # x is at this point the output of the encoder
    enc_out = x

    x = embedding(dec_inp) * sqrt(d_m)
    x = x + pos_encoding(x)
    x = dropout(x)
    mask = causal_mask(x)
    for _ in range(n_dec_layers):
        attn1 = multi_head_attention(x, x, x, mask)
        attn1 = layer_normalization(attn1 + x)

        attn2 = multi_head_attention(attn1, enc_out, enc_out, None)
        attn2 = dropout(attn2)
        attn2 = layer_normalization(attn1 + attn2)

        x = pointwise_ff(attn2)
        x = layer_normalization(attn2 + x)
    return dense(x)

Training Transformer-based architectures can be very expensive, especially for long sentences.[8] Alternative architectures include the Reformer, which reduces the computational load from ${\displaystyle O(N^{2})}$ to ${\displaystyle O(N\ln N)}$, where ${\displaystyle N}$ is the length of the sequence. This is done using locality-sensitive hashing and reversible layers.[9][10]

A benchmark for comparing different xformer architecture was introduced in late 2020.[11]