How is a transformer used in neural networks?

Transformers are a neural network (NN) architecture, or model, that excels at processing sequential data by weighing the importance of different parts of the input sequence. This allows them to capture long-range dependencies and context more effectively than previous architectures, leading to superior performance in natural language processing (NLP) tasks like translation and in computer vision systems.

Self-attention is the differentiator between transformers and previous NN architectures. It’s a way of calculating an “attention score” to identify the relationships between tokens, like words in a sentence or pixels in an image, even if they are not directly adjacent to each other. It enables each token to “pay attention” to every other token to identify and understand interrelationships.

Consider the sentence, “The dog dug a hole until it was deep enough, then he took his bone and buried it.” The word “it” has two meanings based on its position in the sentence. In the first instance, “it” refers to the hole, and in the second, “it” refers to the bone. A transformer model can efficiently capture those contextual nuances.

The technique is called a transformer NN because it uses transformation layers with matrix multiplications to transform the input sequence into the output sequence. It applies transformations to learn relationships within a dataset and produces a more meaningful representation.

Implementing self-attention

The advancement from using self-attention is the ability to mix information from pixels or words that are not directly adjacent and identify longer-range relationships and dependencies. For example, a 3×3 convolutional NN (CNN) mixes information of the nine data points, or tokens, directly around a central data point. That limits the ability to identify longer-range relationships.

The self-attention mechanism in a transformer NN analyzes tokens based on learned properties, not just based on location. That enables a transformer NN to learn and use more complex relationships (Figure 1).

Figure 1. Comparison of CNN learning (left) and transformer learning (right). (Image: Semiconductor Engineering)

How does a transformer NN learn?

Learning in a transformer NN is sometimes called relational learning because it learns relationships between tokens using “Query,” “Key,” and “Value” vectors that enhance the token’s positional encoding.

A token’s query vector is used to determine its alignment with the key vectors of all other tokens. The query vector acts essentially as a “question” that a token poses to the rest of the sequence, allowing the model to dynamically identify and integrate relevant contextual information from other tokens.

The key vector acts as a label for a token that’s used to compare tokens to determine interrelationships. For example, in the sentence, “The dog buried a bone, and she was very happy,” the key vector for “dog” would be compared to the query vector for “she.” The high similarity would lead to a high attention score, ensuring the value vector for “dog” influences the processing of “she,” establishing the pronoun’s reference.

Value vectors of the information that is being processed are combined to form a more contextualized representation of each token in the sequence. Unlike sequential models, the structure of transformer NNs inherently supports processing the entire sequence in parallel since the attention scores for all tokens can be computed simultaneously. That speeds up both training and inference.

The query, key, and value vectors, Q, K, and V in Figure 2, are generated by multiplying the initial token embedding matrix, X, using the W, K, and V weight matrices. The weight matrices in this feedforward NN are learned using self-supervised pretraining. Combined, the Q, K, and V vectors are called the attention layer.

Figure 2. Basic structure of a transformer attention mechanism. (Image: IBM)

Multi-heads are better than one

A multi-head attention structure uses different sets of query, key, and value vectors to focus on different aspects of the dataset. Each head implements the attention process independently. For example, one head might focus on semantic relationships, while another focuses on syntax.

The outputs from the various attention heads are concatenated together, and the combined output is processed through a linear transformation to produce the final output of the combined multi-head attention layer.

Summary

Transformer NNs are optimized for understanding long-range dependencies in data and processing information in parallel. Multi-head attention structures support more nuanced and comprehensive representations of the dataset by simultaneously considering multiple perspectives. That can lead to significantly improved performance in tasks like NLP translations, information summarization, and image recognition.