Transformers

Attention solved a major sequence-modeling problem: every token could look directly at every other token. The next question is architectural.

Can we build an entire neural network around attention?

The answer is the transformer. A transformer is not just attention. It combines self-attention with positional information, feed-forward networks, residual connections, layer normalization, and many repeated blocks.

The Remaining Problem

Attention lets tokens communicate directly, but attention alone does not know word order. The sentences below contain the same words but mean very different things:

A transformer needs both direct token-to-token communication and information about position.

Think of a transformer layer as a meeting. Every token enters with its own meaning. Through self-attention, each token listens to the other tokens it finds relevant. Then a small neural network processes the updated meaning of each token.

After one layer, every token is more contextual. After many layers, each token has absorbed increasingly rich information from the whole sequence.

The Transformer Block

Attention mixes information between tokens. The feed-forward network transforms each token representation after that mixing. Residual connections and normalization keep the stack trainable.

Token Embeddings

Text first becomes tokens, and each token becomes a vector:

A sequence of $n$ tokens becomes a matrix:

Positional Encoding

Self-attention needs position information. The simplest view is that each token embedding receives a position vector:

Common choices include fixed sinusoidal positional encodings, learned positional embeddings, and rotary positional embeddings in many modern language models.

Self-Attention

Each token representation is projected into queries, keys, and values:

The attention operation is:

This produces contextual token representations. For example, ordered can attend to Jessie and burger so its representation includes who did the ordering and what was ordered.

Multi-Head Attention

Language contains many relationships at once: subject to verb, adjective to noun, pronoun to noun, cause to effect, and instruction to answer. Multi-head attention runs several attention mechanisms in parallel so different heads can learn different relationships.

Each head has its own query, key, and value projections. The head outputs are concatenated and projected back to the model dimension.

Feed-Forward Network

After attention mixes information across tokens, a small neural network transforms each token independently:

This distinction matters: attention mixes across tokens, while the feed-forward network processes each token's updated representation.

Residual Connections

Deep networks are easier to train when layers learn corrections:

The shortcut lets information and gradients flow through the network. If a layer is not useful yet, the model can keep the original representation and add only a small change.

Layer Normalization

As representations move through many layers, their scale can become unstable. Layer normalization stabilizes each token representation by normalizing across its feature dimension.

The intuition is simple: keep activations in a reasonable range so optimization does not spiral out of control.

Encoder, Decoder, And Causal Masking

Decoder-only transformers use causal masking. When predicting the next token, a token can attend only to itself and earlier tokens, not future tokens.

Why Transformers Beat RNNs

• They process all tokens in parallel during training.
• Any token can directly attend to any other token.
• Stacked attention layers learn rich contextual representations.

Computational Cost

Full self-attention compares every token with every other token. For sequence length $n$, the attention matrix has $n\times n$ entries:

This quadratic scaling is why long context windows are expensive.

This is a tiny decoder-style transformer block. It includes self-attention, a causal mask, residual connections, layer normalization, and a feed-forward network.

import torch
from torch import nn

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            batch_first=True,
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(self, x):
        seq_len = x.size(1)
        mask = torch.triu(
            torch.ones(seq_len, seq_len, device=x.device),
            diagonal=1,
        ).bool()

        attn_out, attn_weights = self.attn(
            x,
            x,
            x,
            attn_mask=mask,
            need_weights=True,
        )
        x = self.norm1(x + attn_out)

        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        return x, attn_weights

block = TransformerBlock(d_model=64, num_heads=4, d_ff=256)
tokens = torch.randn(8, 12, 64)
output, weights = block(tokens)

print(output.shape)   # (8, 12, 64)
print(weights.shape)  # attention matrix

A full language model stacks many of these blocks, adds token and positional embeddings before them, and adds an output layer that predicts the next token.