The primary goal of langugae models is next-word/next-sequence prediction. The transformer architecture is built on the premise of “attention”, this was developed to solve the precursor’s weakness in modeling long-length text sequences (namely recursive neural networks).

In its most basic format, an RNN is a reduce over a sequence with a carried accumulator (in machine-learning terms, the current state).

For context, reduce(function, iterable, initializer) goes through the iterable from left to right, applying the function at each call. It takes the output of function(accumulator, item) at each call, returning the new accumulator and moves to the next item. The initializer is the starting accumulator/state

In machine learning/stats notation, the forward pass of an RNN is:

from numpy import tanh
from functools import reduce

reduce(lambda h, x: tanh(Whx @ x + Whh @ h + bh), inputs, h0)

h0 is the initial accumulator (state). inputs are the iterable being folded, in this case the sequence of embeddings representing your input string/text. The lambda function passes the projection of the inputs Whx @ x summed to the projection of the previous state Whh @ h and sums that to a bias bh. The tanh nonlinearity squashes the output (which is easier for training an acutal model in terms of stabilizing the function gradients).

Note that @ in Python is matrix multiplication. Unrolling the reduce, it comes down to

h <- h0
h <- tanh(Wxh @ x1 + Whh @ h + bh) where h is h0
h <- tanh(Wxh @ x2 + Whh @ h + bh) where h is the output of the previous step 

Note, I’m using “input” here, but it’s embedded versions of the inputs, so a vectorized representation of input tokens

The two main issues with RNNs are as follows:

1. Sequential dependency: An RNN goes through each cell n times, and it much more difficult to parallelize, since by definition each step must be done in sequence/prior to the next step
2. Path length/long-range dependencies: the prior state from more than 1 step before is not directly accessible, which means long-range dependencies are harder to learn.

What Attention Aims to Solve#

RNNs lose context over longer sequences, and do not do well in modeling context amongst input tokens. The way to make this most obvious is in the classic example of language translation:

In English: “The books that the professor recommended during the very long and tedious lecture last Monday are expensive” translates to “Les livres que le professeur a recommandés … sont chers”. The two long-distance dependencies are subject-ver number agreement livres and sont (not est) and recommandés carries plural, because the direct object livres precedes it.

Confusing French asside, this highlights the non-linear relation/importance between the input tokens and its impact on the outputs. Both of these words end up at the end of the output French target, but the actual controlling token livres appears at the start of the source (about 10 tokens apart), each of which overwrite/update the state of the prior pass as the input sequence is passed.

Calculating The Context Vector#

Attention aims to solve this problem by enriching the input representation called a context vector. Instead of working directly on the embeddedings of input tokens, these embeddings are “enriched” using element-specific context vectors. Each input token would have a context vector corresponding to every other input token in the input sequence, which is calculated as follows:

Step I - Get the Attention Scores#

$$ \omega^{ij} = x^i \cdot x^j $$

This attention score for each input element $j$ is specific to the input element $i$, and is taken as the dot-product of the input embeddings. This results in a row vector $[\omega^{i1},\omega^{i2},…,\omega^{iT}]$ containing the attention scores for each input element with respect to the input element $i$.

Step II - Normalize The Scores into Weights#

This is the application of some nonlinear function to normalize the input scores into weights $\alpha^i$. This ensures that extreme vector values are made tractable so when these weights are used in training, the gradients are more stable with the added benefit of having the weights always being positive (in the case of using the softmax function)

Step III - Create the Context Vector#

This is taken as a sum of all the input embeddings multiplied by the respective attention weights to give the context vector for the element-specific input:

$$ Z^i = \sum_{\text{j to T}} \alpha^{ij} \times x^j $$

Making This Trainable#

The previous section covered the steps for calculating the context vector. However, in order for the model to learn these vectors, it must be made trainable. This is done by introducing three new sets of matrices.

The intuition is as follows:

1. Take a vector for element $i$ and multiply it by a vector for every other element $1$ to $T$ to get the attention scores
2. Take these outputs and normalize to get the attention weights
3. Finally take the attention weights and multiply it by some other learned vector for each input elements

As formalized in the Vaswani paper, these three vectors mentioned are called Query, Key and Value vectors, are produced by multiplying the input embeddings by learned projection matrices $W^Q, W^K, W^V$. Instead of taking the dot-product on the raw embeddings, each embedding is first projected into three separate roles:

• query: the thing I’m looking for
• key: how to find the other things/how the others find the current element
• value: what is actually passed on

Formally, we get the attention scores: $$ \omega^i = Q^i\cdot \hat{K}^j\ \text{For the entire input sequence: }\ \hat{\omega} = \hat{Q}\hat{K}^T\ $$

The attention weights are a normalized version of $\omega^i$.

Finally, the context vector is found by multiplying the value vector by each of the attention weights: $$ Z_2 = \omega^2\times \hat{V} $$

There’s one thing to note in the scaling, prior the softmax function was used as the scaling function; however, the actual scaling function implemented in the paper is: $$ \sigma\left(\frac{QK^T}{\sqrt{d_k}}\right) $$

This ensures that the score distribution is exactly 1 regardless of the size of the matrix. This is done since having one logit significantly larger than the rest, the gradients of the softmax function collapse towards zero, and large-variance logits push softmax into this low-gardient regime causing training to stall.

Ironically, writing the actual trainable module to implement this is super simple in Tensorflow:

import tensorflow as tf

class SelfAttention(tf.keras.layers.Layer):
    def __init__(self, d_in, d_out):
        # these three lines set up our trainable weights 
        # in theory we could also parameterize the initializer, 
        #    or set custom inputs shapes for each
        self.W_query = self.add_weight(shape=(d_in, d_out), initializer="uniform")
        self.W_key = self.add_weight(shape=(d_in, d_out), initializer="uniform")
        self.W_value = self.add_weight(shape=(d_in, d_out), initializer="uniform")

    def call(self, x):
        keys = x @ self.W_key # get the key matrix
        queries = x @ self.W_query # get the query matrix
        values = x @ self.W_value # get the value matrix

        # calculate the attention scores
        #   transpose to ensure element-by-element multplication
        attn_scores = queries @ tf.transnspose(keys)
        attn_weights = tf.nn.softmax(
                attn_scores / tf.shape(keys)[-1] ** 0.5
            , axis=-1)
        return attn_weights @ values # return the context vector

References#

• Vaswani et al., Attention Is All You Need (NeurIPS 2017)
• Pascanu, Mikolov & Bengio, On the Difficulty of Training Recurrent Neural Networks (ICML 2013, arXiv:1211.5063)
• Bahdanau, Cho & Bengio (2015)