Attention

Let's define $q, K, V$ as a query, set of keys, and set of values respectively:

$$q = \begin{pmatrix}\{q_1, q_2, \dots, q_l \}\end{pmatrix} \\[7mm] K = \begin{pmatrix}\{k^{(1)}_1, \dots, k^{(1)}_l \}\\ \vdots \hspace{.25in} \vdots \hspace{.25in} \vdots \\ \{k^{(m)}_1, \dots, k^{(m)}_l \} \end{pmatrix} \\[7mm] V = \begin{pmatrix}\{v^{(1)}_1, \dots, v^{(1)}_l \}\\ \vdots \hspace{.25in} \vdots \hspace{.25in} \vdots \\ \{v^{(m)}_1, \dots, v^{(m)}_l \} \end{pmatrix} \\[7mm] K, V \in \mathbb{R}^{m \times l \times h}$$

where $m$ is batch size and $l$ is sequence length.

We can define attention as:

$$(1) \hspace{.25in} \text{Attention}(q, \hat{K}, \hat{V}) \coloneqq \sum_{t = 0}^{l} \alpha(q,\hat{K})_t \odot \hat{V}_t \rightarrow \mathbb{R}^{h} \\[3mm] \hat{Q}, \hat{K}, \hat{V} \in \mathbb{R}^{l \times h} \\[3mm] q, \hat{K}_t, \hat{V}_t \in \mathbb{R}^h$$

where $h$ is the dimensionality of the value space and $\alpha$ is an arbitrary weighting function, generating scalar values as the attention scores.

Attention, in essence, is taking a query, $q$, and a set of keys, $K$, computing a similarity score denoted by the function $\alpha(\cdot) \rightarrow \mathbb{R}$ which is then multiplied element-wise by value vector, $v$.

$$(2) \hspace{.25in} \alpha(q, \hat{K})_t = \frac{\exp(a(q, \hat{K}_t))}{\sum_{j=0}^l\exp(a(q, \hat{K}_j))} \rightarrow \mathbb{R}$$

where $\alpha$ is constructed as the softmax of an arbitrary similarity function $a$, used to compute the similarity of a given $q$ and $\hat{K}_t$.

We're essentially computing a similarity $a(\cdot)$, and normalizing it to be $\in [0, 1]$, to be interpretable as a probability distribution if we consider all other similarity scores for all other $t$, $\alpha(\cdot)_t$.

The sum $\sum_{t}^l$ shows that for sequence length $l$, we're pooling all $l$ attention weights by multiplying element wise with $\hat{V}_t$ (as we broadcast $\alpha(q, \hat{K})$ to a vector $\in \mathbb{R}^h$), and then summing over $l$. You might as well express this as a dot product between a matrix of attention weights $\vec{\alpha}(\cdot) \in \mathbb{R}^{l \times h}$ and a matrix of value vectors $V^{\top} \in \mathbb{R}^{(h \times l)}$, where we get $\rightarrow \mathbb{R}^{l \times l}$, where $l$ is the sequence length. This is essentially how the attention mechanism becomes parallelizable in large-scale training of neural networks, as you can easily distribute the matrix operations amongst different H200s.

Another property of the summation ( or matrix multiplication, if you follow the above intuition ), is that the resulting attention vector corresponding to the value vector at $t$ will have a magnitude of a fraction of the original $|\hat{V}_t|$ as $\alpha(q, \hat{K})$ is normalized by the $\text{softmax}(\cdot)$, which turns the operation to be a convex combination ensuring that the output at $t$ stays within the convex hull of $\mathcal{V}$, which proves to be useful for stabilizing training.

Intuitively, imagine summing a set of vectors where each vector isn't normalized by a given attenion weight, and instead has unbounded magnitude which increaes as $l$ grows. This could destabilize training as gradients can become seemingly very large. It also allows a model to express the attention vectors given a constraint, which can act as a form of "forcing" the model to learn the right values, as it must focus on distributing attention weights properly amongst different $\hat{V}_t$.

REMARK
In a Seq2Seq model, $V$ and $K$, are the set of all hidden states for the decoder (at a current layer) up to the current time $t' = l$. $q$ is the final hidden state of the encoder at a given layer (at the equivalent layer for the decoder).

Interpreting attention (as described below) can become more intuitive if $q, K, V$ are seen as such, where the decoder hidden state at $t - 1$, as $q$, is the query for the set of encoder hidden state as keys $\hat{K}_t$, to "extract" the similarity score between the two, helping compute how much attention should be paid to the final hidden state of the decoder (value vector) given a query, $q$, at any $t \in [1, l]$, where $l$ is the sequence length at the current time step.

