How Attention Works in Large Language Models

Matt George
Matt George

Fractional CTO · Founder of Injuria.ai

The attention mechanism is the core innovation that enables modern large language models (LLMs) like GPT and Llama. It solves the fundamental problem of how a model can process a sequence of information, like text, while keeping track of which parts are most relevant to each other, regardless of their distance. Before attention, models struggled to connect words at the beginning of a long paragraph with words at the end. Attention provides a dynamic, learned method for creating these connections, making long-range understanding possible and paving the way for the deep language capabilities we see today.

This article explains how the attention mechanism works, from its historical context to the specific mathematical operations that power it. It details the practical costs and tradeoffs for technical leaders and corrects common misconceptions about its function. The goal is to provide a clear, accurate reference for engineers building with or around these models.

The Path to Attention

To appreciate attention, we must first understand the limitations of the models that preceded it. Through the mid-2010s, the dominant architecture for sequence processing was the Recurrent Neural Network[1] (RNN), and its more sophisticated variant, Long Short-Term Memory[2] (LSTM). RNNs process sequences one token at a time, maintaining an internal 'state' or 'memory' that is updated with each new token. This state is then passed to the next step.

This sequential nature created two major problems. First, it was slow. A sentence with 1000 words required 1000 sequential computation steps, a process that could not be easily parallelized on modern hardware like GPUs. Second, RNNs suffered from the vanishing gradient problem[3]. As information traveled through the sequence, its influence on the model's state would diminish. This made it extremely difficult for the model to learn connections between words that were far apart in the text.

The initial breakthrough came with sequence-to-sequence (seq2seq) models, often used for machine translation. A 2014 paper by Sutskever et al.[4] proposed using one LSTM (the encoder) to read an entire source sentence and compress it into a single, fixed-length vector. A second LSTM (the decoder) would then use this vector to generate the translated sentence. The problem was that this single vector became an information bottleneck. It was impossible to summarize a complex, 50-word sentence into one vector without losing crucial details.

Concurrent research that same year addressed this bottleneck. In a paper on machine translation, Bahdanau et al. (2014)[5] proposed the first influential version of an attention mechanism. Instead of forcing the encoder to produce a single context vector, they allowed the decoder to look back at the hidden states from every step of the encoder. At each step of generating an output word, the decoder would use attention to calculate a weighted average of the encoder states, focusing on the parts of the source sentence most relevant to the word it was about to produce. This was the key insight: not all information from the past is equally important, so the model should learn where to 'pay attention'.

The Mechanism: Scaled Dot-Product Attention

The modern form of attention, used in nearly all LLMs, was defined in the seminal 2017 paper "Attention Is All You Need" by Vaswani et al.[6]. This paper introduced the Transformer architecture and a highly optimized form of attention called Scaled Dot-Product Attention. This is the mechanism that powers models today.

The core idea is to determine the relevance of every token in a sequence to every other token. This is achieved by transforming each token's input embedding into three distinct vectors: a Query, a Key, and a Value.

  • Query (Q): Represents the current token's request for information. It's like asking, "Given who I am, what other tokens should I pay attention to?"
  • Key (K): Represents what information a token has to offer. It's like a label or index for a piece of information, saying, "This is the kind of content I represent."
  • Value (V): Represents the actual content of the token. This is the information that will be passed on if the Key is matched by a Query.

These Q, K, and V vectors are not fixed. They are generated by multiplying the input embedding for each token by three separate weight matrices (W_Q, W_K, W_V) that are learned during the model's training process. This means the model learns the best way to project its input tokens into these query, key, and value roles.

Let's walk through the steps for a single token looking for what to pay attention to.

  1. Calculate Scores: To find out how much attention our token's Query should pay to every other token, we calculate the dot product[7] of its Query vector with the Key vector of every other token in the sequence (including itself). A high dot product means the Query and Key are similar or aligned, indicating high relevance. A low dot product means they are not relevant to each other. This is performed for all tokens simultaneously using efficient matrix multiplication[8] (Scores = QKᵀ).

  2. Scale: The dot products can grow very large in magnitude, which pushes the subsequent softmax function into regions where its gradients are extremely small. This can stall the learning process. The Transformer paper's solution is to scale the scores down by dividing them by the square root of the dimension of the key vectors (d_k). For a model with a key dimension of 64, all scores would be divided by 8.

  3. Softmax: The scaled scores are then passed through a softmax function[9]. This converts the scores into a set of positive numbers that sum to 1.0. These numbers are the attention weights. They represent a probability distribution, indicating the percentage of attention the current token should give to every other token in the sequence.

  4. Weighted Sum of Values: Finally, the attention weights are used to create a weighted sum of all the Value vectors in the sequence. Tokens with high attention weights will have their Value vectors contribute significantly to the output, while tokens with low weights will have their Values mostly ignored. The result of this step is a new vector for our token, one that is a rich blend of information from all other tokens, weighted by their learned relevance.

