Self-Attention: Modeling Long-Range Dependencies

Recommended resources

The freely available resources this lesson will direct you to.

Introduction

In our previous lesson, we meticulously derived the formula for scaled dot-product attention, establishing it as a mechanism for computing a context vector based on Query-Key-Value interactions. We focused on the how and why of the equation $Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V$ , particularly the statistical motivation for the scaling factor.

This lesson builds directly on that foundation to introduce one of the most transformative ideas in modern AI: self-attention. We will explore what happens when the queries, keys, and values all originate from the same input sequence. Our primary goal is to understand how this specific application of attention fundamentally solves the long-range dependency problem that constrained earlier architectures like RNNs. By the end of this lesson, you will be able to explain the theoretical advantages of self-attention's non-recurrent, parallelizable structure.

The Problem with Recurrence

Before we delve into the solution, let's briefly revisit the core architectural limitation of Recurrent Neural Networks (RNNs) that we discussed in Module 9. An RNN processes a sequence token by token, maintaining a hidden state that is passed from one time step to the next: $h_{t} = f (h_{t - 1}, x_{t})$ .

This sequential chain imposes two critical constraints:

Long Path Length: To relate two tokens that are far apart in a sequence, information and gradients must traverse all intermediate steps. For a sequence of length $n$ , the maximum path length is $O (n)$ . This long path makes it difficult to learn dependencies between distant tokens, as gradients can vanish or explode when propagated through many successive operations.
Sequential Computation: The calculation for time step $t$ depends on the result from $t - 1$ , making the process inherently sequential and difficult to parallelize over the time dimension.

The following video clip provides a concise recap of these challenges, which set the stage for the development of the Transformer architecture.

Attention is all you need (Transformer) - Model explanation (including math), Inference and Training

This segment from Umar Jamil's tutorial reviews the key problems with RNNs, particularly the difficulty in modeling long-range dependencies.

Watch the segment from this timestamp that discusses the "difficulty in accessing information from long time ago." This will refresh your memory on the specific problem that self-attention was designed to solve.

The fundamental question the Transformer's authors sought to answer was: can we model dependencies between tokens without relying on recurrence and its sequential nature?

Self-Attention: A Sequence Attending to Itself

Self-attention provides an elegant answer to this question. The core idea is to apply the attention mechanism we've already learned, but in a specific configuration where the queries, keys, and values are all derived from the same input sequence.

Given an input sequence represented by a matrix of token embeddings $X \in R^{n \times d_{m o d e l}}$ (where $n$ is sequence length and $d_{m o d e l}$ is the embedding dimension), we first create the Q, K, and V matrices using three distinct, learnable weight matrices $W_{Q}$ , $W_{K}$ , and $W_{V}$ :

Q = X W_{Q} K = X W_{K} V = X W_{V}

The self-attention output is then calculated using the familiar scaled dot-product attention formula:

SelfAttention (X) = softmax (\frac{Q K ^{T}}{d _{k}}) V

Each row of the output matrix is a new, context-aware representation for the corresponding input token. This new representation is a weighted sum of all value vectors in the sequence, where the weights are determined by the similarity between that token's query and every other token's key.

This diagram illustrates the self-attention mechanism. For a given input $x_2$, it generates a query $q_2$, a key $k_2$, and a value $v_2$. The query $q_2$ is then compared with all keys ($k_1, k_2, ..., k_T$) to compute attention scores, which become weights for all the values ($v_1, v_2, ..., v_T$). The final output for this position, $A_2$, is the weighted sum of these values.

This is a powerful concept. Instead of a token only "knowing" about its immediate predecessor's state, it can now directly look at and draw information from every other token in the sequence. To deepen your intuition, the following resource provides a rigorous definition and interpretation.

Attention Mechanisms in Neural Networks A ...

This chapter from a comprehensive monograph on attention provides a formal definition of self-attention and its mathematical properties.

Please read section 3.1, "Definition and Formulation". Pay close attention to how equations 3.1-3.4 define the Q, K, V projections from a single input matrix X and then apply the scaled dot-product formula. The interpretation subsection explains the roles of query, key, and value in this specific context.

