Overview

Suppose we're translating a story. The information we need for the next word may be somewhere in the source text, rather than in the last word we translated. How can a model revisit that information and decide what to use?

When we revisit a passage, we don't give every word equal attention. We focus on what seems relevant to the question at hand. An attention mechanism gives a model a way to learn a similar kind of selective focus. We'll follow one small calculation, then see how transformers use it to build contextual representations.

Outcomes

After completing this lesson, you'll be able to …

• describe the transformer architecture
• trace attention scores, weights, and weighted value vectors
• distinguish learned projection parameters from the representations they produce
• explain what causal and padding masks exclude
• connect positional information and bidirectional context to BERT

Prerequisites

We'll build on tensor shapes and transformations, including rows, columns, and transposing a matrix. The RNN encoder-decoder discussion provides the motivation for revisiting different parts of a sequence. We'll work through the attention calculation in small steps.

Motivation

As we saw previously, an RNN carries a learned "memory" of what it has seen in its hidden state. But carrying information forward in a state isn't the same as being able to revisit a particular part of the input.

Shortcomings of RNNs

Let's consider two reasons we might want something more than a recurrent update alone.

Memory and focus limitations

In a basic RNN encoder-decoder setup, we use the final encoder state to seed the decoder. In other words, we cram information from the source (potentially a very long sequence!) into a single fixed-size representation. What if the decoder could consult the encoder's representations at different positions as it makes each prediction?

That gives us access to more information, but we still need to decide what to use. What should we be paying attention to when making each prediction?

If you were to translate a story, you might revisit the source text at various times, focusing on a different aspect each time (ex. agreement, characters, and actions). Gates help an RNN learn what to retain or discard. An attention mechanism provides another route. It allows information from different positions to be selectively combined in a weighted fashion. Bahdanau et al. [2014] introduced this approach for neural machine translation.

Attention mechanisms are compatible with RNNs, and often they can improve performance. For example, Bahdanau et al. [2014] and Luong et al. [2015] report improvements in neural machine translation. The Bahdanau model combines a recurrent encoder and decoder with attention over source representations. The transformer takes a different step. It uses attention without sequence-aligned recurrent updates.

For a visual guide to attention mechanism, see the Let's Pay Attention Now section of Jay Alammar's Visualizing A Neural Machine Translation Model (Mechanics of Seq2seq Models With Attention).

Slow to train

Within an RNN sequence, we need the prior hidden state to compute the next one. We can process different examples in parallel, but that dependency remains along each sequence. Transformer layers can compute representations across positions in parallel during training. Generating a continuation still proceeds one token at a time, which we'll revisit in the LLM unit.

Attention

There are many formulations of attention, or attention mechanisms. They differ in how they calculate scores and which representations they combine. Here, we'll work through scaled dot product attention, the formulation used in the original transformer.

In scaled dot product attention, we compare each query with the available keys, turn the scores into weights, and use those weights to combine the values. The result is a new representation. It is not a vocabulary prediction.

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Softmax produces the weights first. We then multiply those weights by $V$. This is Equation 1 of Attention is All You Need.

Where do queries, keys, and values come from?

Suppose $X$ contains one representation per input token. A self-attention head uses learned projection matrices^† to produce three different views of those representations:

$$Q=XW^Q,\qquad K=XW^K,\qquad V=XW^V$$

The matrices $W^Q$, $W^K$, and $W^V$ are learned parameters^†. The resulting $Q$, $K$, and $V$ depend on the current input. Change the input and those representations can change, even when the model's parameters stay fixed.

Think of looking up information in a database. We submit a query, compare it with the keys associated with entries, and retrieve the corresponding values. In attention, a query helps determine what information to draw on, a key is what we compare it with, and a value is the information contributed to the result. These are three different roles^†. The analogy has a limit, though. Ordinary soft attention combines values according to their weights, rather than returning just one winning entry.

A small attention example

Let's consider a simple example:

clowns wear shoes

For this example, each word is one token. We'll supply small query, key, and value vectors so that we can inspect the arithmetic. These are illustrative numbers, rather than representations learned from text. Their coordinates don't stand for meanings such as "clownness" or shoe size.

The queries and keys happen to be equal here, but a real head can learn different projections for them.

1. Compare queries with keys

