Attention mechanism

An attention mechanism is a component of machine learning models that allows a model to assign different levels of importance to different elements of an input when producing an output or representation.

Attention mechanisms are widely used in transformers and have applications in natural language processing, computer vision, speech processing, and other machine-learning tasks.

A simplified representation is:

$\text{Input}\to \text{Attention Weights}\to \text{Weighted Information}\to \text{Output}$

The attention mechanism allows a model to focus computationally on information that is more relevant to the current prediction or representation.

Attention mechanism components

A standard attention mechanism uses three types of representations:

• queries;
• keys;
• values.

A simplified representation is:

$Q,K,V\to \text{Attention Mechanism}\to \text{Output}$

where:

• $Q$ = query representations;
• $K$ = key representations;
• $V$ = value representations.

The queries determine what information is being sought, the keys determine how information is matched against the queries, and the values contain the information that is aggregated.

Scaled dot-product attention

A commonly used form of attention is scaled dot-product attention.

It is defined as:

$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$

where:

• $Q$ = query matrix;
• $K$ = key matrix;
• $V$ = value matrix;
• $d_k$ = dimensionality of the key vectors.

The matrix product $Q{K}^T$ produces similarity scores between queries and keys.

The scores are scaled by $\sqrt{d_k}$ and passed through the softmax function to produce attention weights.

The resulting weights are applied to the value matrix.

Attention weights

For a particular query, attention weights can be represented as:

${\alpha }_{ij}=\frac{\exp \left(\frac{q_i{k}_j^T}{\sqrt{d_k}}\right)}{\sum _l\exp \left(\frac{q_i{k}_l^T}{\sqrt{d_k}}\right)}$

where ${\alpha }_{ij}$ represents the attention weight assigned by query $i$ to key-value pair $j$.

The weights satisfy:

$\sum _j{\alpha }_{ij}=1$

and are generally non-negative.

The output corresponding to query $i$ can therefore be represented as:

$o_i=\sum _j{\alpha }_{ij}{v}_j$

Thus, the output is a weighted combination of the value vectors.

Self-attention

Self-attention is a form of attention in which the queries, keys, and values are derived from the same input sequence.

A simplified representation is:

$X\to \{\begin{array}{c}Q=X{W}^Q\\ K=X{W}^K\\ V=X{W}^V\end{array}\to \text{Self-Attention}$

where:

• $X$ = input sequence representation;
• $W^Q$ = query projection matrix;
• $W^K$ = key projection matrix;
• $W^V$ = value projection matrix.

Self-attention allows each element of a sequence to incorporate information from other elements.

For example, in natural language processing, a token can assign greater attention to other tokens that provide relevant contextual information.

Cross-attention

Cross-attention is an attention mechanism in which the queries and the keys and values originate from different representations.

A simplified representation is:

$Q=X{W}^Q$

$K=Y{W}^K,V=Y{W}^V$

followed by:

$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$

Cross-attention allows one representation to selectively retrieve information from another representation.

It is used in architectures such as encoder–decoder transformers, where decoder representations can attend to representations produced by an encoder.

Multi-head attention

Multi-head attention applies multiple attention operations in parallel.

For each attention head:

${\text{head}}_i=\mathrm{Attention}(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$

The outputs are concatenated and projected:

$\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W^O$

where:

• $h$ = number of attention heads;
• $W_i^Q$, $W_i^K$, and $W_i^V$ = learned projection matrices for head $i$;
• $W^O$ = output projection matrix.

Different attention heads can learn different relationships or patterns within the input.

Causal attention

Causal attention restricts the positions that a token can attend to so that information from future positions is not used when predicting the current token.

For an autoregressive model, a simplified condition is:

${\alpha }_{ij}=0\text{for }j>i$

where $i$ represents the current position and $j$ represents an attended position.

This restriction allows the model to predict the next token using only previously available tokens.

Attention masking

Attention masks can be used to control which positions participate in an attention calculation.

A simplified masked attention operation is:

$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^T+M}{\sqrt{d_k}}\right)V$

where $M$ is an attention mask.

Elements of the mask can be assigned values that prevent particular positions from receiving attention.

Masks can be used for purposes such as causal attention or ignoring padding tokens.

Attention in transformers

Attention is a central component of the transformer architecture.

A simplified transformer attention process is:

$\text{Input Representations}\to Q,K,V\to \text{Attention Scores}\to \text{Attention Weights}\to \text{Weighted Values}\to \text{Output Representations}$

The attention mechanism allows transformer models to model relationships between different positions in a sequence without requiring recurrent processing.

Computational complexity

For standard self-attention, a sequence containing $n$ positions produces an attention score matrix with dimensions approximately $n\times n$.

Consequently, the number of pairwise attention relationships grows approximately as:

$O(n^2)$

with sequence length.

This quadratic scaling can create computational and memory challenges for very long sequences.

Various attention variants have therefore been developed to reduce the computational requirements of long-context processing.

Applications

Attention mechanisms are used in a wide range of machine-learning applications, including:

• natural language processing;
• machine translation;
• text generation;
• question answering;
• information retrieval;
• speech processing;
• computer vision;
• image generation;
• multimodal learning;
• sequence-to-sequence modeling.

Advantages

Attention mechanisms provide several useful characteristics:

• they allow models to selectively weight information;
• they can capture relationships between distant elements;
• they can be combined with parallel computation;
• they can be adapted to different input modalities;
• they provide a general mechanism for relating one representation to another.

Limitations

Attention mechanisms also have limitations.

Standard self-attention has computational and memory requirements that increase quadratically with sequence length.

Other limitations can include:

• high computational requirements for large inputs;
• substantial memory requirements;
• sensitivity to the number and dimensionality of attention heads;
• difficulty interpreting attention weights as direct explanations of model decisions;
• increased computational cost when processing very long sequences.

The practical characteristics depend on the architecture and implementation.

See also

• Transformer (machine learning model)
• Self-attention
• Multi-head attention
• Cross-attention
• Deep learning
• Natural language processing
• Large language model
• Softmax function
• Artificial neural network