The Solution to Long-Range Dependencies

How exactly does this architecture solve the problems of recurrence? The answer lies in the matrix multiplication $Q K^{T}$ . This single operation computes the dot product between every query vector (row in $Q$ ) and every key vector (row in $K$ , after transpose). The result is an $n \times n$ matrix of scores, where the entry at $(i, j)$ represents the interaction strength between token $i$ and token $j$ .

This has two profound consequences:

Constant Path Length: The path for information and gradient flow between any two positions $i$ and $j$ in the sequence is no longer dependent on their distance. It is a direct computation of length $O (1)$ . This direct connection allows the model to learn long-range dependencies as easily as short-range ones, completely sidestepping the vanishing gradient problem associated with long recurrent chains.
Parallelizability: The computation for each token's output does not depend on the prior computation of any other token. The entire set of matrix multiplications ( $Q K^{T}$ , followed by $A V$ ) can be executed in parallel, making self-attention vastly more efficient on modern hardware (like GPUs) compared to the inherently sequential nature of RNNs.

The following resource provides an excellent comparison of self-attention, RNNs, and CNNs across these critical architectural dimensions.

11.6. Self-Attention and Positional Encoding

This chapter from "Dive into Deep Learning" clearly compares the architectural properties of different sequence models.

Please read the section "Comparing CNNs, RNNs, and Self-Attention". The table and the accompanying text provide a concise but powerful summary of the trade-offs. Focus on the rows for "Sequential Operations" and "Maximum Path Length" to see the core theoretical advantage of self-attention.

A Consequence: Permutation Equivariance

The "all-to-all" connectivity of self-attention introduces a crucial property: permutation equivariance. Without any additional information, the self-attention mechanism is insensitive to the order of the input sequence. If you were to shuffle the input tokens, the output vectors would be identical, just shuffled in the same order. This is because the attention mechanism computes interactions based on the content of the tokens, not their positions.

From a theoretical standpoint, for any permutation matrix $P$ , it holds that:

SelfAttention (P X) = P \cdot SelfAttention (X)

This means that self-attention, in its pure form, treats the input as a "bag of tokens" or a set, not a sequence. The formal proof of this property relies on the properties of matrix multiplication and is quite elegant.

Attention Mechanisms in Neural Networks A ...

This section of the arXiv paper provides the formal mathematical proof for the permutation equivariance of self-attention.

For a deeper theoretical understanding, review section 3.2.1, "Permutation Equivariance". The derivation walks through how permuting the input matrix X leads to a corresponding permutation of the output, demonstrating the mechanism's inherent order-agnosticism.

This property might seem like a flaw, as order is critical in language and many other sequential tasks. However, it simply means that information about position must be explicitly provided to the model. This is the motivation behind positional encodings, which are vectors added to the input embeddings to give the model a sense of sequence order. We will explore the details of positional encodings when we assemble the full Transformer architecture in a later lesson.

Conclusion

In this lesson, we have defined self-attention and, more importantly, understood its theoretical superiority over recurrent models for capturing long-range dependencies.

Key takeaways from our discussion include:

Self-attention is a specific application of the attention mechanism where the Query, Key, and Value vectors are all projected from the same input sequence.
It creates a direct, $O (1)$ path between any two tokens in a sequence, allowing the model to learn long-range dependencies without the signal degradation or vanishing gradients associated with the $O (n)$ path length in RNNs.
The computation is highly parallelizable, offering significant computational advantages over the sequential nature of recurrence.
A fundamental property of self-attention is permutation equivariance, meaning it is inherently order-agnostic. This necessitates the use of positional encodings to inject sequence order information.

You now grasp how a single self-attention "head" operates. However, the relationships within a sequence can be multifaceted. For example, one token might relate to another syntactically, while simultaneously relating to a different token semantically. Forcing a single attention mechanism to capture all these different relationships could be a bottleneck. In our next lesson, we will explore multi-head attention, a technique that allows the model to learn diverse types of relationships in parallel.