This entire process happens for every token in the sequence in parallel, producing a new, context-aware representation for each one.

A Concrete Example

Imagine a tiny sequence of three tokens: "cat sat here". Let's simplify and assume their input embeddings and the learned Q, K, V vectors are very small. We'll focus on calculating the output for the second token, "sat".

import numpy as np

# Simplified for illustration. Real dimensions are much larger (e.g., 768+).
d_k = 2 # Dimension of Key/Query vectors

# Assume we have Q, K, V vectors for each token.
# Let's focus on the Query for "sat" (token 2) and all Keys.

# Query for "sat"
q_sat = np.array([0, 1])

# Keys for "cat", "sat", "here"
k_cat = np.array([1, 0])
k_sat = np.array([0, 1])
k_here = np.array([1, 1])

# Values for "cat", "sat", "here"
v_cat = np.array([1, 2])
v_sat = np.array([3, 4])
v_here = np.array([5, 6])

# Step 1: Calculate Scores for "sat"'s Query against all Keys
score_cat = np.dot(q_sat, k_cat) # 0*1 + 1*0 = 0
score_sat = np.dot(q_sat, k_sat) # 0*0 + 1*1 = 1
score_here = np.dot(q_sat, k_here) # 0*1 + 1*1 = 1
scores = np.array([0, 1, 1])

# Step 2: Scale
scaled_scores = scores / np.sqrt(d_k) # sqrt(2) is approx 1.414
# scaled_scores = [0 / 1.414, 1 / 1.414, 1 / 1.414] = [0, 0.707, 0.707]

# Step 3: Softmax
def softmax(x):
 return np.exp(x) / np.sum(np.exp(x), axis=0)

attention_weights = softmax(scaled_scores)
# attention_weights is approx [0.26, 0.37, 0.37]
# This means "sat" pays 26% attention to "cat", 37% to itself, and 37% to "here".

# Step 4: Weighted Sum of Values
output_sat = (attention_weights[0] * v_cat +
 attention_weights[1] * v_sat +
 attention_weights[2] * v_here)

# output_sat = (0.26 * [1, 2] + 0.37 * [3, 4] + 0.37 * [5, 6])
# output_sat = ([0.26, 0.52] + [1.11, 1.48] + [1.85, 2.22])
# output_sat = approx [3.22, 4.22]

# This new vector [3.22, 4.22] is the output for "sat".
# It's a context-aware representation that has blended information from the entire sequence.

This new output vector for "sat" now contains contextual information from "cat" and "here", and it becomes the input to the next layer of the model.

Multi-Head Attention

The authors of "Attention Is All You Need" found that a single attention calculation was not enough. A token might need to relate to other tokens in multiple ways simultaneously. For example, in the sentence "The dog chased the cat, it was fast", the word "it" needs to attend to "dog" for pronoun resolution, but "fast" might need to attend to "chased" to understand the action being described.

Multi-Head Attention solves this by running the Scaled Dot-Product Attention process multiple times in parallel. Each parallel run is called an "attention head". Crucially, each head uses its own separate, learned weight matrices (W_Q, W_K, W_V). This allows each head to learn to focus on different types of relationships, one head might learn syntactic dependencies, another might track semantic similarity, and another might focus on positional proximity.

After each head computes its output vector, the results are concatenated together and passed through a final linear projection layer. This combines the insights from all the different heads into a single, unified output vector for each token.

The Transformer Block and Model Architecture

The attention mechanism doesn't operate in a vacuum. It's the centerpiece of a larger component called a Transformer block[10]. A full LLM is typically a deep stack of these blocks.

