Building a Transformer From Scratch

machine learning
transformers
pytorch
A readable, from-scratch GPT-style Transformer in PyTorch — embeddings, causal attention, multi-head, MLP, residual blocks, training and generation.
Author

Utkarsh736

Published

September 9, 2026

Building a Transformer From Scratch

Transformers have become the foundation of modern language models, but using one through a library can make the underlying architecture feel surprisingly abstract.

In this notebook, I build a small GPT-style Transformer from scratch using PyTorch.

The goal isn’t to build a competitive language model. Instead, the goal is to understand what happens inside a Transformer by implementing its core components ourselves: embeddings, positional information, causal self-attention, multi-head attention, MLP layers, residual connections, and the final language modelling objective.

By the end, we’ll have a small autoregressive model that can be trained to predict the next token and generate text.

This is intentionally a simple first version. The implementation prioritizes readability and understanding over architectural fidelity or performance. Later iterations can improve the architecture, training setup, tokenization, and inference.

What we’re building

A GPT-style Transformer can be thought of as a pipeline:

Tokens
   ↓
Token Embeddings + Positional Embeddings
   ↓
Transformer Block × N
   ├── LayerNorm
   ├── Multi-Head Causal Attention
   └── MLP
   ↓
Final LayerNorm
   ↓
Vocabulary Projection
   ↓
Logits

We’ll build each component independently before assembling them into the complete model.

The notebook follows this order:

  1. Configuration
  2. Layer Normalization
  3. Token Embeddings
  4. Positional Embeddings
  5. Causal Masking
  6. Single-Head Attention
  7. Multi-Head Attention
  8. MLP
  9. Transformer Block
  10. Full Transformer
  11. Training
  12. Text Generation

Setup

The implementation uses PyTorch and a deliberately small number of dependencies.

Since the purpose of this notebook is to understand the architecture itself, most of the components are implemented directly rather than relying on high-level Transformer abstractions.

We also set a random seed to make experiments more reproducible and use a GPU when one is available.

import math
import random

import torch
import torch.nn as nn
import torch.nn.functional as F

from dataclasses import dataclass
device = "cuda" if torch.cuda.is_available() else "cpu"

print("PyTorch: ", torch.__version__)
print("Device: ", device)
PyTorch:  2.11.0+cpu
Device:  cpu
SEED = 42

random.seed(SEED)
torch.manual_seed(SEED)

if torch.cuda.is_available():
    torch.cuda.manual_seed(SEED)

Configuration

Before building the model, we define the architectural hyperparameters in one place.

Some of the most important parameters are:

  • vocab_size: number of possible tokens
  • context_length: maximum number of tokens the model can see at once
  • d_model: size of each token’s internal representation
  • n_heads: number of attention heads
  • d_head: dimensionality used by each attention head
  • d_mlp: hidden size of the feed-forward network
  • n_layers: number of Transformer blocks

For this first implementation, the model is intentionally small. The goal is to make the architecture easy to inspect and train rather than to build a powerful language model.

@dataclass
class Config:
  vocab_size: int = 1000
  context_length: int = 128

  d_model: int = 128

  n_heads: int = 4
  d_head: int = 32

  d_mlp: int = 512
  n_layers: int = 4

  dropout: float = 0.0
cfg = Config()

print(cfg)
Config(vocab_size=1000, context_length=128, d_model=128, n_heads=4, d_head=32, d_mlp=512, n_layers=4, dropout=0.0)
assert cfg.d_model == cfg.n_heads * cfg.d_head

Understanding the input

Before implementing the model, it’s useful to establish the shapes flowing through it.

The input to the model is a batch of token IDs:

(batch, sequence)

After embedding, every token becomes a vector:

(batch, sequence, d_model)

The Transformer processes these token representations through multiple layers while preserving this overall shape.

Finally, the model projects each token representation into vocabulary-sized logits:

(batch, sequence, vocab_size)

