Transformer Attention Mechanism in NLP (original) (raw)

Last Updated : 9 Oct, 2025

Transformer model is a type of neural network architecture designed to handle sequential data primarily for tasks such as language translation, text generation and many more. Unlike traditional recurrent neural networks (RNNs) or convolutional neural networks (CNNs), Transformers uses attention mechanism to capture relationships between all words in a sentence regardless of their distance from each other.

The attention mechanism is a technique that allows models to focus on specific parts of the input sequence when producing each element of the output sequence. It assigns different weights to different input elements enabling the model to prioritize certain information over others. This is particularly useful in tasks like language translation where the meaning of a word often depends on its context. In this article we will learn about different types of Transformer’s attention mechanism.

Types of Attention Mechanisms in Transformers

1. Scaled Dot-Product Attention

The Scaled Dot-Product Attention is the fundamental building block of the Transformer's attention mechanism. It involves three main components: queries (Q), keys (K) and values (V). The attention score is computed as the dot product of the query and key vectors, scaled by the square root of the dimension of the key vectors. This score is then passed through a softmax function to obtain the attention weights which are used to compute a weighted sum of the value vectors.

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

where d_k is the dimension of the key vectors.

2. Multi-Head Attention

Multi-Head Attention enhances the model's ability to focus on different parts of the input sequence simultaneously. It involves multiple attention heads each with its own set of query, key and value matrices. The outputs of these heads are concatenated and linearly transformed to produce the final output. This allows the model to capture different features and dependencies in the input sequence.

Formula:

\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \ldots, \text{head}_h) W^O

where each \text{Attention}\big(Q W_i^{Q}, \; K W_i^{K}, \; V W_i^{V}\big) and W^O is the output matrix.

3. Self-Attention

Self-Attention is also known as intra-attention which allows the model to consider different positions of the same sequence when computing the representation of a word. In the context of the Transformer, self-attention is applied in both the encoder and decoder layers. It enables the model to capture long-range dependencies and relationships within the input sequence.

4. Encoder-Decoder Attention

Encoder-Decoder Attention also known as cross-attention, is used in the decoder layers of the Transformer. It allows the decoder to focus on relevant parts of the input sequence (encoded by the encoder) when generating each word of the output sequence. This type of attention ensures that the decoder has access to the entire input sequence, helping it produce more accurate and contextually appropriate translations.

5. Causal or Masked Self-Attention

Causal or Masked Self-Attention is used in the decoder to ensure that the prediction for a given position only depends on the known outputs at positions before it. This is crucial for tasks like language modeling where future tokens should not be visible during training. The attention scores for future tokens are masked out, ensuring that the model cannot look ahead.

Formula:

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

where M is the mask matrix with - \infty in positions that should be masked.

Significance of Transformer Attention Mechanism

The attention mechanism in Transformers offers several advantages:

  1. **Parallel Processing: Unlike RNNs, Transformers can process all words in a sequence simultaneously, significantly reducing training time.
  2. **Long-Range Dependencies: The attention mechanism can capture relationships between distant words, addressing the limitations of traditional models that struggle with long-range dependencies.
  3. **Scalability: Transformers can handle larger datasets and complex tasks due to their scalable architecture.

Implementation of Transformer Attention Mechanism

Step 1: Import Necessary Libraries

First import TensorFlow and other required libraries.

Python `

import tensorflow as tf from tensorflow.keras.layers import Layer, Dense, Dropout, LayerNormalization, Embedding import numpy as np

`

Step 2: Scaled Dot-Product Attention

The Scaled Dot-Product Attention mechanism is the foundational building block of the attention mechanism. It computes the attention scores based on the dot product of the query (Q) and key (K) vectors, scales it by the square root of the dimension of the key vectors and applies a softmax function to calculate the attention weights.

class ScaledDotProductAttention(Layer): def init(self): super(ScaledDotProductAttention, self).init()

def call(self, q, k, v, mask=None):
    matmul_qk = tf.matmul(q, k, transpose_b=True)
    dk = tf.cast(tf.shape(k)[-1], tf.float32)
    scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e9)

    attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
    output = tf.matmul(attention_weights, v)
    return output, attention_weights

`

Step 3: Multi-Head Attention

Define the MultiHeadAttention class. This class uses multiple attention heads to focus on different parts of the sequence simultaneously.

class MultiHeadAttention(Layer): def init(self, d_model, num_heads): super(MultiHeadAttention, self).init() self.num_heads = num_heads self.d_model = d_model assert d_model % num_heads == 0 self.depth = d_model // num_heads

    self.wq = Dense(d_model)
    self.wk = Dense(d_model)
    self.wv = Dense(d_model)
    self.dense = Dense(d_model)

def split_heads(self, x, batch_size):
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])

def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]

    q = self.wq(q)
    k = self.wk(k)
    v = self.wv(v)
    q = self.split_heads(q, batch_size)
    k = self.split_heads(k, batch_size)
    v = self.split_heads(v, batch_size)

    scaled_attention, attention_weights = ScaledDotProductAttention()(q, k, v, mask)

    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
    concat_attention = tf.reshape(
        scaled_attention, (batch_size, -1, self.d_model))
    output = self.dense(concat_attention)

    return output, attention_weights

`

Step 4: Testing the Encoder Layer and Encoder

To ensure that our Scaled Dot-Product Attention and Multi-Head Attention work correctly let’s test them with some random inputs:

q = tf.random.uniform((64, 50, 512)) k = tf.random.uniform((64, 50, 512)) v = tf.random.uniform((64, 50, 512)) mask = None attention_output, attention_weights = ScaledDotProductAttention()(q, k, v, mask)

print("Attention Output Shape:", attention_output.shape) multi_head_attention = MultiHeadAttention(d_model=512, num_heads=8) output, attn_weights = multi_head_attention(v, k, q, mask)

print("Multi-Head Attention Output Shape:", output.shape)

`

**Output:

Attention Output Shape: (64, 50, 512)
Multi-Head Attention Output Shape: (64, 50, 512)

Transformer’s attention mechanism is a key innovation that allows it to outperform traditional models on many NLP tasks. By using different types of attention like Scaled Dot-Product, Multi-Head, Self-Attention, Encoder-Decoder and Causal Attention the model can efficiently capture complex relationships between words in a sequence.