For each time step in the forward pass of a given RNN cell, the decoder hidden state ($q$) at $t-1$, varies at every $t$, allowing $K$ and $V$ to denote how much attention should be paid to $q$ in the forward pass at time step $t$.

Note that there can be two interpretations for time $t$ here ($t$ and $t'$). In the recurrent forward pass for the RNN, $t'$ is defined as the current time step or token in the sequence which the RNN is processing. In the summation for attention, $t$ is defined as an intermediate time step $\in [0, l]$.

Now, in the definition for $\text{Attention}$, we iterate over the sequence of length $l$, extracting different $\hat{V}_t$ and $\hat{K}_t$ at each $t$ to compute how "important" or how much the value vector at should be attended to for a given $q$, (given by the element-wise multiplication of the $t$th attention score, $\alpha(q, \hat{K})_t$ and the value vector at $t$, $\hat{V}_t$) and then sum over the sequence, to encode in a final $\text{Attention}$ vector how much "attention" should be paid to $\hat{V}_t$ given a $q$.

The most straightforward way to compute the similarity, $a(\cdot)$ is the dot product (or the scaled dot product as in scaled dot product attention), as it's scalar output denotes a measure for similarity.

Reminder that $\vec{a}^\top\vec{b} = |\vec{a}| |\vec{b}| \cos \theta$, thereby the larger positive the dot product of the two vectors are, the smaller $\theta$ is and thereby, the more similar the vectors are.

The inverse is true for large negative and a dot product leading to $0$ denotes orthogonality, indicating no similarity but no notion of dissimilarity.

Bahdanau Attention

Neural Machine Translation by Jointly Learning to Align and Translate

Bahdanau Attention, is defined similarly as above, with the difference being that $a(\cdot)$ is defined as:

$$(3) \hspace{.25in} a(\mathbf q, \mathbf k_t) = w^\top \textrm{tanh}( W_q q + W_k k_t) \in \mathbb{R},$$

where $W_q$ and $W_k$ are transformation matrices transforming the vectors $q$ and $k_t$ from their respective dimensionality (query space and value space), to the dimensionality of attention space.

$w^\top$ is the vector which encodes $\tanh(\cdot)$ into a final scalar value, which serves as the attention score logit, prior to the $\text{softmax}$ in eq. $2$.

Note that the query space and the key space don't need to be of the same dimensionality, as $W_q$ and $W_k$ transform both vectors into the same dimensionality transform both vectors into the same dimensionality.

It's difficult to precisely know what $W_q$ and $W_k$ are really doing to both $q$ and $k$, neural networks are black boxes. But more generally, we can see both matrices as encoding the relevant information from both $q$ and $k$ which enables the final attention score to be effectively computed. It all really pieces together when we factor in how gradient descent finds the optimal $W$ to minimize the loss $\mathcal{L}$.

Again, noting the dot product as a metric for similarity between two vectors, perhaps we can interpret as the more similar $q$ or $k$ are to the set of respective row vectors in $W_k, W_q$, the higher magnitude the input to $\tanh$ will have and therefore, the larger the attention score will be if $w$ is also similar in direction to the output of $\tanh$.

Coded out,

import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt

def bahdanau(X, query_space, key_space, value_space, attention_space):
    batch_size, embed_dim = X.size(0), X.size(1)
    
    W1_q, W1_k, W1_v = torch.randn(query_space, embed_dim), torch.randn(key_space, embed_dim), torch.randn(value_space, embed_dim)
    W2_q, W2_k, w2_v = torch.randn(attention_space, query_space), torch.randn(attention_space, key_space), torch.randn(attention_space, 1)
    
    w1_q_mul, w1_k_mul, w1_v_mul = W1_q.unsqueeze(0).expand(batch_size, -1, -1), W1_k.unsqueeze(0).expand(batch_size, -1, -1), W1_v.unsqueeze(0).expand(batch_size, -1, -1)
    w2_q_mul, w2_k_mul, w2_v_mul = W2_q.unsqueeze(0).expand(batch_size, -1, -1), W2_k.unsqueeze(0).expand(batch_size, -1, -1), w2_v.T.unsqueeze(0).expand(batch_size, -1, -1)
    
    Q, K, V = torch.bmm(w1_q_mul, X), torch.bmm(w1_k_mul, X), torch.bmm(w1_v_mul, X)
    A_q, A_k = torch.bmm(w2_q_mul, Q), torch.bmm(w2_k_mul, K)
    attention_logits = torch.bmm(w2_v_mul, torch.tanh(A_q + A_k))
    
    attention_scores = F.softmax(attention_logits, dim=2)
    attention = torch.mul(attention_scores, V)
    attention_pool = torch.sum(attention, dim=2)
    
    return attention, attention_scores, attention_pool

if __name__ == "__main__":
    batch_size, hidden_neurons, seq_len = 2, 10, 15
    query_space, value_space, key_space, attention_space = 18, 20, 19, 30
    
    X = torch.randn(batch_size, hidden_neurons, seq_len)
    attention, attention_scores, attention_pool = bahdanau(X, query_space, key_space, value_space, attention_space)
    
    plt.imshow(attention_scores.view(2, 20), cmap='hot', interpolation='nearest')
    plt.title('Attention. All we need.', {'size': 30})
    plt.show()

Of course, now in modern LLMs, we typically use the scaled dot product, $\frac{q^Tk}{\sqrt{d}}$ where $d$ is the dimensionality of $k$ and $q$.

Multihead Attention

Given a set of queries, keys, and values, it's a good idea to consider a model that can compute attention pooling in multiple different manners which capture different dependencies in a sequence (long-range vs short-range for instance). This is akin to having multiple kernels, $\mathcal{K}$ in a convolutional layer, which are optimized to extract different features of the input feature map. Each set of parameters happens to capture different patterns, in the ConvNet perhaps different edge patterns, in the attention head perhaps different semantic meaning.

Therefore, it's beneficial to allow a model to jointly use different representation spaces of the query space, value space, and key space.

Instead of performing a single attention pooling operation, we can compute $h$ independent attention pooling operations.

Each $h$ attention pooling output is called a "head", hence, multi-head attention.

Assuming a single token.

Given a query $q \in \mathbb{R}^{h_q}$, $k \in \mathbb{R}^{h_k}$, $v \in \mathbb{R}^{h_v}$, we can compute a single attention head as,

$$h_i = f(W_i^qq, W_i^kk, W_i^vv) \in \mathbb{R}^{h_v}$$

where $f$ is a similarity function, $W_i^q \in \mathbb{R}^{(h_a \times h_q)}$, $W_i^k \in \mathbb{R}^{(h_a \times h_k)}$, $W_i^v \in \mathbb{R}^{(h_a \times h_v)}$, $h_a$ is the dimension of the attention space.

Including something such dot product operation ( no scaling, it's simpler notation, less LaTeX to type out ),

$$h_i = \sum_{t=0}^{l}\underbrace{\text{softmax} \left( \frac{(W_i^q q)^\top \cdot (W_i^k k_t)}{\sum_{j=1}^{l} \exp((W_i^q q)^\top \cdot (W_i^k k_j)} \right)}_{\alpha_{i,t}}(W_i^{v}v_t) \in \mathbb{R}^{h_v}$$

where $\alpha_{i, t}$ is the attention weight for the $i$th head at token $t$.

After computing all $h_i$, we concatenate all $h$ into a single vector,

$$H_{\text{cat}} = \text{Concat}(h_1, h_2, \dots, h_n) \in \mathbb{R^{n \cdot h_v}}$$

where $n$ is the count of attention heads.

Finally, we compute through a fully connected layer,

$$H_{out} = W^oH_{\text{cat}} \in \mathbb{R}^{h_e}\\[3mm]$$

where $h_e$ is the original dimensionality of the original embedding space and $W^{(o)} \in \mathbb{R}^{(h_e \times (n\cdot h_v))}$.

Positional Encoding

Assuming we're using attention in a transformer, and naturally, transformers don't have recurrent connections as RNNs, we want a way to embed positional information of a token.

Otherwise, the model wouldn't be able to to tell which token comes first, second, $\dots$, last, and knowing which $t$ $W_i^vv$ to pay attention to the most would become difficult as there is no notion if which word comes first, second, etc in the sequence.

Assume input embeddings of size $h_e$, then we can create positional vectors of same size $h_e$,

$$p_{i, 2j} = \sin\left(\frac{i}{1000^{\frac{2j}{d}}}\right) \\[3mm] p_{i, {2j + 1}} = \cos\left(\frac{i}{1000^{\frac{2j}{d}}}\right)$$

where $i$ is the $i$th row of the input matrix, $X \in \mathbb{R}^{l \times h_e}$ ($h_e$ is embedding dimension and $l$ is sequence length), and $j$ is the $j$th column of the input matrix.

So $\sin$ is computed for every row, every even column, while $\cos$ is computed for every row, odd column.

Ultimately we get a matrix, $\mathcal{P} \in \mathbb{R}^{l \times h_e}$ holding the positional encodings for a given input sequence.

This allows us to effectively encode positional embeddings as we introduce deterministic values for position, given deterministic computation of $\sin(\cdot)$ and $\cos(\cdot)$. Mathematically, this introduces (approximately for larger $d$) orthogonality between each vector, guaranteeing linear independence, which indicates that each row vector in $\mathcal{P}$ is able to capture a unique direction in the row-space.

As a simple example, consider $a = \begin{bmatrix} 1 \\ 0 \end{bmatrix}, b = \begin{bmatrix} 0 \\ 1 \end{bmatrix}$.

These two vectors are orthogonal (orthonormal in fact), and this guarantees linear independence as both vectors capture unique information in different directions (in the $\hat{i}$ and $\hat{j}$ respectively) and you can't express $a$ in terms of $b$.

Positional encodings approximate orthogonality for larger $d$ (<s>curse</s> blessing of dimensionality), and therefore allow us to capture unique information, in this context with regards to position.

Coded out,

class PositionalEncoding(nn.Module):
    
    def __init__(self, dropout_p: float, d_model: int, max_len: int):
        super().__init__()
       
        self.dropout = nn.Dropout(p=dropout_p)

        position = torch.arange(max_len).unsqueeze(1)  
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))

        pe = torch.zeros(size=(max_len, d_model))  
        pe[:, 0::2] = torch.sin(position * div_term) 
        pe[:, 1::2] = torch.cos(position * div_term) 

        pe = pe.unsqueeze(0) 
        self.register_buffer('pe', pe)

    def forward(self, x):
        pe = self.pe[:, :x.size(1)] 
        return pe

Masked Softmax

Assume we compute the raw attention scores $a(\cdot)$ in matrix form, prior to the softmax, as:

$$Z_{a} = \frac{QK^\top}{\sqrt{d}}$$

where $Q \in \mathbb{R}^{l \times h_q}$ and $K \in \mathbb{R}^{l \times h_k}$, and $Z_q \in \mathbb{R}^{l \times l}$ where $l$ is sequence length and $h_q = h_k$.

We transpose $K$, in such that the matrix multiplication yields an $l \times l$ matrix and we get the raw attention scores with respect to the sequence (length) and not some meaningless arbitrary dimension ( with respect to our task ) $h_q$.

Assume $l = 4$, then $\text{softmax}(Z_a)$:

$$Z_a = \begin{bmatrix} .2, .5, .9, .6 \\ .4, .4, .9, .3 \\ .1, .9, .3, .6 \\ .7, .2, .9, .2 \end{bmatrix}\\[10mm] A = \begin{bmatrix} 0.1708 , 0.2305 , 0.3439 , 0.2548 \\ 0.2196 , 0.2196 , 0.3621 , 0.1987 \\ 0.1641 , 0.3651 , 0.2004 , 0.2705 \\ 0.2912 , 0.1766 , 0.3556 , 0.1766 \end{bmatrix}$$

$A$ is size $l \times l$, which can be interpreted as follows. Let's rename our variables for now in context of $A$, where $A$ is now $\mathbb{R}^{i \times j}, i = j$. For every $j$th column vector, it's $i$th (row) value denotes how much "attention" weighting a the $j$th token in the sequence pays attention to all other $i$ tokens (or equivalently $l$) in the sequence.

Note that once we compute the matrix of attention weights, we compute a matrix multiplication with $W^v{V} \in \mathbb{R}^{h_a \times l}$ where $W^v \in \mathbb{R}^{h_a \times h_v}$ and $V$ in this case is $\in \mathbb{R}^{h_v \times l}$

$$A(W^vV)^\top \in \mathbb{R}^{l \times h_a}$$

Remember that in $W^vV$ each column vector corresponds to a single token $t$ in the sequence, given the column size $l$.

Therefore, in the matrix multiplication of $A(W^vV)^\top \in \mathbb{R}^{l \times h_a}$, we're multiplying each $t$ value column vector (or row if we account for the transpose) in the attention space. Note that at time $t$, we have $(W^vV)_t$ but at $t$, we don't expect the given $t$ token to have seen any other tokens at $t + n < l$.

Therefore, for any attention weights in $A$, we need to zero out a portion of the values.

Transformers

Following up on the iterations that led to the formulation for the attention mechanism, Vaswani et al., formulated the Transformer architecture in their seminal paper, Attention is All you Need.

I explain this in depth here.