Thinking in terms of tensor shapes is one of the easiest ways to reason about Transformer implementations. Almost every component in the model can be understood by asking two questions:

  1. What information is this layer transforming?
  2. How does the tensor shape change?
batch_size = 2
seq_len = 10

tokens = torch.randint(
    0,
    cfg.vocab_size,
    (batch_size, seq_len)
)

print(tokens.shape)
torch.Size([2, 10])

Layer Normalization

Neural network activations can change significantly as they pass through many layers. Layer Normalization helps keep these representations numerically well-behaved.

For each token representation, we calculate its mean and variance across the feature dimension:

\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}

We then apply learnable scale and bias parameters:

y = \gamma \hat{x} + \beta

In this implementation, normalization happens across the d_model dimension independently for every token.

Modern GPT-style architectures commonly use Pre-LayerNorm, where normalization happens before the attention and MLP sublayers.

class LayerNorm(nn.Module):
  def __init__(self, d_model, eps=1e-5):
    super().__init__()
    self.gamma = nn.Parameter(torch.ones(d_model))
    self.beta = nn.Parameter(torch.zeros(d_model))
    self.eps = eps

  def forward(self, x):
    mean = torch.mean(x, dim=-1, keepdim=True)
    var = ((x-mean)**2).mean(dim=-1, keepdim=True)

    x_norm = (x - mean) / torch.sqrt(var + self.eps)

    op = self.gamma * x_norm + self.beta

    return op

Token Embeddings

The model receives integer token IDs, but integers themselves don’t contain useful semantic representations.

An embedding layer maps every token ID to a learned vector of size d_model.

Conceptually:

Token ID
   ↓
Embedding Lookup
   ↓
Vector of length d_model

After embedding, our input has the shape:

(B, T, d_model)

These vectors are learned during training, allowing the model to represent tokens in a continuous vector space.