A standard Transformer block contains two main sub-layers:

  1. A Multi-Head Attention mechanism.
  2. A position-wise Feed-Forward Network (a simple two-layer neural network applied independently to each token's representation).

Both of these sub-layers have residual connections and layer normalization around them. This helps stabilize training for very deep models. The output of one block becomes the input to the next, allowing the model to build up progressively more complex and abstract representations of the text.

It is also important to distinguish self-attention from cross-attention. The mechanism described so far is self-attention, where tokens within a single sequence attend to each other. This is the primary mechanism used in decoder-only models like the GPT series[11]. In encoder-decoder models (like the original Transformer), the decoder also employs cross-attention, where Query vectors from the decoder attend to the Key and Value vectors from the final output of the encoder. This is how the model consults the source text when generating its output.

One final piece is critical: attention itself is position-agnostic. The dot product Q * K has no inherent sense of whether a token is at the beginning or end of a sentence. To fix this, Transformers add positional encodings to the input embeddings. These are vectors that provide information about the absolute or relative position of tokens in the sequence, ensuring the model doesn't treat the text as just a bag of words.

The Decision Lens: Costs and Tradeoffs

For technical leaders, understanding the implications of the attention mechanism is crucial for making informed decisions about model selection, infrastructure, and application design.

The Quadratic Cost

The single most important tradeoff of the attention mechanism is its computational complexity. The calculation of the QKᵀ score matrix is the bottleneck. For a sequence of length n, this matrix has n * n = n² elements. This means that both the computational cost and memory requirement of attention scale quadratically with the sequence length. Doubling the context window from 2048 to 4096 tokens doesn't double the cost, it roughly quadruples it.

This O(n²) complexity is the primary reason LLMs have a finite context window. It makes processing very long documents or conversations extremely expensive. A model with a 100,000-token context window is performing calculations over a 100,000 x 100,000 matrix (in parts), which requires specialized hardware and optimization.

For Leaders: When evaluating a use case, ask: "What is the actual context length we need?" Using a massive context window for a task that only requires a few hundred tokens of context is a waste of computational resources and will increase latency and cost. Match the model's capability to the problem's real requirements.

Failure Modes

While powerful, attention is not infallible. Its failure modes are important to understand.

  • Difficulty with Very Long-Range Dependencies: Despite being better than RNNs, attention can still struggle with "needle in a haystack" problems where a single critical detail is located tens of thousands of tokens away from where it's needed. The attention signal can be diffused over such a long distance.
  • Factual Inaccuracy and Hallucination: Attention is a statistical pattern-matching mechanism. It can create highly plausible but factually incorrect outputs by attending to tokens that are thematically related but factually wrong. It finds correlations, not truth.
  • Over-reliance on Local Context: Sometimes, attention mechanisms can be biased towards attending to nearby tokens, even when the most relevant information is further away. Positional encodings help, but this bias can persist.

Common Misconceptions Corrected

Misconception 1: "Attention is how the model 'understands' or 'thinks'."

Correction: Attention is a sophisticated weighting mechanism, not a form of consciousness or human-like understanding. It allows the model to identify and prioritize parts of the input data that are statistically relevant for producing a likely output. It's a mathematical tool for correlation and information synthesis, not a reasoning engine.

Misconception 2: "Visualizing attention weights shows you exactly why the model made a decision."

Correction: Attention weights can be a useful debugging signal, but they are not a reliable or complete explanation of the model's behavior. Research, notably the 2019 paper "Attention is Not Explanation"[12], has shown that the final prediction is not always strongly correlated with the highest attention weights. The entire system, including residual connections, feed-forward layers, and the depth of the network, contributes to the final output in complex ways.

Where the Technology Stands Today

Attention is the undisputed foundation of modern AI. Every state-of-the-art LLM, including OpenAI's GPT series, Google's Gemini, Anthropic's Claude, and Meta's Llama, is built on a Transformer architecture with attention at its core. Its success has been so profound that it has expanded beyond NLP.

  • Vision Transformers (ViT): A 2020 paper[13] showed that by treating an image as a sequence of patches, a standard Transformer with attention could achieve state-of-the-art results in computer vision, challenging the dominance of Convolutional Neural Networks (CNNs).

  • Efficient Attention Variants: The O(n²) cost of standard attention is a major research focus. Numerous variants have been proposed to create more efficient approximations. Methods like Longformer[14] use a combination of sliding window attention and global attention to reduce complexity. More recently, FlashAttention[15] doesn't change the mathematical mechanism but reorders the computation to be I/O-aware, dramatically speeding up the process on GPUs by minimizing slow memory reads and writes.

  • Adoption: The concepts of attention are implemented in all major deep learning frameworks, including PyTorch[16] and TensorFlow[17]. Understanding how it works is no longer an academic exercise but a prerequisite for any engineer working seriously with modern AI models.

The attention mechanism was a paradigm shift. It replaced sequential processing with a parallelizable, highly effective method for learning relationships within data. While not without its costs and limitations, its invention directly enabled the powerful and flexible language models that are reshaping the technological landscape.

Sources (17)
  1. Recurrent neural network - Wikipedia
  2. Long short-term memory - Wikipedia
  3. Vanishing gradient problem - Wikipedia
  4. Sequence to Sequence Learning with Neural Networks (Sutskever et al., 2014)
  5. Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau et al., 2014)
  6. Attention Is All You Need (Vaswani et al., 2017)
  7. Dot product - Wikipedia
  8. Matrix multiplication - Wikipedia
  9. Softmax function - Wikipedia
  10. Transformer (machine learning model) - Wikipedia
  11. Language Models are Few-Shot Learners (GPT-3 Paper)
  12. Attention is Not Explanation (Jain & Wallace, 2019)
  13. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (ViT Paper)
  14. Longformer: The Long-Document Transformer (Beltagy et al., 2020)
  15. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)
  16. PyTorch
  17. TensorFlow