Let's follow wear. Its query is [0, 1]. Comparing it with the three keys gives dot products [0, 1, 1]. Each query and key has two features, so $d_k=2$. Dividing by $\sqrt{2}$ gives scores of approximately [0, 0.707, 0.707].

import math
import torch

tokens = ["clowns", "wear", "shoes"]
Q = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
K = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
V = torch.tensor([[2.0, 0.0], [0.0, 2.0], [1.0, 3.0]])

# One row per query position, one column per key position.
scores = Q @ K.T / math.sqrt(Q.shape[-1])
print(scores.shape)  # (3, 3)
print(scores[1])     # The query for "wear" compared with every key.

We can compute every query's scores in one matrix operation. Rows ask, columns offer information.

2. Turn scores into weights

The scores aren't probabilities yet. Softmax gives us nonnegative weights that sum to one across the available keys for each query.

As we saw in Tensor basics, dim=-1 means the last axis. Here, scores has shape (queries, keys), so the last axis runs over key positions. We normalize across those keys separately for each query, rather than across queries.

# Normalize each query's row across the key positions.
weights = torch.softmax(scores, dim=-1)
print(weights[1])    # Approximately [0.198, 0.401, 0.401].
print(weights.sum(dim=-1))  # Approximately [1, 1, 1].

The wear query gives equal weight to wear and shoes in this example. It still gives some weight to clowns. Soft attention doesn't have to pick just one token.

3. Combine the values

Now we use those weights to combine the value vectors:

$$0.198[2,0]+0.401[0,2]+0.401[1,3]\approx[0.797,2.005]$$

Those weights are rounded for display. Using the full values gives approximately [0.797, 2.006].

# Each query gets a weighted combination of the value vectors.
context = weights @ V
print(context.shape)  # (3, 2): query positions, value features.
print(context[1])     # Approximately [0.797, 2.006].

Notice what changed. The attention weights had one entry per key position. The new representation has one entry per value feature. Its entries need not sum to one.

Attention visual

Which positions may contribute?

The available context depends on the kind of attention we're using.

Where might we use each of these? Let's connect the available context to a task:

These aren't mutually exclusive choices for an entire model. In the original translation transformer [2017], the encoder uses self-attention over the source, while the decoder uses both causal self-attention over the output sequence and encoder-decoder attention over the source representations.

Predict the next token without peeking

Suppose we want a model to continue text. We train it to predict the next token using the text that precedes that token. During generation, we append each newly generated token to the context and predict another.

During training, we already have the complete example clowns wear shoes, but the representation at wear is used to predict shoes. If it could attend to the later shoes token, we'd be giving away the answer! It can use clowns and wear, but it mustn't read shoes before making that prediction. Training a language model to predict tokens from their preceding context is called autoregressive language modeling.

You may also have encountered the name causal language modeling^† for this setup. Here, "causal" refers to predicting from preceding context, without seeing future tokens.

Causal masking blocks later positions before softmax. For the wear query at position 1, shoes at position 2 is excluded. The query can still attend to clowns and wear itself. This is the restriction used in causal decoder self-attention, which we'll revisit in the decoder-only transformers section.

This isn't a rule for every kind of attention. When classifying a complete review, we can use context from both sides. When translating, we can attend to the full source sentence, even though the translation itself is generated one token at a time.

# Restrict decoder self-attention to the current and earlier positions.
# Give later key positions an additive score of negative infinity.
causal_bias = torch.tensor([
    [0.0, float("-inf"), float("-inf")],
    [0.0, 0.0, float("-inf")],
    [0.0, 0.0, 0.0],
])
causal_weights = torch.softmax(scores + causal_bias, dim=-1)
causal_context = causal_weights @ V

print(causal_weights[1])  # Approximately [0.330, 0.670, 0.000].
print(causal_context[1])  # Approximately [0.660, 1.340].

This snippet computes the masked attention representations, rather than a complete text generation model. For next-token training, the prediction made from each position is compared with the following token, not the token already at that position.

A padding mask excludes artificial padding positions, regardless of whether they are earlier or later than the query. These masks solve different problems^†. An implementation may combine them.

Transformer

The transformer introduced by Vaswani et al. [2017] combines attention with position information and ordinary neural network layers. The original architecture has an encoder and a decoder. BERT will use the encoder side of this family of ideas.

Encoding positions