class TokenEmbedding(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.embedding = nn.Embedding(
        cfg.vocab_size,
        cfg.d_model,
    )

  def forward(self, tokens):
    return self.embedding(tokens)
# test snippet
# Expected: [batch, sequence, d_model]

embedding = TokenEmbedding(cfg)
x = embedding(tokens)

print(x.shape)
torch.Size([2, 10, 128])

Positional Embeddings

Self-attention does not inherently know the order of tokens.

Without positional information, a sequence such as:

The cat sat

and a reordered version contain the same set of token representations.

To provide information about position, we learn an embedding for every position in the context window and add it to the token embedding:

x = \text{TokenEmbedding} + \text{PositionEmbedding}

This gives each token representation information about both:

  • what token it is
  • where it appears in the sequence
class PositionalEmbedding(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.embedding = nn.Embedding(
        cfg.context_length,
        cfg.d_model,
    )

  def forward(self, tokens):
    batch_size, seq_len = tokens.shape

    positions = torch.arange(
        seq_len,
        device=tokens.device,
    )
    return self.embedding(positions)
pos_embedding = PositionalEmbedding(cfg)

pos = pos_embedding(tokens)

print(pos.shape)
torch.Size([10, 128])

Causal Masking

A language model predicts the next token.

When predicting a token at position (t), the model should only have access to tokens at positions (0) through (t). It must not be allowed to look into the future.

This is enforced using a causal mask.

For a sequence of length four:

✓ ✗ ✗ ✗
✓ ✓ ✗ ✗
✓ ✓ ✓ ✗
✓ ✓ ✓ ✓

The lower-triangular structure allows each token to attend to itself and previous tokens while masking future tokens.

During attention, the disallowed positions are assigned -inf before the softmax operation. After softmax, those positions receive zero probability.

This is what makes the model autoregressive.

def causal_mask(seq_len, device):
  mask = torch.tril(torch.ones(seq_len, seq_len, device=device))

  return mask.bool()
# Viz
mask = causal_mask(5, device)

print(mask.int())
tensor([[1, 0, 0, 0, 0],
        [1, 1, 0, 0, 0],
        [1, 1, 1, 0, 0],
        [1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1]], dtype=torch.int32)

Single-Head Attention

Self-attention allows each token to dynamically gather information from other tokens in the sequence.

The input representation is projected into three different representations:

  • Query ((Q))
  • Key ((K))
  • Value ((V))

The attention scores are calculated using:

QK^T

These scores measure how strongly one token should attend to another.

The scores are scaled by:

\frac{1}{\sqrt{d_{head}}}

and then passed through a causal mask and softmax:

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

The resulting attention probabilities are then used to compute a weighted combination of the value vectors.

In other words, attention allows each token to decide which previous tokens are relevant when constructing its next representation.

class AttentionHead(nn.Module):
  def __init__(self, cfg):
    super().__init__()
    self.key = nn.Linear(cfg.d_model, cfg.d_head, bias=False)
    self.query = nn.Linear(cfg.d_model, cfg.d_head, bias=False)
    self.value = nn.Linear(cfg.d_model, cfg.d_head, bias=False)

    self.dropout = nn.Dropout(cfg.dropout)

  def forward(self, x):
    B, T, C = x.shape

    k = self.key(x)
    q = self.query(x)
    v = self.value(x)

    attn_scores = q @ k.transpose(-2, -1) /math.sqrt(cfg.d_head)
    mask = causal_mask(T, x.device)
    attn_scores = attn_scores.masked_fill(~mask, float("-inf"))
    attn_probs = F.softmax(attn_scores, dim=-1)

    attn_probs = self.dropout(attn_probs)

    return attn_probs @ v

Multi-Head Attention

A single attention mechanism produces one way of relating tokens to each other.

Multi-head attention allows the model to learn multiple attention patterns simultaneously.

Each attention head has its own Query, Key, and Value projections:

                Input
                  │
        ┌─────────┼─────────┐
        ↓         ↓         ↓
      Head 1    Head 2    Head N
        │         │         │
        └─────────┼─────────┘
                  ↓
             Concatenate
                  ↓
          Output Projection

In this first implementation, the heads are implemented independently using a ModuleList. This is easier to read and understand than a fully vectorized implementation.

A later version can combine all heads into larger matrix operations for improved efficiency.

class MultiHeadAttention(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.heads = nn.ModuleList(
        [AttentionHead(cfg) for _ in range(cfg.n_heads)]
    )

    self.proj = nn.Linear(cfg.n_heads * cfg.d_head, cfg.d_model)

    self.dropout = nn.Dropout(cfg.dropout)

  def forward(self, x):
    out = torch.cat([h(x) for h in self.heads], dim=-1)
    out = self.dropout(self.proj(out))

    return out

The Feed-Forward Network

Attention allows tokens to exchange information with one another.

The MLP, on the other hand, processes each token representation independently.

The representation is first expanded:

d_model → d_mlp

A non-linear GELU activation is applied, and the representation is projected back:

d_mlp → d_model

Conceptually:

Attention → communication between tokens

MLP → computation on each token representation

Both operations are important: attention determines where information comes from, while the MLP transforms that information.

class MLP(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.linear1 = nn.Linear(
        cfg.d_model,
        cfg.d_mlp
    )

    self.linear2 = nn.Linear(
        cfg.d_mlp,
        cfg.d_model,
    )

  def forward(self, x):

    x = self.linear1(x)

    # Apply activation

    x = F.gelu(x)

    x = self.linear2(x)

    return x

Putting the Components Together

A Transformer block combines attention and the MLP with Layer Normalization and residual connections.

The structure used here is:

x
│
├── LayerNorm
│       ↓
│   Attention
│       ↓
└───── (+)
        │
        ↓
    LayerNorm
        │
       MLP
        │
      (+)
        │
      Output

The residual connections allow the model to preserve information from earlier layers while learning incremental transformations.

This architecture is repeated multiple times, allowing the model to progressively build richer representations of the sequence.

class TransformerBlock(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.ln1 = LayerNorm(cfg.d_model)
    self.attn = MultiHeadAttention(cfg)

    self.ln2 = LayerNorm(cfg.d_model)
    self.mlp = MLP(cfg)

  def forward(self, x):

    # Attention + residual
    x = x + self.attn(self.ln1(x))

    # MLP + residual
    x = x + self.mlp(self.ln2(x))

    return x

The Complete Transformer

Now we can assemble the individual components into a complete autoregressive language model:

Token IDs
    ↓
Token Embeddings
    +
Positional Embeddings
    ↓
Transformer Block × N
    ↓
Final LayerNorm
    ↓
Vocabulary Projection
    ↓
Logits

The final linear layer maps the internal representation of each token back into the vocabulary space.

The output is not a token directly. Instead, the model produces a score, or logit, for every possible token in the vocabulary.

class Transformer(nn.Module):
  def __init__(self, cfg):
    super().__init__()

    self.cfg = cfg

    self.token_embedding = TokenEmbedding(cfg)
    self.position_embedding = PositionalEmbedding(cfg)

    self.blocks = nn.ModuleList(
        [TransformerBlock(cfg) for _ in range(cfg.n_layers)]
    )

    self.final_ln = LayerNorm(cfg.d_model)

    self.unembedding = nn.Linear(
        cfg.d_model,
        cfg.vocab_size,
        bias=False,
    )

  def forward(self, tokens):

    # Token embeddings
    x = self.token_embedding(tokens)

    # positional embeddings
    pos = self.position_embedding(tokens)

    x = x + pos

    # transformer blocks
    for block in self.blocks:
      x = block(x)

    # final layer norm
    x = self.final_ln(x)

    # Vocabulary logits
    logits = self.unembedding(x)

    return logits

Model Test

model = Transformer(cfg).to(device)

tokens = torch.randint(
    0,
    cfg.vocab_size,
    (2, 16)
).to(device)

logits = model(tokens)

print(logits.shape)
# Expected: torch.size([2, 16, vocab_size])
torch.Size([2, 16, 1000])

Check parameter count

num_params = sum(p.numel() for p in model.parameters())

print(f"Number of parameters: {num_params}")
Number of parameters: 1064192

Loss Function

Training Objective: Next-Token Prediction

The model is trained using next-token prediction.

Given:

The cat sat on

the input and target sequences are shifted by one position:

Input:   The → cat → sat → on

Target:  cat → sat → on → ...

For every position, the model predicts a probability distribution over the vocabulary.

Cross-entropy loss measures how much probability the model assigns to the correct next token.

During training, we flatten the batch and sequence dimensions so that PyTorch can treat every token prediction as an individual classification problem.

def compute_loss(logits, targets):
  B, T, C = logits.shape

  logits = logits.view(B * T, C)
  targets = targets.view(B * T)

  loss = F.cross_entropy(logits, targets)

  return loss

A Tiny Character-Level Dataset

For this first implementation, I use a very small character-level dataset.

This is not intended to produce high-quality language generation. Instead, it removes additional complexity from the experiment.

A character-level vocabulary allows us to focus on the Transformer architecture without introducing:

  • subword tokenization
  • large datasets
  • data pipelines
  • pretrained vocabularies

The trade-off is that the model has very limited training data and a simplistic representation of language.

text = """
Transformers are neural networks based on attention.
Attention allows tokens to interact with each other.
"""

Vocab

chars = sorted(set(text))

stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for ch, i in stoi.items()}

encode = lambda s: [stoi[c] for c in s]
decode = lambda l: "".join([itos[i] for i in l])

print(encode("hello"))
print(decode(encode("hello")))
[11, 9, 14, 14, 17]
hello

Dataset

data = encode(text)

def get_batch(
    data,
    batch_size,
    context_length,
    device,
):

  starts = torch.randint(
      0,
      len(data) - context_length,
      (batch_size,)
  )

  x = torch.stack(
      [torch.tensor(data[start: start + context_length]) for start in starts]
  )

  y = torch.stack(
      [torch.tensor(data[start + 1: start + context_length + 1]) for start in starts]
  )

  return x.to(device), y.to(device)

Training

The training loop follows the standard neural network workflow:

Sample Batch
     ↓
Forward Pass
     ↓
Calculate Loss
     ↓
Clear Gradients
     ↓
Backpropagation
     ↓
Optimizer Step

The model gradually adjusts its parameters to assign higher probability to the correct next token.

Because the dataset and model are small, this experiment is primarily useful for verifying that the architecture can learn and that gradients flow correctly through the entire network.

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

num_steps = 1000
for step in range(num_steps):
  model.train()

  x, y = get_batch(
      data=data,
      batch_size=16,
      context_length=16,
      device=device,
  )

  logits = model(x)

  loss = compute_loss(logits, y)

  optimizer.zero_grad(set_to_none=True)
  loss.backward()
  optimizer.step()


  if step % 100 == 0:
    print(f"Step {step} | Loss: {loss.item():.4f}")
Step 0 | Loss: 7.2028
Step 100 | Loss: 0.1983
Step 200 | Loss: 0.1069
Step 300 | Loss: 0.1660
Step 400 | Loss: 0.1371
Step 500 | Loss: 0.1377
Step 600 | Loss: 0.1383
Step 700 | Loss: 0.1241
Step 800 | Loss: 0.1439
Step 900 | Loss: 0.1236

Generating Text

After training, we can use the model autoregressively.

The generation process is:

Initial Prompt
      ↓
Predict next token
      ↓
Append predicted token
      ↓
Predict again
      ↓
Repeat

At every step, the model receives the existing context and predicts the next token.

For this first version, generation uses greedy decoding: the token with the highest predicted probability is selected at every step.

This is simple and deterministic, although more advanced generation strategies such as temperature sampling and top-k sampling can produce more varied outputs.

@torch.no_grad()
def generate(model, tokens, max_new_tokens, context_length):

  model.eval()

  for _ in range(max_new_tokens):

    tokens_cond = tokens[:, -context_length:]

    logits = model(tokens_cond)

    next_token_logits = logits[:, -1, :]

    # Greedy Predictions
    next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

    tokens = torch.cat((tokens, next_token), dim=1)

  return tokens
# Encode a short prompt
prompt = "Transformers are"
prompt_ids = torch.tensor([encode(prompt)], dtype=torch.long, device=device)

# Generate
generated = generate(
    model,
    prompt_ids,
    max_new_tokens=100,          # how many new characters to produce
    context_length=cfg.context_length
)

# Decode
print(decode(generated[0].tolist()))
Transformers are n ne arare ne ne ne nereral neural nere ns al ns nere nereuraral nerare neraloral ararareach al ner

Results and Limitations

The model successfully runs end-to-end:

  • tokens are embedded
  • positional information is added
  • causal attention prevents access to future tokens
  • multiple Transformer blocks process the sequence
  • the model learns through next-token prediction
  • generated tokens can be fed back into the model autoregressively

However, the generated text is understandably limited.

This first version uses a very small model and an extremely small character-level dataset. It is designed to demonstrate the mechanics of a Transformer rather than to generate coherent language.

There are also several architectural and training improvements that could make the implementation more capable and efficient.

Next Steps

Possible improvements include:

  • Better tokenization
  • A larger and more meaningful dataset
  • Train/validation splits
  • Improved sampling strategies
  • Dropout and regularization
  • Learning-rate scheduling
  • Vectorized multi-head attention
  • Efficient attention implementations
  • KV caching during generation
  • Comparing the implementation against a real GPT architecture
Back to top