{"slug": "building-a-transformer-from-scratch-part-1-the-embedding-layer", "title": "Building a Transformer from Scratch: Part 1 — The Embedding Layer", "summary": "A developer published a from-scratch PyTorch implementation of the embedding layer used in autoregressive Transformers like GPT-2, combining a learned token lookup table with learned positional embeddings via element-wise addition. The writeup explains why raw integer token IDs must be projected into continuous vectors, why nn.Embedding is a direct row lookup equivalent to one-hot matrix multiplication, and includes a guard clause that raises an error when sequence length exceeds the maximum context window.", "body_md": "In an autoregressive Transformer such as GPT, the work begins before the attention mechanism processes a single tensor. A model cannot operate on raw text, and integer token IDs carry no semantic structure on their own. The embedding layer bridges this gap by converting discrete token IDs into continuous, high-dimensional vectors that encode both the meaning of each token and its position in the sequence.\n\nThis article explains how the embedding layer works and presents a clean PyTorch implementation.\n\n**The Problem Embeddings Solve**\n\nA tokenizer splits raw text into segments and maps each one to an integer index. For example, the word \"transformer\" might map to `41551` (an illustrative value).\n\nPassing these integers directly into a neural network is problematic, because numerical values imply an ordinal relationship that does not exist. Token `41552` is not \"greater than\" token `41551` in any meaningful sense.\n\nThe solution uses two learned lookup tables:\n\n`W_e`)` d_model`. Through backpropagation, tokens that appear in similar linguistic contexts converge toward similar regions of the vector space.`W_p`)` 0, 1, ..., T-1` is therefore mapped to a learned vector of dimension The module below combines token embeddings with learned positional embeddings, following the architecture used in GPT-2.\n\n``` python\nimport torch\nimport torch.nn as nn\n\nclass TransformerEmbeddings(nn.Module):\n    def __init__(self, vocab_size: int, d_model: int, max_seq_len: int):\n        super().__init__()\n        # Token lookup matrix: shape (vocab_size, d_model)\n        self.token_embeddings = nn.Embedding(vocab_size, d_model)\n\n        # Position lookup matrix: shape (max_seq_len, d_model)\n        self.position_embeddings = nn.Embedding(max_seq_len, d_model)\n\n    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:\n        # Expected input shape: (batch_size, seq_len)\n        batch_size, seq_len = input_ids.shape\n\n        # Guard clause: protect against context window overflow\n        if seq_len > self.position_embeddings.num_embeddings:\n            raise ValueError(\n                f\"Sequence length ({seq_len}) exceeds maximum context window \"\n                f\"({self.position_embeddings.num_embeddings})\"\n            )\n\n        # Generate positional indices [0, 1, ..., seq_len - 1]\n        # Shape: (1, seq_len), placed on the same device as the input\n        positions = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0)\n\n        # Retrieve representations\n        tok_emb = self.token_embeddings(input_ids)      # (batch_size, seq_len, d_model)\n        pos_emb = self.position_embeddings(positions)   # (1, seq_len, d_model)\n\n        # Element-wise addition with broadcasting across the batch dimension\n        return tok_emb + pos_emb\n```\n\n`nn.Embedding` Works\n`nn.Embedding` is a trainable weight matrix:\n\n```\nW ∈ R^(vocab_size × d_model)\n```\n\nIt is mathematically equivalent to multiplying a one-hot vector by this matrix (`x_one_hot @ W`), but it is implemented as a direct row lookup. This avoids constructing large, sparse one-hot tensors and removes the associated memory and compute overhead.\n\n`device=input_ids.device` ensures that the position indices are created on the same device (CPU, CUDA, or Apple Silicon MPS) as the input. Omitting it defaults to the CPU and raises a device mismatch error when training on a GPU.`(1, seq_len, d_model)`, while `tok_emb` has shape `(batch_size, seq_len, d_model)`. When the two are added, PyTorch broadcasts the position vectors across every sequence in the batch.\nConcatenating the two vectors would widen each token representation to `2 × d_model`, increasing the parameter count and memory usage of every subsequent linear projection.\n\nElement-wise addition preserves the channel dimension at `d_model`. Conceptually, the positional embedding acts as an offset applied to the token vector, allowing the model to distinguish the same token at different positions.\n\nThe combined embeddings leave this layer with shape `(batch_size, seq_len, d_model)` and feed directly into the first Transformer block. The next article in this series covers the attention mechanism.", "url": "https://wpnews.pro/news/building-a-transformer-from-scratch-part-1-the-embedding-layer", "canonical_source": "https://dev.to/heywinterbell/building-a-transformer-from-scratch-part-1-the-embedding-layer-68k", "published_at": "2026-10-10 11:04:44+00:00", "updated_at": "2026-10-10 11:13:24.823029+00:00", "lang": "en", "topics": ["large-language-models", "neural-networks", "ai-research", "developer-tools"], "entities": ["PyTorch", "GPT-2", "Transformer"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/building-a-transformer-from-scratch-part-1-the-embedding-layer", "markdown": "https://wpnews.pro/news/building-a-transformer-from-scratch-part-1-the-embedding-layer.md", "text": "https://wpnews.pro/news/building-a-transformer-from-scratch-part-1-the-embedding-layer.txt", "jsonld": "https://wpnews.pro/news/building-a-transformer-from-scratch-part-1-the-embedding-layer.jsonld"}}