Building a Transformer from Scratch: Part 1 — The Embedding Layer

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.

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. This article explains how the embedding layer works and presents a clean PyTorch implementation. The Problem Embeddings Solve A 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 . Passing 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. The solution uses two learned lookup tables: 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. python import torch import torch.nn as nn class TransformerEmbeddings nn.Module : def init self, vocab size: int, d model: int, max seq len: int : super . init Token lookup matrix: shape vocab size, d model self.token embeddings = nn.Embedding vocab size, d model Position lookup matrix: shape max seq len, d model self.position embeddings = nn.Embedding max seq len, d model def forward self, input ids: torch.Tensor - torch.Tensor: Expected input shape: batch size, seq len batch size, seq len = input ids.shape Guard clause: protect against context window overflow if seq len self.position embeddings.num embeddings: raise ValueError f"Sequence length {seq len} exceeds maximum context window " f" {self.position embeddings.num embeddings} " Generate positional indices 0, 1, ..., seq len - 1 Shape: 1, seq len , placed on the same device as the input positions = torch.arange 0, seq len, device=input ids.device .unsqueeze 0 Retrieve representations tok emb = self.token embeddings input ids batch size, seq len, d model pos emb = self.position embeddings positions 1, seq len, d model Element-wise addition with broadcasting across the batch dimension return tok emb + pos emb nn.Embedding Works nn.Embedding is a trainable weight matrix: W ∈ R^ vocab size × d model It 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. 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. Concatenating the two vectors would widen each token representation to 2 × d model , increasing the parameter count and memory usage of every subsequent linear projection. Element-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. The 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.