Without information about order, an unmasked self-attention calculation relates the supplied representations without knowing which token came first. clowns wear shoes and shoes wear clowns contain the same words, but tell very different stories.

The original transformer adds positional encodings to token embeddings. Its fixed sine/cosine encoding varies both with position and with feature dimension. Each sine/cosine pair shares a frequency, but different pairs use different frequencies. The model receives the combined representation, which lets it use both token identity and position^†.

A position encoding visual

Our eight-feature encoding is a small example. The original transformer used 512 features [2017]. A position vector has the same size as the token embedding it is added to.

The same formula gives us the values for any position $p$:

$$\operatorname{PE}(p, 2i) = \sin\left(\frac{p}{10000^{2i/d}}\right), \qquad \operatorname{PE}(p, 2i+1) = \cos\left(\frac{p}{10000^{2i/d}}\right)$$

Here, $d$ is the number of features, and $i$ identifies a sine/cosine pair, starting at zero. With $d=8$, the four pairs divide the position by 1, 10, 100, and 1000. Those larger divisors give us slower waves^†.

We can encode more positions by continuing to supply larger values of $p$. A longer sequence doesn't require a longer position vector. Being able to compute those values, however, doesn't guarantee that a model will handle sequences longer than those it saw during training.

A single wave repeats, so periodicity alone isn't enough to distinguish positions. One useful building block is a sine/cosine pair. Together, the two features describe where a position falls within a cycle. Moving forward a fixed number of positions shifts both waves by a consistent amount along their cycles, wherever we start. This gives a model a regular relationship it can use to learn about relative positions, such as one token being a few places after another.

A pair still repeats, though. In the visual, the pairs move through their cycles at different rates. Their values together form a richer positional pattern. The researchers who introduced this idea hypothesized that this regular relationship would help a model learn to attend by relative positions [2017]. That is useful structure, rather than a guarantee that a trained model will use it successfully.

Multi-head attention

Why use more than one head? Each head has its own query, key, and value projections, giving the model different ways to compare and combine information^†. We concatenate the heads' outputs and apply a learned output projection $W^O$.

What else is in a block?

A transformer block contains several components, often called sublayers. Attention is one, and the feedforward network is another.

Attention mixes information across positions. The feedforward sublayer is a multilayer perceptron (MLP) applied separately at each position, using the same weights at every position. It combines features within each representation, rather than mixing information between positions. Residual connections^† provide a direct route around each of these components, and layer normalization helps control the scale of representations.

This follows the original encoder's ordering. Some transformer variants put normalization before the sublayer instead. For now, the important distinction is mixing information between positions, then transforming the information at each position.

Transformer-based models

A number of architectures known as transformer-based models have been proposed. Many of these use only an encoder (ex. BERT [2019] and its variants) or only a decoder (ex. the Generative Pre-trained Transformer [2018] (GPT-$n$) family).

Encoder-only transformers

BERT

Perhaps the most impactful variant of encoder-only transformers is BERT (Bidirectional Encoder Representations from Transformers) [2019] which uses a stack of encoder blocks. In each block, attention for each token is computed over the output of the previous block. Unlike autoregressive language models, BERT is bidirectional. For a token $t_{i}$, the attention calculation involves tokens $t_{0:i-1}$ (the left context) and $t_{i+1:n}$ (the right context).

Features of BERT base

• 12 layers of transformer blocks
• each block contains a multi-head attention layer comprised of 12 attention heads
• A subword vocabulary
• Originally 30,000 tokens generated using the WordPiece algorithm (Schuster and Nakajima, 2012) [2012]
• Gaussian Error Linear Unit (GELU) activation (Hendrycks and Gimpel, 2016)
• Available in PyTorch
• 110 million parameters

Pretraining

BERT uses two self-supervised pretraining tasks to learn general features of language.

Masked Language Modeling (MLM)

Masked Language Modeling is akin to a Cloze test. We select some token positions and ask the model to recover their original tokens using the surrounding context. In the original BERT recipe, 15% of token positions are selected. Of those, 80% are replaced with [MASK], 10% with a random token, and 10% are left unchanged. The prediction task applies to the selected positions, rather than every token. BERT can use both left and right context for these predictions.

Next Sentence Prediction (NSP)

