Understanding BERT's Multi-Head Attention Mechanism

Introduction to BERT

BERT (Bidirectional Encoder Representations from Transformers) is a pre-trained model designed for natural language understanding. It processes text by encoding it with a deep bidirectional context, allowing it to consider both left and right context in all layers. BERT consists of 12 transformer layers for BERT-base and 24 layers for BERT-large, utilizing only the transformer encoder to generate contextual representations of tokens.

Multi-Head Attention in BERT

To understand BERT's multi-head attention mechanism, consider the sentence: "The cat sat on the mat."

Self-Attention Mechanism

In self-attention, each token in the sentence computes how much focus it should give to every other token, including itself, to build its representation. Let's break it down:

Token Embeddings

The sentence is split into tokens: ["The", "cat", "sat", "on", "the", "mat"]. Each token is represented by a vector. For simplicity, let's denote them as T₁, T₂, T₃, T₄, T₅, T₆ respectively.

Creating Q, K, and V Vectors

For each token T_i, three vectors are created: Q_i (Query), K_i (Key), and V_i (Value) through learned linear projections.

Attention Calculation

Dot-Product Scores: Compute the dot product of the query vector of one token with the key vectors of all tokens to get attention scores. For instance, the attention scores for the token "cat" would be computed by taking the dot product of Q_cat with all key vectors K_The, K_cat, K_sat, K_on, K_the, K_mat.

Multi-Head Attention

Multi-head attention allows the model to focus on different parts of the sequence simultaneously, learning various contextual relationships between words. For example:

Head 1: Might focus on syntactic relationships, giving more importance to "sat" (since it's the verb associated with "cat").
Head 2: Might capture semantic relationships, emphasizing "mat" (since "cat" and "mat" are related in a spatial sense).

By employing multiple attention heads, BERT can capture a wide range of linguistic features, enhancing its understanding of language context.

Back to Blog