In the original Next Sentence Prediction (NSP) task, BERT predicts whether the second text segment follows the first in the source document or was sampled from elsewhere. This was intended to help with relationships between segments. It is not an essential ingredient of every successful encoder. RoBERTa, for example, removes this objective.

Special tokens

BERT uses several special tokens. Like the ordinary subword tokens, they have learned embeddings. Here, [SEP] also marks the end of a segment. BERT doesn't add a separate [END] token.

Decoder-only transformers

Some variations of the transformer architecture such as GPT are decoder-only. These models are autoregressive (i.e., single direction models where each prediction only factors in prior predictions). Decoder-only transformers are a form of generative models. Text-based decoder-only transformers generate output one word (fragment) at a time. To predict the next token, the previously generated token is concatenated to the input^†.

Decoder-only transformers (such as those based on the GPT architecture) have received widespread attention for their ability to generate plausible sounding text using zero or few-shot in-context learning (ICL) in the form of prompts. Essentially, the model is fed a task description, optionally some examples of desired output, and an incomplete example. Given this input, the model generates a continuation (or completion) of the prompt. Each new token depends on the preceding context. Generating a long answer requires repeated prediction steps, rather than one encoder pass. That doesn't establish a universal cost ranking between encoder and decoder models. Size, sequence length, implementation, and the task all matter.

Adapting pretrained models

Foundation models

Transformers come in many sizes. Large pretrained models can be expensive to train from scratch, often relying on specialized hardware and large datasets. Fortunately, we can reuse what a model has already learned. A foundation model is trained on broad data and can be adapted to many downstream tasks. The term describes that role, rather than requiring one architecture or a fixed parameter count. For classification, we might add a classification head and fine-tune the model using our task's data.

Fine-tuning

Suppose we have a useful pretrained model, but want it to do better on our task. Full fine-tuning updates all its parameters. This is a valid approach, but it can be expensive for a large model. Fine-tuning can also change useful behavior learned during pretraining, so we need to evaluate more than the training loss. Another option is to freeze most parameters and train only a head or selected layers. That limits what we can change, but reduces the number of parameters we need to update.

Another efficient alternative to fine-tuning a subset of network layers is to train an adapter (a modular network component that is task-specific). An added adapter may introduce extra computation at inference time.

One approach to parameter-efficient fine-tuning is Low-Rank Adaptation of Large Language Models (LoRA)¹. Instead of updating a selected weight matrix directly, we freeze it and learn an update represented by two smaller matrices (see Figure 1). This reduces the number of trainable parameters and the storage needed for task-specific updates. The learned update can be merged into the original weights for inference without adding a separate computation. If we keep the update separate, its computation still has a cost. We'll return to model adaptation in the LLM unit.

Transformer variants

There are many variants of the transformer architecture. Most of the differences can be characterized in terms of a) different tokenization algorithms, b) encoder-only/decoder-only/encoder-decoder design, c) the number of model parameters, d) variations of pre-training tasks, or e) modality (text vs text+images). Below is a summary of several popular transformer-based architectures:

Practice

In the practice activity for this unit, we'll follow the same path: compare queries with keys, normalize each row, combine values, and consider which positions a mask allows. Start by predicting what should change, then use the code or attention visual to check your reasoning. The goal is to understand the flow of information, rather than memorize every box in the architecture diagram.

Resources

Primary sources

• Attention is All You Need (Vaswani et al., 2017)
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin et al., 2019)

Code examples

• The Annotated Transformer
• a walkthrough and PyTorch implementation of the transformer architecture

Visual introductions

• The Illustrated Transformer
• An animated walkthrough of the transformer from Jay Alammar

Textbook chapters

• Section 9.7 of Jurafsky and Martin's Speech and Language Processing
• Chapter 11 of Jurafsky and Martin's Speech and Language Processing

Libraries and implementations

The Hugging Face transformers library is a Python package that provides a unified interface to many popular transformer-based models, including support for fine-tuning these models for a variety of tasks (classification, translation, etc.). Hugging Face also maintains a Model Hub where users can share their models along with the datasets used to train them. The Hugging Face transformers API is compatible with PyTorch, allowing users to easily extend models that incorporate transformers in custom architectures.

Getting started with the Hugging Face transformers library

• Chapters 1-8 of the official Hugging Face course
• NOTE: Hugging Face estimates each chapter taking 6-8 hours to complete
• Natural Language Processing with Transformers (book)

References